5698 lines
249 KiB
Python
5698 lines
249 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||
|
||
"""
|
||
llama-server inference backend for GGUF models.
|
||
|
||
Manages a llama-server subprocess and proxies chat completions
|
||
through its OpenAI-compatible /v1/chat/completions endpoint.
|
||
"""
|
||
|
||
import atexit
|
||
import contextlib
|
||
import json
|
||
import os
|
||
import re
|
||
import struct
|
||
import structlog
|
||
from loggers import get_logger
|
||
import shutil
|
||
import signal
|
||
import socket
|
||
import subprocess
|
||
import sys
|
||
import threading
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Generator, Iterable, List, Optional
|
||
from urllib.parse import urlparse
|
||
|
||
import httpx
|
||
|
||
from core.inference.llama_server_args import (
|
||
parse_cache_override,
|
||
parse_ctx_override,
|
||
resolve_cache_type_kv,
|
||
resolve_requested_ctx,
|
||
)
|
||
from core.tool_healing import (
|
||
_TC_END_TAG_RE,
|
||
_TC_FUNC_CLOSE_RE,
|
||
_TC_FUNC_START_RE,
|
||
_TC_JSON_START_RE,
|
||
_TC_PARAM_CLOSE_RE,
|
||
_TC_PARAM_START_RE,
|
||
_TOOL_ALL_PATS,
|
||
_TOOL_CLOSED_PATS,
|
||
parse_tool_calls_from_text,
|
||
strip_tool_call_markup,
|
||
)
|
||
from utils.native_path_leases import child_env_without_native_path_secret
|
||
from utils.subprocess_compat import (
|
||
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
|
||
)
|
||
from core.inference.tool_call_parser import (
|
||
parse_tool_calls_from_text as _shared_parse_tool_calls_from_text,
|
||
)
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
# ── Pre-compiled patterns for plan-without-action re-prompt ──
|
||
# Forward-looking intent signals that indicate the model is
|
||
# describing what it *will* do rather than giving a final answer.
|
||
_INTENT_SIGNAL = re.compile(
|
||
r"(?i)("
|
||
# Direct intent: "I'll ...", "I will ...", "Let me ...", "I am going to ..."
|
||
# Handles both straight and curly apostrophes.
|
||
# Excludes "I can", "I should", "I want to", "let's" which
|
||
# appear frequently in direct answers / explanations.
|
||
# Negative lookahead drops negated forms ("I will not", "I'll never")
|
||
# so a refusal doesn't trigger a re-prompt.
|
||
r"\b(i['\u2019](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)"
|
||
r"|"
|
||
# Step/plan framing: "First ...", "Step 1:", "Here's my plan"
|
||
r"\b(?:first\b|step \d+:?|here['\u2019]?s (?:my |the |a )?(?:plan|approach))"
|
||
r"|"
|
||
# "Now I" / "Next I" patterns
|
||
r"\b(?:now i|next i)\b"
|
||
r")"
|
||
)
|
||
_MAX_REPROMPTS = 3
|
||
|
||
# Without max_tokens, llama-server defaults to n_predict = n_ctx (up to
|
||
# 262144 for Qwen3.5), producing many-minute zombie decodes when cancel
|
||
# fails. t_max_predict_ms is a wall-clock backstop applied unconditionally,
|
||
# but the llama.cpp README notes it ONLY fires after a newline has been
|
||
# generated -- a model stuck in a long unbroken non-newline sequence is
|
||
# unbounded by it. So we still want a token cap as the front-line limiter.
|
||
#
|
||
# The cap is the model's effective context length when we know it,
|
||
# falling back to a generous floor when metadata is unavailable. 4096 was
|
||
# too low: Qwen3 / gpt-oss reasoning traces routinely exceed it, and any
|
||
# OpenAI-API caller that omits max_tokens (langchain, llama-index, raw
|
||
# curl) sees responses silently truncated mid-sentence.
|
||
_DEFAULT_MAX_TOKENS_FLOOR = 32768
|
||
_DEFAULT_T_MAX_PREDICT_MS = 600_000 # 10 min
|
||
_REPROMPT_MAX_CHARS = 2000
|
||
|
||
# ── Pre-compiled patterns for GGUF shard detection ───────────
|
||
_SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$")
|
||
_SHARD_RE = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$")
|
||
|
||
|
||
# ── Sliding-window-pattern resolver ───────────────────────────
|
||
# Resolves the per-layer SWA mask when a GGUF reports a sliding window
|
||
# but no `sliding_window_pattern` field. Tier order in
|
||
# `_resolve_swa_pattern`: GGUF metadata, on-disk cache, bootstrap dict
|
||
# below, transformers introspection, HF Hub config.json, legacy 1/4
|
||
# fallback. Period N means layer i is SWA iff `(i + 1) % N != 0`,
|
||
# matching transformers. Skipped on purpose: phi3 (no key/val length
|
||
# in GGUF, window >= ctx anyway), qwen2 family (converter strips
|
||
# sliding_window when use_sliding_window=False), mistral v0.1/v0.2
|
||
# (all-SWA can't be expressed as a period).
|
||
_BOOTSTRAP_SWA_DEFAULTS: dict[str, int] = {
|
||
"gemma2": 2, # Gemma2Config.sliding_window_pattern
|
||
"gemma3": 6, # Gemma3TextConfig.sliding_window_pattern
|
||
"gemma3n": 5, # text_config.layer_types: SWA*4 + FULL
|
||
"gpt_oss": 2, # text_config.layer_types: alternating
|
||
"cohere2": 4, # Cohere2Config.sliding_window_pattern
|
||
}
|
||
|
||
# Process-wide cache backed by JSON on disk. Values are int period or
|
||
# list[bool] mask. Lazy-loaded.
|
||
_SWA_CACHE: Optional[dict] = None
|
||
_SWA_CACHE_LOCK = threading.Lock()
|
||
|
||
|
||
def _probe_dns_dead(host: str = "huggingface.co", timeout: float = 2.0) -> bool:
|
||
"""Quick DNS check. Runs on a daemon thread so concurrent sockets
|
||
in the same process are not affected by socket.setdefaulttimeout."""
|
||
result: list[Optional[bool]] = [None]
|
||
|
||
def _probe() -> None:
|
||
try:
|
||
socket.gethostbyname(host)
|
||
result[0] = False
|
||
except Exception:
|
||
result[0] = True
|
||
|
||
t = threading.Thread(target = _probe, daemon = True)
|
||
t.start()
|
||
t.join(timeout)
|
||
# Thread still running -> resolver wedged -> treat as dead.
|
||
return True if result[0] is None else result[0]
|
||
|
||
|
||
@contextlib.contextmanager
|
||
def _hf_offline_if_dns_dead():
|
||
"""Set HF_HUB_OFFLINE for the body of this block only when DNS to
|
||
huggingface.co fails. Restores the env on exit so a transient
|
||
resolver hiccup at the start of one load can't quarantine the whole
|
||
process. Respects an explicit user setting (no-op if already set)."""
|
||
if "HF_HUB_OFFLINE" in os.environ:
|
||
yield False
|
||
return
|
||
if not _probe_dns_dead():
|
||
yield False
|
||
return
|
||
|
||
transformers_was_set = "TRANSFORMERS_OFFLINE" in os.environ
|
||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||
if not transformers_was_set:
|
||
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
||
logger.warning("huggingface.co unreachable; using local HF cache for this load.")
|
||
try:
|
||
yield True
|
||
finally:
|
||
os.environ.pop("HF_HUB_OFFLINE", None)
|
||
if not transformers_was_set:
|
||
os.environ.pop("TRANSFORMERS_OFFLINE", None)
|
||
|
||
|
||
def _swa_cache_path() -> Path:
|
||
home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME")
|
||
base = Path(home) if home else Path.home() / ".unsloth" / "studio"
|
||
return base / "swa_cache.json"
|
||
|
||
|
||
def _load_swa_cache() -> dict:
|
||
global _SWA_CACHE
|
||
with _SWA_CACHE_LOCK:
|
||
if _SWA_CACHE is not None:
|
||
return _SWA_CACHE
|
||
try:
|
||
with open(_swa_cache_path()) as f:
|
||
_SWA_CACHE = json.load(f)
|
||
if not isinstance(_SWA_CACHE, dict):
|
||
_SWA_CACHE = {}
|
||
except (FileNotFoundError, json.JSONDecodeError, OSError):
|
||
_SWA_CACHE = {}
|
||
return _SWA_CACHE
|
||
|
||
|
||
def _save_swa_cache(cache: dict) -> None:
|
||
try:
|
||
path = _swa_cache_path()
|
||
path.parent.mkdir(parents = True, exist_ok = True)
|
||
tmp = path.with_suffix(".json.tmp")
|
||
with open(tmp, "w") as f:
|
||
json.dump(cache, f, indent = 2, sort_keys = True)
|
||
tmp.replace(path)
|
||
except OSError:
|
||
pass
|
||
|
||
|
||
def _period_from_layer_types(layer_types: list) -> Optional[int]:
|
||
"""Smallest period N where `(i+1) % N != 0` matches the SWA mask,
|
||
or None if no fixed period fits."""
|
||
if not layer_types:
|
||
return None
|
||
is_swa = ["full" not in str(t).lower() for t in layer_types]
|
||
n = len(is_swa)
|
||
for N in range(1, n + 1):
|
||
if all(((i + 1) % N != 0) == is_swa[i] for i in range(n)):
|
||
return N
|
||
return None
|
||
|
||
|
||
def _fetch_swa_entry_from_hf(repo_id: str) -> Optional[object]:
|
||
try:
|
||
from huggingface_hub import hf_hub_download
|
||
|
||
cfg_path = hf_hub_download(repo_id, "config.json", repo_type = "model")
|
||
with open(cfg_path) as f:
|
||
cfg = json.load(f)
|
||
except Exception:
|
||
return None
|
||
|
||
src = cfg.get("text_config") if isinstance(cfg.get("text_config"), dict) else cfg
|
||
period = src.get("sliding_window_pattern")
|
||
if isinstance(period, int) and period > 0:
|
||
return period
|
||
lt = src.get("layer_types")
|
||
if isinstance(lt, list) and lt:
|
||
return _period_from_layer_types(lt) or [
|
||
"full" not in str(t).lower() for t in lt
|
||
]
|
||
return None
|
||
|
||
|
||
def _arch_aliases(arch: str) -> tuple:
|
||
# GGUF emits `falcon-h1`; HF model_type is `falcon_h1`. Normalise both ways.
|
||
seen = []
|
||
for a in (arch, arch.replace("-", "_"), arch.replace("_", "-")):
|
||
if a and a not in seen:
|
||
seen.append(a)
|
||
return tuple(seen)
|
||
|
||
|
||
def _swa_entry_from_config_obj(cfg) -> Optional[object]:
|
||
src = getattr(cfg, "text_config", None) or cfg
|
||
period = getattr(src, "sliding_window_pattern", None)
|
||
if isinstance(period, int) and period > 0:
|
||
return period
|
||
lt = getattr(src, "layer_types", None)
|
||
if isinstance(lt, list) and lt:
|
||
return _period_from_layer_types(lt) or [
|
||
"full" not in str(t).lower() for t in lt
|
||
]
|
||
return None
|
||
|
||
|
||
_SWA_PATTERN_SOURCE_RE = re.compile(
|
||
r"sliding_window_pattern\s*(?::\s*[\w\[\], ]*)?\s*=\s*(\d+)"
|
||
)
|
||
|
||
|
||
def _resolve_swa_entry_from_transformers(arch: str) -> Optional[object]:
|
||
"""Default-instantiate the matching Config; on failure, regex-parse
|
||
its source for `sliding_window_pattern = N`."""
|
||
try:
|
||
from transformers.models.auto.configuration_auto import (
|
||
CONFIG_MAPPING,
|
||
CONFIG_MAPPING_NAMES,
|
||
)
|
||
except Exception:
|
||
return None
|
||
|
||
cfg_class = None
|
||
for alias in _arch_aliases(arch):
|
||
if alias in CONFIG_MAPPING_NAMES:
|
||
try:
|
||
cfg_class = CONFIG_MAPPING[alias]
|
||
break
|
||
except Exception:
|
||
cfg_class = None
|
||
if cfg_class is None:
|
||
return None
|
||
|
||
try:
|
||
if (entry := _swa_entry_from_config_obj(cfg_class())) is not None:
|
||
return entry
|
||
except Exception:
|
||
pass
|
||
|
||
import inspect
|
||
|
||
candidates = [cfg_class]
|
||
text_cfg_class = getattr(cfg_class, "sub_configs", {}).get("text_config")
|
||
if text_cfg_class is not None:
|
||
candidates.append(text_cfg_class)
|
||
for cls in candidates:
|
||
try:
|
||
src = inspect.getsource(cls)
|
||
except (OSError, TypeError):
|
||
continue
|
||
if m := _SWA_PATTERN_SOURCE_RE.search(src):
|
||
period = int(m.group(1))
|
||
if period > 0:
|
||
return period
|
||
return None
|
||
|
||
|
||
def _resolve_swa_pattern(
|
||
arch: Optional[str],
|
||
n_layers: Optional[int],
|
||
source_repo_candidates: tuple = (),
|
||
*,
|
||
allow_network: Optional[bool] = None,
|
||
) -> Optional[list]:
|
||
if not arch or not n_layers:
|
||
return None
|
||
if allow_network is None:
|
||
allow_network = os.environ.get("UNSLOTH_STUDIO_OFFLINE", "0") not in (
|
||
"1",
|
||
"true",
|
||
"True",
|
||
"yes",
|
||
)
|
||
|
||
cache = _load_swa_cache()
|
||
|
||
def _entry_to_mask(entry):
|
||
if isinstance(entry, int) and entry > 0:
|
||
return [(i + 1) % entry != 0 for i in range(n_layers)]
|
||
if isinstance(entry, list) and entry:
|
||
return [bool(entry[i % len(entry)]) for i in range(n_layers)]
|
||
return None
|
||
|
||
def _persist(entry):
|
||
with _SWA_CACHE_LOCK:
|
||
cache[arch] = entry
|
||
_save_swa_cache(cache)
|
||
|
||
if (entry := cache.get(arch)) is not None:
|
||
if (mask := _entry_to_mask(entry)) is not None:
|
||
return mask
|
||
|
||
if (entry := _BOOTSTRAP_SWA_DEFAULTS.get(arch)) is not None:
|
||
return _entry_to_mask(entry)
|
||
|
||
entry = _resolve_swa_entry_from_transformers(arch)
|
||
if entry is not None:
|
||
_persist(entry)
|
||
return _entry_to_mask(entry)
|
||
|
||
# Tier 3: live HF fetch (with persistent caching of the result)
|
||
if allow_network:
|
||
for repo_id in source_repo_candidates:
|
||
if not repo_id:
|
||
continue
|
||
entry = _fetch_swa_entry_from_hf(repo_id)
|
||
if entry is not None:
|
||
_persist(entry)
|
||
return _entry_to_mask(entry)
|
||
|
||
return None
|
||
|
||
|
||
def _hf_repo_from_url(url: Optional[str]) -> Optional[str]:
|
||
"""Strip `https://huggingface.co/owner/name(/...)` to `owner/name`."""
|
||
if not url or "huggingface.co/" not in url:
|
||
return None
|
||
tail = url.split("huggingface.co/", 1)[1].rstrip("/")
|
||
parts = tail.split("/")
|
||
if len(parts) < 2:
|
||
return None
|
||
return f"{parts[0]}/{parts[1]}"
|
||
|
||
|
||
# Model size extraction — lazy import to avoid pulling in transformers
|
||
# at module level. See PR description for the full explanation.
|
||
def _extract_model_size_b(model_id: str):
|
||
from utils.models import extract_model_size_b
|
||
|
||
return extract_model_size_b(model_id)
|
||
|
||
|
||
_TOOL_TEMPLATE_MARKERS = (
|
||
"{%- if tools %}",
|
||
"{%- if tools -%}",
|
||
"{% if tools %}",
|
||
"{% if tools -%}",
|
||
'"role" == "tool"',
|
||
"'role' == 'tool'",
|
||
'message.role == "tool"',
|
||
"message.role == 'tool'",
|
||
)
|
||
|
||
|
||
def detect_reasoning_flags(
|
||
chat_template: Optional[str],
|
||
model_identifier: Optional[str] = None,
|
||
*,
|
||
log_source: Optional[str] = None,
|
||
) -> dict:
|
||
"""Classify a chat template's reasoning and tool-calling capabilities.
|
||
|
||
Returns a dict with the same five keys populated by the GGUF sniffer:
|
||
``supports_reasoning``, ``reasoning_style``
|
||
(``"enable_thinking"`` | ``"reasoning_effort"``),
|
||
``reasoning_always_on``, ``supports_preserve_thinking``, and
|
||
``supports_tools``. Used by both the llama-server backend at load
|
||
time and the safetensors/transformers paths in ``routes/inference``
|
||
so the two agree on what the frontend will see.
|
||
"""
|
||
flags = {
|
||
"supports_reasoning": False,
|
||
"reasoning_style": "enable_thinking",
|
||
"reasoning_always_on": False,
|
||
"supports_preserve_thinking": False,
|
||
"supports_tools": False,
|
||
}
|
||
if not chat_template:
|
||
return flags
|
||
tpl = chat_template
|
||
prefix = f"{log_source}: " if log_source else ""
|
||
|
||
if "enable_thinking" in tpl:
|
||
flags["supports_reasoning"] = True
|
||
flags["reasoning_style"] = "enable_thinking"
|
||
logger.info(f"{prefix}model supports reasoning (enable_thinking)")
|
||
elif "reasoning_effort" in tpl:
|
||
# gpt-oss / Harmony templates use reasoning_effort
|
||
# ("low" | "medium" | "high") instead of a boolean.
|
||
flags["supports_reasoning"] = True
|
||
flags["reasoning_style"] = "reasoning_effort"
|
||
logger.info(f"{prefix}model supports reasoning (reasoning_effort)")
|
||
elif "thinking" in tpl:
|
||
# DeepSeek uses 'thinking' instead of 'enable_thinking'
|
||
normalized_id = (model_identifier or "").lower()
|
||
if "deepseek" in normalized_id:
|
||
flags["supports_reasoning"] = True
|
||
logger.info(f"{prefix}model supports reasoning (DeepSeek thinking)")
|
||
|
||
# Hardcoded <think> tags or reasoning_content in the template mean
|
||
# thinking is always on (no toggle to disable it).
|
||
if not flags["supports_reasoning"]:
|
||
if ("<think>" in tpl and "</think>" in tpl) or "reasoning_content" in tpl:
|
||
flags["supports_reasoning"] = True
|
||
flags["reasoning_always_on"] = True
|
||
logger.info(f"{prefix}model always reasons (<think> tags in template)")
|
||
|
||
# preserve_thinking is an independent kwarg on some Qwen templates
|
||
# that keeps historical <think> blocks in prior assistant turns.
|
||
if "preserve_thinking" in tpl:
|
||
flags["supports_preserve_thinking"] = True
|
||
logger.info(f"{prefix}model supports preserve_thinking")
|
||
|
||
if any(marker in tpl for marker in _TOOL_TEMPLATE_MARKERS):
|
||
flags["supports_tools"] = True
|
||
logger.info(f"{prefix}model supports tool calling")
|
||
|
||
return flags
|
||
|
||
|
||
def _is_mtp_model_name(
|
||
model_identifier: Optional[str],
|
||
gguf_path: Optional[str] = None,
|
||
) -> bool:
|
||
"""Name-based MTP detector. Fallback for the metadata signal."""
|
||
for cand in (model_identifier, Path(gguf_path).name if gguf_path else None):
|
||
if cand and "-mtp" in cand.lower():
|
||
return True
|
||
return False
|
||
|
||
|
||
def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool:
|
||
"""User passed --spec-type / --spec-default? llama-server takes a
|
||
single --spec-type (comma-separated to chain), so suppress
|
||
auto-emit when this is true."""
|
||
if not extra_args:
|
||
return False
|
||
for raw in extra_args:
|
||
tok = str(raw)
|
||
if not tok.startswith("--"):
|
||
continue
|
||
flag = tok.split("=", 1)[0]
|
||
if flag in ("--spec-type", "--spec-default"):
|
||
return True
|
||
return False
|
||
|
||
|
||
def _build_ngram_mod_flags(
|
||
caps: Optional[dict],
|
||
n_match: int = 24,
|
||
n_min: int = 48,
|
||
n_max: int = 64,
|
||
) -> list[str]:
|
||
"""Emit the right ngram-mod knob flags for the running llama-server.
|
||
|
||
Post-rename builds expose ``--spec-ngram-mod-n-{match,min,max}``;
|
||
pre-rename builds expose the legacy ``--spec-ngram-size-n`` /
|
||
``--draft-min`` / ``--draft-max``. ``caps`` comes from
|
||
``probe_server_capabilities``; ``ngram_mod_flavor`` tells us which
|
||
set is real (vs a removal-stub entry). Returns ``[]`` when neither
|
||
set is available so the caller can drop ngram-mod entirely.
|
||
"""
|
||
flavor = caps.get("ngram_mod_flavor") if caps else None
|
||
if flavor == "new":
|
||
return [
|
||
"--spec-ngram-mod-n-match",
|
||
str(n_match),
|
||
"--spec-ngram-mod-n-min",
|
||
str(n_min),
|
||
"--spec-ngram-mod-n-max",
|
||
str(n_max),
|
||
]
|
||
if flavor == "legacy":
|
||
# Legacy llama.cpp before the spec arg rename: same knobs lived
|
||
# under --spec-ngram-size-n (lookup length) and the generic
|
||
# --draft-min / --draft-max (ngram size N range).
|
||
return [
|
||
"--spec-ngram-size-n",
|
||
str(n_match),
|
||
"--draft-min",
|
||
str(n_min),
|
||
"--draft-max",
|
||
str(n_max),
|
||
]
|
||
return []
|
||
|
||
|
||
# Canonical Speculative Decoding modes exposed by the Studio chat UI.
|
||
# The dropdown renders five options (auto, mtp, ngram, mtp+ngram, off);
|
||
# the load API also accepts legacy values that the original Switch and
|
||
# external callers emit (default, draft-mtp, ngram-mod, ngram-simple).
|
||
_CANONICAL_SPEC_MODES = {"auto", "mtp", "ngram", "mtp+ngram", "off", "ngram-simple"}
|
||
_LEGACY_SPEC_MODE_MAP = {
|
||
"default": "auto",
|
||
"draft-mtp": "mtp",
|
||
"ngram-mod": "ngram",
|
||
}
|
||
|
||
|
||
def _canonicalize_spec_mode(value):
|
||
"""Map any accepted ``speculative_type`` input onto a canonical mode.
|
||
|
||
Returns one of ``auto``, ``mtp``, ``ngram``, ``mtp+ngram``, ``off``,
|
||
``ngram-simple``, or ``None`` (callers treat ``None`` as ``auto``).
|
||
Unknown strings collapse to ``auto`` so a stale UI value or typo
|
||
falls back to the safe platform-aware path.
|
||
"""
|
||
if value is None:
|
||
return None
|
||
if not isinstance(value, str):
|
||
return None
|
||
stripped = value.strip().lower()
|
||
if not stripped:
|
||
return None
|
||
if stripped in _CANONICAL_SPEC_MODES:
|
||
return stripped
|
||
if stripped in _LEGACY_SPEC_MODE_MAP:
|
||
return _LEGACY_SPEC_MODE_MAP[stripped]
|
||
# llama.cpp comma-chains are emitted by old persisted state e.g.
|
||
# "ngram-mod,draft-mtp"; collapse the most common one explicitly.
|
||
pieces = [p.strip() for p in stripped.split(",") if p.strip()]
|
||
has_mtp = any(p in ("mtp", "draft-mtp") for p in pieces)
|
||
has_ngram = any(p in ("ngram", "ngram-mod") for p in pieces)
|
||
if has_mtp and has_ngram:
|
||
return "mtp+ngram"
|
||
if has_mtp:
|
||
return "mtp"
|
||
if has_ngram:
|
||
return "ngram"
|
||
return "auto"
|
||
|
||
|
||
def _backfill_usage_from_timings(usage, timings):
|
||
"""Synthesize ``usage`` from llama-server's ``timings`` when the
|
||
OpenAI-style usage block is missing or reports zero tokens.
|
||
|
||
The Studio chat UI computes generation t/s from
|
||
``meta.usage.completion_tokens / totalStreamTime``. llama-server
|
||
always populates ``timings.predicted_n`` (true decoded count) and
|
||
``timings.prompt_n``, but the ``usage`` field on the final SSE chunk
|
||
can be absent or zero on some server builds / streaming
|
||
configurations, which makes the UI fall back to wall-clock t/s and
|
||
dilute speculative-decoding speedups.
|
||
"""
|
||
if not timings:
|
||
return usage
|
||
if usage and usage.get("completion_tokens"):
|
||
return usage
|
||
predicted_n = timings.get("predicted_n")
|
||
prompt_n = timings.get("prompt_n")
|
||
if predicted_n is None and prompt_n is None:
|
||
return usage
|
||
out = dict(usage or {})
|
||
if not out.get("completion_tokens") and predicted_n is not None:
|
||
out["completion_tokens"] = predicted_n
|
||
if not out.get("prompt_tokens") and prompt_n is not None:
|
||
out["prompt_tokens"] = prompt_n
|
||
out["total_tokens"] = int(out.get("prompt_tokens") or 0) + int(
|
||
out.get("completion_tokens") or 0
|
||
)
|
||
return out
|
||
|
||
|
||
class LlamaCppBackend:
|
||
"""
|
||
Manages a llama-server subprocess for GGUF model inference.
|
||
|
||
Lifecycle:
|
||
1. load_model() — starts llama-server with the GGUF file
|
||
2. generate_chat_completion() — proxies to /v1/chat/completions, streams back
|
||
3. unload_model() — terminates llama-server subprocess
|
||
"""
|
||
|
||
def __init__(self):
|
||
self._process: Optional[subprocess.Popen] = None
|
||
self._port: Optional[int] = None
|
||
self._model_identifier: Optional[str] = None
|
||
self._gguf_path: Optional[str] = None
|
||
self._hf_repo: Optional[str] = None
|
||
self._hf_variant: Optional[str] = None
|
||
self._is_vision: bool = False
|
||
self._healthy = False
|
||
# Set by _classify_gpu_offload after _wait_for_health.
|
||
self._gpu_offload_active: Optional[bool] = None
|
||
self._context_length: Optional[int] = None
|
||
self._effective_context_length: Optional[int] = None
|
||
self._max_context_length: Optional[int] = None
|
||
self._chat_template: Optional[str] = None
|
||
self._chat_template_override: Optional[str] = None
|
||
self._supports_reasoning: bool = False
|
||
self._reasoning_always_on: bool = False
|
||
self._reasoning_style: str = "enable_thinking"
|
||
self._supports_preserve_thinking: bool = False
|
||
self._supports_tools: bool = False
|
||
self._cache_type_kv: Optional[str] = None
|
||
self._reasoning_default: bool = True
|
||
self._speculative_type: Optional[str] = None
|
||
# Canonical UI-facing mode the user requested: one of
|
||
# ``auto``/``mtp``/``ngram``/``mtp+ngram``/``off``/``ngram-simple``.
|
||
# Round-tripped through the status API so the dropdown reflects
|
||
# the picked mode rather than the resolved internal flag set
|
||
# (auto on a 27B MTP GGUF resolves to draft-mtp but the dropdown
|
||
# should still read "Auto").
|
||
self._requested_spec_mode: Optional[str] = None
|
||
# User-supplied --spec-draft-n-max override (None = platform default).
|
||
self._spec_draft_n_max: Optional[int] = None
|
||
# KV-cache estimation fields (populated by _read_gguf_metadata)
|
||
self._n_layers: Optional[int] = None
|
||
self._n_kv_heads: Optional[int] = None
|
||
self._n_kv_heads_by_layer: Optional[list[int]] = None
|
||
self._n_heads: Optional[int] = None
|
||
self._embedding_length: 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
|
||
self._sliding_window: Optional[int] = None
|
||
self._sliding_window_pattern: Optional[list[bool]] = None
|
||
self._full_attention_interval: Optional[int] = None
|
||
self._kv_lora_rank: Optional[int] = None
|
||
self._key_length_mla: Optional[int] = None
|
||
self._kv_key_length_swa: Optional[int] = None
|
||
self._kv_value_length_swa: Optional[int] = None
|
||
self._ssm_inner_size: Optional[int] = None
|
||
self._ssm_state_size: Optional[int] = None
|
||
# Last N layers reuse KV from earlier layers and don't allocate
|
||
# their own cache (Gemma 3n / Gemma 4: <arch>.attention.shared_kv_layers).
|
||
self._shared_kv_layers: Optional[int] = None
|
||
# MTP head count (llama.cpp #22673); >0 enables --spec-type draft-mtp.
|
||
self._nextn_predict_layers: Optional[int] = None
|
||
self._lock = threading.Lock()
|
||
# Wraps load_model() end-to-end so concurrent loads serialise
|
||
# and never coexist as two llama-server processes (#5401).
|
||
self._serial_load_lock = threading.Lock()
|
||
# Last extra_args / requested n_ctx, preserved across unload so
|
||
# the chat UI's /unload+/load Apply path can inherit them (#5401).
|
||
# ``_extra_args_source`` records the (model_identifier, hf_variant)
|
||
# the stored args came from so the route can refuse cross-model
|
||
# inheritance.
|
||
self._extra_args: Optional[List[str]] = None
|
||
self._extra_args_source: Optional[tuple[str, Optional[str]]] = None
|
||
self._requested_n_ctx: int = 0
|
||
self._stdout_lines: list[str] = []
|
||
self._stdout_thread: Optional[threading.Thread] = None
|
||
# llama-server tee log (see _drain_stdout / _kill_process).
|
||
self._llama_log_fh = None
|
||
self._llama_log_path: Optional[Path] = None
|
||
self._cancel_event = threading.Event()
|
||
self._api_key: Optional[str] = None
|
||
# True once a probe has completed; cleared on transient failure.
|
||
self._is_audio: bool = False
|
||
self._audio_type: Optional[str] = None
|
||
self._audio_probed: bool = False
|
||
# Monotonic timestamp set in _kill_process; read by load_model
|
||
# to decide whether to wait for the VRAM reclaim to finish.
|
||
self._last_kill_monotonic: float = 0.0
|
||
|
||
self._kill_orphaned_servers()
|
||
atexit.register(self._cleanup)
|
||
|
||
# ── Properties ────────────────────────────────────────────────
|
||
|
||
@property
|
||
def is_loaded(self) -> bool:
|
||
return self._process is not None and self._healthy
|
||
|
||
@property
|
||
def is_active(self) -> bool:
|
||
"""True if a llama-server process exists (loading or loaded)."""
|
||
return self._process is not None
|
||
|
||
@property
|
||
def base_url(self) -> str:
|
||
return f"http://127.0.0.1:{self._port}"
|
||
|
||
@property
|
||
def model_identifier(self) -> Optional[str]:
|
||
return self._model_identifier
|
||
|
||
@property
|
||
def is_vision(self) -> bool:
|
||
return self._is_vision
|
||
|
||
@property
|
||
def hf_variant(self) -> Optional[str]:
|
||
return self._hf_variant
|
||
|
||
@property
|
||
def extra_args(self) -> Optional[List[str]]:
|
||
"""Extra llama-server flags from the last load. Copy; None = never
|
||
set, [] = explicitly cleared. Used by the route for inheritance."""
|
||
return list(self._extra_args) if self._extra_args is not None else None
|
||
|
||
@property
|
||
def requested_n_ctx(self) -> int:
|
||
"""n_ctx the last load was invoked with (not the effective cap).
|
||
0 means Auto. Used by the route to detect Auto-vs-explicit flips."""
|
||
return self._requested_n_ctx
|
||
|
||
@property
|
||
def extra_args_source(self) -> Optional[tuple[str, Optional[str]]]:
|
||
"""(model_identifier, hf_variant) the stored extra_args came from.
|
||
``None`` if no extras have ever been recorded. Used by the route
|
||
to refuse cross-model inheritance (#5401)."""
|
||
return self._extra_args_source
|
||
|
||
@property
|
||
def context_length(self) -> Optional[int]:
|
||
"""Return the effective context length the server is running at."""
|
||
return self._effective_context_length or self._context_length
|
||
|
||
@property
|
||
def max_context_length(self) -> Optional[int]:
|
||
"""Return the largest context that fits on this hardware at load time.
|
||
|
||
This is the "safe zone" threshold the UI renders warnings
|
||
against. For a model whose weights fit on some GPU subset, it
|
||
is the binary-search cap from ``_fit_context_to_vram`` for that
|
||
subset. For a model whose weights exceed 90% of every GPU
|
||
subset, it is the 4096 fallback -- the spec's default when the
|
||
model will not fit. The UI slider ceiling is
|
||
``native_context_length``; dragging above ``max_context_length``
|
||
triggers the "might be slower" warning.
|
||
"""
|
||
return self._max_context_length or self._context_length
|
||
|
||
@property
|
||
def native_context_length(self) -> Optional[int]:
|
||
"""Return the model's native context length from GGUF metadata."""
|
||
return self._context_length
|
||
|
||
def load_progress(self) -> Optional[dict]:
|
||
"""Return live model-load progress, or None if not loading.
|
||
|
||
While llama-server is warming up, its process is typically in
|
||
kernel state D (disk sleep) mmap'ing the weight shards into
|
||
page cache before pushing layers to VRAM. During that window
|
||
``/api/inference/status`` only reports ``loading``, which gives
|
||
the UI nothing to display besides a spinner that looks stuck
|
||
for minutes on large MoE models.
|
||
|
||
This method samples ``/proc/<pid>/status VmRSS`` against the
|
||
sum of the GGUF shard sizes so the UI can render a real bar
|
||
and compute rate / ETA. Returns ``None`` when no load is in
|
||
flight (no process, or process already healthy).
|
||
|
||
Shape::
|
||
|
||
{
|
||
"phase": "mmap" | "ready",
|
||
"bytes_loaded": int, # VmRSS of the llama-server
|
||
"bytes_total": int, # sum of shard file sizes
|
||
"fraction": float, # bytes_loaded / bytes_total, 0..1
|
||
}
|
||
|
||
Linux-only in the current implementation. On macOS/Windows the
|
||
equivalent would be a different API; this returns ``None`` on
|
||
platforms where ``/proc/<pid>/status`` is unavailable.
|
||
"""
|
||
proc = self._process
|
||
if proc is None:
|
||
return None
|
||
pid = proc.pid
|
||
if pid is None:
|
||
return None
|
||
|
||
# Sum up shard sizes (primary + any extras sitting alongside).
|
||
bytes_total = 0
|
||
gguf_path = self._gguf_path
|
||
if gguf_path:
|
||
primary = Path(gguf_path)
|
||
try:
|
||
if primary.is_file():
|
||
bytes_total += primary.stat().st_size
|
||
except OSError:
|
||
pass
|
||
# Extra shards live alongside the primary with the same prefix
|
||
# before the shard index (e.g. ``-00001-of-00004.gguf``).
|
||
try:
|
||
parent = primary.parent
|
||
stem = primary.name
|
||
m = _SHARD_RE.match(stem)
|
||
prefix = m.group(1) if m else None
|
||
if prefix and parent.is_dir():
|
||
for sibling in parent.iterdir():
|
||
if (
|
||
sibling.is_file()
|
||
and sibling.name.startswith(prefix)
|
||
and sibling.name != stem
|
||
and sibling.suffix == ".gguf"
|
||
):
|
||
try:
|
||
bytes_total += sibling.stat().st_size
|
||
except OSError:
|
||
pass
|
||
except OSError:
|
||
pass
|
||
|
||
# Read VmRSS from /proc/<pid>/status. Kilobytes on Linux.
|
||
bytes_loaded = 0
|
||
try:
|
||
with open(f"/proc/{pid}/status", "r", encoding = "utf-8") as f:
|
||
for line in f:
|
||
if line.startswith("VmRSS:"):
|
||
kb = int(line.split()[1])
|
||
bytes_loaded = kb * 1024
|
||
break
|
||
except (FileNotFoundError, PermissionError, ValueError, OSError):
|
||
return None
|
||
|
||
phase = "ready" if self._healthy else "mmap"
|
||
fraction = 0.0
|
||
if bytes_total > 0:
|
||
fraction = min(1.0, bytes_loaded / bytes_total)
|
||
return {
|
||
"phase": phase,
|
||
"bytes_loaded": bytes_loaded,
|
||
"bytes_total": bytes_total,
|
||
"fraction": round(fraction, 4),
|
||
}
|
||
|
||
@property
|
||
def chat_template(self) -> Optional[str]:
|
||
return self._chat_template
|
||
|
||
@property
|
||
def chat_template_override(self) -> Optional[str]:
|
||
return self._chat_template_override
|
||
|
||
@property
|
||
def supports_reasoning(self) -> bool:
|
||
return self._supports_reasoning
|
||
|
||
@property
|
||
def reasoning_always_on(self) -> bool:
|
||
return self._reasoning_always_on
|
||
|
||
@property
|
||
def reasoning_style(self) -> str:
|
||
return self._reasoning_style
|
||
|
||
@property
|
||
def supports_preserve_thinking(self) -> bool:
|
||
return self._supports_preserve_thinking
|
||
|
||
@property
|
||
def reasoning_default(self) -> bool:
|
||
return self._reasoning_default
|
||
|
||
def _reasoning_kwargs(self, enable_thinking: bool) -> dict:
|
||
if self._reasoning_style == "reasoning_effort":
|
||
return {"reasoning_effort": "high" if enable_thinking else "low"}
|
||
return {"enable_thinking": enable_thinking}
|
||
|
||
def _request_reasoning_kwargs(
|
||
self,
|
||
enable_thinking: Optional[bool],
|
||
reasoning_effort: Optional[str] = None,
|
||
preserve_thinking: Optional[bool] = None,
|
||
) -> Optional[dict]:
|
||
"""Build chat_template_kwargs from per-request reasoning fields.
|
||
|
||
Produces a merged dict covering the active model's reasoning style
|
||
(``enable_thinking`` or ``reasoning_effort``) plus the independent
|
||
``preserve_thinking`` kwarg when the template supports it.
|
||
"""
|
||
kwargs: dict = {}
|
||
# Always-on reasoning models hardcode <think> tags in their template
|
||
# and do not consume enable_thinking / reasoning_effort -- skip.
|
||
if self._supports_reasoning and not self._reasoning_always_on:
|
||
if self._reasoning_style == "reasoning_effort":
|
||
if reasoning_effort in ("low", "medium", "high"):
|
||
kwargs["reasoning_effort"] = reasoning_effort
|
||
elif enable_thinking is not None:
|
||
kwargs["reasoning_effort"] = "high" if enable_thinking else "low"
|
||
else:
|
||
if enable_thinking is not None:
|
||
kwargs["enable_thinking"] = enable_thinking
|
||
if self._supports_preserve_thinking and preserve_thinking is not None:
|
||
kwargs["preserve_thinking"] = preserve_thinking
|
||
return kwargs or None
|
||
|
||
@property
|
||
def supports_tools(self) -> bool:
|
||
return self._supports_tools
|
||
|
||
@property
|
||
def cache_type_kv(self) -> Optional[str]:
|
||
return self._cache_type_kv
|
||
|
||
@property
|
||
def speculative_type(self) -> Optional[str]:
|
||
return self._speculative_type
|
||
|
||
@property
|
||
def requested_spec_mode(self) -> Optional[str]:
|
||
"""Canonical UI-facing mode the user requested (see field doc)."""
|
||
return self._requested_spec_mode
|
||
|
||
@property
|
||
def spec_draft_n_max(self) -> Optional[int]:
|
||
"""User --spec-draft-n-max override active on the load, or None
|
||
when the platform default (6 GPU / 3 CPU) is in effect."""
|
||
return self._spec_draft_n_max
|
||
|
||
# ── Binary discovery ──────────────────────────────────────────
|
||
|
||
@staticmethod
|
||
def _find_llama_server_binary() -> Optional[str]:
|
||
"""
|
||
Locate the llama-server binary.
|
||
|
||
Search order:
|
||
1. LLAMA_SERVER_PATH environment variable (direct path to binary)
|
||
1b. UNSLOTH_LLAMA_CPP_PATH env var (custom llama.cpp install dir)
|
||
2. ~/.unsloth/llama.cpp/llama-server (make build, root dir)
|
||
3. ~/.unsloth/llama.cpp/build/bin/llama-server (cmake build, Linux)
|
||
4. ~/.unsloth/llama.cpp/build/bin/Release/llama-server.exe (cmake build, Windows)
|
||
5. ./llama.cpp/llama-server (legacy: make build, root dir)
|
||
6. ./llama.cpp/build/bin/llama-server (legacy: cmake in-tree build)
|
||
7. llama-server on PATH (system install)
|
||
8. ./bin/llama-server (legacy: extracted binary)
|
||
"""
|
||
binary_name = "llama-server.exe" if sys.platform == "win32" else "llama-server"
|
||
|
||
# 1. Env var — direct path to binary
|
||
env_path = os.environ.get("LLAMA_SERVER_PATH")
|
||
if env_path and Path(env_path).is_file():
|
||
return env_path
|
||
|
||
# 1b. UNSLOTH_LLAMA_CPP_PATH — custom llama.cpp install directory
|
||
custom_llama_cpp = os.environ.get("UNSLOTH_LLAMA_CPP_PATH")
|
||
if custom_llama_cpp:
|
||
custom_dir = Path(custom_llama_cpp)
|
||
# Root dir (make builds)
|
||
root_bin = custom_dir / binary_name
|
||
if root_bin.is_file():
|
||
return str(root_bin)
|
||
# build/bin/ (cmake builds on Linux)
|
||
cmake_bin = custom_dir / "build" / "bin" / binary_name
|
||
if cmake_bin.is_file():
|
||
return str(cmake_bin)
|
||
# build/bin/Release/ (cmake builds on Windows)
|
||
if sys.platform == "win32":
|
||
win_bin = custom_dir / "build" / "bin" / "Release" / binary_name
|
||
if win_bin.is_file():
|
||
return str(win_bin)
|
||
|
||
# 2-4. Match installer layout: env-mode -> $STUDIO_HOME/llama.cpp;
|
||
# default/HOME-redirect -> ~/.unsloth/llama.cpp (sibling of studio).
|
||
legacy_llama = Path.home() / ".unsloth" / "llama.cpp"
|
||
try:
|
||
from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433
|
||
|
||
_resolved_sr = _sr()
|
||
_legacy_studio = Path.home() / ".unsloth" / "studio"
|
||
try:
|
||
_is_legacy = _resolved_sr.resolve() == _legacy_studio.resolve()
|
||
except (OSError, ValueError):
|
||
_is_legacy = _resolved_sr == _legacy_studio
|
||
if _is_legacy:
|
||
search_roots = [legacy_llama]
|
||
else:
|
||
# why: _kill_orphaned_servers excludes the legacy root in custom
|
||
# mode; discovery must match so we never spawn a server we then
|
||
# refuse to clean up. UNSLOTH_LLAMA_CPP_PATH (handled earlier)
|
||
# is the explicit way to share a build across roots.
|
||
search_roots = [_resolved_sr / "llama.cpp"]
|
||
except (ImportError, OSError, ValueError):
|
||
search_roots = [legacy_llama]
|
||
_seen_roots: set[str] = set()
|
||
_unique_roots: list[Path] = []
|
||
for r in search_roots:
|
||
k = str(r)
|
||
if k not in _seen_roots:
|
||
_seen_roots.add(k)
|
||
_unique_roots.append(r)
|
||
for unsloth_home in _unique_roots:
|
||
home_root = unsloth_home / binary_name
|
||
if home_root.is_file():
|
||
return str(home_root)
|
||
home_linux = unsloth_home / "build" / "bin" / binary_name
|
||
if home_linux.is_file():
|
||
return str(home_linux)
|
||
if sys.platform == "win32":
|
||
home_win = unsloth_home / "build" / "bin" / "Release" / binary_name
|
||
if home_win.is_file():
|
||
return str(home_win)
|
||
|
||
# 5–6. Legacy: in-tree build (older setup.sh / setup.ps1 versions)
|
||
project_root = Path(__file__).resolve().parents[4]
|
||
# Root dir (make builds)
|
||
root_path = project_root / "llama.cpp" / binary_name
|
||
if root_path.is_file():
|
||
return str(root_path)
|
||
# build/bin/ (cmake builds)
|
||
build_path = project_root / "llama.cpp" / "build" / "bin" / binary_name
|
||
if build_path.is_file():
|
||
return str(build_path)
|
||
if sys.platform == "win32":
|
||
win_path = (
|
||
project_root / "llama.cpp" / "build" / "bin" / "Release" / binary_name
|
||
)
|
||
if win_path.is_file():
|
||
return str(win_path)
|
||
|
||
# 7. System PATH
|
||
system_path = shutil.which("llama-server")
|
||
if system_path:
|
||
return system_path
|
||
|
||
# 8. Legacy: extracted to bin/
|
||
bin_path = project_root / "bin" / binary_name
|
||
if bin_path.is_file():
|
||
return str(bin_path)
|
||
|
||
return None
|
||
|
||
# ── llama-server capability probe ─────────────────────────────
|
||
|
||
# Cached on (path, mtime); `unsloth studio update` bumps mtime.
|
||
_capability_cache: dict[tuple[str, int], dict[str, object]] = {}
|
||
|
||
@classmethod
|
||
def probe_server_capabilities(
|
||
cls, binary: Optional[str] = None
|
||
) -> dict[str, object]:
|
||
"""Parse `llama-server --help` for feature flags. Returns
|
||
{found, mtp_token, supports_mtp, ngram_mod_flavor,
|
||
supports_ngram_mod, spec_draft_n_max_flag}.
|
||
|
||
``ngram_mod_flavor`` is ``"new"`` when the binary exposes the
|
||
post-rename ``--spec-ngram-mod-n-match / -n-min / -n-max`` as
|
||
real args, ``"legacy"`` when only the pre-rename
|
||
``--spec-ngram-size-n / --draft-min / --draft-max`` are real
|
||
(the rename ships with stub removal entries for the legacy
|
||
names; we tell stubs apart by the "argument has been removed"
|
||
description), or ``None`` if neither set is usable.
|
||
|
||
``spec_draft_n_max_flag`` is the actual flag name the binary
|
||
accepts: ``--spec-draft-n-max`` on post-rename builds, or
|
||
``--draft-max`` on legacy. ``None`` means n_max cannot be set.
|
||
"""
|
||
bin_path = binary or cls._find_llama_server_binary()
|
||
if not bin_path or not Path(bin_path).is_file():
|
||
return {
|
||
"found": False,
|
||
"mtp_token": None,
|
||
"supports_mtp": False,
|
||
"ngram_mod_flavor": None,
|
||
"supports_ngram_mod": False,
|
||
"spec_draft_n_max_flag": None,
|
||
}
|
||
try:
|
||
mtime = int(Path(bin_path).stat().st_mtime)
|
||
except OSError:
|
||
mtime = 0
|
||
cache_key = (bin_path, mtime)
|
||
cached = cls._capability_cache.get(cache_key)
|
||
if cached is not None:
|
||
return cached
|
||
|
||
mtp_token: Optional[str] = None
|
||
ngram_mod_flavor: Optional[str] = None
|
||
spec_draft_n_max_flag: Optional[str] = None
|
||
try:
|
||
result = subprocess.run(
|
||
[bin_path, "--help"],
|
||
capture_output = True,
|
||
text = True,
|
||
timeout = 10,
|
||
check = False,
|
||
)
|
||
help_text = (result.stdout or "") + "\n" + (result.stderr or "")
|
||
# Split into per-flag blocks: each --flag line plus its
|
||
# indented continuation lines, so the "argument has been
|
||
# removed" description sits with its flag.
|
||
blocks: dict[str, str] = {}
|
||
current_flags: list[str] = []
|
||
current_desc: list[str] = []
|
||
for line in help_text.splitlines():
|
||
stripped = line.strip()
|
||
if stripped.startswith("-") and not line.startswith(" "):
|
||
# New flag line; flush previous.
|
||
if current_flags:
|
||
desc = " ".join(current_desc)
|
||
for f in current_flags:
|
||
blocks[f] = desc
|
||
current_flags = []
|
||
current_desc = [stripped]
|
||
# Extract long-form flag tokens from the DECLARATION
|
||
# prefix only (comma-separated aliases). Stop at the
|
||
# first token that isn't itself a flag, so flag
|
||
# references inside descriptions are ignored.
|
||
for tok in re.split(r"[,\s]+", stripped):
|
||
if tok.startswith("--") and re.match(
|
||
r"--[A-Za-z][A-Za-z0-9_-]*$", tok
|
||
):
|
||
current_flags.append(tok)
|
||
elif tok.startswith("-") and len(tok) > 1:
|
||
# short alias like -fa; keep scanning aliases.
|
||
continue
|
||
else:
|
||
# First non-flag token marks end of decl.
|
||
break
|
||
else:
|
||
current_desc.append(stripped)
|
||
if current_flags:
|
||
desc = " ".join(current_desc)
|
||
for f in current_flags:
|
||
blocks[f] = desc
|
||
|
||
def _is_real(flag: str) -> bool:
|
||
"""True if the flag exists AND is not a removal stub."""
|
||
desc = blocks.get(flag)
|
||
if desc is None:
|
||
return False
|
||
return "argument has been removed" not in desc
|
||
|
||
# MTP token detection from --spec-type line.
|
||
spec_line = ""
|
||
for line in help_text.splitlines():
|
||
if "--spec-type" in line:
|
||
spec_line = line
|
||
break
|
||
# PR #22673 used draft-mtp; later renamed to mtp.
|
||
if "draft-mtp" in spec_line:
|
||
mtp_token = "draft-mtp"
|
||
elif re.search(r"[|,\[]mtp[|,\]]", spec_line):
|
||
mtp_token = "mtp"
|
||
|
||
# ngram-mod flag flavor. Post-rename builds advertise both
|
||
# the new args (real) and the legacy ones (stubs); pre-rename
|
||
# builds only have the legacy ones as real.
|
||
new_ngram_real = (
|
||
_is_real("--spec-ngram-mod-n-match")
|
||
and _is_real("--spec-ngram-mod-n-min")
|
||
and _is_real("--spec-ngram-mod-n-max")
|
||
)
|
||
legacy_ngram_real = (
|
||
_is_real("--spec-ngram-size-n")
|
||
and _is_real("--draft-max")
|
||
and _is_real("--draft-min")
|
||
)
|
||
if new_ngram_real:
|
||
ngram_mod_flavor = "new"
|
||
elif legacy_ngram_real:
|
||
ngram_mod_flavor = "legacy"
|
||
|
||
# n_max flag: prefer post-rename, fall back to legacy.
|
||
if _is_real("--spec-draft-n-max"):
|
||
spec_draft_n_max_flag = "--spec-draft-n-max"
|
||
elif _is_real("--draft-max"):
|
||
spec_draft_n_max_flag = "--draft-max"
|
||
except (OSError, subprocess.SubprocessError) as exc:
|
||
logger.debug(f"llama-server --help probe failed: {exc}")
|
||
|
||
info = {
|
||
"found": True,
|
||
"mtp_token": mtp_token,
|
||
"supports_mtp": mtp_token is not None,
|
||
"ngram_mod_flavor": ngram_mod_flavor,
|
||
"supports_ngram_mod": ngram_mod_flavor is not None,
|
||
"spec_draft_n_max_flag": spec_draft_n_max_flag,
|
||
}
|
||
cls._capability_cache[cache_key] = info
|
||
return info
|
||
|
||
# ── GPU allocation ────────────────────────────────────────────
|
||
|
||
@staticmethod
|
||
def _get_gguf_size_bytes(model_path: str) -> int:
|
||
"""Get total GGUF size in bytes, including split shards."""
|
||
main = Path(model_path)
|
||
total = main.stat().st_size
|
||
|
||
# Check for split shards (e.g., model-00001-of-00003.gguf)
|
||
m = _SHARD_FULL_RE.match(main.name)
|
||
if m:
|
||
prefix, _, num_total = m.group(1), m.group(2), m.group(3)
|
||
sibling_pat = re.compile(
|
||
r"^"
|
||
+ re.escape(prefix)
|
||
+ r"-\d{5}-of-"
|
||
+ re.escape(num_total)
|
||
+ r"\.gguf$"
|
||
)
|
||
for sibling in main.parent.iterdir():
|
||
if sibling != main and sibling_pat.match(sibling.name):
|
||
total += sibling.stat().st_size
|
||
|
||
return total
|
||
|
||
@staticmethod
|
||
def _amd_apu_wants_unified_memory() -> bool:
|
||
"""True only for AMD unified-memory APUs (gfx1150/gfx1151), where
|
||
GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM.
|
||
False for discrete AMD, NVIDIA, CPU and macOS (the env hurts discrete
|
||
GPUs). ROCm reuses torch.cuda.*; the gcnArchName suffix is stripped."""
|
||
try:
|
||
import torch
|
||
|
||
if getattr(torch.version, "hip", None) is None:
|
||
return False
|
||
if not (hasattr(torch, "cuda") and torch.cuda.is_available()):
|
||
return False
|
||
for _i in range(torch.cuda.device_count()):
|
||
try:
|
||
_arch = (
|
||
getattr(torch.cuda.get_device_properties(_i), "gcnArchName", "")
|
||
or ""
|
||
)
|
||
except Exception:
|
||
continue
|
||
if _arch.split(":")[0].strip().lower() in {"gfx1150", "gfx1151"}:
|
||
return True
|
||
except Exception:
|
||
return False
|
||
return False
|
||
|
||
@staticmethod
|
||
def _get_gpu_free_memory() -> list[tuple[int, int]]:
|
||
"""Query free memory per GPU.
|
||
|
||
Order:
|
||
1. ``nvidia-smi`` (NVIDIA CUDA hosts) -- respects
|
||
``CUDA_VISIBLE_DEVICES``.
|
||
2. ``torch.cuda.mem_get_info`` -- universal fallback that
|
||
works on AMD ROCm too because the HIP runtime
|
||
reuses the entire ``torch.cuda.*`` namespace. Covers the
|
||
AMD case for issue #5106 (nvidia-smi-only probe silently
|
||
returned [] on AMD hosts) and also rescues NVIDIA hosts
|
||
where ``nvidia-smi`` is missing from PATH.
|
||
|
||
Returns list of (gpu_index, free_mib) sorted by index. Empty
|
||
list if no supported GPU is reachable.
|
||
"""
|
||
# ── NVIDIA via nvidia-smi ────────────────────────────────────
|
||
try:
|
||
result = subprocess.run(
|
||
[
|
||
"nvidia-smi",
|
||
"--query-gpu=index,memory.free",
|
||
"--format=csv,noheader,nounits",
|
||
],
|
||
capture_output = True,
|
||
text = True,
|
||
timeout = 10,
|
||
env = child_env_without_native_path_secret(),
|
||
**_windows_hidden_subprocess_kwargs(),
|
||
)
|
||
if result.returncode == 0:
|
||
allowed: Optional[set[int]] = None
|
||
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||
if cvd is not None:
|
||
try:
|
||
# `if x.strip()` filters trailing-comma masks like
|
||
# "0,1," which would otherwise raise ValueError on
|
||
# an empty token. An explicitly empty mask (CVD="")
|
||
# yields an empty `allowed` set so all GPUs are
|
||
# filtered out, matching the codebase convention.
|
||
allowed = set(
|
||
int(x.strip()) for x in cvd.split(",") if x.strip()
|
||
)
|
||
except ValueError:
|
||
pass
|
||
gpus: list[tuple[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))
|
||
# Match the docstring's sort-by-id guarantee. nvidia-smi
|
||
# almost always returns sorted output, but driver order
|
||
# is not formally guaranteed.
|
||
gpus.sort(key = lambda g: g[0])
|
||
if gpus:
|
||
return gpus
|
||
except Exception as e:
|
||
logger.debug(f"nvidia-smi probe failed: {e}")
|
||
|
||
# ── Torch fallback (covers AMD ROCm and missing nvidia-smi) ──
|
||
try:
|
||
import torch
|
||
|
||
if not hasattr(torch, "cuda") or not torch.cuda.is_available():
|
||
return []
|
||
if not hasattr(torch.cuda, "mem_get_info"):
|
||
return []
|
||
# torch.cuda enumerates GPUs RELATIVE to the visibility mask.
|
||
# On NVIDIA builds the mask is CUDA_VISIBLE_DEVICES; on AMD
|
||
# ROCm builds it is HIP_VISIBLE_DEVICES (or ROCR_VISIBLE_DEVICES
|
||
# if HIP is unset). Downstream we feed these IDs back into the
|
||
# llama-server subprocess as CVD, so we must translate visible
|
||
# ordinals back to physical indices first; otherwise launching
|
||
# with ``CUDA_VISIBLE_DEVICES=2,3`` would get rewritten to
|
||
# ``CUDA_VISIBLE_DEVICES=0,1`` and target the wrong GPUs.
|
||
physical_ids: Optional[list[int]] = None
|
||
# Match the codebase convention in
|
||
# ``utils/hardware/hardware.py::_get_parent_visible_gpu_spec``:
|
||
# treat an explicitly empty mask (``HIP_VISIBLE_DEVICES=""``)
|
||
# as "set to no GPUs" rather than falling through to the next
|
||
# var. ``or`` would coerce empty string to falsy and silently
|
||
# promote the wrong source.
|
||
if getattr(torch.version, "hip", None) is not None:
|
||
hip_v = os.environ.get("HIP_VISIBLE_DEVICES")
|
||
rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES")
|
||
cvd = (
|
||
hip_v
|
||
if hip_v is not None
|
||
else rocr_v
|
||
if rocr_v is not None
|
||
else os.environ.get("CUDA_VISIBLE_DEVICES")
|
||
)
|
||
else:
|
||
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||
if cvd is not None:
|
||
try:
|
||
# Empty mask (CVD="") yields an empty list so the
|
||
# below loop produces no GPUs, consistent with the
|
||
# nvidia-smi path and utils/hardware/hardware.py.
|
||
physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()]
|
||
except ValueError:
|
||
physical_ids = None
|
||
gpus = []
|
||
for ordinal in range(torch.cuda.device_count()):
|
||
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)))
|
||
# Match the nvidia-smi path's docstring guarantee of sorted-by-id.
|
||
return sorted(gpus, key = lambda g: g[0])
|
||
except Exception as e:
|
||
logger.debug(f"torch GPU probe failed: {e}")
|
||
return []
|
||
|
||
# Skip the wait when the last kill is older than this; the GPU
|
||
# driver has already reclaimed the prior process's allocations.
|
||
_VRAM_SETTLE_WINDOW_S: float = 15.0
|
||
|
||
@staticmethod
|
||
def _wait_for_vram_settle(
|
||
max_wait: float = 2.0,
|
||
interval: float = 0.25,
|
||
tolerance_mib: int = 256,
|
||
since_kill: float = 0.0,
|
||
) -> None:
|
||
"""Poll ``_get_gpu_free_memory`` until free VRAM stabilises.
|
||
|
||
The GPU driver reclaims a dead process's allocations
|
||
asynchronously, so sampling free memory in the kill-to-spawn
|
||
window reads artificially low and pushes ``_select_gpus`` /
|
||
``_fit_context_to_vram`` toward needless CPU offload -- on a
|
||
tight VRAM card this is the Apply-reload OOM that bare-shell
|
||
launches with the same flags never see.
|
||
|
||
Short-circuits on cold start (``since_kill`` zero) or stale
|
||
kill (older than ``_VRAM_SETTLE_WINDOW_S``); also on CPU-only
|
||
hosts (empty probe), probe exceptions, and GPU-set changes.
|
||
``max_wait`` is a wall-clock bound that includes probe time,
|
||
so a wedged ``nvidia-smi`` cannot extend the reload.
|
||
"""
|
||
now = time.monotonic()
|
||
if since_kill <= 0.0:
|
||
return
|
||
if now - since_kill > LlamaCppBackend._VRAM_SETTLE_WINDOW_S:
|
||
return
|
||
deadline = now + max_wait
|
||
|
||
def _probe_or_none():
|
||
if time.monotonic() >= deadline:
|
||
return None
|
||
try:
|
||
return LlamaCppBackend._get_gpu_free_memory()
|
||
except Exception:
|
||
return None
|
||
|
||
prev = _probe_or_none()
|
||
if prev is None or not prev:
|
||
return
|
||
while time.monotonic() < deadline:
|
||
remaining = deadline - time.monotonic()
|
||
if remaining <= 0:
|
||
return
|
||
# Clip the nap so a near-zero ``max_wait`` is respected.
|
||
time.sleep(min(interval, remaining))
|
||
curr = _probe_or_none()
|
||
if curr is None or not curr or len(curr) != len(prev):
|
||
return
|
||
prev_map = dict(prev)
|
||
stable = True
|
||
for idx, free in curr:
|
||
if idx not in prev_map:
|
||
stable = False
|
||
break
|
||
prev_free = prev_map[idx]
|
||
# Adaptive: 2 % of the larger sample dominates the
|
||
# 256 MiB floor on large-VRAM cards.
|
||
per_gpu_tol = max(tolerance_mib, int(max(free, prev_free) * 0.02))
|
||
if abs(free - prev_free) >= per_gpu_tol:
|
||
stable = False
|
||
break
|
||
if stable:
|
||
return
|
||
prev = curr
|
||
|
||
# Free-VRAM fraction at which Studio pins the GPU directly instead
|
||
# of deferring to ``--fit on``. 5% headroom covers CUDA context +
|
||
# compute buffers; 0.90 was too conservative and dropped 91-94%
|
||
# fits to CPU offload (#5106). The fork's --fit on still catches
|
||
# the truly-too-large case.
|
||
_GPU_PIN_VRAM_FRACTION = 0.95
|
||
|
||
@staticmethod
|
||
def _windows_pip_nvidia_dll_dirs(prefix: str) -> list[str]:
|
||
"""Return DLL dirs from pip-installed CUDA wheels under
|
||
``<prefix>/Lib/site-packages/`` so llama-server.exe can load
|
||
``cudart64_X.dll`` / ``cublas64_X.dll`` without a system CUDA
|
||
toolkit. Mirrors the Linux ``nvidia/cu*/lib`` LD_LIBRARY_PATH
|
||
block, with parity for the Windows-specific wheel layouts seen
|
||
in the wild. Covered patterns:
|
||
* ``nvidia/<pkg>/bin`` -- legacy modular wheels
|
||
(``nvidia-cuda-runtime-cu12``, ``nvidia-cublas-cu12``, etc.).
|
||
* ``nvidia/<pkg>/bin/x86_64`` and ``.../bin/x64`` -- current
|
||
CUDA 13 wheel layout used by the unsuffixed
|
||
``nvidia-cuda-runtime`` / ``nvidia-cublas`` packages, which
|
||
ship under ``nvidia/cu13/bin/x86_64/`` (#5106).
|
||
* ``nvidia/<pkg>/Library/bin`` (and arch subdirs) -- conda-
|
||
style wheel repacks.
|
||
* ``torch/lib`` -- PyTorch's own CUDA-bundled Windows wheel,
|
||
which can ship ``cudart64_*.dll`` directly here instead of
|
||
as separate ``nvidia-*`` wheels. The install-side helper
|
||
``python_runtime_dirs`` in ``install_llama_prebuilt.py``
|
||
covers this path for the same reason.
|
||
|
||
Walks the tree with ``Path.iterdir`` rather than ``glob.glob``
|
||
so the resolver is safe against Windows paths containing
|
||
``[`` or ``]`` (valid in usernames; would otherwise be
|
||
interpreted as a glob character class and silently miss
|
||
existing dirs)."""
|
||
site_packages = Path(prefix) / "Lib" / "site-packages"
|
||
out: list[str] = []
|
||
seen: set[str] = set()
|
||
|
||
def _add(path: Path) -> None:
|
||
if not path.is_dir():
|
||
return
|
||
key = os.path.normcase(os.path.abspath(str(path)))
|
||
if key in seen:
|
||
return
|
||
seen.add(key)
|
||
out.append(str(path))
|
||
|
||
nvidia_root = site_packages / "nvidia"
|
||
if nvidia_root.is_dir():
|
||
for pkg_dir in nvidia_root.iterdir():
|
||
if not pkg_dir.is_dir():
|
||
continue
|
||
# Order matters for PATH search: arch-specific subdirs
|
||
# first so the explicit cudart64_X.dll location wins
|
||
# over a sibling ``bin`` that might be empty.
|
||
for sub in (
|
||
pkg_dir / "bin" / "x86_64",
|
||
pkg_dir / "bin" / "x64",
|
||
pkg_dir / "bin",
|
||
pkg_dir / "Library" / "bin" / "x86_64",
|
||
pkg_dir / "Library" / "bin" / "x64",
|
||
pkg_dir / "Library" / "bin",
|
||
):
|
||
_add(sub)
|
||
_add(site_packages / "torch" / "lib")
|
||
return out
|
||
|
||
@staticmethod
|
||
def _build_windows_path_dirs(
|
||
binary_dir: str, prefix: str, cuda_path: str
|
||
) -> list[str]:
|
||
"""Ordered PATH entries the win32 branch of start_llama_server
|
||
prepends so llama-server.exe resolves cudart / cublas DLLs:
|
||
binary_dir, pip nvidia wheels, CUDA_PATH/bin, CUDA_PATH/bin/x64.
|
||
Extracted so test_windows_gpu_detection_mock asserts against
|
||
production logic, not a hand-copy. #5106."""
|
||
path_dirs = [binary_dir]
|
||
path_dirs.extend(LlamaCppBackend._windows_pip_nvidia_dll_dirs(prefix))
|
||
if cuda_path:
|
||
cuda_bin = os.path.join(cuda_path, "bin")
|
||
if os.path.isdir(cuda_bin):
|
||
path_dirs.append(cuda_bin)
|
||
cuda_bin_x64 = os.path.join(cuda_path, "bin", "x64")
|
||
if os.path.isdir(cuda_bin_x64):
|
||
path_dirs.append(cuda_bin_x64)
|
||
return path_dirs
|
||
|
||
@staticmethod
|
||
def _select_gpus(
|
||
model_size_bytes: int,
|
||
gpus: list[tuple[int, int]],
|
||
) -> tuple[Optional[list[int]], bool]:
|
||
"""Pick GPU(s) for a model based on estimated VRAM and free memory.
|
||
|
||
``model_size_bytes`` should include both model weights and estimated
|
||
KV cache. The ``_GPU_PIN_VRAM_FRACTION`` threshold provides headroom
|
||
for compute buffers, CUDA context, and other runtime overhead.
|
||
|
||
Returns (gpu_indices, use_fit):
|
||
- ([1], False) model fits on 1 GPU at the headroom threshold
|
||
- ([1, 2], False) model needs 2 GPUs
|
||
- (None, True) model too large, let --fit handle it
|
||
"""
|
||
if not gpus:
|
||
return None, True
|
||
|
||
model_size_mib = model_size_bytes / (1024 * 1024)
|
||
usable_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION
|
||
|
||
# Sort GPUs by free memory descending
|
||
ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
|
||
|
||
# Try fitting on 1 GPU at the usable-VRAM threshold.
|
||
if ranked[0][1] * usable_fraction >= model_size_mib:
|
||
return [ranked[0][0]], False
|
||
|
||
# Try fitting on N GPUs (accumulate free memory from most-free)
|
||
cumulative = 0
|
||
selected = []
|
||
for idx, free_mib in ranked:
|
||
selected.append(idx)
|
||
cumulative += free_mib * usable_fraction
|
||
if cumulative >= model_size_mib:
|
||
return sorted(selected), False
|
||
|
||
# Model is too large even for all GPUs, let --fit handle it
|
||
logger.debug(
|
||
"Model does not fit in available GPU memory, falling back to --fit",
|
||
model_size_mib = round(model_size_mib, 2),
|
||
ranked_gpus = ranked,
|
||
)
|
||
return None, True
|
||
|
||
# ── KV cache VRAM estimation ─────────────────────────────────────
|
||
|
||
def _can_estimate_kv(self) -> bool:
|
||
"""True if we have enough GGUF metadata to estimate KV cache size."""
|
||
if self._n_layers is None:
|
||
return False
|
||
# MLA: kv_lora_rank is sufficient (K-only cache)
|
||
if self._kv_lora_rank is not None:
|
||
return True
|
||
# New-style: need both explicit key AND value dimensions
|
||
if self._kv_key_length is not None and self._kv_value_length is not None:
|
||
return True
|
||
# Legacy: need embedding_length + a head count (scalar or per-layer).
|
||
return self._embedding_length is not None and (
|
||
self._n_kv_heads is not None
|
||
or self._n_heads is not None
|
||
or self._n_kv_heads_by_layer is not None
|
||
)
|
||
|
||
def _kv_heads_for_layer(self, layer_idx: int, fallback: int) -> int:
|
||
if self._n_kv_heads_by_layer is not None and layer_idx < len(
|
||
self._n_kv_heads_by_layer
|
||
):
|
||
return self._n_kv_heads_by_layer[layer_idx]
|
||
return fallback
|
||
|
||
def _estimate_kv_cache_bytes(
|
||
self,
|
||
n_ctx: int,
|
||
cache_type_kv: Optional[str] = None,
|
||
*,
|
||
swa_full: bool = False,
|
||
n_parallel: int = 1,
|
||
kv_unified: bool = True,
|
||
ctx_checkpoints: int = 0,
|
||
) -> int:
|
||
"""Estimate KV cache VRAM for a given context length.
|
||
|
||
Uses 5-path architecture-aware estimation:
|
||
1. MLA -- compressed KV latent + RoPE, K-only (no separate V)
|
||
2. Hybrid -- only attention layers need KV (Mamba layers don't)
|
||
3. SWA -- sliding-window layers cache min(ctx, window) tokens
|
||
4. GQA -- standard full KV with explicit key/value dimensions
|
||
5. Legacy -- fallback using embed // n_heads
|
||
|
||
Server-flag knobs (mirror llama-server's CLI):
|
||
swa_full -- ``--swa-full``: force SWA layers to cache the
|
||
full ``n_ctx`` (collapses path 3 to path 4
|
||
sizing for the SWA layers).
|
||
n_parallel -- ``--parallel``: number of server slots.
|
||
Verified empirically against llama-server:
|
||
non-SWA layers stay constant (cells split
|
||
across slots), SWA layers scale linearly
|
||
(per-slot window).
|
||
kv_unified -- ``--kv-unified`` (default on): retained for
|
||
API forward-compat. Currently a no-op for
|
||
memory math because the unified buffer total
|
||
matches per-slot buffers in measured cases.
|
||
ctx_checkpoints -- ``--ctx-checkpoints``: SWA snapshot count per
|
||
slot (PR #15293). Each snapshot stores one
|
||
sliding-window of state per SWA layer.
|
||
|
||
Returns 0 if metadata is insufficient for estimation.
|
||
"""
|
||
if not self._can_estimate_kv() or n_ctx <= 0:
|
||
return 0
|
||
|
||
n_layers = self._n_layers # type: ignore[assignment]
|
||
# Gemma 3n / Gemma 4 reuse KV from earlier layers in the last
|
||
# ``shared_kv_layers`` blocks -- those don't allocate their own
|
||
# cache. Floor at 1 so a misconfigured GGUF can't zero out KV.
|
||
shared = self._shared_kv_layers or 0
|
||
n_layers_kv = max(1, n_layers - shared)
|
||
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)
|
||
|
||
slots = max(1, n_parallel)
|
||
|
||
# Path 1: MLA (DeepSeek-V2/V3, GLM-4.7, GLM-5, Kimi-K2.5)
|
||
# MLA stores one compressed KV latent per token/layer (shared across heads).
|
||
# V is reconstructed from the latent on the fly -- no separate V cache.
|
||
# key_length = kv_lora_rank + rope_dim (the full compressed representation).
|
||
# MLA GGUFs set head_count_kv=1; default to 1 if absent to avoid
|
||
# falling back to n_heads (e.g., 128 for DeepSeek-V3) which would 128x.
|
||
if self._kv_lora_rank is not None:
|
||
n_kv_mla = self._n_kv_heads or 1
|
||
rope_dim = self._key_length_mla or 64
|
||
key_len = self._kv_key_length or (self._kv_lora_rank + rope_dim)
|
||
return int(n_layers_kv * n_ctx * n_kv_mla * key_len * bpe)
|
||
|
||
key_len = self._kv_key_length
|
||
val_len = self._kv_value_length
|
||
|
||
# Path 2: Hybrid Mamba/Attention (Qwen3.5-27B, Qwen3.5-35B-A3B)
|
||
# Only 1 in N layers is attention; the rest are Mamba (no KV cache).
|
||
if (
|
||
self._ssm_inner_size is not None
|
||
and self._full_attention_interval is not None
|
||
):
|
||
fai = self._full_attention_interval
|
||
n_attn = -(-n_layers // fai) if fai > 0 else n_layers # ceiling division
|
||
if key_len is not None and val_len is not None:
|
||
return int(n_attn * n_ctx * n_kv * (key_len + val_len) * bpe)
|
||
head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator]
|
||
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 is filled in 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 empirically
|
||
# against ``llama-server``):
|
||
# * non-SWA layers: total cells = n_ctx, partitioned across
|
||
# slots -> total memory CONSTANT in slots.
|
||
# * SWA layers: per-slot cells = 2 * sliding_window
|
||
# (capped at n_ctx and at per_slot_ctx
|
||
# when ctx is split among many slots) ->
|
||
# total memory grows LINEARLY in slots.
|
||
# ``--swa-full`` forces full n_ctx for SWA layers instead.
|
||
# ``--ctx-checkpoints N`` adds N snapshots per SWA layer per slot.
|
||
if (
|
||
self._sliding_window is not None
|
||
and self._sliding_window > 0
|
||
and key_len is not None
|
||
and val_len is not None
|
||
):
|
||
swa = self._sliding_window
|
||
per_slot_ctx = max(1, n_ctx // slots)
|
||
# ``--swa-full`` makes SWA layers cache the full context just
|
||
# like non-SWA: cells get partitioned across slots, so per-slot
|
||
# cells = per_slot_ctx and the slots*per-slot product collapses
|
||
# back to the constant ``n_ctx`` total. Otherwise SWA caches
|
||
# 2*sliding_window per slot, clamped at the 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
|
||
val_len_swa = self._kv_value_length_swa or val_len
|
||
if self._sliding_window_pattern is not None:
|
||
global_bytes = 0.0 # constant across slots
|
||
swa_bytes_per_slot = 0.0 # multiplied by slots
|
||
checkpoint_extra_per_slot = 0.0
|
||
# Iterate only over layers that allocate their own KV;
|
||
# the trailing ``shared`` layers reuse earlier caches.
|
||
for layer_idx in range(n_layers_kv):
|
||
layer_n_kv = self._kv_heads_for_layer(layer_idx, n_kv)
|
||
is_swa = (
|
||
layer_idx < len(self._sliding_window_pattern)
|
||
and self._sliding_window_pattern[layer_idx]
|
||
)
|
||
if is_swa:
|
||
swa_bytes_per_slot += (
|
||
swa_cells_per_slot
|
||
* layer_n_kv
|
||
* (key_len_swa + val_len_swa)
|
||
* bpe
|
||
)
|
||
if ctx_checkpoints > 0 and not swa_full:
|
||
checkpoint_extra_per_slot += (
|
||
ctx_checkpoints
|
||
* swa
|
||
* layer_n_kv
|
||
* (key_len_swa + val_len_swa)
|
||
* bpe
|
||
)
|
||
else:
|
||
global_bytes += n_ctx * layer_n_kv * (key_len + val_len) * bpe
|
||
return int(
|
||
global_bytes
|
||
+ slots * (swa_bytes_per_slot + checkpoint_extra_per_slot)
|
||
)
|
||
n_global = max(1, n_layers_kv // 4)
|
||
n_swa = n_layers_kv - n_global
|
||
kv_per_token = n_kv * (key_len + val_len) * bpe
|
||
kv_per_token_swa = n_kv * (key_len_swa + val_len_swa) * bpe
|
||
global_bytes = n_global * n_ctx * kv_per_token
|
||
swa_bytes_per_slot = n_swa * swa_cells_per_slot * kv_per_token_swa
|
||
checkpoint_extra_per_slot = (
|
||
ctx_checkpoints * n_swa * swa * kv_per_token_swa
|
||
if ctx_checkpoints > 0 and not swa_full
|
||
else 0.0
|
||
)
|
||
return int(
|
||
global_bytes + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot)
|
||
)
|
||
|
||
# Path 4: Standard GQA with explicit key/value dimensions
|
||
if key_len is not None and val_len is not None:
|
||
return int(n_layers_kv * n_ctx * n_kv * (key_len + val_len) * bpe)
|
||
|
||
# Path 5: Legacy fallback (old GGUFs without explicit dimensions)
|
||
head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator]
|
||
return int(2 * n_kv * head_dim * n_layers_kv * n_ctx * bpe)
|
||
|
||
def _fit_context_to_vram(
|
||
self,
|
||
requested_ctx: int,
|
||
available_mib: int,
|
||
model_size_bytes: int,
|
||
cache_type_kv: Optional[str] = None,
|
||
min_ctx: int = 4096,
|
||
*,
|
||
swa_full: bool = False,
|
||
n_parallel: int = 1,
|
||
kv_unified: bool = True,
|
||
ctx_checkpoints: int = 0,
|
||
kv_on_gpu: bool = True,
|
||
mtp_engaged: bool = False,
|
||
) -> 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, while pinning conservatively just defers to
|
||
--fit on. If the weights alone don't fit, returns
|
||
``requested_ctx`` unchanged.
|
||
|
||
``kv_on_gpu`` mirrors ``--kv-offload`` (default on). When False
|
||
the KV cache lives in CPU RAM and doesn't compete with weights
|
||
for VRAM; the requested context is honored verbatim. The other
|
||
keyword args mirror ``_estimate_kv_cache_bytes``.
|
||
|
||
``mtp_engaged`` reserves extra VRAM for the MTP draft model's
|
||
KV cache + compute graph buffers. llama.cpp's MTP path keeps a
|
||
secondary cache sized off the target's KV; on tight VRAM tiers
|
||
(e.g. 32 GB) auto-fit at native context would otherwise spill
|
||
and force llama-server into a slower partial-offload path.
|
||
"""
|
||
if not self._can_estimate_kv():
|
||
logger.debug(
|
||
"Skipping context fit because KV cache metadata is unavailable",
|
||
requested_ctx = requested_ctx,
|
||
available_mib = available_mib,
|
||
)
|
||
return requested_ctx
|
||
|
||
# KV lives off-GPU: no VRAM accounting needed for the cache itself.
|
||
if not kv_on_gpu:
|
||
return requested_ctx
|
||
|
||
kv_kwargs = dict(
|
||
swa_full = swa_full,
|
||
n_parallel = n_parallel,
|
||
kv_unified = kv_unified,
|
||
ctx_checkpoints = ctx_checkpoints,
|
||
)
|
||
|
||
# MTP needs a tighter budget; drop from 0.90 to 0.85.
|
||
budget_frac = 0.85 if mtp_engaged else 0.90
|
||
budget_bytes = available_mib * 1024 * 1024 * budget_frac
|
||
model_footprint = model_size_bytes
|
||
|
||
# Check if requested context already fits
|
||
kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs)
|
||
if model_footprint + kv <= budget_bytes:
|
||
return requested_ctx
|
||
|
||
# Model weights alone exceed budget -- can't help by reducing ctx.
|
||
# Return requested_ctx unchanged; --fit will handle VRAM management.
|
||
if model_footprint >= budget_bytes:
|
||
logger.debug(
|
||
"Model footprint exceeds GPU budget before KV cache",
|
||
requested_ctx = requested_ctx,
|
||
available_mib = available_mib,
|
||
model_size_gb = round(model_footprint / (1024**3), 2),
|
||
)
|
||
return requested_ctx
|
||
|
||
# Binary search for max context that fits
|
||
remaining = budget_bytes - model_footprint
|
||
effective_min = min(min_ctx, requested_ctx)
|
||
lo, hi = effective_min, requested_ctx
|
||
best = effective_min
|
||
while lo <= hi:
|
||
mid = (lo + hi) // 2
|
||
kv = self._estimate_kv_cache_bytes(mid, cache_type_kv, **kv_kwargs)
|
||
if kv <= remaining:
|
||
best = mid
|
||
lo = mid + 1
|
||
else:
|
||
hi = mid - 1
|
||
|
||
# Round down to nearest 256 for alignment, but never exceed requested_ctx
|
||
best = (best // 256) * 256
|
||
best = max(effective_min, best)
|
||
best = min(best, requested_ctx)
|
||
return best
|
||
|
||
# ── Variant fallback ────────────────────────────────────────────
|
||
|
||
@staticmethod
|
||
def _find_smallest_fitting_variant(
|
||
hf_repo: str,
|
||
free_bytes: int,
|
||
hf_token: Optional[str] = None,
|
||
) -> Optional[tuple[str, int]]:
|
||
"""Find the smallest GGUF variant (including all shards) that fits.
|
||
|
||
Groups split shards by variant prefix and sums their sizes.
|
||
For example, UD-Q4_K_XL with 9 shards of 50 GB each = 450 GB total.
|
||
|
||
Returns (first_shard_filename, total_size_bytes) or None if nothing fits.
|
||
"""
|
||
try:
|
||
from huggingface_hub import get_paths_info, list_repo_files
|
||
|
||
files = list_repo_files(hf_repo, token = hf_token)
|
||
gguf_files = [
|
||
f for f in files if f.endswith(".gguf") and "mmproj" not in f.lower()
|
||
]
|
||
if not gguf_files:
|
||
return None
|
||
|
||
# Get sizes for all GGUF files
|
||
path_infos = list(get_paths_info(hf_repo, gguf_files, token = hf_token))
|
||
size_map = {p.path: (p.size or 0) for p in path_infos}
|
||
|
||
# Group files by variant: shards share a prefix before -NNNNN-of-NNNNN
|
||
variants: dict[str, list[str]] = {}
|
||
for f in gguf_files:
|
||
m = _SHARD_RE.match(f)
|
||
key = m.group(1) if m else f
|
||
variants.setdefault(key, []).append(f)
|
||
|
||
# Sum shard sizes per variant, track the first shard (for download)
|
||
variant_sizes: list[tuple[str, int, list[str]]] = []
|
||
for key, shard_files in variants.items():
|
||
total = sum(size_map.get(f, 0) for f in shard_files)
|
||
first = sorted(shard_files)[0]
|
||
variant_sizes.append((first, total, shard_files))
|
||
|
||
# Sort by total size ascending and pick the smallest that fits
|
||
variant_sizes.sort(key = lambda x: x[1])
|
||
for first_file, total_size, _ in variant_sizes:
|
||
if total_size > 0 and total_size <= free_bytes:
|
||
return first_file, total_size
|
||
|
||
return None
|
||
except Exception:
|
||
return None
|
||
|
||
# ── Port allocation ───────────────────────────────────────────
|
||
|
||
@staticmethod
|
||
def _find_free_port() -> int:
|
||
"""Find an available TCP port."""
|
||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||
s.bind(("127.0.0.1", 0))
|
||
return s.getsockname()[1]
|
||
|
||
# ── Stdout drain (prevents pipe deadlock on Windows) ─────────
|
||
|
||
def _drain_stdout(self):
|
||
"""
|
||
Read lines from the subprocess stdout in a background thread.
|
||
|
||
This prevents a pipe-buffer deadlock on Windows where the default
|
||
pipe buffer is only ~4 KB. Without draining, llama-server blocks
|
||
on writes and never becomes healthy.
|
||
|
||
Each line is also teed to ``self._llama_log_fh`` when set so a
|
||
post-mortem (especially in CI) has the full subprocess output
|
||
even if the crash predates the drain-thread join in
|
||
``_wait_for_health``.
|
||
"""
|
||
try:
|
||
for line in self._process.stdout:
|
||
line = line.rstrip()
|
||
if line:
|
||
self._stdout_lines.append(line)
|
||
logger.debug(f"[llama-server] {line}")
|
||
fh = getattr(self, "_llama_log_fh", None)
|
||
if fh is not None:
|
||
try:
|
||
fh.write(line + "\n")
|
||
fh.flush()
|
||
except (ValueError, OSError):
|
||
# Log file closed under us; tee silently.
|
||
pass
|
||
except (ValueError, OSError):
|
||
# Pipe closed — process is terminating
|
||
pass
|
||
|
||
# GGUF KV type sizes for fast skipping
|
||
_GGUF_TYPE_SIZE = {
|
||
0: 1,
|
||
1: 1,
|
||
2: 2,
|
||
3: 2,
|
||
4: 4,
|
||
5: 4,
|
||
6: 4,
|
||
7: 1,
|
||
10: 8,
|
||
11: 8,
|
||
12: 8,
|
||
}
|
||
|
||
@staticmethod
|
||
def _gguf_skip_value(f, vtype: int) -> None:
|
||
"""Skip a GGUF KV value without reading it."""
|
||
sz = LlamaCppBackend._GGUF_TYPE_SIZE.get(vtype)
|
||
if sz is not None:
|
||
f.seek(sz, 1)
|
||
elif vtype == 8: # STRING
|
||
slen = struct.unpack("<Q", f.read(8))[0]
|
||
f.seek(slen, 1)
|
||
elif vtype == 9: # ARRAY
|
||
atype = struct.unpack("<I", f.read(4))[0]
|
||
alen = struct.unpack("<Q", f.read(8))[0]
|
||
elem_sz = LlamaCppBackend._GGUF_TYPE_SIZE.get(atype)
|
||
if elem_sz is not None:
|
||
f.seek(elem_sz * alen, 1)
|
||
elif atype == 8:
|
||
for _ in range(alen):
|
||
slen = struct.unpack("<Q", f.read(8))[0]
|
||
f.seek(slen, 1)
|
||
else:
|
||
for _ in range(alen):
|
||
LlamaCppBackend._gguf_skip_value(f, atype)
|
||
|
||
@staticmethod
|
||
def _gguf_read_array_value(f, atype: int, alen: int) -> Optional[list]:
|
||
if atype == 4: # UINT32
|
||
return [struct.unpack("<I", f.read(4))[0] for _ in range(alen)]
|
||
if atype == 5: # INT32
|
||
return [struct.unpack("<i", f.read(4))[0] for _ in range(alen)]
|
||
if atype == 7: # BOOL
|
||
return [struct.unpack("<?", f.read(1))[0] for _ in range(alen)]
|
||
|
||
for _ in range(alen):
|
||
LlamaCppBackend._gguf_skip_value(f, atype)
|
||
return None
|
||
|
||
def _read_gguf_metadata(self, gguf_path: str) -> None:
|
||
"""Read context_length, architecture params, and chat_template from a GGUF header.
|
||
|
||
Parses only the KV pairs we need (~30ms even for multi-GB files).
|
||
For split GGUFs, metadata is always in shard 1.
|
||
"""
|
||
# Reset metadata from any previously loaded model so stale flags
|
||
# (eg _supports_reasoning) do not carry over when switching models.
|
||
self._context_length = None
|
||
self._chat_template = None
|
||
self._supports_reasoning = False
|
||
self._reasoning_always_on = False
|
||
self._reasoning_style = "enable_thinking"
|
||
self._reasoning_default = True
|
||
self._supports_preserve_thinking = False
|
||
self._supports_tools = False
|
||
self._n_layers = None
|
||
self._n_kv_heads = None
|
||
self._n_kv_heads_by_layer = None
|
||
self._n_heads = None
|
||
self._embedding_length = None
|
||
self._kv_key_length = None
|
||
self._kv_value_length = None
|
||
self._sliding_window = None
|
||
self._sliding_window_pattern = None
|
||
self._full_attention_interval = None
|
||
self._kv_lora_rank = None
|
||
self._key_length_mla = None
|
||
self._kv_key_length_swa = None
|
||
self._kv_value_length_swa = None
|
||
self._ssm_inner_size = None
|
||
self._ssm_state_size = None
|
||
self._shared_kv_layers = None
|
||
self._nextn_predict_layers = None
|
||
|
||
try:
|
||
WANTED = {
|
||
"general.architecture",
|
||
"tokenizer.chat_template",
|
||
# Source-repo hints for the SWA resolver's HF fallback.
|
||
"general.source.huggingface.repository",
|
||
"general.source.url",
|
||
"general.source.repo_url",
|
||
"general.base_model.0.repo_url",
|
||
"general.base_model.0.organization",
|
||
"general.base_model.0.name",
|
||
"general.basename",
|
||
"general.organization",
|
||
"general.size_label",
|
||
"general.finetune",
|
||
}
|
||
# Additional arch-specific keys are added dynamically once
|
||
# we know the architecture name.
|
||
arch_keys: dict[str, str] = {} # gguf_key -> attribute name
|
||
arch = None
|
||
sliding_window_pattern_period: Optional[int] = None
|
||
general: dict[str, str] = {}
|
||
|
||
with open(gguf_path, "rb") as f:
|
||
magic = struct.unpack("<I", f.read(4))[0]
|
||
if magic != 0x46554747: # b"GGUF" as little-endian u32
|
||
return
|
||
_version = struct.unpack("<I", f.read(4))[0]
|
||
_tensor_count, kv_count = struct.unpack("<QQ", f.read(16))
|
||
|
||
for _ in range(kv_count):
|
||
# Tolerate truncated input (e.g., a partial header
|
||
# fetched via HTTP byte-range): bail out gracefully
|
||
# so the resolver fallback still runs on whatever
|
||
# we did manage to parse.
|
||
try:
|
||
key_len_bytes = f.read(8)
|
||
if len(key_len_bytes) < 8:
|
||
break
|
||
key_len = struct.unpack("<Q", key_len_bytes)[0]
|
||
key_bytes = f.read(key_len)
|
||
if len(key_bytes) < key_len:
|
||
break
|
||
key = key_bytes.decode("utf-8")
|
||
vtype_bytes = f.read(4)
|
||
if len(vtype_bytes) < 4:
|
||
break
|
||
vtype = struct.unpack("<I", vtype_bytes)[0]
|
||
except (struct.error, UnicodeDecodeError):
|
||
break
|
||
|
||
try:
|
||
if key in WANTED or key in arch_keys:
|
||
if vtype == 8: # STRING
|
||
slen = struct.unpack("<Q", f.read(8))[0]
|
||
val_s = f.read(slen).decode("utf-8")
|
||
if (
|
||
key.startswith("general.")
|
||
and key != "general.architecture"
|
||
):
|
||
general[key] = val_s
|
||
if key == "general.architecture":
|
||
arch = val_s
|
||
arch_keys = {
|
||
f"{arch}.context_length": "context_length",
|
||
f"{arch}.block_count": "n_layers",
|
||
f"{arch}.attention.head_count_kv": "n_kv_heads",
|
||
f"{arch}.attention.head_count": "n_heads",
|
||
f"{arch}.embedding_length": "embedding_length",
|
||
f"{arch}.attention.key_length": "kv_key_length",
|
||
f"{arch}.attention.value_length": "kv_value_length",
|
||
f"{arch}.attention.sliding_window": "sliding_window",
|
||
f"{arch}.attention.sliding_window_pattern": "sliding_window_pattern",
|
||
f"{arch}.full_attention_interval": "full_attention_interval",
|
||
f"{arch}.attention.kv_lora_rank": "kv_lora_rank",
|
||
f"{arch}.attention.key_length_mla": "key_length_mla",
|
||
f"{arch}.attention.key_length_swa": "kv_key_length_swa",
|
||
f"{arch}.attention.value_length_swa": "kv_value_length_swa",
|
||
f"{arch}.attention.shared_kv_layers": "shared_kv_layers",
|
||
f"{arch}.ssm.inner_size": "ssm_inner_size",
|
||
f"{arch}.ssm.state_size": "ssm_state_size",
|
||
f"{arch}.nextn_predict_layers": "nextn_predict_layers",
|
||
}
|
||
elif key == "tokenizer.chat_template":
|
||
self._chat_template = val_s
|
||
elif vtype in (4, 10): # UINT32 or UINT64
|
||
val_i = (
|
||
struct.unpack("<I", f.read(4))[0]
|
||
if vtype == 4
|
||
else struct.unpack("<Q", f.read(8))[0]
|
||
)
|
||
attr = arch_keys.get(key)
|
||
if attr:
|
||
if attr == "sliding_window_pattern":
|
||
sliding_window_pattern_period = val_i
|
||
else:
|
||
setattr(self, f"_{attr}", val_i)
|
||
elif vtype == 9: # ARRAY
|
||
atype = struct.unpack("<I", f.read(4))[0]
|
||
alen = struct.unpack("<Q", f.read(8))[0]
|
||
val_a = self._gguf_read_array_value(f, atype, alen)
|
||
attr = arch_keys.get(key)
|
||
if attr == "n_kv_heads" and val_a is not None:
|
||
self._n_kv_heads_by_layer = [int(x) for x in val_a]
|
||
if self._n_kv_heads is None and val_a:
|
||
self._n_kv_heads = max(int(x) for x in val_a)
|
||
elif (
|
||
attr == "sliding_window_pattern"
|
||
and val_a is not None
|
||
):
|
||
self._sliding_window_pattern = [
|
||
bool(x) for x in val_a
|
||
]
|
||
sliding_window_pattern_period = None
|
||
else:
|
||
self._gguf_skip_value(f, vtype)
|
||
else:
|
||
self._gguf_skip_value(f, vtype)
|
||
except (struct.error, UnicodeDecodeError):
|
||
# Truncated input (e.g., HTTP byte-range fetch
|
||
# of just the GGUF header); break so the
|
||
# resolver fallback still runs on what we have.
|
||
break
|
||
|
||
# Expand a scalar period straight from the GGUF first.
|
||
if (
|
||
self._sliding_window_pattern is None
|
||
and sliding_window_pattern_period
|
||
and self._n_layers
|
||
):
|
||
self._sliding_window_pattern = [
|
||
(i + 1) % sliding_window_pattern_period != 0
|
||
for i in range(self._n_layers)
|
||
]
|
||
|
||
# Otherwise hand off to the resolver (cache / bootstrap /
|
||
# transformers / HF). See `_resolve_swa_pattern`.
|
||
if (
|
||
self._sliding_window_pattern is None
|
||
and self._sliding_window
|
||
and self._n_layers
|
||
):
|
||
hf_repo_candidates = (
|
||
general.get("general.source.huggingface.repository"),
|
||
_hf_repo_from_url(general.get("general.source.url")),
|
||
_hf_repo_from_url(general.get("general.source.repo_url")),
|
||
_hf_repo_from_url(general.get("general.base_model.0.repo_url")),
|
||
(
|
||
f"{general['general.base_model.0.organization']}/"
|
||
f"{general['general.base_model.0.name']}".replace(" ", "-")
|
||
if general.get("general.base_model.0.organization")
|
||
and general.get("general.base_model.0.name")
|
||
else None
|
||
),
|
||
(
|
||
f"{general['general.organization']}/"
|
||
f"{general['general.basename']}".replace(" ", "-")
|
||
if general.get("general.organization")
|
||
and general.get("general.basename")
|
||
else None
|
||
),
|
||
)
|
||
self._sliding_window_pattern = _resolve_swa_pattern(
|
||
arch,
|
||
self._n_layers,
|
||
hf_repo_candidates,
|
||
)
|
||
|
||
if self._context_length:
|
||
logger.info(f"GGUF metadata: context_length={self._context_length}")
|
||
if self._chat_template:
|
||
logger.info(
|
||
f"GGUF metadata: chat_template={len(self._chat_template)} chars"
|
||
)
|
||
# Detect thinking/reasoning support from chat template
|
||
flags = detect_reasoning_flags(
|
||
self._chat_template,
|
||
self._model_identifier,
|
||
log_source = "GGUF metadata",
|
||
)
|
||
self._supports_reasoning = flags["supports_reasoning"]
|
||
self._reasoning_style = flags["reasoning_style"]
|
||
self._reasoning_always_on = flags["reasoning_always_on"]
|
||
self._supports_preserve_thinking = flags["supports_preserve_thinking"]
|
||
self._supports_tools = flags["supports_tools"]
|
||
except Exception as e:
|
||
logger.warning(f"Failed to read GGUF metadata: {e}")
|
||
|
||
# ── HF download (no lock held) ───────────────────────────────
|
||
|
||
def _download_gguf(
|
||
self,
|
||
*,
|
||
hf_repo: str,
|
||
hf_variant: Optional[str] = None,
|
||
hf_token: Optional[str] = None,
|
||
) -> str:
|
||
"""Download GGUF file(s) from HuggingFace. Returns local path.
|
||
|
||
Runs WITHOUT self._lock so that unload_model() can set
|
||
_cancel_event at any time. Checks _cancel_event between
|
||
each shard download.
|
||
"""
|
||
try:
|
||
from huggingface_hub import hf_hub_download
|
||
except ImportError:
|
||
raise RuntimeError(
|
||
"huggingface_hub is required for HF model loading. "
|
||
"Install it with: pip install huggingface_hub"
|
||
)
|
||
|
||
# Determine the filename from the variant
|
||
gguf_filename = None
|
||
gguf_extra_shards: list[str] = []
|
||
if hf_variant:
|
||
try:
|
||
from huggingface_hub import list_repo_files
|
||
|
||
files = list_repo_files(hf_repo, token = hf_token)
|
||
variant_lower = hf_variant.lower()
|
||
boundary = re.compile(
|
||
r"(?<![a-zA-Z0-9])" + re.escape(variant_lower) + r"(?![a-zA-Z0-9])"
|
||
)
|
||
gguf_files = sorted(
|
||
f
|
||
for f in files
|
||
if f.endswith(".gguf") and boundary.search(f.lower())
|
||
)
|
||
if gguf_files:
|
||
gguf_filename = gguf_files[0]
|
||
m = _SHARD_FULL_RE.match(gguf_filename)
|
||
if m:
|
||
prefix = m.group(1)
|
||
total = m.group(3)
|
||
sibling_pat = re.compile(
|
||
r"^"
|
||
+ re.escape(prefix)
|
||
+ r"-\d{5}-of-"
|
||
+ re.escape(total)
|
||
+ r"\.gguf$"
|
||
)
|
||
gguf_extra_shards = [
|
||
f for f in gguf_files[1:] if sibling_pat.match(f)
|
||
]
|
||
except Exception as e:
|
||
logger.warning(f"Could not list repo files: {e}")
|
||
|
||
# Offline: resolve variant -> filename from the local HF cache.
|
||
# The heuristic below assumes filenames echo the repo name,
|
||
# which breaks for e.g. Qwen3.6-27B-MTP-GGUF (no "MTP" in file).
|
||
# Match against the rel path (not just basename) so subdir
|
||
# layouts like ``BF16/foo.gguf`` are findable.
|
||
if not gguf_filename:
|
||
try:
|
||
from utils.models.model_config import _iter_hf_cache_snapshots
|
||
|
||
boundary = re.compile(
|
||
r"(?<![a-zA-Z0-9])"
|
||
+ re.escape(hf_variant.lower())
|
||
+ r"(?![a-zA-Z0-9])"
|
||
)
|
||
for snap in _iter_hf_cache_snapshots(hf_repo):
|
||
matches = sorted(
|
||
p.relative_to(snap).as_posix()
|
||
for p in snap.rglob("*.gguf")
|
||
if "mmproj" not in p.name.lower()
|
||
and boundary.search(p.relative_to(snap).as_posix().lower())
|
||
)
|
||
if not matches:
|
||
continue
|
||
gguf_filename = matches[0]
|
||
m = _SHARD_FULL_RE.match(Path(gguf_filename).name)
|
||
if m:
|
||
prefix = m.group(1)
|
||
total = m.group(3)
|
||
sibling_pat = re.compile(
|
||
r"^"
|
||
+ re.escape(prefix)
|
||
+ r"-\d{5}-of-"
|
||
+ re.escape(total)
|
||
+ r"\.gguf$"
|
||
)
|
||
gguf_extra_shards = [
|
||
f
|
||
for f in matches[1:]
|
||
if sibling_pat.match(Path(f).name)
|
||
]
|
||
logger.info(
|
||
"Resolved variant %s -> %s from local HF cache",
|
||
hf_variant,
|
||
gguf_filename,
|
||
)
|
||
break
|
||
except Exception as e:
|
||
logger.debug(f"Offline cache lookup for variant failed: {e}")
|
||
|
||
if not gguf_filename:
|
||
repo_name = hf_repo.split("/")[-1].replace("-GGUF", "")
|
||
gguf_filename = f"{repo_name}-{hf_variant}.gguf"
|
||
|
||
# Check disk space and fall back to a smaller variant if needed
|
||
all_gguf_files = [gguf_filename] + gguf_extra_shards
|
||
try:
|
||
from huggingface_hub import get_paths_info, try_to_load_from_cache
|
||
|
||
path_infos = list(get_paths_info(hf_repo, all_gguf_files, token = hf_token))
|
||
total_bytes = sum((p.size or 0) for p in path_infos)
|
||
|
||
# Subtract bytes already present in the HF cache so we only
|
||
# preflight against what we actually have to download. Without
|
||
# this, re-loading a cached large model (e.g. MiniMax-M2.7-GGUF
|
||
# at 131 GB) fails cold whenever free disk is below the full
|
||
# weight footprint, even though nothing needs downloading.
|
||
already_cached_bytes = 0
|
||
for p in path_infos:
|
||
if not p.size:
|
||
continue
|
||
try:
|
||
cached_path = try_to_load_from_cache(hf_repo, p.path)
|
||
except Exception:
|
||
cached_path = None
|
||
if isinstance(cached_path, str) and os.path.exists(cached_path):
|
||
try:
|
||
on_disk = os.path.getsize(cached_path)
|
||
except OSError:
|
||
on_disk = 0
|
||
# Count as satisfied only when the full blob is present.
|
||
if on_disk >= p.size:
|
||
already_cached_bytes += p.size
|
||
|
||
total_download_bytes = max(0, total_bytes - already_cached_bytes)
|
||
|
||
if total_download_bytes > 0:
|
||
cache_dir = os.environ.get(
|
||
"HF_HUB_CACHE",
|
||
str(Path.home() / ".cache" / "huggingface" / "hub"),
|
||
)
|
||
Path(cache_dir).mkdir(parents = True, exist_ok = True)
|
||
free_bytes = shutil.disk_usage(cache_dir).free
|
||
|
||
total_gb = total_download_bytes / (1024**3)
|
||
free_gb = free_bytes / (1024**3)
|
||
cached_gb = already_cached_bytes / (1024**3)
|
||
|
||
logger.info(
|
||
f"GGUF download: {total_gb:.1f} GB needed "
|
||
f"({cached_gb:.1f} GB already cached), "
|
||
f"{free_gb:.1f} GB free on disk"
|
||
)
|
||
|
||
if total_download_bytes > free_bytes:
|
||
smaller = self._find_smallest_fitting_variant(
|
||
hf_repo,
|
||
free_bytes,
|
||
hf_token,
|
||
)
|
||
if smaller:
|
||
fallback_file, fallback_size = smaller
|
||
logger.info(
|
||
f"Selected variant too large ({total_gb:.1f} GB), "
|
||
f"falling back to {fallback_file} ({fallback_size / (1024**3):.1f} GB)"
|
||
)
|
||
gguf_filename = fallback_file
|
||
_m = _SHARD_RE.match(gguf_filename)
|
||
_prefix = _m.group(1) if _m else None
|
||
if _prefix:
|
||
gguf_extra_shards = sorted(
|
||
f
|
||
for f in all_gguf_files
|
||
if f.startswith(_prefix)
|
||
and f != gguf_filename
|
||
and "mmproj" not in f.lower()
|
||
)
|
||
else:
|
||
gguf_extra_shards = []
|
||
else:
|
||
raise RuntimeError(
|
||
f"Not enough disk space to download any variant. "
|
||
f"Only {free_gb:.1f} GB free in {cache_dir}"
|
||
)
|
||
except RuntimeError:
|
||
raise
|
||
except Exception as e:
|
||
logger.warning(f"Could not check disk space: {e}")
|
||
|
||
gguf_label = f"{hf_repo}/{gguf_filename}" + (
|
||
f" (+{len(gguf_extra_shards)} shards)" if gguf_extra_shards else ""
|
||
)
|
||
logger.info(f"Resolving GGUF: {gguf_label}")
|
||
try:
|
||
if self._cancel_event.is_set():
|
||
raise RuntimeError("Cancelled")
|
||
dl_start = time.monotonic()
|
||
local_path = hf_hub_download(
|
||
repo_id = hf_repo,
|
||
filename = gguf_filename,
|
||
token = hf_token,
|
||
)
|
||
for shard in gguf_extra_shards:
|
||
if self._cancel_event.is_set():
|
||
raise RuntimeError("Cancelled")
|
||
logger.info(f"Resolving GGUF shard: {shard}")
|
||
hf_hub_download(
|
||
repo_id = hf_repo,
|
||
filename = shard,
|
||
token = hf_token,
|
||
)
|
||
except RuntimeError as e:
|
||
if "Cancelled" in str(e):
|
||
raise
|
||
raise RuntimeError(
|
||
f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}"
|
||
)
|
||
except Exception as e:
|
||
raise RuntimeError(
|
||
f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}"
|
||
)
|
||
|
||
dl_elapsed = time.monotonic() - dl_start
|
||
if dl_elapsed < 2.0:
|
||
logger.info(f"GGUF resolved from cache: {local_path}")
|
||
else:
|
||
logger.info(f"GGUF downloaded in {dl_elapsed:.1f}s: {local_path}")
|
||
return local_path
|
||
|
||
def _download_mmproj(
|
||
self,
|
||
*,
|
||
hf_repo: str,
|
||
hf_token: Optional[str] = None,
|
||
) -> Optional[str]:
|
||
"""Download the mmproj (vision projection) file from a GGUF repo.
|
||
|
||
Prefers mmproj-F16.gguf, falls back to any mmproj*.gguf file.
|
||
Returns the local path, or None if no mmproj file exists.
|
||
"""
|
||
|
||
def _pick_mmproj(candidates: list[str]) -> Optional[str]:
|
||
mmproj_files = sorted(
|
||
f
|
||
for f in candidates
|
||
if f.lower().endswith(".gguf") and "mmproj" in Path(f).name.lower()
|
||
)
|
||
if not mmproj_files:
|
||
return None
|
||
for f in mmproj_files:
|
||
if f.lower().endswith("-f16.gguf"):
|
||
return f
|
||
return mmproj_files[0]
|
||
|
||
target: Optional[str] = None
|
||
try:
|
||
from huggingface_hub import list_repo_files
|
||
|
||
target = _pick_mmproj(list_repo_files(hf_repo, token = hf_token))
|
||
except Exception as e:
|
||
logger.debug(f"Could not list repo files for mmproj: {e}")
|
||
|
||
# Offline: resolve mmproj from the local HF cache snapshot, same
|
||
# shape as _download_gguf's offline fallback above.
|
||
if target is None:
|
||
try:
|
||
from utils.models.model_config import _iter_hf_cache_snapshots
|
||
|
||
for snap in _iter_hf_cache_snapshots(hf_repo):
|
||
rel_files = [
|
||
p.relative_to(snap).as_posix() for p in snap.rglob("*.gguf")
|
||
]
|
||
target = _pick_mmproj(rel_files)
|
||
if target is not None:
|
||
logger.info("Resolved mmproj %s from local HF cache", target)
|
||
break
|
||
except Exception as e:
|
||
logger.debug(f"Offline cache lookup for mmproj failed: {e}")
|
||
|
||
if target is None:
|
||
return None
|
||
|
||
try:
|
||
from huggingface_hub import hf_hub_download
|
||
|
||
logger.info(f"Downloading mmproj: {hf_repo}/{target}")
|
||
local_path = hf_hub_download(
|
||
repo_id = hf_repo,
|
||
filename = target,
|
||
token = hf_token,
|
||
)
|
||
return local_path
|
||
except Exception as e:
|
||
logger.warning(f"Could not download mmproj: {e}")
|
||
return None
|
||
|
||
def _resolve_launch_mmproj_path(
|
||
self,
|
||
*,
|
||
model_path: str,
|
||
mmproj_path: Optional[str],
|
||
) -> Optional[str]:
|
||
"""Return mmproj_path iff it exists on disk AND matches the model family.
|
||
|
||
Returns None if mmproj_path is None, missing on disk, or family-mismatched.
|
||
"""
|
||
if not mmproj_path:
|
||
return None
|
||
|
||
mmproj = Path(mmproj_path)
|
||
if not mmproj.is_file():
|
||
logger.warning(f"mmproj file not found: {mmproj_path}")
|
||
return None
|
||
|
||
from utils.models.model_config import mmproj_matches_model_family
|
||
|
||
if not mmproj_matches_model_family(model_path, str(mmproj)):
|
||
logger.warning(
|
||
f"mmproj does not match model family: model={Path(model_path).name} "
|
||
f"mmproj={mmproj.name}"
|
||
)
|
||
return None
|
||
|
||
return str(mmproj)
|
||
|
||
# ── Lifecycle ─────────────────────────────────────────────────
|
||
|
||
# GGUF ``general.architecture`` values for diffusion / image models.
|
||
# llama.cpp proper has no such architectures, so loading one as a chat
|
||
# model dies with "unknown model architecture: '<arch>'". These match
|
||
# the patched stable-diffusion.cpp / ComfyUI-GGUF enums (LLM_ARCH_FLUX,
|
||
# LLM_ARCH_QWEN_IMAGE, ...). Unsloth publishes FLUX and Qwen-Image GGUFs
|
||
# under https://huggingface.co/collections/unsloth/unsloth-diffusion-ggufs.
|
||
# Matched exactly (not as a substring) so a chat arch merely containing a
|
||
# short token like "wan"/"sd1" (e.g. "taiwan") is not misrouted to Images.
|
||
_DIFFUSION_ARCHES = frozenset(
|
||
(
|
||
"qwen_image",
|
||
"flux",
|
||
"sd1",
|
||
"sdxl",
|
||
"sd3",
|
||
"aura",
|
||
"hidream",
|
||
"cosmos",
|
||
"ltxv",
|
||
"hyvid",
|
||
"wan",
|
||
"lumina2",
|
||
)
|
||
)
|
||
|
||
@staticmethod
|
||
def _classify_llama_start_failure(
|
||
output: str,
|
||
gguf_path: Optional[str],
|
||
model_identifier: Optional[str],
|
||
) -> str:
|
||
"""Explain *why* llama-server failed to start, from its output.
|
||
|
||
Several distinct failures all otherwise collapse into the same
|
||
opaque "invalid GGUF or out of memory" message. The worst case is
|
||
a diffusion / image GGUF (FLUX, Qwen-Image, ...) loaded as a chat
|
||
model: the file is perfectly valid and there is plenty of memory,
|
||
but llama.cpp has no such architecture, so the user is told to free
|
||
memory that was never the problem (issue #5842). Pick the most
|
||
specific message the captured output supports.
|
||
"""
|
||
lowered = (output or "").lower()
|
||
|
||
# Detect Ollama source up front so the arch branch can keep the
|
||
# Ollama hint instead of the generic "unsupported arch" message.
|
||
gguf = gguf_path or ""
|
||
is_ollama = (
|
||
".studio_links" in gguf
|
||
or os.sep + "ollama_links" + os.sep in gguf
|
||
or os.sep + ".cache" + os.sep + "ollama" + os.sep in gguf
|
||
or (model_identifier or "").startswith("ollama/")
|
||
)
|
||
|
||
# "unknown model architecture: '<arch>'": diffusion -> Images page,
|
||
# Ollama -> Ollama hint, else a precise "unsupported" message. Exact
|
||
# match so chat archs are never misrouted.
|
||
arch_match = re.search(r"unknown model architecture:\s*'([^']+)'", lowered)
|
||
if arch_match:
|
||
arch = arch_match.group(1)
|
||
if arch in LlamaCppBackend._DIFFUSION_ARCHES:
|
||
return (
|
||
f"'{arch}' is a diffusion (image-generation) GGUF, which "
|
||
"llama-server cannot run as a chat/completion model. Use "
|
||
"Studio's Images page to generate with local diffusion "
|
||
"GGUFs such as FLUX and Qwen-Image."
|
||
)
|
||
if is_ollama:
|
||
return (
|
||
"Some Ollama models do not work with llama.cpp. Try a "
|
||
"different model, or use this model directly through "
|
||
"Ollama instead."
|
||
)
|
||
return (
|
||
f"llama.cpp does not support this GGUF's model architecture "
|
||
f"('{arch}'). The file is valid, but this model type cannot "
|
||
"be run with llama-server."
|
||
)
|
||
|
||
# Other Ollama compat failures that do not name an arch. Only when
|
||
# the output shows a GGUF compat issue, not OOM / missing binaries.
|
||
if is_ollama:
|
||
gguf_compat_hints = (
|
||
"key not found",
|
||
"unknown model architecture",
|
||
"failed to load model",
|
||
)
|
||
if any(h in lowered for h in gguf_compat_hints):
|
||
return (
|
||
"Some Ollama models do not work with llama.cpp. Try a "
|
||
"different model, or use this model directly through "
|
||
"Ollama instead."
|
||
)
|
||
|
||
# Fallback: genuinely unknown failure (OOM, missing binary, ...).
|
||
return (
|
||
"llama-server failed to start. "
|
||
"Check that the GGUF file is valid and you have enough memory."
|
||
)
|
||
|
||
def load_model(
|
||
self,
|
||
*,
|
||
# Local mode: pass a path to a .gguf file
|
||
gguf_path: Optional[str] = None,
|
||
# Vision projection (mmproj) for local vision models
|
||
mmproj_path: Optional[str] = None,
|
||
# HF mode: let llama-server download via -hf "repo:quant"
|
||
hf_repo: Optional[str] = None,
|
||
hf_variant: Optional[str] = None,
|
||
hf_token: Optional[str] = None,
|
||
# Common
|
||
model_identifier: str,
|
||
is_vision: bool = False,
|
||
n_ctx: int = 4096,
|
||
chat_template_override: Optional[str] = None,
|
||
cache_type_kv: Optional[str] = None,
|
||
speculative_type: Optional[str] = None,
|
||
spec_draft_n_max: Optional[int] = None,
|
||
n_threads: Optional[int] = None,
|
||
n_gpu_layers: Optional[int] = None, # Accepted for caller compat, unused
|
||
n_parallel: int = 1,
|
||
extra_args: Optional[List[str]] = None,
|
||
) -> bool:
|
||
"""
|
||
Start llama-server with a GGUF model.
|
||
|
||
Two modes:
|
||
- Local: ``gguf_path="/path/to/model.gguf"`` → uses ``-m``
|
||
- HF: ``hf_repo="unsloth/gemma-3-4b-it-GGUF", hf_variant="Q4_K_M"`` → uses ``-hf``
|
||
|
||
In HF mode, llama-server handles downloading, caching, and
|
||
auto-loading mmproj files for vision models.
|
||
|
||
Returns True if server started and health check passed.
|
||
"""
|
||
# Serialise the whole load so concurrent /load calls never
|
||
# leave two llama-server processes alive (#5401 / #5161). Does
|
||
# not block /unload, /status, /load-progress.
|
||
with self._serial_load_lock:
|
||
# Duplicate /load that raced past the route-level check
|
||
# (the first one hadn't published _healthy=True yet). If the
|
||
# live server already satisfies this request, do nothing.
|
||
if self._already_in_target_state(
|
||
gguf_path = gguf_path,
|
||
model_identifier = model_identifier,
|
||
hf_variant = hf_variant,
|
||
n_ctx = n_ctx,
|
||
cache_type_kv = cache_type_kv,
|
||
speculative_type = speculative_type,
|
||
spec_draft_n_max = spec_draft_n_max,
|
||
chat_template_override = chat_template_override,
|
||
extra_args = extra_args,
|
||
is_vision = is_vision,
|
||
):
|
||
logger.info(
|
||
f"load_model: backend already in target state for "
|
||
f"'{model_identifier}', skipping reload"
|
||
)
|
||
# Retry probe only if a prior attempt didn't complete.
|
||
if not self._audio_probed:
|
||
try:
|
||
detected = self._detect_audio_type_strict()
|
||
self._audio_probed = True
|
||
except Exception as exc:
|
||
logger.debug("Fast-path audio probe failed: %s", exc)
|
||
detected = None
|
||
if detected in ("snac", "bicodec", "dac"):
|
||
with self._lock:
|
||
if not self._healthy:
|
||
return False
|
||
try:
|
||
self.init_audio_codec(detected)
|
||
self._is_audio = True
|
||
self._audio_type = detected
|
||
except Exception as exc:
|
||
logger.warning(
|
||
"Failed to init audio codec '%s': %s",
|
||
detected,
|
||
exc,
|
||
)
|
||
self._audio_probed = False
|
||
return False
|
||
elif detected:
|
||
# csm / whisper / audio_vlm: track type but keep
|
||
# _is_audio False -- GGUF TTS routing only fires
|
||
# for snac/bicodec/dac.
|
||
with self._lock:
|
||
if not self._healthy:
|
||
return False
|
||
self._audio_type = detected
|
||
if not self._healthy:
|
||
return False
|
||
return True
|
||
|
||
self._cancel_event.clear()
|
||
|
||
# ── Phase 1: kill old process (under lock, fast) ──────────
|
||
with self._lock:
|
||
self._kill_process()
|
||
|
||
binary = self._find_llama_server_binary()
|
||
if not binary:
|
||
raise RuntimeError(
|
||
"llama-server binary not found. "
|
||
"Run setup.sh to build it, install llama.cpp, "
|
||
"or set LLAMA_SERVER_PATH environment variable."
|
||
)
|
||
|
||
# ── Phase 2: download (NO lock held, so cancel can proceed) ──
|
||
# Scope HF_HUB_OFFLINE to the download block only when DNS is
|
||
# dead; cleanup runs even on exception so a transient hiccup
|
||
# at the start of one load cannot quarantine future loads.
|
||
if hf_repo:
|
||
with _hf_offline_if_dns_dead():
|
||
model_path = self._download_gguf(
|
||
hf_repo = hf_repo,
|
||
hf_variant = hf_variant,
|
||
hf_token = hf_token,
|
||
)
|
||
# Auto-download mmproj for vision models
|
||
if is_vision and not mmproj_path:
|
||
mmproj_path = self._download_mmproj(
|
||
hf_repo = hf_repo,
|
||
hf_token = hf_token,
|
||
)
|
||
elif gguf_path:
|
||
if not Path(gguf_path).is_file():
|
||
raise FileNotFoundError(f"GGUF file not found: {gguf_path}")
|
||
model_path = gguf_path
|
||
else:
|
||
raise ValueError("Either gguf_path or hf_repo must be provided")
|
||
|
||
# Set identifier early so _read_gguf_metadata can use it for DeepSeek detection
|
||
self._model_identifier = model_identifier
|
||
|
||
# Read GGUF metadata (context_length, chat_template) -- fast, header only
|
||
self._read_gguf_metadata(model_path)
|
||
|
||
# Check cancel after download
|
||
if self._cancel_event.is_set():
|
||
logger.info("Load cancelled after download phase")
|
||
return False
|
||
|
||
# Outside ``self._lock`` so /unload, /cancel, /status are
|
||
# not blocked. ``unload_model`` also records the kill, so
|
||
# the frontend /unload+/load Apply path engages the wait
|
||
# here even though no in-process kill happened.
|
||
self._wait_for_vram_settle(since_kill = self._last_kill_monotonic)
|
||
|
||
# ── Phase 3: start llama-server (under lock) ──────────────
|
||
with self._lock:
|
||
# Re-check cancel inside lock
|
||
if self._cancel_event.is_set():
|
||
logger.info("Load cancelled before server start")
|
||
return False
|
||
|
||
self._port = self._find_free_port()
|
||
|
||
# Select GPU(s) based on model size + estimated KV cache.
|
||
# Seed safe defaults before GPU probing so the except path
|
||
# still has valid state to publish.
|
||
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)
|
||
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"
|
||
)
|
||
effective_ctx = (
|
||
requested_ctx if requested_ctx > 0 else (self._context_length or 0)
|
||
)
|
||
max_available_ctx = self._context_length or effective_ctx
|
||
gpus: list[tuple[int, int]] = []
|
||
try:
|
||
model_size = self._get_gguf_size_bytes(model_path)
|
||
gpus = self._get_gpu_free_memory()
|
||
|
||
# Resolve effective context: 0 means let llama-server use the
|
||
# model's native length. Only expand to a known native length
|
||
# if metadata is available; otherwise preserve 0 as a sentinel.
|
||
if requested_ctx > 0:
|
||
effective_ctx = requested_ctx
|
||
elif self._context_length is not None:
|
||
effective_ctx = self._context_length
|
||
else:
|
||
effective_ctx = 0
|
||
original_ctx = effective_ctx
|
||
# Default UI ceiling to the model's native context length.
|
||
# GPU/VRAM-fit logic below may shrink this if hardware is limited.
|
||
max_available_ctx = self._context_length or effective_ctx
|
||
|
||
# Will MTP engage on this load? If so, the auto-fit
|
||
# budget needs to reserve extra VRAM for the draft
|
||
# model's KV cache + compute graph. Mirrors the
|
||
# canonical-mode resolver in _build_speculative_flags:
|
||
# forced mtp / mtp+ngram always engage; auto only
|
||
# engages on an MTP GGUF >= 3B (sub-3B auto falls
|
||
# back to ngram-mod which doesn't need headroom);
|
||
# ngram / ngram-simple / off never engage MTP.
|
||
_mtp_canonical = _canonicalize_spec_mode(speculative_type)
|
||
_mtp_effective = _mtp_canonical or "auto"
|
||
_mtp_size_for_fit = _extract_model_size_b(model_identifier)
|
||
_mtp_sub_3b_for_fit = (
|
||
_mtp_size_for_fit is not None and _mtp_size_for_fit < 3.0
|
||
)
|
||
_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)
|
||
)
|
||
and not _mtp_sub_3b_for_fit
|
||
)
|
||
)
|
||
)
|
||
|
||
# Auto-cap context to fit in GPU VRAM and select GPUs.
|
||
#
|
||
# Two policies depending on whether the user set n_ctx:
|
||
#
|
||
# Explicit n_ctx (user chose a context length):
|
||
# Honor it. Try the full requested context with _select_gpus
|
||
# (which uses as many GPUs as needed). Only cap if it doesn't
|
||
# fit on any GPU combination.
|
||
#
|
||
# Auto n_ctx=0 (model's native context):
|
||
# Prefer fewer GPUs with reduced context over more GPUs,
|
||
# since multi-GPU is slower and the user didn't ask for a
|
||
# specific context length.
|
||
gpu_indices, use_fit = None, True
|
||
explicit_ctx = requested_ctx > 0
|
||
|
||
if gpus and self._can_estimate_kv() and effective_ctx > 0:
|
||
# Compute the largest hardware-aware cap from the model's
|
||
# native context across all usable GPU subsets (for UI
|
||
# 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
|
||
)
|
||
best_cap = 0
|
||
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)
|
||
capped = self._fit_context_to_vram(
|
||
native_ctx_for_cap,
|
||
pool_mib,
|
||
model_size,
|
||
cache_type_kv,
|
||
n_parallel = n_parallel,
|
||
mtp_engaged = _mtp_will_engage,
|
||
)
|
||
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 * 0.90:
|
||
best_cap = max(best_cap, capped)
|
||
if best_cap > 0:
|
||
max_available_ctx = best_cap
|
||
else:
|
||
# Weights exceed 90% of every GPU subset's free
|
||
# memory, so there is no fitting context. Anchor
|
||
# the UI's "safe zone" threshold at 4096 (the
|
||
# spec's default when the model cannot fit) so
|
||
# the ctx slider shows the "might be slower"
|
||
# warning as soon as the user drags above the
|
||
# fallback default instead of never.
|
||
max_available_ctx = min(4096, native_ctx_for_cap)
|
||
|
||
if explicit_ctx:
|
||
# Honor the user's requested context verbatim. If it
|
||
# fits, pin GPUs and skip --fit; if it doesn't, ship
|
||
# -c <user_ctx> --fit on and let llama-server flex
|
||
# -ngl (CPU layer offload). The UI is expected to
|
||
# have surfaced the "might be slower" warning before
|
||
# the user submitted a ctx above the fit ceiling.
|
||
requested_total = (
|
||
model_size
|
||
+ self._estimate_kv_cache_bytes(
|
||
effective_ctx, cache_type_kv, n_parallel = n_parallel
|
||
)
|
||
)
|
||
gpu_indices, use_fit = self._select_gpus(
|
||
requested_total, gpus
|
||
)
|
||
# No silent shrink: effective_ctx stays == requested_ctx.
|
||
else:
|
||
# Auto context: prefer fewer GPUs, cap context
|
||
# to fit. Same headroom threshold as
|
||
# _select_gpus (#5106).
|
||
ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
|
||
pin_fraction = self._GPU_PIN_VRAM_FRACTION
|
||
for n_gpus in range(1, len(ranked) + 1):
|
||
subset = ranked[:n_gpus]
|
||
pool_mib = sum(free for _, free in subset)
|
||
capped = self._fit_context_to_vram(
|
||
effective_ctx,
|
||
pool_mib,
|
||
model_size,
|
||
cache_type_kv,
|
||
n_parallel = n_parallel,
|
||
mtp_engaged = _mtp_will_engage,
|
||
)
|
||
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:
|
||
effective_ctx = capped
|
||
gpu_indices = sorted(idx for idx, _ in subset)
|
||
use_fit = False
|
||
break
|
||
else:
|
||
# Native ctx doesn't fit. Drop to 4096 and
|
||
# re-check before deferring to --fit on:
|
||
# a model that overflows at 131k may pin
|
||
# comfortably with a 4096 KV cache (#5106).
|
||
effective_ctx = min(4096, effective_ctx)
|
||
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:
|
||
gpu_indices = sorted(
|
||
idx for idx, _ in subset
|
||
)
|
||
use_fit = False
|
||
break
|
||
|
||
elif gpus:
|
||
# Can't estimate KV -- fall back to file-size-only check.
|
||
# Without KV estimation we cannot prove a hardware cap, so
|
||
# keep the ceiling at the native context (already the default).
|
||
logger.debug(
|
||
"Falling back to file-size-only GPU selection",
|
||
model_size_gb = round(model_size / (1024**3), 2),
|
||
)
|
||
gpu_indices, use_fit = self._select_gpus(model_size, gpus)
|
||
if use_fit and not explicit_ctx:
|
||
# Weights don't fit on any subset. Default the UI to
|
||
# 4096 so the slider doesn't land on an unusable native
|
||
# context. --fit on will flex -ngl at runtime.
|
||
effective_ctx = (
|
||
min(4096, effective_ctx) if effective_ctx > 0 else 4096
|
||
)
|
||
|
||
if effective_ctx < original_ctx:
|
||
kv_est = self._estimate_kv_cache_bytes(
|
||
effective_ctx, cache_type_kv, n_parallel = n_parallel
|
||
)
|
||
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)"
|
||
)
|
||
|
||
kv_cache_bytes = self._estimate_kv_cache_bytes(
|
||
effective_ctx, cache_type_kv, n_parallel = n_parallel
|
||
)
|
||
logger.info(
|
||
f"GGUF size: {model_size / (1024**3):.1f} GB, "
|
||
f"est. KV cache: {kv_cache_bytes / (1024**3):.1f} GB, "
|
||
f"context: {effective_ctx}, "
|
||
f"GPUs free: {gpus}, selected: {gpu_indices}, fit: {use_fit}"
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"GPU selection failed ({e}), using --fit on")
|
||
gpu_indices, use_fit = None, True
|
||
effective_ctx = requested_ctx # fall back to original
|
||
|
||
launch_mmproj_path = self._resolve_launch_mmproj_path(
|
||
model_path = model_path,
|
||
mmproj_path = mmproj_path,
|
||
)
|
||
# Need both a resolved mmproj AND the config vision flag; a stray
|
||
# mmproj passing the family-name heuristic must not flip a non-VLM
|
||
# GGUF into vision mode.
|
||
effective_is_vision = bool(launch_mmproj_path) and bool(is_vision)
|
||
if is_vision and not effective_is_vision:
|
||
logger.warning(
|
||
"Vision-capable GGUF loaded without a usable mmproj; "
|
||
"image input will be disabled for this session"
|
||
)
|
||
|
||
cmd = [
|
||
binary,
|
||
"-m",
|
||
model_path,
|
||
"--port",
|
||
str(self._port),
|
||
"-c",
|
||
str(effective_ctx) if effective_ctx > 0 else "0",
|
||
"--parallel",
|
||
str(n_parallel),
|
||
"--flash-attn",
|
||
"on", # Force flash attention for speed
|
||
# Error out at n_ctx instead of silently rotating the KV cache; frontend catches it and points the user at "Context Length".
|
||
"--no-context-shift",
|
||
]
|
||
|
||
if use_fit:
|
||
cmd.extend(["--fit", "on"])
|
||
elif gpu_indices is not None:
|
||
# Model fits on selected GPU(s) -- offload all layers
|
||
cmd.extend(["-ngl", "-1"])
|
||
|
||
# -1 = llama.cpp auto-detect (physical cores). Pass explicitly so we
|
||
# do not inherit llama-server's internal default, which has historically
|
||
# varied (hardware concurrency incl. hyperthreads on some builds).
|
||
cmd.extend(
|
||
["--threads", str(n_threads if n_threads is not None else -1)]
|
||
)
|
||
|
||
# Always enable Jinja chat template rendering for proper template support
|
||
cmd.extend(["--jinja"])
|
||
|
||
# KV cache data type
|
||
_valid_cache_types = {
|
||
"f16",
|
||
"bf16",
|
||
"q8_0",
|
||
"q4_0",
|
||
"q4_1",
|
||
"q5_0",
|
||
"q5_1",
|
||
"iq4_nl",
|
||
"f32",
|
||
}
|
||
if cache_type_kv and cache_type_kv in _valid_cache_types:
|
||
cmd.extend(
|
||
[
|
||
"--cache-type-k",
|
||
cache_type_kv,
|
||
"--cache-type-v",
|
||
cache_type_kv,
|
||
]
|
||
)
|
||
self._cache_type_kv = cache_type_kv
|
||
logger.info(f"KV cache type: {cache_type_kv}")
|
||
else:
|
||
self._cache_type_kv = None
|
||
|
||
# Speculative decoding (n-gram self-speculation, zero VRAM cost)
|
||
# ngram-mod: ~16 MB shared hash pool, constant memory/complexity,
|
||
# variable draft lengths. Helps most when the model repeats
|
||
# existing text (code refactoring, summarization, reasoning).
|
||
# For general chat with low repetition, overhead is ~5 ms.
|
||
#
|
||
# Benchmarks from upstream llama.cpp speculative-decoding PRs:
|
||
# Scenario | Without | With | Speedup
|
||
# gpt-oss-120b code refactor | 181 t/s | 446 t/s | 2.5x
|
||
# Qwen3-235B offloaded | 12 t/s | 21 t/s | 1.8x
|
||
# gpt-oss-120b repeat (92% accept)| 181 t/s | 814 t/s | 4.5x
|
||
#
|
||
# Params from llama.cpp server README:
|
||
# --spec-ngram-mod-n-match 24 (lookup length)
|
||
# --spec-ngram-mod-n-min 48 --spec-ngram-mod-n-max 64
|
||
# (MoEs need long drafts; dense models can reduce these)
|
||
# ref: https://github.com/ggml-org/llama.cpp/blob/master/docs/speculative.md
|
||
# ref: https://github.com/ggml-org/llama.cpp/pull/19164
|
||
# ref: https://github.com/ggml-org/llama.cpp/pull/18471
|
||
# draft-mtp: MTP heads on Unsloth's *-MTP GGUFs
|
||
# (llama.cpp #22673). Auto-enabled via nextn_predict_layers,
|
||
# fallback to -MTP in name. GPU: MTP-only. CPU/Mac: chain
|
||
# with ngram-mod. See unsloth.ai/docs/models/qwen3.6#mtp-guide.
|
||
spec_flags = self._build_speculative_flags(
|
||
speculative_type = speculative_type,
|
||
spec_draft_n_max = spec_draft_n_max,
|
||
extra_args = extra_args,
|
||
model_identifier = model_identifier,
|
||
model_path = model_path,
|
||
gpus = bool(gpus),
|
||
binary = binary,
|
||
)
|
||
cmd.extend(spec_flags)
|
||
|
||
# Apply custom chat template override if provided
|
||
self._chat_template_override = chat_template_override
|
||
if chat_template_override:
|
||
import tempfile
|
||
|
||
flags = detect_reasoning_flags(
|
||
chat_template_override,
|
||
self._model_identifier,
|
||
log_source = "GGUF chat template override",
|
||
)
|
||
self._supports_reasoning = flags["supports_reasoning"]
|
||
self._reasoning_style = flags["reasoning_style"]
|
||
self._reasoning_always_on = flags["reasoning_always_on"]
|
||
self._supports_preserve_thinking = flags[
|
||
"supports_preserve_thinking"
|
||
]
|
||
self._supports_tools = flags["supports_tools"]
|
||
|
||
self._chat_template_file = tempfile.NamedTemporaryFile(
|
||
mode = "w",
|
||
suffix = ".jinja",
|
||
delete = False,
|
||
prefix = "unsloth_chat_template_",
|
||
)
|
||
self._chat_template_file.write(chat_template_override)
|
||
self._chat_template_file.close()
|
||
cmd.extend(["--chat-template-file", self._chat_template_file.name])
|
||
logger.info(
|
||
f"Using custom chat template file: {self._chat_template_file.name}"
|
||
)
|
||
|
||
# For reasoning models, set default thinking mode.
|
||
# Qwen3.5/3.6 models below 9B (0.8B, 2B, 4B) disable thinking by default.
|
||
# Only 9B and larger enable thinking.
|
||
# Always-on templates ignore the kwarg entirely, so skip.
|
||
if self._supports_reasoning and not self._reasoning_always_on:
|
||
thinking_default = True
|
||
mid = (model_identifier or "").lower()
|
||
if "qwen3.5" in mid or "qwen3.6" in mid:
|
||
size_val = _extract_model_size_b(mid)
|
||
if size_val is not None and size_val < 9:
|
||
thinking_default = False
|
||
self._reasoning_default = thinking_default
|
||
reasoning_kw = self._reasoning_kwargs(thinking_default)
|
||
cmd.extend(
|
||
[
|
||
"--chat-template-kwargs",
|
||
json.dumps(reasoning_kw),
|
||
]
|
||
)
|
||
logger.info(f"Reasoning model: {reasoning_kw} by default")
|
||
|
||
if launch_mmproj_path and effective_is_vision:
|
||
cmd.extend(["--mmproj", launch_mmproj_path])
|
||
logger.info(f"Using mmproj for vision: {launch_mmproj_path}")
|
||
|
||
# Option C: add --api-key for direct client access when enabled
|
||
import os as _os
|
||
import secrets as _secrets
|
||
|
||
if _os.getenv("UNSLOTH_DIRECT_STREAM", "0") == "1":
|
||
self._api_key = _secrets.token_urlsafe(32)
|
||
cmd.extend(["--api-key", self._api_key])
|
||
logger.info(
|
||
"llama-server started with --api-key for direct streaming"
|
||
)
|
||
else:
|
||
self._api_key = None
|
||
|
||
# User-supplied pass-through args go last so llama.cpp's
|
||
# last-wins flag parsing lets the user override Studio's
|
||
# auto-set tier-2 flags (e.g. --cache-type-k, --spec-type).
|
||
# The route layer has already validated this list against
|
||
# the managed-flag denylist via validate_extra_args().
|
||
if extra_args:
|
||
cmd.extend(str(a) for a in extra_args)
|
||
logger.info(
|
||
f"Appending user extra args to llama-server: {list(extra_args)}"
|
||
)
|
||
|
||
_log_cmd = list(cmd)
|
||
if "--api-key" in _log_cmd:
|
||
_ki = _log_cmd.index("--api-key") + 1
|
||
if _ki < len(_log_cmd):
|
||
_log_cmd[_ki] = "<redacted>"
|
||
logger.info(f"Starting llama-server: {' '.join(_log_cmd)}")
|
||
|
||
# Set library paths so llama-server can find its shared libs and CUDA DLLs
|
||
import os
|
||
import sys
|
||
|
||
env = child_env_without_native_path_secret()
|
||
binary_dir = str(Path(binary).parent)
|
||
|
||
# AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use
|
||
# shared system RAM. setdefault so a user value wins.
|
||
if self._amd_apu_wants_unified_memory():
|
||
env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1")
|
||
logger.info(
|
||
"AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1"
|
||
)
|
||
|
||
if sys.platform == "win32":
|
||
# See _build_windows_path_dirs for ordering. #5106.
|
||
path_dirs = self._build_windows_path_dirs(
|
||
binary_dir,
|
||
sys.prefix,
|
||
os.environ.get("CUDA_PATH", ""),
|
||
)
|
||
existing_path = env.get("PATH", "")
|
||
env["PATH"] = ";".join(path_dirs) + ";" + existing_path
|
||
|
||
# ROCm: the llama.cpp prebuilt bundles its own rocblas.dll
|
||
# but NOT the Tensile kernel library files it needs
|
||
# (rocblas/library/TensileLibrary*.dat + *.hsaco). The
|
||
# bundled DLL searches relative to its own location by
|
||
# default (i.e. <binary_dir>/rocblas/library/) which does
|
||
# not exist, causing a silent crash on the first GEMM.
|
||
# ROCBLAS_TENSILE_LIBPATH overrides that search to point at
|
||
# the ROCm installation where the kernel files actually are.
|
||
_hip_path = os.environ.get(
|
||
"HIP_PATH", os.environ.get("ROCM_PATH", "")
|
||
)
|
||
if _hip_path:
|
||
_rocblas_lib = os.path.join(
|
||
_hip_path, "bin", "rocblas", "library"
|
||
)
|
||
if os.path.isdir(_rocblas_lib):
|
||
env.setdefault("ROCBLAS_TENSILE_LIBPATH", _rocblas_lib)
|
||
else:
|
||
# Linux: set LD_LIBRARY_PATH for shared libs next to the binary
|
||
# and CUDA runtime libs (libcudart, libcublas, etc.)
|
||
import platform
|
||
|
||
lib_dirs = [binary_dir]
|
||
_arch = platform.machine() # x86_64, aarch64, etc.
|
||
|
||
# Pip-installed nvidia CUDA runtime libs (e.g. torch's
|
||
# bundled cuda-bindings). The prebuilt llama.cpp binary
|
||
# links against libcudart.so.13 / libcublas.so.13 which
|
||
# live here, not in /usr/local/cuda.
|
||
import glob as _glob
|
||
|
||
for _nv_pattern in [
|
||
os.path.join(
|
||
sys.prefix,
|
||
"lib",
|
||
"python*",
|
||
"site-packages",
|
||
"nvidia",
|
||
"cu*",
|
||
"lib",
|
||
),
|
||
os.path.join(
|
||
sys.prefix,
|
||
"lib",
|
||
"python*",
|
||
"site-packages",
|
||
"nvidia",
|
||
"cudnn",
|
||
"lib",
|
||
),
|
||
os.path.join(
|
||
sys.prefix,
|
||
"lib",
|
||
"python*",
|
||
"site-packages",
|
||
"nvidia",
|
||
"nvjitlink",
|
||
"lib",
|
||
),
|
||
]:
|
||
for _nv_dir in _glob.glob(_nv_pattern):
|
||
if os.path.isdir(_nv_dir):
|
||
lib_dirs.append(_nv_dir)
|
||
|
||
for cuda_lib in [
|
||
"/usr/local/cuda/lib64",
|
||
f"/usr/local/cuda/targets/{_arch}-linux/lib",
|
||
# Fallback CUDA compat paths (e.g. binary built with
|
||
# CUDA 12 on a system where default /usr/local/cuda
|
||
# points to CUDA 13+).
|
||
"/usr/local/cuda-12/lib64",
|
||
"/usr/local/cuda-12.8/lib64",
|
||
f"/usr/local/cuda-12/targets/{_arch}-linux/lib",
|
||
f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib",
|
||
]:
|
||
if os.path.isdir(cuda_lib):
|
||
lib_dirs.append(cuda_lib)
|
||
existing_ld = env.get("LD_LIBRARY_PATH", "")
|
||
new_ld = ":".join(lib_dirs)
|
||
env["LD_LIBRARY_PATH"] = (
|
||
f"{new_ld}:{existing_ld}" if existing_ld else new_ld
|
||
)
|
||
|
||
# Pin to selected GPU(s). On ROCm, llama-server (and any torch
|
||
# in the subprocess) honors HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES;
|
||
# narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child seeing
|
||
# the full HIP/ROCR set the parent inherited.
|
||
if gpu_indices is not None:
|
||
pinned = ",".join(str(i) for i in gpu_indices)
|
||
env["CUDA_VISIBLE_DEVICES"] = pinned
|
||
try:
|
||
import torch as _torch
|
||
|
||
if getattr(_torch.version, "hip", None) is not None:
|
||
env["HIP_VISIBLE_DEVICES"] = pinned
|
||
env["ROCR_VISIBLE_DEVICES"] = pinned
|
||
except Exception as e:
|
||
logger.debug(
|
||
"Failed to set ROCm visibility env vars for child: %s", e
|
||
)
|
||
|
||
# Defensive kill: if a concurrent load slipped past Phase 1
|
||
# (because its `self._process` was None at the time) and
|
||
# already stored a Popen handle here, drop that orphan
|
||
# before we overwrite the reference. See issue #5161.
|
||
self._kill_process()
|
||
|
||
self._stdout_lines = []
|
||
# Tee llama-server output to a dedicated log file so a
|
||
# post-mortem in CI (or after a remote-debug session)
|
||
# has the full subprocess trail even when the parent
|
||
# only stored the last 50 lines. Path lives under the
|
||
# studio home so it ships in the same place all other
|
||
# Studio logs live.
|
||
self._llama_log_fh = None
|
||
try:
|
||
log_dir = _swa_cache_path().parent / "logs" / "llama-server"
|
||
log_dir.mkdir(parents = True, exist_ok = True)
|
||
self._llama_log_path = (
|
||
log_dir / f"llama-{int(time.time())}-port-{self._port}.log"
|
||
)
|
||
self._llama_log_fh = open(
|
||
self._llama_log_path,
|
||
"w",
|
||
encoding = "utf-8",
|
||
buffering = 1,
|
||
)
|
||
logger.info(f"llama-server stdout/stderr -> {self._llama_log_path}")
|
||
except OSError as e:
|
||
# Best-effort; never block the load on logging.
|
||
logger.debug(f"Could not open llama-server log file: {e}")
|
||
self._llama_log_path = None
|
||
self._process = subprocess.Popen(
|
||
cmd,
|
||
stdout = subprocess.PIPE,
|
||
stderr = subprocess.STDOUT,
|
||
text = True,
|
||
env = env,
|
||
**_windows_hidden_subprocess_kwargs(),
|
||
)
|
||
|
||
# Start background thread to drain stdout and prevent pipe deadlock
|
||
self._stdout_thread = threading.Thread(
|
||
target = self._drain_stdout, daemon = True, name = "llama-stdout"
|
||
)
|
||
self._stdout_thread.start()
|
||
|
||
# Store the resolved on-disk path, not the caller's kwarg. In
|
||
# HF mode the caller passes gguf_path=None and the real path
|
||
# (``model_path``) is what llama-server is actually mmap'ing.
|
||
# Downstream consumers (load_progress, log lines, etc.) need
|
||
# the path that exists on disk.
|
||
self._gguf_path = model_path
|
||
self._hf_repo = hf_repo
|
||
# For local GGUF files, extract variant from filename if not provided
|
||
if hf_variant:
|
||
self._hf_variant = hf_variant
|
||
elif gguf_path:
|
||
try:
|
||
from utils.models.model_config import _extract_quant_label
|
||
|
||
self._hf_variant = _extract_quant_label(gguf_path)
|
||
except Exception:
|
||
self._hf_variant = None
|
||
else:
|
||
self._hf_variant = None
|
||
self._is_vision = effective_is_vision
|
||
self._model_identifier = model_identifier
|
||
|
||
# Store the effective (possibly capped) context separately.
|
||
# Do NOT overwrite _context_length -- it holds the model's native
|
||
# context length from GGUF metadata and is used for display/info.
|
||
self._effective_context_length = (
|
||
effective_ctx if effective_ctx > 0 else self._context_length
|
||
)
|
||
self._max_context_length = (
|
||
max_available_ctx
|
||
if max_available_ctx > 0
|
||
else self._effective_context_length
|
||
)
|
||
|
||
# Wait for llama-server to become healthy
|
||
if not self._wait_for_health(timeout = 600.0):
|
||
self._kill_process()
|
||
raise RuntimeError(
|
||
self._classify_llama_start_failure(
|
||
"\n".join(self._stdout_lines[-50:]),
|
||
gguf_path,
|
||
self._model_identifier,
|
||
)
|
||
)
|
||
|
||
self._healthy = True
|
||
|
||
# Commit caller intent only after _healthy=True so a
|
||
# failed startup can't poison the next inheritance check.
|
||
# None keeps prior, [] clears, list sets. Source records
|
||
# the caller's hf_variant (None for local files) so the
|
||
# route's same_source check stays symmetric.
|
||
if extra_args is not None:
|
||
self._extra_args = list(extra_args)
|
||
self._extra_args_source = (model_identifier, hf_variant)
|
||
self._requested_n_ctx = int(n_ctx)
|
||
|
||
# Catch silent CPU fallback when GPU was intended (#5106).
|
||
self._gpu_offload_active = self._classify_gpu_offload(
|
||
gpu_indices is not None or use_fit, gpus or []
|
||
)
|
||
if self._gpu_offload_active is False:
|
||
logger.warning(
|
||
"llama-server appears to have loaded the model entirely "
|
||
"on CPU even though Studio detected at least one GPU. "
|
||
"This usually means the prebuilt binary's GPU backend "
|
||
"failed to load -- on Windows, cudart64_X.dll / "
|
||
"cublas64_X.dll could not be resolved. Reinstall the "
|
||
"Studio llama.cpp prebuilt or install a matching CUDA "
|
||
"toolkit (issue unslothai/unsloth#5106).",
|
||
)
|
||
|
||
logger.info(
|
||
f"llama-server ready on port {self._port} "
|
||
f"for model '{model_identifier}'"
|
||
)
|
||
|
||
# Probe outside _lock (interruptible by /unload); init inside.
|
||
self._is_audio = False
|
||
self._audio_type = None
|
||
self._audio_probed = False
|
||
try:
|
||
detected = self._detect_audio_type_strict()
|
||
self._audio_probed = True
|
||
except Exception as exc:
|
||
logger.debug("Audio probe failed: %s", exc)
|
||
detected = None
|
||
if detected in ("snac", "bicodec", "dac"):
|
||
with self._lock:
|
||
if not self._healthy:
|
||
return False
|
||
try:
|
||
self.init_audio_codec(detected)
|
||
self._is_audio = True
|
||
self._audio_type = detected
|
||
except Exception as exc:
|
||
# Surface as HTTP 500 -- matches pre-PR contract.
|
||
logger.warning(
|
||
"Failed to init audio codec '%s': %s",
|
||
detected,
|
||
exc,
|
||
)
|
||
self._audio_probed = False
|
||
return False
|
||
elif detected:
|
||
# csm / whisper / audio_vlm: track type but keep _is_audio
|
||
# False -- GGUF TTS routing only fires for snac/bicodec/dac.
|
||
with self._lock:
|
||
if not self._healthy:
|
||
return False
|
||
self._audio_type = detected
|
||
|
||
if not self._healthy:
|
||
return False
|
||
return True
|
||
|
||
def _build_speculative_flags(
|
||
self,
|
||
*,
|
||
speculative_type: Optional[str],
|
||
spec_draft_n_max: Optional[int],
|
||
extra_args: Optional[List[str]],
|
||
model_identifier: str,
|
||
model_path: Optional[str],
|
||
gpus: bool,
|
||
binary: Optional[str],
|
||
) -> List[str]:
|
||
"""Return the llama-server flag list for the requested spec mode.
|
||
|
||
Side effects: sets ``self._speculative_type`` (resolved internal
|
||
emit), ``self._requested_spec_mode`` (canonical UI mode for the
|
||
status round-trip), and ``self._spec_draft_n_max`` (user override
|
||
only; None when the platform default applies).
|
||
|
||
Speculative decoding (n-gram self-speculation, zero VRAM cost):
|
||
ngram-mod uses a ~16 MB shared hash pool, constant memory /
|
||
complexity, variable draft lengths. Helps most when the model
|
||
repeats existing text (code refactor, summarisation, reasoning).
|
||
For general chat with low repetition, overhead is ~5 ms.
|
||
|
||
Benchmarks from upstream llama.cpp speculative-decoding PRs:
|
||
Scenario | Without | With | Speedup
|
||
gpt-oss-120b code refactor | 181 t/s | 446 t/s | 2.5x
|
||
Qwen3-235B offloaded | 12 t/s | 21 t/s | 1.8x
|
||
gpt-oss-120b repeat (92% accept)| 181 t/s | 814 t/s | 4.5x
|
||
|
||
Sub-3B dense MTP regresses vs spec-off because the draft head's
|
||
per-token cost exceeds the acceptance savings at this scale.
|
||
Q4_K_XL clean bench (each prompt once after an unrelated warmup)
|
||
on B200 + x86 CPU:
|
||
0.8B GPU: draft-mtp n=2 = 0.58x vs OFF; ngram-only = 1.10x
|
||
2B GPU: draft-mtp n=2 = 0.82x vs OFF; OFF or ngram = 1.00x
|
||
0.8B CPU: chained n=2 = 0.86x vs OFF; ngram-only = 1.19x
|
||
2B CPU: chained n=2 = 0.83x vs OFF; ngram-only = 1.01x
|
||
4B+ GPU/CPU: spec on is a net win (1.08x-1.46x).
|
||
Auto falls back to ngram-mod (zero-VRAM, near-zero idle cost on
|
||
diverse content); forced MTP variants engage anyway and just log
|
||
a warning per the user's choice.
|
||
"""
|
||
flags: List[str] = []
|
||
# Reset; emit branches re-set on the resolved emission.
|
||
self._spec_draft_n_max = None
|
||
self._speculative_type = None
|
||
|
||
# Canonical UI-facing requested mode: auto / mtp / ngram /
|
||
# mtp+ngram / off / ngram-simple. Legacy values are mapped via
|
||
# _canonicalize_spec_mode (default->auto, draft-mtp->mtp,
|
||
# ngram-mod->ngram, "ngram-mod,draft-mtp"->mtp+ngram).
|
||
canonical_mode = _canonicalize_spec_mode(speculative_type)
|
||
is_mtp_model = bool(self._nextn_predict_layers) or (
|
||
_is_mtp_model_name(model_identifier, model_path)
|
||
)
|
||
user_owns_spec_type = _extra_args_set_spec_type(extra_args)
|
||
_mtp_size_b = _extract_model_size_b(model_identifier)
|
||
_mtp_too_small = _mtp_size_b is not None and _mtp_size_b < 3.0
|
||
|
||
if user_owns_spec_type:
|
||
# User --spec-type in extra_args wins outright; suppress
|
||
# auto-emit so we don't emit a duplicate / conflicting
|
||
# spec block. Record requested mode as None.
|
||
self._requested_spec_mode = None
|
||
return flags
|
||
|
||
effective_mode = canonical_mode or "auto"
|
||
self._requested_spec_mode = effective_mode
|
||
|
||
def _resolved_draft_n_max() -> int:
|
||
# User override wins; else platform default (the B200 / x86
|
||
# clean-sweep sweet spot from PR #5582 is n=2 GPU, n=3 CPU;
|
||
# raising past 3 starts to regress on essay-style
|
||
# low-acceptance prompts).
|
||
if spec_draft_n_max is not None:
|
||
n = int(spec_draft_n_max)
|
||
self._spec_draft_n_max = n
|
||
return n
|
||
return 2 if gpus else 3
|
||
|
||
def _emit_mtp(*, chain_ngram: bool) -> bool:
|
||
"""Append --spec-type mtp[/draft-mtp][,ngram-mod] + n-max."""
|
||
caps = self.probe_server_capabilities(binary)
|
||
mtp_token = caps.get("mtp_token") if caps else None
|
||
if not mtp_token:
|
||
logger.warning(
|
||
"Requested MTP speculative decoding but "
|
||
"llama-server lacks --spec-type mtp/draft-mtp; "
|
||
"run `unsloth studio update`. Loading without "
|
||
"speculative decoding."
|
||
)
|
||
return False
|
||
draft_n_max = _resolved_draft_n_max()
|
||
n_max_flag = caps.get("spec_draft_n_max_flag") or "--spec-draft-n-max"
|
||
if chain_ngram:
|
||
ngram_knobs = _build_ngram_mod_flags(caps)
|
||
if ngram_knobs:
|
||
spec_value = f"ngram-mod,{mtp_token}"
|
||
else:
|
||
logger.warning(
|
||
"llama-server lacks ngram-mod tuning "
|
||
"flags; loading MTP only (no ngram chain)"
|
||
)
|
||
spec_value = mtp_token
|
||
flags.extend(
|
||
[
|
||
"--spec-type",
|
||
spec_value,
|
||
n_max_flag,
|
||
str(draft_n_max),
|
||
]
|
||
)
|
||
flags.extend(ngram_knobs)
|
||
else:
|
||
flags.extend(
|
||
[
|
||
"--spec-type",
|
||
mtp_token,
|
||
n_max_flag,
|
||
str(draft_n_max),
|
||
]
|
||
)
|
||
self._speculative_type = "draft-mtp"
|
||
chain_label = "chained ngram-mod" if chain_ngram else "MTP-only"
|
||
logger.info(f"Spec decoding: {mtp_token} ({chain_label})")
|
||
return True
|
||
|
||
def _emit_ngram_mod() -> bool:
|
||
"""Append --spec-type ngram-mod + flag-set knobs."""
|
||
ngram_caps = self.probe_server_capabilities(binary)
|
||
ngram_knobs = _build_ngram_mod_flags(ngram_caps)
|
||
flags.extend(["--spec-type", "ngram-mod"])
|
||
if not ngram_knobs:
|
||
logger.warning(
|
||
"llama-server lacks ngram-mod tuning "
|
||
"flags; loading without --spec-ngram-mod-* knobs"
|
||
)
|
||
flags.extend(ngram_knobs)
|
||
self._speculative_type = "ngram-mod"
|
||
logger.info("Spec decoding: ngram-mod")
|
||
return True
|
||
|
||
if effective_mode == "off":
|
||
return flags # nothing to emit
|
||
if effective_mode == "ngram-simple":
|
||
flags.extend(["--spec-type", "ngram-simple"])
|
||
self._speculative_type = "ngram-simple"
|
||
return flags
|
||
if effective_mode == "ngram":
|
||
_emit_ngram_mod()
|
||
return flags
|
||
if effective_mode == "mtp":
|
||
if _mtp_too_small:
|
||
logger.warning(
|
||
f"Forcing MTP on a {_mtp_size_b:.1f}B model; "
|
||
"the bench shows draft-mtp regresses below 3B. "
|
||
"Engaging anyway (user override)."
|
||
)
|
||
elif not is_mtp_model:
|
||
logger.warning(
|
||
"Forcing MTP on a non-MTP GGUF; llama-server may "
|
||
"fall back to spec-off if no nextn head is present. "
|
||
"Engaging anyway (user override)."
|
||
)
|
||
_emit_mtp(chain_ngram = False)
|
||
return flags
|
||
if effective_mode == "mtp+ngram":
|
||
if _mtp_too_small:
|
||
logger.warning(
|
||
f"Forcing MTP+Ngram on a {_mtp_size_b:.1f}B model; "
|
||
"the bench shows the chain regresses below 3B. "
|
||
"Engaging anyway (user override)."
|
||
)
|
||
elif not is_mtp_model:
|
||
logger.warning(
|
||
"Forcing MTP+Ngram on a non-MTP GGUF; llama-server "
|
||
"may fall back to ngram-only if no nextn head is "
|
||
"present. Engaging anyway (user override)."
|
||
)
|
||
_emit_mtp(chain_ngram = True)
|
||
return flags
|
||
|
||
# effective_mode == "auto": today's promotion path. llama.cpp
|
||
# #22673: MTP is compatible with mmproj, so there's no vision gate.
|
||
if is_mtp_model and not _mtp_too_small:
|
||
# GPU: MTP-only. CPU/Mac: chain ngram-mod + MTP.
|
||
_emit_mtp(chain_ngram = not gpus)
|
||
elif is_mtp_model and _mtp_too_small:
|
||
# Sub-3B fallback: drop the MTP draft head, keep ngram-mod
|
||
# when the binary supports it.
|
||
_small_caps = self.probe_server_capabilities(binary)
|
||
if _small_caps.get("supports_ngram_mod"):
|
||
logger.info(
|
||
f"MTP GGUF detected but model size {_mtp_size_b:.1f}B "
|
||
"is below the 3B speedup threshold; using ngram-mod "
|
||
"only (zero-VRAM, no draft head). Override via "
|
||
"--spec-type or the Studio Speculative Decoding "
|
||
"dropdown."
|
||
)
|
||
_emit_ngram_mod()
|
||
else:
|
||
logger.info(
|
||
f"MTP GGUF detected but model size {_mtp_size_b:.1f}B "
|
||
"is below the 3B speedup threshold and the bundled "
|
||
"llama-server does not advertise ngram-mod; "
|
||
"auto-disabling speculative decoding."
|
||
)
|
||
else:
|
||
# Non-MTP model: let llama-server choose its default strategy.
|
||
flags.append("--spec-default")
|
||
self._speculative_type = "default"
|
||
return flags
|
||
|
||
def _already_in_target_state(
|
||
self,
|
||
*,
|
||
model_identifier: str,
|
||
hf_variant: Optional[str],
|
||
n_ctx: int,
|
||
cache_type_kv: Optional[str],
|
||
speculative_type: Optional[str],
|
||
chat_template_override: Optional[str],
|
||
extra_args: Optional[List[str]],
|
||
is_vision: bool,
|
||
gguf_path: Optional[str] = None,
|
||
spec_draft_n_max: Optional[int] = None,
|
||
) -> bool:
|
||
"""True iff the live server already satisfies these load kwargs.
|
||
|
||
Mirrors ``routes/inference.py:_request_matches_loaded_settings``
|
||
but compares raw kwargs so ``load_model`` can short-circuit a
|
||
duplicate /load that raced past the route-level check (#5401).
|
||
"""
|
||
if not self.is_loaded:
|
||
return False
|
||
if (self._model_identifier or "").lower() != (model_identifier or "").lower():
|
||
return False
|
||
# Direct-file loads pass hf_variant=None while the backend
|
||
# stores an extracted filename label; compare paths instead
|
||
# to keep the guard symmetric.
|
||
if gguf_path is not None and self._gguf_path:
|
||
try:
|
||
if Path(self._gguf_path).resolve() != Path(gguf_path).resolve():
|
||
return False
|
||
except OSError:
|
||
return False
|
||
elif (self._hf_variant or "").lower() != (hf_variant or "").lower():
|
||
return False
|
||
if self._requested_n_ctx != int(n_ctx):
|
||
return False
|
||
|
||
def _norm(value):
|
||
if value is None:
|
||
return None
|
||
if isinstance(value, str):
|
||
stripped = value.strip().lower()
|
||
return stripped or None
|
||
return value
|
||
|
||
if _norm(self._cache_type_kv) != _norm(cache_type_kv):
|
||
return False
|
||
|
||
# Compare on the canonical UI-facing mode the user requested.
|
||
# When extra_args carries --spec-type, the route-layer code paths
|
||
# bypass the dropdown anyway and the backend stores
|
||
# _requested_spec_mode = None; the request mirrors that by
|
||
# canonicalising to None.
|
||
if _extra_args_set_spec_type(extra_args):
|
||
req_mode = None
|
||
else:
|
||
req_mode = _canonicalize_spec_mode(speculative_type) or "auto"
|
||
backend_mode = self._requested_spec_mode
|
||
if req_mode != backend_mode:
|
||
return False
|
||
|
||
# spec_draft_n_max only matters when an MTP variant is actually
|
||
# engaged. Compare on the resolved spec rather than the requested
|
||
# mode so an Auto request that auto-promoted to draft-mtp under
|
||
# the hood still bounces a reload when the user changes n_max.
|
||
if (
|
||
self._speculative_type == "draft-mtp"
|
||
and spec_draft_n_max is not None
|
||
and int(spec_draft_n_max) != (self._spec_draft_n_max or 0)
|
||
):
|
||
return False
|
||
|
||
if (self._chat_template_override or None) != (chat_template_override or None):
|
||
return False
|
||
|
||
# extra_args=None means "no opinion" (inherit semantics handled
|
||
# at the route layer); only an explicit list forces equality.
|
||
if extra_args is not None:
|
||
current = list(self._extra_args) if self._extra_args is not None else []
|
||
if list(extra_args) != current:
|
||
return False
|
||
return True
|
||
|
||
def _classify_gpu_offload(
|
||
self,
|
||
expected_gpu: bool,
|
||
detected_gpus: list[tuple[int, int]],
|
||
) -> Optional[bool]:
|
||
"""True if a GPU model buffer was allocated, False if only CPU
|
||
buffers landed despite GPU intent, None when there's no signal
|
||
(no GPU detected, no buffer-size lines, etc.)."""
|
||
if not detected_gpus or not expected_gpu:
|
||
return None
|
||
# llama-server logs one ``... model buffer size = N MiB`` line
|
||
# per backend buffer; CUDA0 / ROCm0 / Metal / Vulkan0 /
|
||
# OpenCL0 / SYCL0 are GPU, CPU / CPU_Mapped are not.
|
||
gpu_markers = ("CUDA", "ROCm", "Metal", "Vulkan", "OpenCL", "SYCL")
|
||
saw_buffer_line = False
|
||
saw_gpu_buffer = False
|
||
for line in self._stdout_lines:
|
||
if "model buffer size" not in line:
|
||
continue
|
||
saw_buffer_line = True
|
||
if any(marker in line for marker in gpu_markers):
|
||
saw_gpu_buffer = True
|
||
break
|
||
if not saw_buffer_line:
|
||
return None
|
||
return saw_gpu_buffer
|
||
|
||
def unload_model(self) -> bool:
|
||
"""Terminate the llama-server subprocess and cancel any in-flight download."""
|
||
self._cancel_event.set()
|
||
with self._lock:
|
||
self._kill_process()
|
||
logger.info(f"Unloaded GGUF model: {self._model_identifier}")
|
||
self._model_identifier = None
|
||
self._gguf_path = None
|
||
self._hf_repo = None
|
||
self._hf_variant = None
|
||
self._is_vision = False
|
||
self._is_audio = False
|
||
self._audio_type = None
|
||
self._audio_probed = False
|
||
self._port = None
|
||
self._healthy = False
|
||
self._context_length = None
|
||
self._effective_context_length = None
|
||
self._max_context_length = None
|
||
self._chat_template = None
|
||
self._chat_template_override = None
|
||
self._supports_reasoning = False
|
||
self._reasoning_always_on = False
|
||
self._reasoning_style = "enable_thinking"
|
||
self._reasoning_default = True
|
||
self._supports_preserve_thinking = False
|
||
self._supports_tools = False
|
||
self._cache_type_kv = None
|
||
self._speculative_type = None
|
||
self._requested_spec_mode = None
|
||
self._spec_draft_n_max = None
|
||
self._n_layers = None
|
||
self._n_kv_heads = None
|
||
self._n_kv_heads_by_layer = None
|
||
self._n_heads = None
|
||
self._embedding_length = None
|
||
self._kv_key_length = None
|
||
self._kv_value_length = None
|
||
self._sliding_window = None
|
||
self._sliding_window_pattern = None
|
||
self._full_attention_interval = None
|
||
self._kv_lora_rank = None
|
||
self._key_length_mla = None
|
||
self._kv_key_length_swa = None
|
||
self._kv_value_length_swa = None
|
||
self._ssm_inner_size = None
|
||
self._ssm_state_size = None
|
||
self._shared_kv_layers = None
|
||
self._nextn_predict_layers = None
|
||
# Clean up temp chat template file
|
||
if hasattr(self, "_chat_template_file") and self._chat_template_file:
|
||
try:
|
||
import os
|
||
|
||
os.unlink(self._chat_template_file.name)
|
||
except Exception:
|
||
pass
|
||
self._chat_template_file = None
|
||
# Free audio codec GPU memory
|
||
if LlamaCppBackend._codec_mgr is not None:
|
||
LlamaCppBackend._codec_mgr.unload()
|
||
LlamaCppBackend._codec_mgr = None
|
||
import torch
|
||
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
return True
|
||
|
||
def _kill_process(self):
|
||
"""Terminate the subprocess if running."""
|
||
if self._process is None:
|
||
return
|
||
try:
|
||
self._process.terminate()
|
||
self._process.wait(timeout = 5)
|
||
except subprocess.TimeoutExpired:
|
||
logger.warning("llama-server did not exit on SIGTERM, sending SIGKILL")
|
||
self._process.kill()
|
||
self._process.wait(timeout = 5)
|
||
except Exception as e:
|
||
logger.warning(f"Error killing llama-server process: {e}")
|
||
finally:
|
||
self._process = None
|
||
# Clear healthy so a /load arriving during the replacement
|
||
# server's warm-up window cannot short-circuit against the
|
||
# previous server's health (#5401).
|
||
self._healthy = False
|
||
# Drives _wait_for_vram_settle in the next load_model;
|
||
# set in finally so both in-process and frontend
|
||
# /unload+/load Apply paths record the kill.
|
||
self._last_kill_monotonic = time.monotonic()
|
||
if self._stdout_thread is not None:
|
||
self._stdout_thread.join(timeout = 2)
|
||
self._stdout_thread = None
|
||
fh = getattr(self, "_llama_log_fh", None)
|
||
if fh is not None:
|
||
try:
|
||
fh.close()
|
||
except Exception:
|
||
pass
|
||
self._llama_log_fh = None
|
||
|
||
@staticmethod
|
||
def _kill_orphaned_servers():
|
||
"""Kill orphaned llama-server processes started by studio.
|
||
|
||
Only kills processes whose resolved binary lives under a known
|
||
Studio install directory (or matches an exact env-var override)
|
||
to avoid terminating unrelated llama-server instances.
|
||
|
||
Mirrors every location that _find_llama_server_binary() can
|
||
return from so that orphans from any supported install path
|
||
are still cleaned up.
|
||
|
||
Uses psutil for cross-platform support (Linux, macOS, Windows).
|
||
Falls back to pgrep + /proc/<pid>/exe on Linux when psutil is
|
||
not installed.
|
||
"""
|
||
try:
|
||
# -- Build the ownership allowlist --------------------------------
|
||
# Two kinds of matches:
|
||
# exact_binaries -- env var overrides (exact path match only)
|
||
# install_roots -- directory trees that are Studio-owned
|
||
# (binary must be *under* one of these)
|
||
install_roots: list[Path] = []
|
||
|
||
# Env-mode custom root (mirrors _find_llama_server_binary).
|
||
_is_custom_root = False
|
||
try:
|
||
from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433
|
||
|
||
_resolved_sr = _sr()
|
||
_legacy_studio = Path.home() / ".unsloth" / "studio"
|
||
try:
|
||
_is_custom_root = _resolved_sr.resolve() != _legacy_studio.resolve()
|
||
except (OSError, ValueError):
|
||
_is_custom_root = _resolved_sr != _legacy_studio
|
||
if _is_custom_root:
|
||
install_roots.append(_resolved_sr / "llama.cpp")
|
||
except (ImportError, OSError, ValueError):
|
||
pass
|
||
|
||
# Primary install dir (default mode only). Env-mode skips this so
|
||
# a custom-root Studio cannot kill a concurrent default-install
|
||
# Studio's llama-server (same OS user, different install).
|
||
if not _is_custom_root:
|
||
install_roots.append(Path.home() / ".unsloth" / "llama.cpp")
|
||
|
||
# Legacy in-tree build dirs (older setup.sh versions)
|
||
project_root = Path(__file__).resolve().parents[4]
|
||
install_roots.append(project_root / "llama.cpp")
|
||
|
||
# Legacy: extracted binary
|
||
install_roots.append(project_root / "bin")
|
||
|
||
# UNSLOTH_LLAMA_CPP_PATH env var (custom install dir)
|
||
custom_dir = os.environ.get("UNSLOTH_LLAMA_CPP_PATH")
|
||
if custom_dir:
|
||
install_roots.append(Path(custom_dir))
|
||
|
||
# LLAMA_SERVER_PATH env var (exact binary path)
|
||
exact_binaries: list[Path] = []
|
||
env_binary = os.environ.get("LLAMA_SERVER_PATH")
|
||
if env_binary:
|
||
try:
|
||
exact_binaries.append(Path(env_binary).resolve())
|
||
except OSError:
|
||
pass
|
||
|
||
# Resolve all roots so is_relative_to works reliably
|
||
resolved_roots: list[Path] = []
|
||
for root in install_roots:
|
||
try:
|
||
resolved_roots.append(root.resolve())
|
||
except OSError:
|
||
pass
|
||
|
||
my_pid = os.getpid()
|
||
|
||
# -- Enumerate processes -------------------------------------------
|
||
# Prefer psutil (cross-platform). Fall back to pgrep + /proc on
|
||
# Linux when psutil is not installed.
|
||
try:
|
||
import psutil
|
||
|
||
has_psutil = True
|
||
except ImportError:
|
||
has_psutil = False
|
||
|
||
if has_psutil:
|
||
for proc in psutil.process_iter(["pid", "name", "exe"]):
|
||
try:
|
||
if proc.info["pid"] == my_pid:
|
||
continue
|
||
|
||
name = proc.info.get("name") or ""
|
||
if not name.lower().startswith("llama-server"):
|
||
continue
|
||
|
||
exe = proc.info.get("exe")
|
||
if not exe:
|
||
continue
|
||
|
||
exe_path = Path(exe).resolve()
|
||
|
||
# Check ownership: exact binary match OR binary is
|
||
# under a known install root (proper ancestry, not
|
||
# substring).
|
||
is_ours = exe_path in exact_binaries or any(
|
||
exe_path.is_relative_to(root) for root in resolved_roots
|
||
)
|
||
if not is_ours:
|
||
continue
|
||
|
||
proc.kill()
|
||
logger.info(
|
||
f"Killed orphaned llama-server process "
|
||
f"(pid={proc.info['pid']})"
|
||
)
|
||
except (
|
||
psutil.NoSuchProcess,
|
||
psutil.AccessDenied,
|
||
psutil.ZombieProcess,
|
||
):
|
||
pass
|
||
else:
|
||
# -- Fallback: pgrep + /proc/<pid>/exe (Linux only) -----------
|
||
if sys.platform != "linux":
|
||
return
|
||
result = subprocess.run(
|
||
["pgrep", "-a", "-f", "llama-server"],
|
||
capture_output = True,
|
||
text = True,
|
||
timeout = 5,
|
||
env = child_env_without_native_path_secret(),
|
||
)
|
||
if result.returncode != 0:
|
||
return
|
||
|
||
for line in result.stdout.strip().splitlines():
|
||
parts = line.strip().split(None, 1)
|
||
if len(parts) < 2:
|
||
continue
|
||
pid = int(parts[0])
|
||
if pid == my_pid:
|
||
continue
|
||
|
||
# Resolve the actual executable. /proc/<pid>/exe is a
|
||
# symlink to the real binary and avoids all cmdline-
|
||
# parsing ambiguities (spaces in paths, argv rewriting).
|
||
# Fall back to the first cmdline token when /proc is
|
||
# unavailable.
|
||
proc_exe = Path(f"/proc/{pid}/exe")
|
||
try:
|
||
binary = proc_exe.resolve(strict = True)
|
||
except (OSError, ValueError):
|
||
cmdline = parts[1]
|
||
token = cmdline.split()[0] if cmdline.strip() else ""
|
||
if not token:
|
||
continue
|
||
binary = Path(token).resolve(strict = False)
|
||
|
||
owned = binary in exact_binaries or any(
|
||
binary.is_relative_to(root) for root in resolved_roots
|
||
)
|
||
if not owned:
|
||
continue
|
||
|
||
try:
|
||
os.kill(pid, signal.SIGKILL)
|
||
logger.info(f"Killed orphaned llama-server process (pid={pid})")
|
||
except ProcessLookupError:
|
||
pass
|
||
except PermissionError:
|
||
pass
|
||
except Exception:
|
||
logger.warning("Error during orphan server cleanup", exc_info = True)
|
||
|
||
def _cleanup(self):
|
||
"""atexit handler to ensure llama-server is terminated."""
|
||
self._kill_process()
|
||
|
||
def _wait_for_health(self, timeout: float = 120.0, interval: float = 0.5) -> bool:
|
||
"""
|
||
Poll llama-server's /health endpoint until it responds 200.
|
||
|
||
Also monitors subprocess for early exit/crash.
|
||
"""
|
||
deadline = time.monotonic() + timeout
|
||
url = f"http://127.0.0.1:{self._port}/health"
|
||
|
||
while time.monotonic() < deadline:
|
||
# Check if process crashed
|
||
if self._process.poll() is not None:
|
||
# Give the drain thread a moment to collect final output
|
||
if self._stdout_thread is not None:
|
||
self._stdout_thread.join(timeout = 2)
|
||
output = "\n".join(self._stdout_lines[-50:])
|
||
logger.error(
|
||
f"llama-server exited with code {self._process.returncode}. "
|
||
f"Output: {output[:2000]}"
|
||
)
|
||
return False
|
||
|
||
try:
|
||
resp = httpx.get(url, timeout = 2.0)
|
||
if resp.status_code == 200:
|
||
return True
|
||
except (
|
||
httpx.ConnectError,
|
||
httpx.TimeoutException,
|
||
# ReadError covers TCP RST mid-read while llama-server is
|
||
# still binding the port (Windows: WinError 10054). The
|
||
# crash-detection branch above catches a real exit; this
|
||
# one keeps a transient socket close from masking it.
|
||
httpx.ReadError,
|
||
httpx.RemoteProtocolError,
|
||
httpx.WriteError,
|
||
):
|
||
pass
|
||
|
||
time.sleep(interval)
|
||
|
||
logger.error(f"llama-server health check timed out after {timeout}s")
|
||
return False
|
||
|
||
# ── Message building (OpenAI format) ──────────────────────────
|
||
|
||
@staticmethod
|
||
def _parse_tool_calls_from_text(content: str) -> list[dict]:
|
||
"""Thin wrapper around the shared parser in tool_call_parser
|
||
so safetensors and llama_cpp pick up the same fixes."""
|
||
return _shared_parse_tool_calls_from_text(content)
|
||
|
||
@staticmethod
|
||
def _build_openai_messages(
|
||
messages: list[dict],
|
||
image_b64: Optional[str] = None,
|
||
) -> list[dict]:
|
||
"""
|
||
Build OpenAI-format messages, optionally injecting an image_url
|
||
content part into the last user message for vision models.
|
||
|
||
If no image is provided, returns messages as-is.
|
||
"""
|
||
if not image_b64:
|
||
return messages
|
||
|
||
# Find the last user message and convert to multimodal content parts
|
||
result = [msg.copy() for msg in messages]
|
||
last_user_idx = None
|
||
for i, msg in enumerate(result):
|
||
if msg["role"] == "user":
|
||
last_user_idx = i
|
||
|
||
if last_user_idx is not None:
|
||
text_content = result[last_user_idx].get("content", "")
|
||
result[last_user_idx]["content"] = [
|
||
{"type": "text", "text": text_content},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:image/png;base64,{image_b64}",
|
||
},
|
||
},
|
||
]
|
||
|
||
return result
|
||
|
||
# ── Generation (proxy to llama-server) ────────────────────────
|
||
|
||
@staticmethod
|
||
def _iter_text_cancellable(
|
||
response: "httpx.Response",
|
||
cancel_event: Optional[threading.Event] = None,
|
||
) -> Generator[str, None, None]:
|
||
"""Iterate over an httpx streaming response with cancel support.
|
||
|
||
Checks cancel_event between chunks and on ReadTimeout. The
|
||
cancel watcher in _stream_with_retry also calls response.close()
|
||
on cancel, which unblocks iter_text() once the response exists.
|
||
During normal streaming llama-server sends tokens frequently,
|
||
so the cancel check between chunks is the primary mechanism.
|
||
"""
|
||
text_iter = response.iter_text()
|
||
while True:
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
response.close()
|
||
return
|
||
try:
|
||
chunk = next(text_iter)
|
||
yield chunk
|
||
except StopIteration:
|
||
return
|
||
except httpx.ReadTimeout:
|
||
# No data within the timeout window -- just loop back
|
||
# and re-check cancel_event.
|
||
continue
|
||
|
||
@staticmethod
|
||
@contextlib.contextmanager
|
||
def _stream_with_retry(
|
||
client: "httpx.Client",
|
||
url: str,
|
||
payload: dict,
|
||
cancel_event: Optional[threading.Event] = None,
|
||
headers: Optional[dict] = None,
|
||
):
|
||
"""Open an httpx streaming POST with cancel support.
|
||
|
||
Sends the request once with a long read timeout (120 s) so
|
||
prompt processing (prefill) can finish without triggering a
|
||
retry storm. The previous 0.5 s timeout caused duplicate POST
|
||
requests every half second, forcing llama-server to restart
|
||
processing each time.
|
||
|
||
A background watcher thread provides cancel by closing the
|
||
response when cancel_event is set. Limitation: httpx does not
|
||
allow interrupting a blocked read from another thread before
|
||
the response object exists, so cancel during the initial
|
||
header wait (prefill phase) only takes effect once headers
|
||
arrive. After that, response.close() unblocks reads promptly.
|
||
In practice llama-server prefill is 1-5 s for typical prompts,
|
||
during which cancel is deferred -- still much better than the
|
||
old retry storm which made prefill slower.
|
||
"""
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
raise GeneratorExit
|
||
|
||
# Background watcher: close the response if cancel is requested.
|
||
# Only effective after response headers arrive (httpx limitation).
|
||
_cancel_closed = threading.Event()
|
||
_response_ref: list = [None]
|
||
|
||
def _cancel_watcher():
|
||
while not _cancel_closed.is_set():
|
||
if cancel_event.wait(timeout = 0.3):
|
||
# Cancel requested. Keep polling until the response object
|
||
# exists so we can close it, or until the main thread
|
||
# finishes on its own (_cancel_closed is set in finally).
|
||
while not _cancel_closed.is_set():
|
||
r = _response_ref[0]
|
||
if r is not None:
|
||
try:
|
||
r.close()
|
||
return
|
||
except Exception as e:
|
||
logger.debug(
|
||
f"Error closing response in cancel watcher: {e}"
|
||
)
|
||
# Response not created yet -- wait briefly and retry
|
||
_cancel_closed.wait(timeout = 0.1)
|
||
return
|
||
|
||
watcher = None
|
||
if cancel_event is not None:
|
||
watcher = threading.Thread(
|
||
target = _cancel_watcher, daemon = True, name = "prefill-cancel"
|
||
)
|
||
watcher.start()
|
||
|
||
try:
|
||
# Long read timeout so prefill (prompt processing) can finish
|
||
# without triggering a retry storm. Cancel during both
|
||
# prefill and streaming is handled by the watcher thread
|
||
# which closes the response, unblocking any httpx read.
|
||
prefill_timeout = httpx.Timeout(
|
||
connect = 30,
|
||
read = 120.0,
|
||
write = 10,
|
||
pool = 10,
|
||
)
|
||
with client.stream(
|
||
"POST",
|
||
url,
|
||
json = payload,
|
||
timeout = prefill_timeout,
|
||
headers = headers,
|
||
) as response:
|
||
_response_ref[0] = response
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
raise GeneratorExit
|
||
yield response
|
||
return
|
||
except (httpx.ReadError, httpx.RemoteProtocolError, httpx.CloseError):
|
||
# Response was closed by the cancel watcher
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
raise GeneratorExit
|
||
raise
|
||
finally:
|
||
_cancel_closed.set()
|
||
|
||
def generate_chat_completion(
|
||
self,
|
||
messages: list[dict],
|
||
image_b64: Optional[str] = None,
|
||
temperature: float = 0.6,
|
||
top_p: float = 0.95,
|
||
top_k: int = 20,
|
||
min_p: float = 0.01,
|
||
max_tokens: Optional[int] = None,
|
||
repetition_penalty: float = 1.0,
|
||
presence_penalty: float = 0.0,
|
||
stop: Optional[list[str]] = None,
|
||
cancel_event: Optional[threading.Event] = None,
|
||
enable_thinking: Optional[bool] = None,
|
||
reasoning_effort: Optional[str] = None,
|
||
preserve_thinking: Optional[bool] = None,
|
||
) -> Generator[str | dict, None, None]:
|
||
"""
|
||
Send a chat completion request to llama-server and stream tokens back.
|
||
|
||
Uses /v1/chat/completions — llama-server handles chat template
|
||
application and vision (multimodal image_url parts) natively.
|
||
|
||
Yields cumulative text (matching InferenceBackend's convention).
|
||
"""
|
||
if not self.is_loaded:
|
||
raise RuntimeError("llama-server is not loaded")
|
||
|
||
openai_messages = self._build_openai_messages(messages, image_b64)
|
||
|
||
payload = {
|
||
"messages": openai_messages,
|
||
"stream": True,
|
||
"temperature": temperature,
|
||
"top_p": top_p,
|
||
"top_k": top_k if top_k >= 0 else 0,
|
||
"min_p": min_p,
|
||
"repeat_penalty": repetition_penalty,
|
||
"presence_penalty": presence_penalty,
|
||
}
|
||
# Pass enable_thinking / reasoning_effort / preserve_thinking per-request
|
||
_reasoning_kw = self._request_reasoning_kwargs(
|
||
enable_thinking, reasoning_effort, preserve_thinking
|
||
)
|
||
if _reasoning_kw is not None:
|
||
payload["chat_template_kwargs"] = _reasoning_kw
|
||
# Default cap to the model's effective context length when known,
|
||
# otherwise the conservative floor. The wall-clock backstop below
|
||
# keeps a stuck model from running indefinitely either way.
|
||
payload["max_tokens"] = (
|
||
max_tokens
|
||
if max_tokens is not None
|
||
else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR)
|
||
)
|
||
payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS
|
||
if stop:
|
||
payload["stop"] = stop
|
||
payload["stream_options"] = {"include_usage": True}
|
||
|
||
url = f"{self.base_url}/v1/chat/completions"
|
||
cumulative = ""
|
||
in_thinking = False
|
||
_stream_done = False
|
||
_metadata_usage = None
|
||
_metadata_timings = None
|
||
|
||
try:
|
||
# _stream_with_retry uses a 120 s read timeout so prefill
|
||
# can finish. Cancel during streaming is handled by the
|
||
# watcher thread (closes the response on cancel_event).
|
||
stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10)
|
||
_auth_headers = (
|
||
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
||
)
|
||
with httpx.Client(
|
||
timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0)
|
||
) as client:
|
||
with self._stream_with_retry(
|
||
client,
|
||
url,
|
||
payload,
|
||
cancel_event,
|
||
headers = _auth_headers,
|
||
) as response:
|
||
if response.status_code != 200:
|
||
error_body = response.read().decode()
|
||
raise RuntimeError(
|
||
f"llama-server returned {response.status_code}: {error_body}"
|
||
)
|
||
|
||
buffer = ""
|
||
has_content_tokens = False
|
||
reasoning_text = ""
|
||
for raw_chunk in self._iter_text_cancellable(
|
||
response, cancel_event
|
||
):
|
||
buffer += raw_chunk
|
||
while "\n" in buffer:
|
||
line, buffer = buffer.split("\n", 1)
|
||
line = line.strip()
|
||
|
||
if not line:
|
||
continue
|
||
if line == "data: [DONE]":
|
||
if in_thinking:
|
||
if has_content_tokens:
|
||
# Real thinking + content: close the tag
|
||
cumulative += "</think>"
|
||
yield cumulative
|
||
else:
|
||
# Only reasoning_content, no content tokens:
|
||
# the model put its entire reply in reasoning
|
||
# (e.g. Qwen3 always-think mode). Show it
|
||
# as the main response, not as a thinking block.
|
||
cumulative = reasoning_text
|
||
yield cumulative
|
||
_stream_done = True
|
||
break # exit inner while
|
||
if not line.startswith("data: "):
|
||
continue
|
||
|
||
try:
|
||
data = json.loads(line[6:])
|
||
# Capture server timings/usage from final chunks
|
||
_chunk_timings = data.get("timings")
|
||
if _chunk_timings:
|
||
_metadata_timings = _chunk_timings
|
||
_chunk_usage = data.get("usage")
|
||
if _chunk_usage:
|
||
_metadata_usage = _chunk_usage
|
||
choices = data.get("choices", [])
|
||
if choices:
|
||
delta = choices[0].get("delta", {})
|
||
|
||
# Handle reasoning/thinking tokens
|
||
# llama-server sends these as "reasoning_content"
|
||
# Wrap in <think> tags for the frontend parser
|
||
reasoning = delta.get("reasoning_content", "")
|
||
if reasoning:
|
||
reasoning_text += reasoning
|
||
if not in_thinking:
|
||
cumulative += "<think>"
|
||
in_thinking = True
|
||
cumulative += reasoning
|
||
yield cumulative
|
||
|
||
token = delta.get("content", "")
|
||
if token:
|
||
has_content_tokens = True
|
||
if in_thinking:
|
||
cumulative += "</think>"
|
||
in_thinking = False
|
||
cumulative += token
|
||
yield cumulative
|
||
except json.JSONDecodeError:
|
||
logger.debug(
|
||
f"Skipping malformed SSE line: {line[:100]}"
|
||
)
|
||
if _stream_done:
|
||
break # exit outer for
|
||
if _metadata_usage or _metadata_timings:
|
||
_metadata_usage = _backfill_usage_from_timings(
|
||
_metadata_usage, _metadata_timings
|
||
)
|
||
yield {
|
||
"type": "metadata",
|
||
"usage": _metadata_usage,
|
||
"timings": _metadata_timings,
|
||
}
|
||
|
||
except httpx.ConnectError:
|
||
raise RuntimeError("Lost connection to llama-server")
|
||
except Exception as e:
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
return
|
||
raise
|
||
|
||
# ── Tool-calling agentic loop ──────────────────────────────
|
||
|
||
def generate_chat_completion_with_tools(
|
||
self,
|
||
messages: list[dict],
|
||
tools: list[dict],
|
||
temperature: float = 0.6,
|
||
top_p: float = 0.95,
|
||
top_k: int = 20,
|
||
min_p: float = 0.01,
|
||
max_tokens: Optional[int] = None,
|
||
repetition_penalty: float = 1.0,
|
||
presence_penalty: float = 0.0,
|
||
stop: Optional[list[str]] = None,
|
||
cancel_event: Optional[threading.Event] = None,
|
||
enable_thinking: Optional[bool] = None,
|
||
reasoning_effort: Optional[str] = None,
|
||
preserve_thinking: Optional[bool] = None,
|
||
max_tool_iterations: int = 25,
|
||
auto_heal_tool_calls: bool = True,
|
||
tool_call_timeout: int = 300,
|
||
session_id: Optional[str] = None,
|
||
confirm_tool_calls: bool = False,
|
||
) -> Generator[dict, None, None]:
|
||
"""
|
||
Agentic loop: let the model call tools, execute them, and continue.
|
||
|
||
Yields dicts with:
|
||
{"type": "status", "text": "Searching: ..."/"Reading: ..."} -- tool status updates
|
||
{"type": "content", "text": "token"} -- streamed content tokens (cumulative)
|
||
{"type": "reasoning", "text": "token"} -- streamed reasoning tokens (cumulative)
|
||
"""
|
||
from core.inference.tools import execute_tool
|
||
from state.tool_approvals import (
|
||
TOOL_REJECTED_MESSAGE,
|
||
begin_tool_decision,
|
||
new_approval_id,
|
||
wait_tool_decision,
|
||
)
|
||
|
||
if not self.is_loaded:
|
||
raise RuntimeError("llama-server is not loaded")
|
||
|
||
conversation = list(messages)
|
||
url = f"{self.base_url}/v1/chat/completions"
|
||
_accumulated_completion_tokens = 0
|
||
_accumulated_predicted_ms = 0.0
|
||
_accumulated_predicted_n = 0
|
||
|
||
def _strip_tool_markup(text: str, *, final: bool = False) -> str:
|
||
if not auto_heal_tool_calls:
|
||
return text
|
||
return strip_tool_call_markup(text, final = final)
|
||
|
||
# XML prefixes that signal a tool call in content.
|
||
# Empty when auto_heal is disabled so the buffer never
|
||
# speculatively holds content for XML detection.
|
||
_TOOL_XML_SIGNALS = (
|
||
("<tool_call>", "<function=") if auto_heal_tool_calls else ()
|
||
)
|
||
_MAX_BUFFER_CHARS = 32
|
||
|
||
# ── Duplicate tool-call detection ────────────────────────
|
||
# Track recent (tool_name, arguments) hashes to detect loops
|
||
# where the model repeats the exact same call. Retries after
|
||
# a transient failure are allowed (only block when the previous
|
||
# identical call succeeded).
|
||
_tool_call_history: list[tuple[str, bool]] = [] # (key, failed)
|
||
|
||
# ── Re-prompt on plan-without-action ─────────────────
|
||
# When the model describes what it intends to do (forward-looking
|
||
# language) without actually calling a tool, re-prompt once.
|
||
# Only triggers on responses that signal intent/planning -- a
|
||
# direct answer like "4" or "Hello!" will not match.
|
||
# Pattern is compiled once at module level (_INTENT_SIGNAL).
|
||
_reprompt_count = 0
|
||
|
||
# Reserve extra iterations for re-prompts so they don't
|
||
# consume the caller's tool-call budget. Only add the
|
||
# extra slot when tool iterations are actually allowed.
|
||
_extra = _MAX_REPROMPTS if max_tool_iterations > 0 else 0
|
||
for iteration in range(max_tool_iterations + _extra):
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
return
|
||
|
||
# Build payload -- stream: True so we detect tool signals
|
||
# in the first 1-2 chunks without a non-streaming penalty.
|
||
payload = {
|
||
"messages": conversation,
|
||
"stream": True,
|
||
"stream_options": {"include_usage": True},
|
||
"temperature": temperature,
|
||
"top_p": top_p,
|
||
"top_k": top_k if top_k >= 0 else 0,
|
||
"min_p": min_p,
|
||
"repeat_penalty": repetition_penalty,
|
||
"presence_penalty": presence_penalty,
|
||
"tools": tools,
|
||
"tool_choice": "auto",
|
||
}
|
||
_reasoning_kw = self._request_reasoning_kwargs(
|
||
enable_thinking, reasoning_effort, preserve_thinking
|
||
)
|
||
if _reasoning_kw is not None:
|
||
payload["chat_template_kwargs"] = _reasoning_kw
|
||
payload["max_tokens"] = (
|
||
max_tokens
|
||
if max_tokens is not None
|
||
else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR)
|
||
)
|
||
payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS
|
||
if stop:
|
||
payload["stop"] = stop
|
||
|
||
try:
|
||
_auth_headers = (
|
||
{"Authorization": f"Bearer {self._api_key}"}
|
||
if self._api_key
|
||
else None
|
||
)
|
||
|
||
# ── Speculative buffer state machine ──────────────────
|
||
# BUFFERING: accumulating content, checking for tool signals
|
||
# STREAMING: no tool detected, yielding tokens to caller
|
||
# DRAINING: tool signal found, silently consuming rest
|
||
_S_BUFFERING = 0
|
||
_S_STREAMING = 1
|
||
_S_DRAINING = 2
|
||
|
||
detect_state = _S_BUFFERING
|
||
content_buffer = "" # Raw content held during BUFFERING
|
||
content_accum = "" # All content tokens (for tool parsing)
|
||
reasoning_accum = ""
|
||
cumulative_display = "" # Cumulative text yielded (with <think>)
|
||
in_thinking = False
|
||
has_content_tokens = False
|
||
tool_calls_acc = {} # Structured delta.tool_calls fragments
|
||
has_structured_tc = False
|
||
_iter_usage = None
|
||
_iter_timings = None
|
||
_stream_done = False
|
||
_last_emitted = ""
|
||
|
||
stream_timeout = httpx.Timeout(
|
||
connect = 10,
|
||
read = 0.5,
|
||
write = 10,
|
||
pool = 10,
|
||
)
|
||
with httpx.Client(
|
||
timeout = stream_timeout,
|
||
limits = httpx.Limits(max_keepalive_connections = 0),
|
||
) as client:
|
||
with self._stream_with_retry(
|
||
client,
|
||
url,
|
||
payload,
|
||
cancel_event,
|
||
headers = _auth_headers,
|
||
) as response:
|
||
if response.status_code != 200:
|
||
error_body = response.read().decode()
|
||
raise RuntimeError(
|
||
f"llama-server returned {response.status_code}: "
|
||
f"{error_body}"
|
||
)
|
||
|
||
raw_buf = ""
|
||
for raw_chunk in self._iter_text_cancellable(
|
||
response,
|
||
cancel_event,
|
||
):
|
||
raw_buf += raw_chunk
|
||
while "\n" in raw_buf:
|
||
line, raw_buf = raw_buf.split("\n", 1)
|
||
line = line.strip()
|
||
|
||
if not line:
|
||
continue
|
||
if line == "data: [DONE]":
|
||
# Flush thinking state for STREAMING
|
||
if detect_state == _S_STREAMING and in_thinking:
|
||
if has_content_tokens:
|
||
cumulative_display += "</think>"
|
||
yield {
|
||
"type": "content",
|
||
"text": _strip_tool_markup(
|
||
cumulative_display,
|
||
final = True,
|
||
),
|
||
}
|
||
else:
|
||
cumulative_display = reasoning_accum
|
||
yield {
|
||
"type": "content",
|
||
"text": cumulative_display,
|
||
}
|
||
_stream_done = True
|
||
break # exit inner while
|
||
if not line.startswith("data: "):
|
||
continue
|
||
|
||
try:
|
||
chunk_data = json.loads(line[6:])
|
||
_ct = chunk_data.get("timings")
|
||
if _ct:
|
||
_iter_timings = _ct
|
||
_cu = chunk_data.get("usage")
|
||
if _cu:
|
||
_iter_usage = _cu
|
||
|
||
choices = chunk_data.get("choices", [])
|
||
if not choices:
|
||
continue
|
||
|
||
delta = choices[0].get("delta", {})
|
||
|
||
# ── Structured tool_calls ──
|
||
tc_deltas = delta.get("tool_calls")
|
||
if tc_deltas:
|
||
# Once visible content has been
|
||
# emitted, do not reclassify this
|
||
# turn as a tool call.
|
||
if _last_emitted:
|
||
continue
|
||
has_structured_tc = True
|
||
detect_state = _S_DRAINING
|
||
for tc_d in tc_deltas:
|
||
idx = tc_d.get("index", 0)
|
||
if idx not in tool_calls_acc:
|
||
tool_calls_acc[idx] = {
|
||
"id": tc_d.get("id", f"call_{idx}"),
|
||
"type": "function",
|
||
"function": {
|
||
"name": "",
|
||
"arguments": "",
|
||
},
|
||
}
|
||
elif tc_d.get("id"):
|
||
# Update ID if real one
|
||
# arrives on a later delta
|
||
tool_calls_acc[idx]["id"] = tc_d["id"]
|
||
func = tc_d.get("function", {})
|
||
if func.get("name"):
|
||
tool_calls_acc[idx]["function"][
|
||
"name"
|
||
] += func["name"]
|
||
if func.get("arguments"):
|
||
tool_calls_acc[idx]["function"][
|
||
"arguments"
|
||
] += func["arguments"]
|
||
continue
|
||
|
||
# ── Reasoning tokens ──
|
||
# Only yield in STREAMING state. In BUFFERING
|
||
# and DRAINING, accumulate silently so we don't
|
||
# corrupt the consumer's prev_text tracker
|
||
# (routes/inference.py never resets prev_text
|
||
# between tool iterations).
|
||
reasoning = delta.get("reasoning_content", "")
|
||
if reasoning:
|
||
reasoning_accum += reasoning
|
||
if detect_state == _S_STREAMING:
|
||
if not in_thinking:
|
||
cumulative_display += "<think>"
|
||
in_thinking = True
|
||
cumulative_display += reasoning
|
||
yield {
|
||
"type": "content",
|
||
"text": cumulative_display,
|
||
}
|
||
|
||
# ── Content tokens ──
|
||
token = delta.get("content", "")
|
||
if token:
|
||
has_content_tokens = True
|
||
content_accum += token
|
||
|
||
if detect_state == _S_DRAINING:
|
||
pass # accumulate silently
|
||
|
||
elif detect_state == _S_STREAMING:
|
||
if in_thinking:
|
||
cumulative_display += "</think>"
|
||
in_thinking = False
|
||
cumulative_display += token
|
||
cleaned = _strip_tool_markup(
|
||
cumulative_display,
|
||
)
|
||
if len(cleaned) > len(_last_emitted):
|
||
_last_emitted = cleaned
|
||
yield {
|
||
"type": "content",
|
||
"text": cleaned,
|
||
}
|
||
|
||
elif detect_state == _S_BUFFERING:
|
||
content_buffer += token
|
||
stripped_buf = content_buffer.lstrip()
|
||
if not stripped_buf:
|
||
continue
|
||
|
||
# Check tool signal prefixes
|
||
is_prefix = False
|
||
is_match = False
|
||
for sig in _TOOL_XML_SIGNALS:
|
||
if stripped_buf.startswith(sig):
|
||
is_match = True
|
||
break
|
||
if sig.startswith(stripped_buf):
|
||
is_prefix = True
|
||
break
|
||
|
||
if is_match:
|
||
detect_state = _S_DRAINING
|
||
elif (
|
||
is_prefix
|
||
and len(stripped_buf)
|
||
< _MAX_BUFFER_CHARS
|
||
):
|
||
pass # keep buffering
|
||
else:
|
||
# Not a tool -- flush buffer
|
||
detect_state = _S_STREAMING
|
||
# Flush any reasoning accumulated
|
||
# during BUFFERING phase
|
||
if reasoning_accum:
|
||
cumulative_display += "<think>"
|
||
cumulative_display += (
|
||
reasoning_accum
|
||
)
|
||
cumulative_display += "</think>"
|
||
cumulative_display += content_buffer
|
||
cleaned = _strip_tool_markup(
|
||
cumulative_display,
|
||
)
|
||
if len(cleaned) > len(_last_emitted):
|
||
_last_emitted = cleaned
|
||
yield {
|
||
"type": "content",
|
||
"text": cleaned,
|
||
}
|
||
|
||
except json.JSONDecodeError:
|
||
logger.debug(
|
||
f"Skipping malformed SSE line: {line[:100]}"
|
||
)
|
||
if _stream_done:
|
||
break # exit outer for
|
||
|
||
# ── Resolve BUFFERING at stream end ──
|
||
if detect_state == _S_BUFFERING:
|
||
stripped_buf = content_buffer.lstrip()
|
||
if (
|
||
stripped_buf
|
||
and auto_heal_tool_calls
|
||
and any(s in stripped_buf for s in _TOOL_XML_SIGNALS)
|
||
):
|
||
detect_state = _S_DRAINING
|
||
elif content_accum or reasoning_accum:
|
||
detect_state = _S_STREAMING
|
||
if content_buffer:
|
||
# Flush any reasoning accumulated first
|
||
if reasoning_accum:
|
||
cumulative_display += "<think>"
|
||
cumulative_display += reasoning_accum
|
||
cumulative_display += "</think>"
|
||
cumulative_display += content_buffer
|
||
yield {
|
||
"type": "content",
|
||
"text": _strip_tool_markup(
|
||
cumulative_display,
|
||
final = True,
|
||
),
|
||
}
|
||
elif reasoning_accum and not has_content_tokens:
|
||
# Reasoning-only response (no content tokens):
|
||
# show reasoning as plain text, matching
|
||
# the final streaming pass behavior for
|
||
# models that put everything in reasoning.
|
||
cumulative_display = reasoning_accum
|
||
yield {
|
||
"type": "content",
|
||
"text": cumulative_display,
|
||
}
|
||
else:
|
||
return
|
||
|
||
# ── STREAMING path: no tool call ──
|
||
if detect_state == _S_STREAMING:
|
||
# Safety net: check for XML tool signals in content.
|
||
# The route layer resets prev_text on tool_start, so
|
||
# post-tool synthesis streams correctly even if
|
||
# content was already emitted before the tool XML.
|
||
_safety_tc = None
|
||
if auto_heal_tool_calls and any(
|
||
s in content_accum for s in _TOOL_XML_SIGNALS
|
||
):
|
||
_safety_tc = self._parse_tool_calls_from_text(
|
||
content_accum,
|
||
)
|
||
if not _safety_tc:
|
||
# ── Re-prompt on plan-without-action ──
|
||
# If the model described what it intends to do
|
||
# (forward-looking language) without calling any
|
||
# tool, nudge it to act. Only fires once per
|
||
# request and only on short responses that
|
||
# contain intent signals -- a direct answer
|
||
# like "4" or "Hello!" won't trigger this.
|
||
# Use content if available, otherwise fall back
|
||
# to reasoning text (reasoning-only stalls).
|
||
_stripped = content_accum.strip()
|
||
if not _stripped:
|
||
_stripped = reasoning_accum.strip()
|
||
if (
|
||
tools
|
||
and _reprompt_count < _MAX_REPROMPTS
|
||
and 0 < len(_stripped) < _REPROMPT_MAX_CHARS
|
||
and _INTENT_SIGNAL.search(_stripped)
|
||
):
|
||
_reprompt_count += 1
|
||
logger.info(
|
||
f"Re-prompt {_reprompt_count}/{_MAX_REPROMPTS}: "
|
||
f"model responded without calling tools "
|
||
f"({len(_stripped)} chars)"
|
||
)
|
||
conversation.append(
|
||
{
|
||
"role": "assistant",
|
||
"content": _stripped,
|
||
}
|
||
)
|
||
conversation.append(
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
"STOP. Do NOT write code or explain. "
|
||
"You MUST call a tool NOW. "
|
||
"Call web_search or python immediately."
|
||
),
|
||
}
|
||
)
|
||
# Accumulate tokens and timing from this iteration
|
||
_fu_r = (
|
||
_backfill_usage_from_timings(_iter_usage, _iter_timings)
|
||
or {}
|
||
)
|
||
_accumulated_completion_tokens += _fu_r.get(
|
||
"completion_tokens", 0
|
||
)
|
||
_it_r = _iter_timings or {}
|
||
_accumulated_predicted_ms += _it_r.get("predicted_ms", 0)
|
||
_accumulated_predicted_n += _it_r.get("predicted_n", 0)
|
||
yield {"type": "status", "text": ""}
|
||
continue
|
||
|
||
# Content was already streamed. Yield metadata.
|
||
yield {"type": "status", "text": ""}
|
||
_fu = (
|
||
_backfill_usage_from_timings(_iter_usage, _iter_timings)
|
||
or {}
|
||
)
|
||
_fc = _fu.get("completion_tokens", 0)
|
||
_fp = _fu.get("prompt_tokens", 0)
|
||
_tc = _fc + _accumulated_completion_tokens
|
||
if (
|
||
_iter_usage
|
||
or _iter_timings
|
||
or _accumulated_completion_tokens
|
||
):
|
||
_mt = dict(_iter_timings) if _iter_timings else {}
|
||
if _accumulated_predicted_ms or _accumulated_predicted_n:
|
||
_mt["predicted_ms"] = (
|
||
_mt.get("predicted_ms", 0)
|
||
+ _accumulated_predicted_ms
|
||
)
|
||
_tn = (
|
||
_mt.get("predicted_n", 0) + _accumulated_predicted_n
|
||
)
|
||
_mt["predicted_n"] = _tn
|
||
_tms = _mt["predicted_ms"]
|
||
if _tms > 0:
|
||
_mt["predicted_per_second"] = _tn / (_tms / 1000.0)
|
||
yield {
|
||
"type": "metadata",
|
||
"usage": {
|
||
"prompt_tokens": _fp,
|
||
"completion_tokens": _tc,
|
||
"total_tokens": _fp + _tc,
|
||
},
|
||
"timings": _mt,
|
||
}
|
||
return
|
||
|
||
# Safety net caught tool XML -- treat as tool call
|
||
tool_calls = _safety_tc
|
||
content_text = _strip_tool_markup(
|
||
content_accum,
|
||
final = True,
|
||
)
|
||
logger.info(
|
||
f"Safety net: parsed {len(tool_calls)} tool call(s) "
|
||
f"from streamed content"
|
||
)
|
||
else:
|
||
# ── DRAINING path: assemble tool_calls ──
|
||
tool_calls = None
|
||
content_text = content_accum
|
||
if has_structured_tc:
|
||
# Filter out incomplete fragments (e.g. from
|
||
# truncation by max_tokens or disconnect).
|
||
tool_calls = [
|
||
tool_calls_acc[i]
|
||
for i in sorted(tool_calls_acc)
|
||
if (
|
||
tool_calls_acc[i]
|
||
.get("function", {})
|
||
.get("name", "")
|
||
.strip()
|
||
)
|
||
] or None
|
||
if (
|
||
not tool_calls
|
||
and auto_heal_tool_calls
|
||
and any(s in content_accum for s in _TOOL_XML_SIGNALS)
|
||
):
|
||
tool_calls = self._parse_tool_calls_from_text(
|
||
content_accum,
|
||
)
|
||
if tool_calls and not has_structured_tc:
|
||
content_text = _strip_tool_markup(
|
||
content_text,
|
||
final = True,
|
||
)
|
||
if tool_calls:
|
||
logger.info(
|
||
f"Parsed {len(tool_calls)} tool call(s) from "
|
||
f"{'structured delta' if has_structured_tc else 'content text'}"
|
||
)
|
||
if not tool_calls:
|
||
# DRAINING but no tool calls (false positive).
|
||
# Merge accumulated metrics from prior tool
|
||
# iterations so they are not silently dropped.
|
||
yield {"type": "status", "text": ""}
|
||
if content_accum:
|
||
# Strip leaked tool-call XML before yielding
|
||
content_accum = _strip_tool_markup(
|
||
content_accum, final = True
|
||
)
|
||
if content_accum:
|
||
yield {"type": "content", "text": content_accum}
|
||
_fu = (
|
||
_backfill_usage_from_timings(_iter_usage, _iter_timings)
|
||
or {}
|
||
)
|
||
_fc = _fu.get("completion_tokens", 0)
|
||
_fp = _fu.get("prompt_tokens", 0)
|
||
_tc = _fc + _accumulated_completion_tokens
|
||
if (
|
||
_iter_usage
|
||
or _iter_timings
|
||
or _accumulated_completion_tokens
|
||
):
|
||
_mt = dict(_iter_timings) if _iter_timings else {}
|
||
if _accumulated_predicted_ms or _accumulated_predicted_n:
|
||
_mt["predicted_ms"] = (
|
||
_mt.get("predicted_ms", 0)
|
||
+ _accumulated_predicted_ms
|
||
)
|
||
_tn = (
|
||
_mt.get("predicted_n", 0) + _accumulated_predicted_n
|
||
)
|
||
_mt["predicted_n"] = _tn
|
||
_tms = _mt["predicted_ms"]
|
||
if _tms > 0:
|
||
_mt["predicted_per_second"] = _tn / (_tms / 1000.0)
|
||
yield {
|
||
"type": "metadata",
|
||
"usage": {
|
||
"prompt_tokens": _fp,
|
||
"completion_tokens": _tc,
|
||
"total_tokens": _fp + _tc,
|
||
},
|
||
"timings": _mt,
|
||
}
|
||
return
|
||
|
||
# ── Execute tool calls ──
|
||
_accumulated_completion_tokens += (
|
||
_backfill_usage_from_timings(_iter_usage, _iter_timings) or {}
|
||
).get("completion_tokens", 0)
|
||
_it = _iter_timings or {}
|
||
_accumulated_predicted_ms += _it.get("predicted_ms", 0)
|
||
_accumulated_predicted_n += _it.get("predicted_n", 0)
|
||
|
||
assistant_msg = {"role": "assistant", "content": content_text}
|
||
if tool_calls:
|
||
assistant_msg["tool_calls"] = tool_calls
|
||
conversation.append(assistant_msg)
|
||
|
||
for tc in tool_calls or []:
|
||
func = tc.get("function", {})
|
||
tool_name = func.get("name", "")
|
||
raw_args = func.get("arguments", {})
|
||
|
||
if isinstance(raw_args, str):
|
||
try:
|
||
arguments = json.loads(raw_args)
|
||
except (json.JSONDecodeError, ValueError):
|
||
if auto_heal_tool_calls:
|
||
arguments = {"query": raw_args}
|
||
else:
|
||
arguments = {"raw": raw_args}
|
||
else:
|
||
arguments = raw_args
|
||
|
||
if tool_name == "web_search":
|
||
_ws_url = (arguments.get("url") or "").strip()
|
||
if _ws_url:
|
||
_parsed = urlparse(_ws_url)
|
||
if _parsed.scheme in ("http", "https") and _parsed.hostname:
|
||
_ws_host = _parsed.hostname
|
||
if _ws_host.startswith("www."):
|
||
_ws_host = _ws_host[4:]
|
||
status_text = f"Reading: {_ws_host}"
|
||
else:
|
||
status_text = "Reading page..."
|
||
else:
|
||
status_text = f"Searching: {arguments.get('query', '')}"
|
||
elif tool_name == "python":
|
||
preview = (
|
||
(arguments.get("code") or "").strip().split("\n")[0][:60]
|
||
)
|
||
status_text = (
|
||
f"Running Python: {preview}"
|
||
if preview
|
||
else "Running Python..."
|
||
)
|
||
elif tool_name == "terminal":
|
||
cmd_preview = (arguments.get("command") or "")[:60]
|
||
status_text = (
|
||
f"Running: {cmd_preview}"
|
||
if cmd_preview
|
||
else "Running command..."
|
||
)
|
||
else:
|
||
status_text = f"Calling: {tool_name}"
|
||
yield {"type": "status", "text": status_text}
|
||
|
||
# ── Duplicate call detection ──────────────
|
||
# str(dict) is stable here: arguments always comes from
|
||
# json.loads on the same model output within one request,
|
||
# so insertion order is deterministic (Python 3.7+).
|
||
_tc_key = tool_name + str(arguments)
|
||
_prev = _tool_call_history[-1] if _tool_call_history else None
|
||
_is_duplicate = bool(_prev) and _prev[0] == _tc_key and not _prev[1]
|
||
# Guard against the model emitting a tool not in the
|
||
# per-request advertised set: filtered MCP names, a
|
||
# built-in the caller opted out of, or a stale name
|
||
# from a prior turn. Mirrors the safetensors loop's
|
||
# allowed_tool_names check.
|
||
_allowed = {
|
||
(t.get("function") or {}).get("name")
|
||
for t in (tools or [])
|
||
if (t.get("function") or {}).get("name")
|
||
}
|
||
_is_disabled = bool(_allowed) and tool_name not in _allowed
|
||
# Only gate calls that would actually run: duplicate or
|
||
# disabled calls are short-circuited below and never
|
||
# execute, so prompting for them would be noise.
|
||
# Registering the slot before tool_start closes the race
|
||
# where the confirmation could arrive before the waiter.
|
||
_needs_confirm = (
|
||
confirm_tool_calls and not _is_duplicate and not _is_disabled
|
||
)
|
||
_approval_id = new_approval_id() if _needs_confirm else ""
|
||
_decision_slot = (
|
||
begin_tool_decision(session_id, _approval_id)
|
||
if _needs_confirm
|
||
else None
|
||
)
|
||
|
||
yield {
|
||
"type": "tool_start",
|
||
"tool_name": tool_name,
|
||
"tool_call_id": tc.get("id", ""),
|
||
"arguments": arguments,
|
||
"approval_id": _approval_id,
|
||
"awaiting_confirmation": _needs_confirm,
|
||
}
|
||
|
||
_denied = False
|
||
if _is_duplicate:
|
||
result = (
|
||
"You already made this exact call. "
|
||
"Do not repeat the same tool call. "
|
||
"Try a different approach: fetch a URL "
|
||
"from previous results, use Python to "
|
||
"process data you already have, or "
|
||
"provide your final answer now."
|
||
)
|
||
elif _is_disabled:
|
||
result = (
|
||
f"Error: tool '{tool_name}' is not enabled "
|
||
"for this request. Use one of the enabled "
|
||
"tools or provide a final answer."
|
||
)
|
||
else:
|
||
_denied = (
|
||
_decision_slot is not None
|
||
and wait_tool_decision(
|
||
_decision_slot,
|
||
_approval_id,
|
||
cancel_event = cancel_event,
|
||
)
|
||
== "deny"
|
||
)
|
||
if _denied:
|
||
result = TOOL_REJECTED_MESSAGE
|
||
else:
|
||
_effective_timeout = (
|
||
None if tool_call_timeout >= 9999 else tool_call_timeout
|
||
)
|
||
result = execute_tool(
|
||
tool_name,
|
||
arguments,
|
||
cancel_event = cancel_event,
|
||
timeout = _effective_timeout,
|
||
session_id = session_id,
|
||
)
|
||
|
||
yield {
|
||
"type": "tool_end",
|
||
"tool_name": tool_name,
|
||
"tool_call_id": tc.get("id", ""),
|
||
"result": result,
|
||
}
|
||
|
||
# Nudge model to try a different approach on errors
|
||
_error_prefixes = (
|
||
"Error",
|
||
"Search failed",
|
||
"Execution error",
|
||
"Blocked:",
|
||
"Exit code",
|
||
"Failed to fetch",
|
||
"Failed to resolve",
|
||
"No query provided",
|
||
)
|
||
_is_error = isinstance(result, str) and result.lstrip().startswith(
|
||
_error_prefixes
|
||
)
|
||
# A user-denied call never executed, so it must not count
|
||
# toward duplicate detection — otherwise re-issuing and
|
||
# approving the same call would be rejected as a duplicate.
|
||
if not _denied:
|
||
_tool_call_history.append((_tc_key, _is_error))
|
||
# Strip image sentinel before feeding result to the LLM
|
||
# (the full result with sentinel is still yielded via
|
||
# tool_end so the frontend can extract image paths).
|
||
_result_content = result
|
||
if "\n__IMAGES__:" in _result_content:
|
||
_result_content = _result_content.rsplit("\n__IMAGES__:", 1)[0]
|
||
if _is_error:
|
||
_result_content = (
|
||
_result_content + "\n\nThe tool call encountered an issue. "
|
||
"Please try a different approach or rephrase your request."
|
||
)
|
||
|
||
tool_msg = {
|
||
"role": "tool",
|
||
"name": tool_name,
|
||
"content": _result_content,
|
||
}
|
||
tool_call_id = tc.get("id")
|
||
if tool_call_id:
|
||
tool_msg["tool_call_id"] = tool_call_id
|
||
conversation.append(tool_msg)
|
||
|
||
# Clear tool status badge before next generation iteration
|
||
yield {"type": "status", "text": ""}
|
||
# Continue the loop to let model respond with context
|
||
continue
|
||
|
||
except httpx.ConnectError:
|
||
raise RuntimeError("Lost connection to llama-server")
|
||
except Exception as e:
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
return
|
||
raise
|
||
|
||
# ── Tool iteration cap reached -- synthesize final answer ──
|
||
# The model used all iterations without producing a final text
|
||
# response. Inject a nudge so the final streaming pass produces
|
||
# a useful answer instead of continuing to request tools.
|
||
if max_tool_iterations > 0:
|
||
conversation.append(
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
"You have used all available tool calls. Based on "
|
||
"everything you have found so far, provide your final "
|
||
"answer now. Do not call any more tools."
|
||
),
|
||
}
|
||
)
|
||
|
||
# Clear status
|
||
yield {"type": "status", "text": ""}
|
||
|
||
# Final streaming pass with the full conversation context
|
||
stream_payload = {
|
||
"messages": conversation,
|
||
"stream": True,
|
||
"temperature": temperature,
|
||
"top_p": top_p,
|
||
"top_k": top_k if top_k >= 0 else 0,
|
||
"min_p": min_p,
|
||
"repeat_penalty": repetition_penalty,
|
||
"presence_penalty": presence_penalty,
|
||
}
|
||
_reasoning_kw = self._request_reasoning_kwargs(
|
||
enable_thinking, reasoning_effort, preserve_thinking
|
||
)
|
||
if _reasoning_kw is not None:
|
||
stream_payload["chat_template_kwargs"] = _reasoning_kw
|
||
stream_payload["max_tokens"] = (
|
||
max_tokens
|
||
if max_tokens is not None
|
||
else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR)
|
||
)
|
||
stream_payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS
|
||
if stop:
|
||
stream_payload["stop"] = stop
|
||
stream_payload["stream_options"] = {"include_usage": True}
|
||
|
||
cumulative = ""
|
||
_last_emitted = ""
|
||
in_thinking = False
|
||
has_content_tokens = False
|
||
reasoning_text = ""
|
||
_metadata_usage = None
|
||
_metadata_timings = None
|
||
_stream_done = False
|
||
|
||
try:
|
||
stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10)
|
||
_auth_headers = (
|
||
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
||
)
|
||
with httpx.Client(
|
||
timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0)
|
||
) as client:
|
||
with self._stream_with_retry(
|
||
client,
|
||
url,
|
||
stream_payload,
|
||
cancel_event,
|
||
headers = _auth_headers,
|
||
) as response:
|
||
if response.status_code != 200:
|
||
error_body = response.read().decode()
|
||
raise RuntimeError(
|
||
f"llama-server returned {response.status_code}: {error_body}"
|
||
)
|
||
|
||
buffer = ""
|
||
for raw_chunk in self._iter_text_cancellable(
|
||
response, cancel_event
|
||
):
|
||
buffer += raw_chunk
|
||
while "\n" in buffer:
|
||
line, buffer = buffer.split("\n", 1)
|
||
line = line.strip()
|
||
|
||
if not line:
|
||
continue
|
||
if line == "data: [DONE]":
|
||
if in_thinking:
|
||
if has_content_tokens:
|
||
cumulative += "</think>"
|
||
yield {
|
||
"type": "content",
|
||
"text": _strip_tool_markup(
|
||
cumulative, final = True
|
||
),
|
||
}
|
||
else:
|
||
cumulative = reasoning_text
|
||
yield {"type": "content", "text": cumulative}
|
||
_stream_done = True
|
||
break # exit inner while
|
||
if not line.startswith("data: "):
|
||
continue
|
||
|
||
try:
|
||
chunk_data = json.loads(line[6:])
|
||
# Capture server timings/usage from final chunks
|
||
_chunk_timings = chunk_data.get("timings")
|
||
if _chunk_timings:
|
||
_metadata_timings = _chunk_timings
|
||
_chunk_usage = chunk_data.get("usage")
|
||
if _chunk_usage:
|
||
_metadata_usage = _chunk_usage
|
||
choices = chunk_data.get("choices", [])
|
||
if choices:
|
||
delta = choices[0].get("delta", {})
|
||
|
||
reasoning = delta.get("reasoning_content", "")
|
||
if reasoning:
|
||
reasoning_text += reasoning
|
||
if not in_thinking:
|
||
cumulative += "<think>"
|
||
in_thinking = True
|
||
cumulative += reasoning
|
||
yield {"type": "content", "text": cumulative}
|
||
|
||
token = delta.get("content", "")
|
||
if token:
|
||
has_content_tokens = True
|
||
if in_thinking:
|
||
cumulative += "</think>"
|
||
in_thinking = False
|
||
cumulative += token
|
||
cleaned = _strip_tool_markup(cumulative)
|
||
# Only emit when cleaned text grows (monotonic).
|
||
if len(cleaned) > len(_last_emitted):
|
||
_last_emitted = cleaned
|
||
yield {"type": "content", "text": cleaned}
|
||
except json.JSONDecodeError:
|
||
logger.debug(
|
||
f"Skipping malformed SSE line: {line[:100]}"
|
||
)
|
||
if _stream_done:
|
||
break # exit outer for
|
||
_final_usage = _metadata_usage or {}
|
||
_final_completion = _final_usage.get("completion_tokens", 0)
|
||
_final_prompt = _final_usage.get("prompt_tokens", 0)
|
||
_total_completion = (
|
||
_final_completion + _accumulated_completion_tokens
|
||
)
|
||
if _metadata_usage or _metadata_timings:
|
||
_merged_timings = (
|
||
dict(_metadata_timings) if _metadata_timings else {}
|
||
)
|
||
if _accumulated_predicted_ms or _accumulated_predicted_n:
|
||
_merged_timings["predicted_ms"] = (
|
||
_merged_timings.get("predicted_ms", 0)
|
||
+ _accumulated_predicted_ms
|
||
)
|
||
_total_predicted_n = (
|
||
_merged_timings.get("predicted_n", 0)
|
||
+ _accumulated_predicted_n
|
||
)
|
||
_merged_timings["predicted_n"] = _total_predicted_n
|
||
_total_predicted_ms = _merged_timings["predicted_ms"]
|
||
if _total_predicted_ms > 0:
|
||
_merged_timings["predicted_per_second"] = (
|
||
_total_predicted_n / (_total_predicted_ms / 1000.0)
|
||
)
|
||
yield {
|
||
"type": "metadata",
|
||
"usage": {
|
||
"prompt_tokens": _final_prompt,
|
||
"completion_tokens": _total_completion,
|
||
"total_tokens": _final_prompt + _total_completion,
|
||
},
|
||
"timings": _merged_timings,
|
||
}
|
||
|
||
except httpx.ConnectError:
|
||
raise RuntimeError("Lost connection to llama-server")
|
||
except Exception as e:
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
return
|
||
raise
|
||
|
||
# ── TTS support ────────────────────────────────────────────
|
||
|
||
def detect_audio_type(self) -> Optional[str]:
|
||
"""Detect audio/TTS codec; swallows errors (use _strict variant to distinguish)."""
|
||
try:
|
||
return self._detect_audio_type_strict()
|
||
except Exception as e:
|
||
logger.debug(f"Audio type detection failed: {e}")
|
||
return None
|
||
|
||
def _detect_audio_type_strict(self) -> Optional[str]:
|
||
"""Codec name on match, None on definitive non-audio, raises on transport/JSON errors."""
|
||
if not self.is_loaded:
|
||
return None
|
||
_auth_headers = (
|
||
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
||
)
|
||
with httpx.Client(timeout = 10, headers = _auth_headers) as client:
|
||
|
||
def _detok(tid: int) -> str:
|
||
# Non-200 means "marker not in vocab" -- keep probing.
|
||
# Transport / JSON errors still raise.
|
||
r = client.post(f"{self.base_url}/detokenize", json = {"tokens": [tid]})
|
||
if r.status_code != 200:
|
||
return ""
|
||
return r.json().get("content", "")
|
||
|
||
def _tok(text: str) -> list[int]:
|
||
r = client.post(
|
||
f"{self.base_url}/tokenize",
|
||
json = {"content": text, "add_special": False},
|
||
)
|
||
if r.status_code != 200:
|
||
return []
|
||
return r.json().get("tokens", [])
|
||
|
||
# Check codec-specific tokens (not generic ones that may exist in non-audio models)
|
||
if "<custom_token_" in _detok(128258) and "<custom_token_" in _detok(
|
||
128259
|
||
):
|
||
return "snac"
|
||
if len(_tok("<|AUDIO|>")) == 1 and len(_tok("<|audio_eos|>")) == 1:
|
||
return "csm"
|
||
if len(_tok("<|startoftranscript|>")) == 1:
|
||
return "whisper"
|
||
if len(_tok("<audio_soft_token>")) == 1:
|
||
return "audio_vlm"
|
||
if (
|
||
len(_tok("<|bicodec_semantic_0|>")) == 1
|
||
and len(_tok("<|bicodec_global_0|>")) == 1
|
||
):
|
||
return "bicodec"
|
||
if len(_tok("<|c1_0|>")) == 1 and len(_tok("<|c2_0|>")) == 1:
|
||
return "dac"
|
||
return None
|
||
|
||
# Prompt format per codec: (template, stop_tokens, needs_token_ids)
|
||
# Matches prompts in InferenceBackend._generate_snac/bicodec/dac
|
||
_TTS_PROMPTS = {
|
||
"snac": (
|
||
"<custom_token_3>{text}<|eot_id|><custom_token_4>",
|
||
["<custom_token_2>"],
|
||
True,
|
||
),
|
||
"bicodec": (
|
||
"<|task_tts|><|start_content|>{text}<|end_content|><|start_global_token|>",
|
||
["<|im_end|>", "</s>"],
|
||
False,
|
||
),
|
||
"dac": (
|
||
"<|im_start|>\n<|text_start|>{text}<|text_end|>\n<|audio_start|><|global_features_start|>\n",
|
||
["<|im_end|>", "<|audio_end|>"],
|
||
False,
|
||
),
|
||
}
|
||
|
||
_codec_mgr = None # Shared AudioCodecManager instance
|
||
|
||
def init_audio_codec(self, audio_type: str) -> None:
|
||
"""Load the audio codec at model load time (mirrors non-GGUF path)."""
|
||
import torch
|
||
from core.inference.audio_codecs import AudioCodecManager
|
||
|
||
if LlamaCppBackend._codec_mgr is None:
|
||
LlamaCppBackend._codec_mgr = AudioCodecManager()
|
||
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
model_repo_path = None
|
||
|
||
# BiCodec needs a repo with BiCodec/ weights — download canonical SparkTTS
|
||
if audio_type == "bicodec":
|
||
from huggingface_hub import snapshot_download
|
||
import os
|
||
|
||
repo_path = snapshot_download(
|
||
"unsloth/Spark-TTS-0.5B", local_dir = "Spark-TTS-0.5B"
|
||
)
|
||
model_repo_path = os.path.abspath(repo_path)
|
||
|
||
LlamaCppBackend._codec_mgr.load_codec(
|
||
audio_type, device, model_repo_path = model_repo_path
|
||
)
|
||
logger.info(f"Loaded audio codec for GGUF TTS: {audio_type}")
|
||
|
||
def generate_audio_response(
|
||
self,
|
||
text: str,
|
||
audio_type: str,
|
||
temperature: float = 0.6,
|
||
top_p: float = 0.95,
|
||
top_k: int = 50,
|
||
min_p: float = 0.0,
|
||
max_new_tokens: int = 2048,
|
||
repetition_penalty: float = 1.1,
|
||
) -> tuple:
|
||
"""
|
||
Generate TTS audio via llama-server /completion + codec decoding.
|
||
Returns (wav_bytes, sample_rate).
|
||
"""
|
||
if audio_type not in self._TTS_PROMPTS:
|
||
raise RuntimeError(f"GGUF TTS does not support '{audio_type}' codec.")
|
||
|
||
tpl, stop, need_ids = self._TTS_PROMPTS[audio_type]
|
||
|
||
payload: dict = {
|
||
"prompt": tpl.format(text = text),
|
||
"stream": False,
|
||
"n_predict": max_new_tokens,
|
||
"temperature": temperature,
|
||
"top_p": top_p,
|
||
"top_k": top_k if top_k >= 0 else 0,
|
||
"min_p": min_p,
|
||
"repeat_penalty": repetition_penalty,
|
||
}
|
||
if stop:
|
||
payload["stop"] = stop
|
||
if need_ids:
|
||
payload["n_probs"] = 1
|
||
|
||
_auth_headers = (
|
||
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
||
)
|
||
with httpx.Client(
|
||
timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers
|
||
) as client:
|
||
resp = client.post(f"{self.base_url}/completion", json = payload)
|
||
if resp.status_code != 200:
|
||
raise RuntimeError(
|
||
f"llama-server returned {resp.status_code}: {resp.text}"
|
||
)
|
||
|
||
data = resp.json()
|
||
token_ids = (
|
||
[p["id"] for p in data.get("completion_probabilities", []) if "id" in p]
|
||
if need_ids
|
||
else None
|
||
)
|
||
|
||
import torch
|
||
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
return LlamaCppBackend._codec_mgr.decode(
|
||
audio_type, device, token_ids = token_ids, text = data.get("content", "")
|
||
)
|