8848a310dfAdds a self-contained RAG stack to Studio: knowledge bases with chunked indexing, hybrid (dense + lexical) retrieval, and an automatic first-pass context inject into chat. Embeddings run through a local llama-server GGUF backend (default unsloth/bge-small-en-v1.5-GGUF) with a sentence-transformers fallback. The chat tool loop gains a search_knowledge_base tool, a per-turn re-search cap, and source citation, layered on top of the shared ToolLoopController.
5516 lines
244 KiB
Python
5516 lines
244 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
|
|
|
|
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 (
|
|
RAG_MAX_SEARCHES_PER_TURN,
|
|
RAG_SEARCH_CAP_NUDGE,
|
|
TOOL_XML_SIGNALS,
|
|
parse_tool_calls_from_text as _shared_parse_tool_calls_from_text,
|
|
)
|
|
from core.inference.tool_loop_controller import (
|
|
ToolLoopController,
|
|
tool_event_provenance,
|
|
)
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
# ── Pre-compiled patterns for plan-without-action re-prompt ──
|
|
# Forward-looking intent signals: the model is describing what it *will*
|
|
# do rather than giving a final answer.
|
|
_INTENT_SIGNAL = re.compile(
|
|
r"(?i)("
|
|
# Direct intent ("I'll ...", "Let me ...", straight + curly apostrophes).
|
|
# Excludes "I can"/"I should"/"I want to"/"let's" (common in answers).
|
|
# Negative lookahead drops negated forms ("I will not") 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 = 1
|
|
|
|
# Without max_tokens, llama-server defaults n_predict = n_ctx (up to 262144 for
|
|
# Qwen3.5), causing many-minute zombie decodes when cancel fails.
|
|
# t_max_predict_ms is a wall-clock backstop but per the llama.cpp README only
|
|
# fires after a newline, so we keep a token cap as the front-line limiter.
|
|
# The cap is the effective context length when known, else this floor. 4096 was
|
|
# too low: Qwen3 / gpt-oss reasoning traces and max_tokens-omitting OpenAI-API
|
|
# callers (langchain, llama-index, curl) got truncated mid-sentence.
|
|
_DEFAULT_MAX_TOKENS_FLOOR = 32768
|
|
_DEFAULT_T_MAX_PREDICT_MS = 600_000 # 10 min
|
|
_REPROMPT_MAX_CHARS = 2000
|
|
_FORCED_REPEAT_PLAN_SIGNAL = re.compile(
|
|
r"\b(?:i\s+will|i'll|let\s+me|going\s+to|need\s+to|call|use|run|search|fetch|render)\b",
|
|
re.I,
|
|
)
|
|
_FINAL_ANSWER_SIGNAL = re.compile(
|
|
r"\b(?:final\s+answer|answer\s*:|here\s+is|here's|in\s+summary|result\s*:)\b",
|
|
re.I,
|
|
)
|
|
|
|
|
|
def _is_short_intent_without_action(text: str) -> bool:
|
|
stripped = text.strip()
|
|
return 0 < len(stripped) < _REPROMPT_MAX_CHARS and _INTENT_SIGNAL.search(stripped) is not None
|
|
|
|
|
|
def _should_suppress_forced_no_tool_output(text: str) -> bool:
|
|
"""Suppress only repeated forced-turn planning text, not final answers."""
|
|
stripped = text.strip()
|
|
if not stripped or len(stripped) >= _REPROMPT_MAX_CHARS:
|
|
return False
|
|
if _FINAL_ANSWER_SIGNAL.search(stripped):
|
|
return False
|
|
return _FORCED_REPEAT_PLAN_SIGNAL.search(stripped) is not None
|
|
|
|
|
|
# ── 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 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 on a daemon thread, so concurrent sockets aren't
|
|
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 -> dead.
|
|
return True if result[0] is None else result[0]
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def _hf_offline_if_dns_dead():
|
|
"""Set HF_HUB_OFFLINE for this block only when DNS to huggingface.co fails;
|
|
restores env on exit so a transient hiccup can't quarantine the process.
|
|
No-op if the user already set it."""
|
|
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, else None."""
|
|
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 (result persistently cached)
|
|
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(/...)` -> `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]}"
|
|
|
|
|
|
# Lazy import to avoid pulling transformers in at module level.
|
|
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 the same five keys as the GGUF sniffer: ``supports_reasoning``,
|
|
``reasoning_style`` (``"enable_thinking"`` | ``"reasoning_effort"``),
|
|
``reasoning_always_on``, ``supports_preserve_thinking``,
|
|
``supports_tools``. Used by both the llama-server backend at load time
|
|
and the safetensors/transformers paths in ``routes/inference`` so they
|
|
agree on what the frontend sees.
|
|
"""
|
|
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 use reasoning_effort
|
|
# ("low" | "medium" | "high"), not 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', not '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).
|
|
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: 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 one
|
|
--spec-type (comma-separated to chain), so suppress auto-emit."""
|
|
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 legacy ``--spec-ngram-size-n`` /
|
|
``--draft-min`` / ``--draft-max``. ``caps`` comes from
|
|
``probe_server_capabilities``; ``ngram_mod_flavor`` says which set is
|
|
real (vs a removal-stub). Returns ``[]`` when neither 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":
|
|
# Pre-rename llama.cpp: same knobs lived under --spec-ngram-size-n
|
|
# (lookup length) and generic --draft-min / --draft-max (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.
|
|
# Dropdown renders five (auto, mtp, ngram, mtp+ngram, off); the load API
|
|
# also accepts legacy values 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 ``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]
|
|
# Old persisted state emits llama.cpp comma-chains 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 final SSE chunk's ``usage`` can be absent
|
|
or zero on some server builds / streaming configs, making 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() — start llama-server with the GGUF file
|
|
2. generate_chat_completion() — proxy to /v1/chat/completions, stream back
|
|
3. unload_model() — terminate the 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
|
|
# (auto/mtp/ngram/mtp+ngram/off/ngram-simple). Round-tripped through the
|
|
# status API so the dropdown reflects the picked mode, not the resolved
|
|
# flag set (auto on a 27B MTP GGUF resolves to draft-mtp but reads "Auto").
|
|
self._requested_spec_mode: Optional[str] = None
|
|
# User --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 earlier layers' KV 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
|
|
# Audio INPUT capability (distinct from _is_audio, which is TTS output).
|
|
self._has_audio_input: bool = False
|
|
self._mmproj_has_audio: bool = False # clip.has_audio_encoder, set at load
|
|
# 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 (a 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.
|
|
|
|
The UI's "safe zone" warning threshold: the ``_fit_context_to_vram``
|
|
binary-search cap for the best GPU subset, or the 4096 fallback if the
|
|
weights exceed 90% of every subset. The slider ceiling is
|
|
``native_context_length``; dragging above this triggers the 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.
|
|
|
|
During warm-up llama-server mmaps weight shards into page cache before
|
|
pushing layers to VRAM, a window where status only reports ``loading``
|
|
and the UI spinner looks stuck for minutes on large MoEs. Samples
|
|
``/proc/<pid>/status VmRSS`` against the sum of GGUF shard sizes for a
|
|
real progress bar. Returns ``None`` when no load is in flight.
|
|
|
|
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; returns ``None`` 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 shard sizes (primary + any extras 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 share the primary's prefix before the shard index.
|
|
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.
|
|
|
|
Merges 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 and don't 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 dir
|
|
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 on Linux)
|
|
cmake_bin = custom_dir / "build" / "bin" / binary_name
|
|
if cmake_bin.is_file():
|
|
return str(cmake_bin)
|
|
# build/bin/Release/ (cmake 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:
|
|
# _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)
|
|
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``: ``"new"`` when the post-rename
|
|
``--spec-ngram-mod-n-match / -n-min / -n-max`` are real args;
|
|
``"legacy"`` when only the pre-rename
|
|
``--spec-ngram-size-n / --draft-min / --draft-max`` are real (the
|
|
rename ships stub removal entries for legacy names, told apart by
|
|
the "argument has been removed" description); ``None`` if neither
|
|
set is usable.
|
|
|
|
``spec_draft_n_max_flag``: the flag the binary accepts --
|
|
``--spec-draft-n-max`` post-rename, ``--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 + its indented
|
|
# continuation), 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 non-flag token 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 from the --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 new
|
|
# args (real) and legacy ones (stubs); pre-rename builds only
|
|
# have 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:
|
|
"""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 elsewhere (the env hurts discrete GPUs). ROCm reuses torch.cuda.*;
|
|
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 (HIP runtime reuses the ``torch.cuda.*``
|
|
namespace). Covers the AMD case for issue #5106 (nvidia-smi
|
|
probe returned [] on AMD) and NVIDIA hosts missing
|
|
``nvidia-smi`` from PATH.
|
|
|
|
Returns list of (gpu_index, free_mib) sorted by index; empty if no
|
|
supported GPU is reachable.
|
|
"""
|
|
# ── 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 ("0,1,").
|
|
# Empty mask (CVD="") yields an empty set -> all GPUs
|
|
# filtered out, per 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 (driver order isn't).
|
|
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. We
|
|
# feed these IDs back into the subprocess as CVD, so visible ordinals
|
|
# must be translated to physical indices first; otherwise CVD=2,3
|
|
# gets rewritten to 0,1 and targets the wrong GPUs.
|
|
physical_ids: Optional[list[int]] = None
|
|
# Match utils/hardware/hardware.py::_get_parent_visible_gpu_spec:
|
|
# treat an empty mask (HIP_VISIBLE_DEVICES="") as "no GPUs" rather
|
|
# than falling through. ``or`` would coerce "" to the wrong source.
|
|
if getattr(torch.version, "hip", None) is not None:
|
|
hip_v = os.environ.get("HIP_VISIBLE_DEVICES")
|
|
rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES")
|
|
cvd = (
|
|
hip_v
|
|
if hip_v is not None
|
|
else rocr_v
|
|
if rocr_v is not None
|
|
else os.environ.get("CUDA_VISIBLE_DEVICES")
|
|
)
|
|
else:
|
|
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
|
|
if cvd is not None:
|
|
try:
|
|
# Empty mask (CVD="") yields an empty list -> no GPUs,
|
|
# consistent with the nvidia-smi path.
|
|
physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()]
|
|
except ValueError:
|
|
physical_ids = None
|
|
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 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 driver reclaims a dead process's allocations asynchronously, so
|
|
sampling free memory in the kill-to-spawn window reads artificially low
|
|
and pushes GPU selection toward needless CPU offload (the Apply-reload
|
|
OOM bare-shell launches never see).
|
|
|
|
Short-circuits on cold start, stale kill (older than
|
|
``_VRAM_SETTLE_WINDOW_S``), CPU-only hosts, probe exceptions, and GPU-set
|
|
changes. ``max_wait`` bounds wall-clock time so a wedged ``nvidia-smi``
|
|
can't 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.
|
|
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 dropped 91-94% fits to CPU offload (#5106).
|
|
_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, covering the
|
|
Windows wheel layouts seen in the wild:
|
|
* ``nvidia/<pkg>/bin`` -- legacy modular wheels.
|
|
* ``nvidia/<pkg>/bin/x86_64`` and ``.../bin/x64`` -- CUDA 13 layout
|
|
for unsuffixed packages (#5106).
|
|
* ``nvidia/<pkg>/Library/bin`` (and arch subdirs) -- conda repacks.
|
|
* ``torch/lib`` -- PyTorch's CUDA-bundled wheel can ship
|
|
``cudart64_*.dll`` here; mirrors install_llama_prebuilt.py.
|
|
|
|
Walks with ``Path.iterdir`` not ``glob.glob`` so it's safe against
|
|
Windows paths containing ``[`` or ``]`` (valid in usernames)."""
|
|
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
|
|
# Arch-specific subdirs first so the explicit cudart64_X.dll
|
|
# location wins over an empty sibling ``bin``.
|
|
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 prepended so llama-server.exe resolves cudart /
|
|
cublas DLLs: binary_dir, pip nvidia wheels, CUDA_PATH/bin, .../bin/x64.
|
|
Extracted so test_windows_gpu_detection_mock tests the real logic. #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 from estimated VRAM and free memory.
|
|
|
|
``model_size_bytes`` should include weights and estimated KV cache.
|
|
``_GPU_PIN_VRAM_FRACTION`` provides headroom for compute buffers,
|
|
CUDA context, and other runtime overhead.
|
|
|
|
Returns (gpu_indices, use_fit):
|
|
- ([1], False) fits on 1 GPU at the headroom threshold
|
|
- ([1, 2], False) needs 2 GPUs
|
|
- (None, True) 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 1 GPU at the usable-VRAM threshold.
|
|
if ranked[0][1] * usable_fraction >= model_size_mib:
|
|
return [ranked[0][0]], False
|
|
|
|
# Try 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
|
|
|
|
# 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 suffices (K-only cache).
|
|
if self._kv_lora_rank is not None:
|
|
return True
|
|
# New-style: need 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.
|
|
|
|
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 full
|
|
``n_ctx`` (collapses path 3 to path 4 for them).
|
|
n_parallel -- ``--parallel`` slots: non-SWA layers stay constant
|
|
(cells split across slots), SWA layers scale linearly.
|
|
kv_unified -- ``--kv-unified`` (default on): no-op for memory math;
|
|
kept for API forward-compat.
|
|
ctx_checkpoints -- ``--ctx-checkpoints`` (PR #15293): N SWA snapshots
|
|
per slot, one sliding-window of state per SWA layer.
|
|
|
|
Returns 0 if metadata is insufficient.
|
|
"""
|
|
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 earlier KV in the last ``shared_kv_layers``
|
|
# blocks (no cache). Floor at 1 so a bad 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)
|
|
# One compressed KV latent per token/layer (shared across heads); V is
|
|
# reconstructed from it, no separate V cache. key_length = kv_lora_rank
|
|
# + rope_dim. 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 128x's.
|
|
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 filled by the resolver at parse time; if absent, falls through
|
|
# to the legacy 1/4-global heuristic below. Per-layer-type --parallel N
|
|
# accounting (verified against llama-server):
|
|
# * non-SWA layers: total cells = n_ctx split across slots -> CONSTANT.
|
|
# * SWA layers: per-slot cells = 2*sliding_window (capped at n_ctx
|
|
# and per_slot_ctx) -> grows LINEARLY in slots.
|
|
# --swa-full forces full n_ctx for SWA layers; --ctx-checkpoints N adds
|
|
# N snapshots per SWA layer per slot.
|
|
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 caches full context like non-SWA (per-slot cells =
|
|
# per_slot_ctx, collapsing to constant n_ctx total); otherwise SWA
|
|
# caches 2*sliding_window per slot, clamped at per-slot ctx.
|
|
swa_cells_per_slot = per_slot_ctx if swa_full else min(n_ctx, 2 * swa, per_slot_ctx)
|
|
key_len_swa = self._kv_key_length_swa or key_len
|
|
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
|
|
# Only layers that allocate their own KV; 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). If the weights alone don't fit, returns ``requested_ctx``.
|
|
|
|
``kv_on_gpu`` mirrors ``--kv-offload`` (default on); when False the KV
|
|
cache lives in CPU RAM and the requested context is honored verbatim.
|
|
Other keyword args mirror ``_estimate_kv_cache_bytes``.
|
|
|
|
``mtp_engaged`` reserves extra VRAM for the MTP draft model's KV cache +
|
|
compute buffers, else tight tiers (e.g. 32 GB) spill to a slower 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
|
|
|
|
# Already fits?
|
|
kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs)
|
|
if model_footprint + kv <= budget_bytes:
|
|
return requested_ctx
|
|
|
|
# Weights alone exceed budget -- reducing ctx can't help; --fit handles it.
|
|
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, never above 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 (e.g.
|
|
UD-Q4_K_XL with 9 shards of 50 GB each = 450 GB total).
|
|
|
|
Returns (first_shard_filename, total_size_bytes) or None.
|
|
"""
|
|
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
|
|
|
|
# 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 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))
|
|
|
|
# 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 subprocess stdout lines in a background thread.
|
|
|
|
Prevents a pipe-buffer deadlock on Windows (~4 KB buffer): 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 has the full 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 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 so stale flags (e.g. _supports_reasoning) don't
|
|
# 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",
|
|
}
|
|
# Arch-specific keys added dynamically once we know the arch.
|
|
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 from an
|
|
# HTTP byte-range fetch): bail out so the resolver
|
|
# fallback runs on whatever we parsed.
|
|
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 header
|
|
# fetch); break so the resolver fallback 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 unload_model() can set _cancel_event at
|
|
any time; checks it 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"
|
|
)
|
|
|
|
# Resolve 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; 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 in the HF cache so we only preflight
|
|
# against what we must 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
|
|
# 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, else any mmproj*.gguf. Returns the local
|
|
path, or None if none 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 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.
|
|
|
|
None if mmproj_path is None, missing, 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 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. Unsloth publishes FLUX and
|
|
# Qwen-Image GGUFs under
|
|
# https://huggingface.co/collections/unsloth/unsloth-diffusion-ggufs.
|
|
# Matched exactly (not a substring) so a chat arch containing "wan"/"sd1"
|
|
# (e.g. "taiwan") isn't 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 otherwise collapse into the same opaque
|
|
"invalid GGUF or out of memory" message. Worst case: a diffusion GGUF
|
|
loaded as a chat model -- valid file, plenty of memory, but llama.cpp
|
|
has no such architecture, so the user is told to free memory that was
|
|
never the problem (#5842). Pick the most specific message we can.
|
|
"""
|
|
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 aren't 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 don't 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, # 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="...-GGUF", hf_variant="Q4_K_M"`` → uses ``-hf``
|
|
|
|
Returns True if the server started and the health check passed.
|
|
"""
|
|
# Serialise the whole load so concurrent /load calls never leave two
|
|
# llama-server processes alive (#5401 / #5161). Doesn't block /unload.
|
|
with self._serial_load_lock:
|
|
# Duplicate /load that raced past the route check: do nothing if the
|
|
# live server already satisfies this request.
|
|
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 finish.
|
|
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
|
|
# Re-derive after a retried probe (_mmproj_has_audio persists).
|
|
from utils.models.model_config import is_audio_input_type
|
|
|
|
self._has_audio_input = bool(is_audio_input_type(self._audio_type)) or bool(
|
|
self._mmproj_has_audio
|
|
)
|
|
if not self._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
|
|
# can't 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 (DeepSeek).
|
|
self._model_identifier = model_identifier
|
|
|
|
# Read GGUF metadata (context_length, chat_template); header-only.
|
|
self._read_gguf_metadata(model_path)
|
|
|
|
if self._cancel_event.is_set():
|
|
logger.info("Load cancelled after download phase")
|
|
return False
|
|
|
|
# Outside ``self._lock`` so /unload, /cancel, /status aren't
|
|
# blocked. ``unload_model`` also records the kill, so the
|
|
# frontend /unload+/load Apply path engages the wait here even
|
|
# without an in-process kill.
|
|
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) from model size + estimated KV cache. Seed
|
|
# safe defaults before probing so the except path 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 exists; else keep 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 native context length;
|
|
# GPU/VRAM-fit logic below may shrink it on limited HW.
|
|
max_available_ctx = self._context_length or effective_ctx
|
|
|
|
# Will MTP engage on this load? If so, auto-fit reserves
|
|
# extra VRAM for the draft model. Mirrors
|
|
# _build_speculative_flags' resolver: forced mtp / mtp+ngram
|
|
# always engage; auto only on an MTP GGUF >= 3B; 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 GPU VRAM and select GPUs. Two
|
|
# policies by whether the user set n_ctx:
|
|
# Explicit n_ctx: honor it. Try the full context with
|
|
# _select_gpus (as many GPUs as needed); cap only if it
|
|
# doesn't fit on any combination.
|
|
# Auto n_ctx=0 (native): prefer fewer GPUs with reduced
|
|
# context, since multi-GPU is slower.
|
|
gpu_indices, use_fit = None, True
|
|
explicit_ctx = requested_ctx > 0
|
|
|
|
if gpus and self._can_estimate_kv() and effective_ctx > 0:
|
|
# Largest hardware-aware cap from the native context
|
|
# across all usable GPU subsets (for UI bounds),
|
|
# independent of the 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, so no
|
|
# context fits. Anchor the UI "safe zone" at 4096
|
|
# so the slider warns above the fallback.
|
|
max_available_ctx = min(4096, native_ctx_for_cap)
|
|
|
|
if explicit_ctx:
|
|
# Honor the requested context verbatim. If it fits,
|
|
# pin GPUs and skip --fit; else ship -c <ctx> --fit
|
|
# on and let llama-server flex -ngl (CPU offload).
|
|
requested_total = model_size + self._estimate_kv_cache_bytes(
|
|
effective_ctx, cache_type_kv, n_parallel = n_parallel
|
|
)
|
|
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 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 --fit on: a model overflowing
|
|
# at 131k may pin fine with a 4096 KV (#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 -- file-size-only check; keep the
|
|
# ceiling at 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 UI to 4096
|
|
# so the slider isn't on an unusable native ctx.
|
|
effective_ctx = min(4096, effective_ctx) if effective_ctx > 0 else 4096
|
|
|
|
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"
|
|
)
|
|
|
|
# Audio input straight from the mmproj (clip.has_audio_encoder),
|
|
# independent of token names.
|
|
self._mmproj_has_audio = False
|
|
if launch_mmproj_path:
|
|
try:
|
|
from utils.models.gguf_metadata import (
|
|
read_mmproj_audio_capability,
|
|
)
|
|
self._mmproj_has_audio = bool(
|
|
read_mmproj_audio_capability(launch_mmproj_path)
|
|
)
|
|
except Exception as e:
|
|
logger.debug(f"mmproj audio-capability read failed: {e}")
|
|
|
|
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:
|
|
# Fits on selected GPU(s) -- offload all layers
|
|
cmd.extend(["-ngl", "-1"])
|
|
|
|
# -1 = llama.cpp auto-detect (physical cores). Pass explicitly
|
|
# so we don't inherit llama-server's internal default, which
|
|
# has varied (hardware concurrency incl. hyperthreads on some
|
|
# builds).
|
|
cmd.extend(["--threads", str(n_threads if n_threads is not None else -1)])
|
|
|
|
# Enable Jinja chat template rendering
|
|
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. See _build_speculative_flags for the
|
|
# mode resolution, benchmarks, and llama.cpp references.
|
|
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}")
|
|
|
|
# Default thinking mode for reasoning models. Qwen3.5/3.6 below
|
|
# 9B disable thinking by default; 9B+ enable it. Always-on
|
|
# templates ignore the kwarg, 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: --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 pass-through args go last so llama.cpp's last-wins parsing
|
|
# lets the user override Studio's auto-set flags. Already
|
|
# validated by the route 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)}")
|
|
|
|
# Library paths so llama-server finds 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":
|
|
# Ordering: see _build_windows_path_dirs. #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 prebuilt bundles rocblas.dll but NOT the Tensile
|
|
# kernel files (rocblas/library/*.dat + *.hsaco); the DLL
|
|
# searches <binary_dir>/rocblas/library/ which doesn't exist
|
|
# -> silent crash on the first GEMM. ROCBLAS_TENSILE_LIBPATH
|
|
# repoints that search at the ROCm install.
|
|
_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: LD_LIBRARY_PATH for shared libs next to the binary
|
|
# plus 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. The prebuilt
|
|
# binary links 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 where default /usr/local/cuda is 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, narrowing only
|
|
# CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full
|
|
# HIP/ROCR set, so set those too.
|
|
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: drop an orphan Popen a concurrent load may
|
|
# have stored before we overwrite the reference (#5161).
|
|
self._kill_process()
|
|
|
|
self._stdout_lines = []
|
|
# Tee llama-server output to a dedicated log file so a
|
|
# post-mortem has the full trail even when the parent only kept
|
|
# the last 50 lines. Path is under the studio home.
|
|
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(),
|
|
)
|
|
|
|
# Background thread to drain stdout (prevents 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 gguf_path is None and ``model_path`` is what
|
|
# llama-server mmap's, which downstream consumers need.
|
|
self._gguf_path = model_path
|
|
self._hf_repo = hf_repo
|
|
# For local GGUF files, extract variant from filename if absent
|
|
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 (the native length for display).
|
|
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 start
|
|
# can't poison the next inheritance check. None keeps prior, []
|
|
# clears, list sets. Source records hf_variant for the route's
|
|
# same_source check.
|
|
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
|
|
self._has_audio_input = 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
|
|
|
|
# Audio input = token probe (audio_vlm/whisper) OR mmproj encoder.
|
|
from utils.models.model_config import is_audio_input_type
|
|
|
|
self._has_audio_input = bool(is_audio_input_type(self._audio_type)) or bool(
|
|
self._mmproj_has_audio
|
|
)
|
|
|
|
if not self._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):
|
|
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 low-repetition chat, 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
|
|
Refs: https://github.com/ggml-org/llama.cpp/blob/master/docs/speculative.md
|
|
https://github.com/ggml-org/llama.cpp/pull/19164
|
|
https://github.com/ggml-org/llama.cpp/pull/18471
|
|
MTP guide: unsloth.ai/docs/models/qwen3.6#mtp-guide
|
|
|
|
Sub-3B dense MTP regresses vs spec-off: 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 (legacy values mapped via
|
|
# _canonicalize_spec_mode).
|
|
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 wins outright; suppress auto-emit to avoid a
|
|
# duplicate spec block.
|
|
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;
|
|
# past 3 regresses 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": the 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 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 requested mode. With --spec-type in
|
|
# extra_args the backend stores None; mirror that here.
|
|
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 engaged. Compare
|
|
# on the resolved spec so an Auto request promoted to draft-mtp still
|
|
# bounces a reload when n_max changes.
|
|
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 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; CUDA/ROCm/Metal/Vulkan/OpenCL/SYCL are GPU, CPU* 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 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._has_audio_input = False
|
|
self._mmproj_has_audio = 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 during the replacement's warm-up can't
|
|
# 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 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 dir (or matches an exact env-var override), to avoid
|
|
terminating unrelated llama-server instances. Mirrors every location
|
|
_find_llama_server_binary() can return, so orphans from any
|
|
supported install path are cleaned up.
|
|
|
|
Uses psutil for cross-platform support (Linux, macOS, Windows);
|
|
falls back to pgrep + /proc/<pid>/exe on Linux when psutil is
|
|
absent.
|
|
"""
|
|
try:
|
|
# -- Build the ownership allowlist --------------------------------
|
|
# exact_binaries -- env var overrides (exact path match).
|
|
# install_roots -- Studio-owned dir trees (binary must be under one).
|
|
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 can't kill a default-install Studio's server.
|
|
if not _is_custom_root:
|
|
install_roots.append(Path.home() / ".unsloth" / "llama.cpp")
|
|
|
|
# Legacy in-tree build dirs (older setup.sh)
|
|
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 absent.
|
|
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()
|
|
|
|
# Ownership: exact match OR binary under a known root.
|
|
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
|
|
|
|
# /proc/<pid>/exe symlinks the real binary, avoiding
|
|
# cmdline-parsing ambiguities; 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 until 200; also detect early exit/crash."""
|
|
deadline = time.monotonic() + timeout
|
|
url = f"http://127.0.0.1:{self._port}/health"
|
|
|
|
while time.monotonic() < deadline:
|
|
# Process crashed?
|
|
if self._process.poll() is not None:
|
|
# Let the drain thread 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 still binding the port
|
|
# (Windows: WinError 10054); the crash branch catches real exits.
|
|
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, *, allow_incomplete: bool = True) -> 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,
|
|
allow_incomplete = allow_incomplete,
|
|
)
|
|
|
|
@staticmethod
|
|
def _build_openai_messages(messages: list[dict], image_b64: Optional[str] = None) -> list[dict]:
|
|
"""Build OpenAI-format messages, optionally injecting an image_url part
|
|
into the last user message for vision models. As-is if no image."""
|
|
if not image_b64:
|
|
return messages
|
|
|
|
# Convert the last user message 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 an httpx streaming response with cancel support.
|
|
|
|
Checks cancel_event between chunks and on ReadTimeout; the
|
|
_stream_with_retry watcher also closes the response on cancel.
|
|
"""
|
|
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 -- 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 once with a long read timeout (120 s) so prefill finishes without
|
|
a retry storm (the old 0.5 s timeout caused duplicate POSTs every half
|
|
second). A watcher thread cancels by closing the response. httpx can't
|
|
interrupt a blocked read before the response exists, so cancel during
|
|
the header wait (1-5 s prefill) is deferred until headers arrive.
|
|
"""
|
|
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. Poll until the response object exists
|
|
# so we can close it, or until the main thread finishes
|
|
# (_cancel_closed 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 can finish without a retry storm.
|
|
# Cancel during prefill and streaming is handled by the watcher
|
|
# thread closing 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,
|
|
seed: Optional[int] = None,
|
|
) -> Generator[str | dict, None, None]:
|
|
"""
|
|
Send a chat completion to llama-server and stream tokens back.
|
|
|
|
Uses /v1/chat/completions -- llama-server applies the chat template
|
|
and handles 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,
|
|
}
|
|
# Per-request enable_thinking / reasoning_effort / preserve_thinking
|
|
_reasoning_kw = self._request_reasoning_kwargs(
|
|
enable_thinking, reasoning_effort, preserve_thinking
|
|
)
|
|
if _reasoning_kw is not None:
|
|
payload["chat_template_kwargs"] = _reasoning_kw
|
|
# Cap to the effective context length when known, else the floor.
|
|
# The wall-clock backstop below stops a stuck model regardless.
|
|
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
|
|
if seed is not None:
|
|
payload["seed"] = seed
|
|
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
|
|
_metadata_finish_reason = 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:
|
|
# model put its whole reply in reasoning
|
|
# (e.g. Qwen3 always-think). Show it as
|
|
# the main response, not a thinking block.
|
|
cumulative = reasoning_text
|
|
yield cumulative
|
|
_stream_done = True
|
|
break # exit inner while
|
|
if not line.startswith("data: "):
|
|
continue
|
|
|
|
try:
|
|
data = json.loads(line[6:])
|
|
# Capture server timings/usage from final chunks.
|
|
_chunk_timings = data.get("timings")
|
|
if _chunk_timings:
|
|
_metadata_timings = _chunk_timings
|
|
_chunk_usage = data.get("usage")
|
|
if _chunk_usage:
|
|
_metadata_usage = _chunk_usage
|
|
choices = data.get("choices", [])
|
|
if choices:
|
|
delta = choices[0].get("delta", {})
|
|
_fr = choices[0].get("finish_reason")
|
|
if _fr:
|
|
_metadata_finish_reason = _fr
|
|
|
|
# Reasoning/thinking tokens: llama-server
|
|
# sends these as "reasoning_content"; wrap
|
|
# in <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 or _metadata_finish_reason:
|
|
_metadata_usage = _backfill_usage_from_timings(
|
|
_metadata_usage, _metadata_timings
|
|
)
|
|
yield {
|
|
"type": "metadata",
|
|
# Never None: a finish-only metadata event (no usage,
|
|
# no timings) would otherwise crash consumers that do
|
|
# usage.get(...) on the non-streaming paths.
|
|
"usage": _metadata_usage or {},
|
|
"timings": _metadata_timings,
|
|
"finish_reason": _metadata_finish_reason,
|
|
}
|
|
|
|
except httpx.ConnectError:
|
|
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,
|
|
rag_scope: Optional[dict] = None,
|
|
seed: Optional[int] = None,
|
|
disable_parallel_tool_use: bool = False,
|
|
) -> Generator[dict, None, None]:
|
|
"""
|
|
Agentic loop: let the model call tools, execute them, and continue.
|
|
|
|
Yields dicts:
|
|
{"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 build_rag_autoinject, execute_tool
|
|
|
|
if not self.is_loaded:
|
|
raise RuntimeError("llama-server is not loaded")
|
|
|
|
conversation = list(messages)
|
|
|
|
# Forced first-pass RAG so a doc question doesn't lose to web_search. Emits
|
|
# the same tool card + citations a real call would.
|
|
_auto = build_rag_autoinject(conversation, rag_scope)
|
|
if _auto:
|
|
for _ev in _auto["events"]:
|
|
yield _ev
|
|
conversation.extend(_auto["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,
|
|
force: bool = False,
|
|
) -> str:
|
|
if not (auto_heal_tool_calls or force):
|
|
return text
|
|
return strip_tool_call_markup(text, final = final)
|
|
|
|
def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str:
|
|
if not (auto_heal_tool_calls or force):
|
|
return text
|
|
for pat in _TOOL_ALL_PATS:
|
|
text = pat.sub("", text)
|
|
return text
|
|
|
|
tool_controller = ToolLoopController(
|
|
tools = tools,
|
|
auto_heal_tool_calls = auto_heal_tool_calls,
|
|
)
|
|
|
|
def _tool_succeeded(tool_name: str) -> bool:
|
|
key_prefix = f"{tool_name}:"
|
|
return any(
|
|
record.executed and not record.is_error and record.key.startswith(key_prefix)
|
|
for record in tool_controller.history
|
|
)
|
|
|
|
_MAX_BUFFER_CHARS = 32
|
|
_append_budget_exhausted_nudge = True
|
|
# RAG: cap knowledge-base searches per assistant turn. The controller is
|
|
# tool-agnostic, so this gate stays in the loop.
|
|
_kb_search_count = 0
|
|
|
|
# ── Re-prompt on plan-without-action ─────────────────
|
|
# When the model describes what it intends to do (forward-looking
|
|
# language) without calling a tool, re-prompt once. Only triggers on
|
|
# responses signaling intent/planning -- a direct answer like "4" or
|
|
# "Hello!" won't match. Pattern compiled at module level
|
|
# (_INTENT_SIGNAL).
|
|
_reprompt_count = 0
|
|
_forced_tool_call_pending = False
|
|
|
|
# Reserve extra iterations for re-prompts so they don't consume the
|
|
# caller's tool-call budget; only when tool iterations are 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
|
|
|
|
active_tools = tool_controller.active_tools()
|
|
if not active_tools:
|
|
_append_budget_exhausted_nudge = False
|
|
break
|
|
_tool_xml_signals = TOOL_XML_SIGNALS if active_tools else ()
|
|
|
|
# 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": active_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
|
|
if seed is not None:
|
|
payload["seed"] = seed
|
|
|
|
try:
|
|
_auth_headers = (
|
|
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
|
)
|
|
|
|
# ── Speculative buffer state machine ──────────────────
|
|
# BUFFERING: accumulate content, check for tool signals
|
|
# STREAMING: no tool detected, yield tokens to caller
|
|
# DRAINING: tool signal found, silently consume 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 yielded text (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
|
|
_iter_finish_reason = None
|
|
_stream_done = False
|
|
_last_emitted = ""
|
|
provisional_render_html_tool_call_ids = set()
|
|
_suppress_visible_output = _forced_tool_call_pending
|
|
|
|
stream_timeout = httpx.Timeout(
|
|
connect = 10,
|
|
read = 0.5,
|
|
write = 10,
|
|
pool = 10,
|
|
)
|
|
with httpx.Client(
|
|
timeout = stream_timeout,
|
|
limits = httpx.Limits(max_keepalive_connections = 0),
|
|
) as client:
|
|
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>"
|
|
if not _suppress_visible_output:
|
|
yield {
|
|
"type": "content",
|
|
"text": _strip_tool_markup(
|
|
cumulative_display,
|
|
final = True,
|
|
),
|
|
}
|
|
else:
|
|
cumulative_display = reasoning_accum
|
|
if not _suppress_visible_output:
|
|
yield {
|
|
"type": "content",
|
|
"text": cumulative_display,
|
|
}
|
|
_stream_done = True
|
|
break # exit inner while
|
|
if not line.startswith("data: "):
|
|
continue
|
|
|
|
try:
|
|
chunk_data = json.loads(line[6:])
|
|
_ct = chunk_data.get("timings")
|
|
if _ct:
|
|
_iter_timings = _ct
|
|
_cu = chunk_data.get("usage")
|
|
if _cu:
|
|
_iter_usage = _cu
|
|
|
|
choices = chunk_data.get("choices", [])
|
|
if not choices:
|
|
continue
|
|
|
|
delta = choices[0].get("delta", {})
|
|
_fr = choices[0].get("finish_reason")
|
|
if _fr:
|
|
_iter_finish_reason = _fr
|
|
|
|
# ── Structured tool_calls ──
|
|
tc_deltas = delta.get("tool_calls")
|
|
if tc_deltas:
|
|
# llama-server can emit visible assistant
|
|
# preface content before native structured
|
|
# tool_calls. Preserve content_accum as
|
|
# the assistant pre-tool text and still
|
|
# drain/execute the structured call.
|
|
has_structured_tc = True
|
|
detect_state = _S_DRAINING
|
|
for tc_d in tc_deltas:
|
|
idx = tc_d.get("index", 0)
|
|
if idx not in tool_calls_acc:
|
|
tool_calls_acc[idx] = {
|
|
"id": tc_d.get("id", f"call_{idx}"),
|
|
"type": "function",
|
|
"function": {
|
|
"name": "",
|
|
"arguments": "",
|
|
},
|
|
}
|
|
elif tc_d.get("id"):
|
|
# Update ID if a real one
|
|
# arrives on a later delta.
|
|
tool_calls_acc[idx]["id"] = tc_d["id"]
|
|
func = tc_d.get("function", {})
|
|
if func.get("name"):
|
|
tool_calls_acc[idx]["function"]["name"] += func[
|
|
"name"
|
|
]
|
|
if func.get("arguments"):
|
|
tool_calls_acc[idx]["function"]["arguments"] += (
|
|
func["arguments"]
|
|
)
|
|
current_name = tool_calls_acc[idx]["function"].get(
|
|
"name", ""
|
|
)
|
|
fallback_id = f"call_{idx}"
|
|
current_id = tool_calls_acc[idx].get("id", fallback_id)
|
|
already_started = (
|
|
current_id in provisional_render_html_tool_call_ids
|
|
)
|
|
has_real_id = current_id != fallback_id
|
|
if (
|
|
current_name == "render_html"
|
|
and not _tool_succeeded("render_html")
|
|
and any(
|
|
(
|
|
(tool.get("function") or {}).get("name")
|
|
== "render_html"
|
|
)
|
|
for tool in active_tools
|
|
)
|
|
and not already_started
|
|
and not provisional_render_html_tool_call_ids
|
|
and has_real_id
|
|
):
|
|
provisional_render_html_tool_call_ids.add(
|
|
current_id
|
|
)
|
|
yield {
|
|
"type": "tool_start",
|
|
"tool_name": "render_html",
|
|
"tool_call_id": current_id,
|
|
"arguments": {},
|
|
"provenance": tool_event_provenance(
|
|
provisional = True,
|
|
),
|
|
}
|
|
continue
|
|
|
|
# ── Reasoning tokens ──
|
|
# Yield only in STREAMING. In BUFFERING and
|
|
# DRAINING, accumulate silently so we don't
|
|
# corrupt the consumer's prev_text tracker
|
|
# (routes/inference.py never resets it
|
|
# between tool iterations).
|
|
reasoning = delta.get("reasoning_content", "")
|
|
if reasoning:
|
|
reasoning_accum += reasoning
|
|
if detect_state == _S_STREAMING:
|
|
if not in_thinking:
|
|
cumulative_display += "<think>"
|
|
in_thinking = True
|
|
cumulative_display += reasoning
|
|
if not _suppress_visible_output:
|
|
yield {
|
|
"type": "content",
|
|
"text": cumulative_display,
|
|
}
|
|
|
|
# ── Content tokens ──
|
|
token = delta.get("content", "")
|
|
if token:
|
|
has_content_tokens = True
|
|
content_accum += token
|
|
|
|
if detect_state == _S_DRAINING:
|
|
pass # accumulate silently
|
|
|
|
elif detect_state == _S_STREAMING:
|
|
if in_thinking:
|
|
cumulative_display += "</think>"
|
|
in_thinking = False
|
|
cumulative_display += token
|
|
cleaned = _strip_tool_markup_streaming(
|
|
cumulative_display
|
|
)
|
|
if len(cleaned) > len(_last_emitted):
|
|
_last_emitted = cleaned
|
|
if not _suppress_visible_output:
|
|
yield {
|
|
"type": "content",
|
|
"text": cleaned,
|
|
}
|
|
|
|
elif detect_state == _S_BUFFERING:
|
|
content_buffer += token
|
|
stripped_buf = content_buffer.lstrip()
|
|
if not stripped_buf:
|
|
continue
|
|
|
|
# Check tool signal prefixes.
|
|
is_prefix = False
|
|
is_match = False
|
|
for sig in _tool_xml_signals:
|
|
if stripped_buf.startswith(sig):
|
|
is_match = True
|
|
break
|
|
if sig.startswith(stripped_buf):
|
|
is_prefix = True
|
|
break
|
|
|
|
if is_match:
|
|
# Tool signal -- flush any visible
|
|
# prefix before DRAINING so the
|
|
# route sends it before tool_start.
|
|
if reasoning_accum:
|
|
cumulative_display += "<think>"
|
|
cumulative_display += reasoning_accum
|
|
cumulative_display += "</think>"
|
|
cumulative_display += content_buffer
|
|
cleaned = _strip_tool_markup_streaming(
|
|
cumulative_display,
|
|
force = True,
|
|
)
|
|
if len(cleaned) > len(_last_emitted):
|
|
_last_emitted = cleaned
|
|
if not _suppress_visible_output:
|
|
yield {
|
|
"type": "content",
|
|
"text": cleaned,
|
|
}
|
|
detect_state = _S_DRAINING
|
|
elif (
|
|
is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS
|
|
):
|
|
pass # keep buffering
|
|
else:
|
|
# Not a tool -- flush buffer
|
|
detect_state = _S_STREAMING
|
|
# Flush reasoning accumulated
|
|
# during BUFFERING.
|
|
if reasoning_accum:
|
|
cumulative_display += "<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
|
|
if not _suppress_visible_output:
|
|
yield {
|
|
"type": "content",
|
|
"text": cleaned,
|
|
}
|
|
|
|
except json.JSONDecodeError:
|
|
logger.debug(f"Skipping malformed SSE line: {line[:100]}")
|
|
if _stream_done:
|
|
break # exit outer for
|
|
|
|
# ── Resolve BUFFERING at stream end ──
|
|
if detect_state == _S_BUFFERING:
|
|
stripped_buf = content_buffer.lstrip()
|
|
if stripped_buf 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 reasoning first.
|
|
if reasoning_accum:
|
|
cumulative_display += "<think>"
|
|
cumulative_display += reasoning_accum
|
|
cumulative_display += "</think>"
|
|
cumulative_display += content_buffer
|
|
if not _suppress_visible_output:
|
|
yield {
|
|
"type": "content",
|
|
"text": _strip_tool_markup(
|
|
cumulative_display,
|
|
final = True,
|
|
),
|
|
}
|
|
elif reasoning_accum and not has_content_tokens:
|
|
# Reasoning-only response: show reasoning as plain
|
|
# text, matching the final streaming pass for
|
|
# models that put everything in reasoning.
|
|
cumulative_display = reasoning_accum
|
|
if not _suppress_visible_output:
|
|
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 emitted
|
|
# before the tool XML.
|
|
_safety_tc = None
|
|
if any(s in content_accum for s in _tool_xml_signals):
|
|
_safety_tc = self._parse_tool_calls_from_text(
|
|
content_accum,
|
|
allow_incomplete = auto_heal_tool_calls,
|
|
)
|
|
if not _safety_tc:
|
|
# ── Re-prompt on plan-without-action ──
|
|
# If the model described its intent (forward-looking
|
|
# language) without calling a tool, nudge it to act.
|
|
# Fires at most once per request, only on short
|
|
# responses with intent signals -- "4" or "Hello!"
|
|
# won't trigger it. Use content if available, else
|
|
# fall back to reasoning text (reasoning-only stalls).
|
|
_stripped = content_accum.strip()
|
|
if not _stripped:
|
|
_stripped = reasoning_accum.strip()
|
|
_render_html_already_done_intent = _tool_succeeded(
|
|
"render_html"
|
|
) and re.search(
|
|
r"(?i)\brender[_\s-]?html\b",
|
|
_stripped,
|
|
)
|
|
if (
|
|
auto_heal_tool_calls
|
|
and active_tools
|
|
and not _render_html_already_done_intent
|
|
and _reprompt_count < _MAX_REPROMPTS
|
|
and _is_short_intent_without_action(_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,
|
|
}
|
|
)
|
|
available_tool_names = [
|
|
(tool.get("function") or {}).get("name")
|
|
for tool in active_tools
|
|
if isinstance(tool, dict) and isinstance(tool.get("function"), dict)
|
|
]
|
|
available_tool_names = [name for name in available_tool_names if name]
|
|
tool_hint = " or ".join(available_tool_names) or "an available tool"
|
|
_forced_tool_call_pending = True
|
|
conversation.append(
|
|
{
|
|
"role": "user",
|
|
"content": (
|
|
"You have access to enabled tools. If a tool is needed to satisfy "
|
|
"the user's request or complete the action you described, call "
|
|
f"{tool_hint} now. If no tool is needed, provide the final answer "
|
|
"and follow the user's requested format."
|
|
),
|
|
}
|
|
)
|
|
# 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
|
|
|
|
if _forced_tool_call_pending:
|
|
_forced_tool_call_pending = False
|
|
if not _should_suppress_forced_no_tool_output(_stripped):
|
|
if cumulative_display:
|
|
forced_visible_text = _strip_tool_markup(
|
|
cumulative_display,
|
|
final = True,
|
|
)
|
|
elif content_accum:
|
|
forced_visible_text = _strip_tool_markup(
|
|
content_accum,
|
|
final = True,
|
|
)
|
|
else:
|
|
forced_visible_text = reasoning_accum
|
|
if forced_visible_text:
|
|
yield {
|
|
"type": "content",
|
|
"text": forced_visible_text,
|
|
}
|
|
|
|
# Content was already streamed. Yield metadata.
|
|
yield {"type": "status", "text": ""}
|
|
_fu = _backfill_usage_from_timings(_iter_usage, _iter_timings) or {}
|
|
_fc = _fu.get("completion_tokens", 0)
|
|
_fp = _fu.get("prompt_tokens", 0)
|
|
_tc = _fc + _accumulated_completion_tokens
|
|
if _iter_usage or _iter_timings or _accumulated_completion_tokens:
|
|
_mt = dict(_iter_timings) if _iter_timings else {}
|
|
if _accumulated_predicted_ms or _accumulated_predicted_n:
|
|
_mt["predicted_ms"] = (
|
|
_mt.get("predicted_ms", 0) + _accumulated_predicted_ms
|
|
)
|
|
_tn = _mt.get("predicted_n", 0) + _accumulated_predicted_n
|
|
_mt["predicted_n"] = _tn
|
|
_tms = _mt["predicted_ms"]
|
|
if _tms > 0:
|
|
_mt["predicted_per_second"] = _tn / (_tms / 1000.0)
|
|
yield {
|
|
"type": "metadata",
|
|
"usage": {
|
|
"prompt_tokens": _fp,
|
|
"completion_tokens": _tc,
|
|
"total_tokens": _fp + _tc,
|
|
},
|
|
"timings": _mt,
|
|
"finish_reason": _iter_finish_reason,
|
|
}
|
|
return
|
|
|
|
# Safety net caught tool XML -- treat as tool call.
|
|
tool_calls = _safety_tc
|
|
content_text = _strip_tool_markup(
|
|
content_accum,
|
|
final = True,
|
|
force = 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:
|
|
# Drop incomplete fragments (e.g. from max_tokens
|
|
# truncation 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 any(s in content_accum for s in _tool_xml_signals):
|
|
tool_calls = self._parse_tool_calls_from_text(
|
|
content_accum,
|
|
allow_incomplete = auto_heal_tool_calls,
|
|
)
|
|
if tool_calls and not has_structured_tc:
|
|
content_text = _strip_tool_markup(
|
|
content_text,
|
|
final = True,
|
|
force = 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 aren't 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,
|
|
"finish_reason": _iter_finish_reason,
|
|
}
|
|
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)
|
|
|
|
# disable_parallel_tool_use: execute only the first tool call
|
|
# this turn. Truncate before building assistant_msg so the
|
|
# conversation stays consistent and extra calls are never executed.
|
|
if disable_parallel_tool_use and tool_calls and len(tool_calls) > 1:
|
|
tool_calls = tool_calls[:1]
|
|
|
|
assistant_msg: dict = {"role": "assistant", "content": content_text}
|
|
assistant_appended = False
|
|
|
|
for tc in tool_calls or []:
|
|
func = tc.get("function", {})
|
|
tool_name = func.get("name", "")
|
|
provisional_render_html_match = (
|
|
tool_name == "render_html"
|
|
and tc.get("id") in provisional_render_html_tool_call_ids
|
|
)
|
|
decision = tool_controller.prepare_call(
|
|
tc,
|
|
forced = _forced_tool_call_pending,
|
|
provisional = provisional_render_html_match,
|
|
)
|
|
|
|
if not decision.should_execute:
|
|
if content_text and not assistant_appended:
|
|
conversation.append(assistant_msg)
|
|
assistant_appended = True
|
|
completion = tool_controller.record_noop(decision)
|
|
conversation.append(completion.model_message())
|
|
if _forced_tool_call_pending:
|
|
_forced_tool_call_pending = False
|
|
logger.info(
|
|
"Suppressed local GGUF tool call as internal no-op: "
|
|
f"action={decision.action} tool={decision.tool_name}"
|
|
)
|
|
break
|
|
|
|
if not assistant_appended:
|
|
assistant_msg["tool_calls"] = [decision.as_assistant_tool_call()]
|
|
conversation.append(assistant_msg)
|
|
assistant_appended = True
|
|
else:
|
|
assistant_msg.setdefault("tool_calls", []).append(
|
|
decision.as_assistant_tool_call()
|
|
)
|
|
|
|
yield {"type": "status", "text": decision.status_text}
|
|
yield decision.tool_start_event()
|
|
|
|
_effective_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout
|
|
# RAG: cap paraphrased KB re-searches that slip past the dup guard.
|
|
if (
|
|
decision.tool_name == "search_knowledge_base"
|
|
and _kb_search_count >= RAG_MAX_SEARCHES_PER_TURN
|
|
):
|
|
result = RAG_SEARCH_CAP_NUDGE
|
|
else:
|
|
result = execute_tool(
|
|
decision.tool_name,
|
|
decision.arguments,
|
|
cancel_event = cancel_event,
|
|
timeout = _effective_timeout,
|
|
session_id = session_id,
|
|
rag_scope = rag_scope,
|
|
)
|
|
if decision.tool_name == "search_knowledge_base":
|
|
_kb_search_count += 1
|
|
completion = tool_controller.record_result(decision, result)
|
|
yield completion.tool_end_event()
|
|
conversation.append(completion.tool_message())
|
|
|
|
if _forced_tool_call_pending:
|
|
_forced_tool_call_pending = False
|
|
|
|
# Clear tool status badge before next generation/final pass.
|
|
yield {"type": "status", "text": ""}
|
|
if tool_controller.force_final_answer or not tool_controller.active_tools():
|
|
_append_budget_exhausted_nudge = False
|
|
break
|
|
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 a final text response. Nudge
|
|
# the final streaming pass to produce a useful answer instead of
|
|
# continuing to request tools.
|
|
if max_tool_iterations > 0 and _append_budget_exhausted_nudge:
|
|
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
|
|
if seed is not None:
|
|
stream_payload["seed"] = seed
|
|
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
|
|
_metadata_finish_reason = 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", {})
|
|
_fr = choices[0].get("finish_reason")
|
|
if _fr:
|
|
_metadata_finish_reason = _fr
|
|
|
|
reasoning = delta.get("reasoning_content", "")
|
|
if reasoning:
|
|
reasoning_text += reasoning
|
|
if not in_thinking:
|
|
cumulative += "<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)
|
|
# Emit only when cleaned text grows (monotonic).
|
|
if len(cleaned) > len(_last_emitted):
|
|
_last_emitted = cleaned
|
|
yield {"type": "content", "text": cleaned}
|
|
except json.JSONDecodeError:
|
|
logger.debug(f"Skipping malformed SSE line: {line[:100]}")
|
|
if _stream_done:
|
|
break # exit outer for
|
|
_final_usage = _metadata_usage or {}
|
|
_final_completion = _final_usage.get("completion_tokens", 0)
|
|
_final_prompt = _final_usage.get("prompt_tokens", 0)
|
|
_total_completion = _final_completion + _accumulated_completion_tokens
|
|
if _metadata_usage or _metadata_timings or _metadata_finish_reason:
|
|
_merged_timings = dict(_metadata_timings) if _metadata_timings else {}
|
|
if _accumulated_predicted_ms or _accumulated_predicted_n:
|
|
_merged_timings["predicted_ms"] = (
|
|
_merged_timings.get("predicted_ms", 0) + _accumulated_predicted_ms
|
|
)
|
|
_total_predicted_n = (
|
|
_merged_timings.get("predicted_n", 0) + _accumulated_predicted_n
|
|
)
|
|
_merged_timings["predicted_n"] = _total_predicted_n
|
|
_total_predicted_ms = _merged_timings["predicted_ms"]
|
|
if _total_predicted_ms > 0:
|
|
_merged_timings["predicted_per_second"] = _total_predicted_n / (
|
|
_total_predicted_ms / 1000.0
|
|
)
|
|
yield {
|
|
"type": "metadata",
|
|
"usage": {
|
|
"prompt_tokens": _final_prompt,
|
|
"completion_tokens": _total_completion,
|
|
"total_tokens": _final_prompt + _total_completion,
|
|
},
|
|
"timings": _merged_timings,
|
|
"finish_reason": _metadata_finish_reason,
|
|
}
|
|
|
|
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
|
|
|
|
# ── Prompt token counting ──────────────────────────────────
|
|
|
|
def count_chat_tokens(
|
|
self,
|
|
messages,
|
|
system = None,
|
|
tools = None,
|
|
strict: bool = False,
|
|
) -> int:
|
|
"""Count prompt tokens for a chat request via llama-server.
|
|
|
|
Non-strict callers keep the historical best-effort behavior and receive
|
|
0 when a count cannot be determined. Strict callers (public count_tokens
|
|
endpoints) get an exception instead of a successful-looking zero when
|
|
tokenizer/template calls fail or a multimodal prompt would fall back to a
|
|
text-only approximation.
|
|
"""
|
|
if not self.is_loaded:
|
|
if strict:
|
|
raise RuntimeError("llama-server is not loaded")
|
|
return 0
|
|
|
|
def _has_non_text_content(content) -> bool:
|
|
if isinstance(content, list):
|
|
for block in content:
|
|
if isinstance(block, str):
|
|
continue
|
|
if not isinstance(block, dict):
|
|
return True
|
|
if block.get("type") == "text" and isinstance(block.get("text"), str):
|
|
continue
|
|
if isinstance(block.get("text"), str):
|
|
continue
|
|
return True
|
|
return False
|
|
|
|
def _has_non_text_prompt_parts() -> bool:
|
|
if _has_non_text_content(system):
|
|
return True
|
|
for msg in messages or []:
|
|
if isinstance(msg, dict) and _has_non_text_content(msg.get("content", "")):
|
|
return True
|
|
return False
|
|
|
|
def _block_text(content) -> str:
|
|
if isinstance(content, str):
|
|
return content
|
|
if isinstance(content, list):
|
|
parts = []
|
|
for block in content:
|
|
if isinstance(block, dict):
|
|
if block.get("type") == "text" and isinstance(block.get("text"), str):
|
|
parts.append(block["text"])
|
|
elif isinstance(block.get("text"), str):
|
|
parts.append(block["text"])
|
|
elif isinstance(block, str):
|
|
parts.append(block)
|
|
return "".join(parts)
|
|
return ""
|
|
|
|
# Normalize system into a leading message / plain text.
|
|
system_text = ""
|
|
if isinstance(system, str):
|
|
system_text = system
|
|
elif isinstance(system, list):
|
|
system_text = _block_text(system)
|
|
|
|
try:
|
|
_auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
|
with httpx.Client(timeout = 10, headers = _auth_headers) as client:
|
|
|
|
def _tokenize(text: str) -> int:
|
|
r = client.post(
|
|
f"{self.base_url}/tokenize",
|
|
json = {"content": text, "add_special": True},
|
|
)
|
|
if r.status_code != 200:
|
|
if strict:
|
|
raise RuntimeError("llama-server tokenizer failed")
|
|
return 0
|
|
tokens = r.json().get("tokens", [])
|
|
if not isinstance(tokens, list):
|
|
if strict:
|
|
raise RuntimeError("llama-server tokenizer returned invalid tokens")
|
|
return 0
|
|
return len(tokens)
|
|
|
|
# 1. Try /apply-template to render the real chat prompt.
|
|
template_messages = list(messages) if messages else []
|
|
if system_text:
|
|
template_messages = [
|
|
{"role": "system", "content": system_text}
|
|
] + template_messages
|
|
apply_template_failed = False
|
|
try:
|
|
# llama-server's /apply-template renders tool declarations
|
|
# into the prompt when ``tools`` is supplied, so pass them
|
|
# through — otherwise tool-schema tokens go uncounted.
|
|
template_body = {"messages": template_messages}
|
|
if tools:
|
|
template_body["tools"] = tools
|
|
resp = client.post(
|
|
f"{self.base_url}/apply-template",
|
|
json = template_body,
|
|
)
|
|
if resp.status_code == 200:
|
|
prompt = resp.json().get("prompt", "")
|
|
if isinstance(prompt, str):
|
|
return _tokenize(prompt)
|
|
apply_template_failed = True
|
|
except Exception:
|
|
apply_template_failed = True
|
|
|
|
if strict and apply_template_failed and _has_non_text_prompt_parts():
|
|
raise RuntimeError(
|
|
"cannot fall back to text-only token counting for multimodal messages"
|
|
)
|
|
|
|
# 2. Fallback: concatenate plain text and tokenize. Append a
|
|
# serialized form of the tools so they still contribute to the
|
|
# count when /apply-template is unavailable.
|
|
parts = []
|
|
if system_text:
|
|
parts.append(system_text)
|
|
for msg in messages or []:
|
|
if isinstance(msg, dict):
|
|
parts.append(_block_text(msg.get("content", "")))
|
|
if tools:
|
|
try:
|
|
parts.append(json.dumps(tools, ensure_ascii = False))
|
|
except Exception:
|
|
pass
|
|
return _tokenize("\n".join(p for p in parts if p))
|
|
except Exception:
|
|
if strict:
|
|
raise
|
|
return 0
|
|
|
|
# ── TTS support ────────────────────────────────────────────
|
|
|
|
def detect_audio_type(self) -> Optional[str]:
|
|
"""Detect audio/TTS codec; swallows errors (use _strict 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 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", [])
|
|
|
|
# Codec-specific tokens (not generic ones that non-audio models may have)
|
|
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"
|
|
# Gemma 3n: <audio_soft_token>; Gemma 4: <|audio|> (not csm's <|AUDIO|>).
|
|
if len(_tok("<audio_soft_token>")) == 1 or len(_tok("<|audio|>")) == 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 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 the 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 decode.
|
|
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", "")
|
|
)
|