From fbcd3fa511e86798ed988514d3f193bd26ff122b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 10 Jul 2026 03:07:39 -0700 Subject: [PATCH 1/9] CI: retry transient HTTP timeouts in Studio smoke probes (#7052) * CI: retry transient HTTP timeouts in Studio smoke probes The post() helper in the Studio inference smoke workflows does a single urlopen with a 240s timeout against the local Studio server. On shared runners this sporadically hits TimeoutError while the server is stalled, failing the whole job for a transport hiccup; the same flake has recurred across unrelated PRs on Linux and Windows (JSON/images and tool-calling jobs) and passes on rerun. Retry the probe up to 3 times on transport-level failures only (TimeoutError, ConnectionError, non-HTTP URLError), 15s apart. HTTP status errors still surface immediately, so genuine server failures are unaffected. post_sse() is left unchanged: it has a 600s budget and has not flaked. * CI: retry only short probes so worst case fits the job budget Some json-images calls pass timeout=600; three attempts there could spend 30 minutes in one step and hit the job's timeout-minutes instead of failing with the Python error. Retry (3 attempts) only when timeout <= 300s, which covers the observed flaky 180-240s probes; longer probes keep the pre-PR single attempt. * CI: give long smoke probes one capped retry Round two of bounding the retries: timeout>300s probes previously got a single attempt, so a transient stall in the 600s JSON-mode probes still failed on first occurrence. Give them one retry with the attempt timeout capped at 300s. Worst cases stay inside timeout-minutes: 240s probes 12.5 min, one 600s probe 15.25 min, the Windows JSON job's two long probes 30.5 min against its 35 minute budget. --- .github/workflows/studio-inference-smoke.yml | 44 +++++++++++++++++-- .../workflows/studio-mac-inference-smoke.yml | 44 +++++++++++++++++-- .../studio-windows-inference-smoke.yml | 44 +++++++++++++++++-- 3 files changed, 120 insertions(+), 12 deletions(-) diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index aebf90380a..f540c11da4 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -444,6 +444,8 @@ jobs: python - <<'PY' import json import os + import time + import urllib.error import urllib.request BASE = os.environ["BASE_URL"] @@ -464,8 +466,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) def post_sse(path, body, *, timeout = 600): """POST a streaming request and accumulate the assistant @@ -938,6 +956,8 @@ jobs: import base64 import json import os + import time + import urllib.error import urllib.request from openai import OpenAI from anthropic import Anthropic @@ -956,8 +976,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) # ── 1. response_format = json_object (JSON mode) ───────────── # llama.cpp's HTTP server supports OpenAI-compatible JSON diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index d562294d42..03c0a8580d 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -430,6 +430,8 @@ jobs: python - <<'PY' import json import os + import time + import urllib.error import urllib.request BASE = os.environ["BASE_URL"] @@ -450,8 +452,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) def post_sse(path, body, *, timeout = 600): """POST a streaming request and accumulate the assistant @@ -825,6 +843,8 @@ jobs: import base64 import json import os + import time + import urllib.error import urllib.request from openai import OpenAI from anthropic import Anthropic @@ -848,8 +868,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) # ── 1. response_format = json_object (JSON mode) ───────────── # llama.cpp's HTTP server supports OpenAI-compatible JSON diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index dbb0f9ea6f..0453c9212a 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -634,6 +634,8 @@ jobs: python - <<'PY' import json import os + import time + import urllib.error import urllib.request BASE = os.environ["BASE_URL"] @@ -656,8 +658,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) def post_sse(path, body, *, timeout = 600): body = {**body, "stream": True} @@ -1063,6 +1081,8 @@ jobs: import base64 import json import os + import time + import urllib.error import urllib.request from openai import OpenAI from anthropic import Anthropic @@ -1082,8 +1102,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) # ── 1. response_format = json_object (JSON mode) ───────────── status, data = post("/v1/chat/completions", { From 33119c9bf73a6655f800168160ae6a9effa32472 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Fri, 10 Jul 2026 10:55:11 -0700 Subject: [PATCH 2/9] fix: guard remove_special_tokens against tokenizers without a BOS token (#7048) --- .../test_remove_special_tokens_no_bos.py | 46 +++++++++++++++++++ unsloth/chat_templates.py | 5 +- 2 files changed, 49 insertions(+), 2 deletions(-) create mode 100644 tests/python/test_remove_special_tokens_no_bos.py diff --git a/tests/python/test_remove_special_tokens_no_bos.py b/tests/python/test_remove_special_tokens_no_bos.py new file mode 100644 index 0000000000..94c5ea3027 --- /dev/null +++ b/tests/python/test_remove_special_tokens_no_bos.py @@ -0,0 +1,46 @@ +import ast +from pathlib import Path + + +def _load_remove_special_tokens(): + # Extract remove_special_tokens without importing unsloth (importing unsloth + # needs unsloth_zoo / a GPU). The function is pure Python and uses no imports, + # so it execs cleanly in an empty namespace. + source = Path(__file__).parents[2] / "unsloth" / "chat_templates.py" + tree = ast.parse(source.read_text(encoding = "utf-8")) + funcs = [ + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "remove_special_tokens" + ] + namespace = {} + module = ast.Module(body = funcs, type_ignores = []) + ast.fix_missing_locations(module) + exec(compile(module, str(source), "exec"), namespace) + return namespace["remove_special_tokens"] + + +class _StubTokenizer: + def __init__(self, bos_token): + self.bos_token = bos_token + + +def test_no_bos_tokenizer_does_not_crash(): + # Tokenizers such as Qwen2 / Qwen2.5, GPT-2, Falcon and GPT-NeoX have no BOS + # token, so tokenizer.bos_token is None. remove_special_tokens must leave the + # prompt untouched instead of raising + # "TypeError: startswith first arg must be str or a tuple of str, not NoneType". + remove_special_tokens = _load_remove_special_tokens() + assert remove_special_tokens(_StubTokenizer(None), "Hello world") == "Hello world" + + +def test_double_bos_is_stripped(): + # A tokenizer with a BOS token still has a single leading BOS removed. + remove_special_tokens = _load_remove_special_tokens() + assert remove_special_tokens(_StubTokenizer(""), "Hello world") == "Hello world" + + +def test_prompt_without_leading_bos_unchanged(): + # A BOS-bearing tokenizer leaves a prompt that does not start with BOS alone. + remove_special_tokens = _load_remove_special_tokens() + assert remove_special_tokens(_StubTokenizer(""), "Hello world") == "Hello world" diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index dd1e433471..2d3674fb04 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -2103,8 +2103,9 @@ def get_chat_template( def remove_special_tokens(tokenizer, prompt): # Removes double BOS token - if prompt.startswith(tokenizer.bos_token): - prompt = prompt[len(tokenizer.bos_token):] + bos_token = getattr(tokenizer, "bos_token", None) + if bos_token is not None and prompt.startswith(bos_token): + prompt = prompt[len(bos_token):] return prompt From fef37cb25b35e36b12c1d9c7d020aa0fb402ce89 Mon Sep 17 00:00:00 2001 From: Apoze <158856608+Apoze@users.noreply.github.com> Date: Fri, 10 Jul 2026 22:05:48 +0200 Subject: [PATCH 3/9] Studio: queue local GGUF OpenAI-compatible requests before llama-server (#7047) --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- .../backend/core/inference/llama_admission.py | 368 +++++ studio/backend/core/inference/llama_cpp.py | 23 + studio/backend/routes/inference.py | 1257 ++++++++++++++- studio/backend/tests/test_llama_admission.py | 320 ++++ ...test_llama_cpp_effective_parallel_slots.py | 53 + .../tests/test_openai_tool_passthrough.py | 1396 +++++++++++++++++ .../test_stream_cancel_registration_timing.py | 24 +- 7 files changed, 3408 insertions(+), 33 deletions(-) create mode 100644 studio/backend/core/inference/llama_admission.py create mode 100644 studio/backend/tests/test_llama_admission.py create mode 100644 studio/backend/tests/test_llama_cpp_effective_parallel_slots.py diff --git a/studio/backend/core/inference/llama_admission.py b/studio/backend/core/inference/llama_admission.py new file mode 100644 index 0000000000..b6a939c87b --- /dev/null +++ b/studio/backend/core/inference/llama_admission.py @@ -0,0 +1,368 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Admission control for local llama-server generation requests. + +The helpers in this module deliberately know nothing about FastAPI, SSE, or the +OpenAI-compatible route shape. They only coordinate how many upstream generation +requests may be active for one llama-server backend and provide a cancellable +FIFO queue for excess requests. +""" + +from __future__ import annotations + +import asyncio +import os +import threading +from collections import deque +from dataclasses import dataclass +from typing import Deque, Optional + + +ADMISSION_CONTROL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL" +ADMISSION_QUEUE_TIMEOUT_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT" +ADMISSION_KEEPALIVE_INTERVAL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL" +ADMISSION_MAX_QUEUE_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE" + +DEFAULT_ADMISSION_ENABLED = True +DEFAULT_ADMISSION_QUEUE_TIMEOUT_S = None +DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S = 5.0 +DEFAULT_ADMISSION_MAX_QUEUE = 64 + + +@dataclass(frozen = True) +class LlamaAdmissionConfig: + enabled: bool = DEFAULT_ADMISSION_ENABLED + queue_timeout_s: Optional[float] = DEFAULT_ADMISSION_QUEUE_TIMEOUT_S + keepalive_interval_s: float = DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S + max_queue: Optional[int] = DEFAULT_ADMISSION_MAX_QUEUE + + +@dataclass(frozen = True) +class LlamaAdmissionSnapshot: + key: str + capacity: int + active: int + queued: int + + +class LlamaAdmissionError(Exception): + def __init__( + self, + message: str, + *, + snapshot: Optional[LlamaAdmissionSnapshot] = None, + ): + super().__init__(message) + self.snapshot = snapshot + + +class LlamaAdmissionQueueFull(LlamaAdmissionError): + pass + + +class LlamaAdmissionTimeout(LlamaAdmissionError): + pass + + +class LlamaAdmissionCancelled(LlamaAdmissionError): + pass + + +def _bool_env(name: str, default: bool) -> bool: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + value = value.strip().lower() + if value in {"1", "true", "yes", "on"}: + return True + if value in {"0", "false", "no", "off"}: + return False + return default + + +def _optional_positive_float_env(name: str, default: Optional[float]) -> Optional[float]: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + try: + parsed = float(value.strip()) + except ValueError: + return default + return parsed if parsed > 0 else None + + +def _positive_float_env(name: str, default: float) -> float: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + try: + parsed = float(value.strip()) + except ValueError: + return default + return parsed if parsed > 0 else default + + +def _optional_positive_int_env(name: str, default: Optional[int]) -> Optional[int]: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + try: + parsed = int(value.strip()) + except ValueError: + return default + return parsed if parsed > 0 else None + + +def llama_admission_config_from_env() -> LlamaAdmissionConfig: + return LlamaAdmissionConfig( + enabled = _bool_env(ADMISSION_CONTROL_ENV, DEFAULT_ADMISSION_ENABLED), + queue_timeout_s = _optional_positive_float_env( + ADMISSION_QUEUE_TIMEOUT_ENV, + DEFAULT_ADMISSION_QUEUE_TIMEOUT_S, + ), + keepalive_interval_s = _positive_float_env( + ADMISSION_KEEPALIVE_INTERVAL_ENV, + DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S, + ), + max_queue = _optional_positive_int_env( + ADMISSION_MAX_QUEUE_ENV, + DEFAULT_ADMISSION_MAX_QUEUE, + ), + ) + + +@dataclass +class _Waiter: + loop: asyncio.AbstractEventLoop + future: asyncio.Future + cancelled: bool = False + granted_lease: Optional["LlamaAdmissionLease"] = None + + +class LlamaAdmissionLease: + def __init__(self, queue: Optional["LlamaAdmissionQueue"]): + self._queue = queue + self._released = False + self._release_lock = threading.Lock() + + def release(self) -> None: + queue = None + with self._release_lock: + if self._released: + return + self._released = True + queue = self._queue + if queue is not None: + queue.release() + + async def __aenter__(self) -> "LlamaAdmissionLease": + return self + + async def __aexit__(self, *_args) -> None: + self.release() + + +class LlamaAdmissionReservation: + def __init__( + self, + *, + queue: Optional["LlamaAdmissionQueue"], + lease: Optional[LlamaAdmissionLease] = None, + waiter: Optional[_Waiter] = None, + snapshot: Optional[LlamaAdmissionSnapshot] = None, + ): + self._queue = queue + self._lease = lease + self._waiter = waiter + self.snapshot = snapshot + + @property + def is_cancelled(self) -> bool: + return self._lease is None and self._waiter is None + + def lease_nowait(self) -> Optional[LlamaAdmissionLease]: + if self._lease is not None: + return self._lease + if self._waiter is None or not self._waiter.future.done(): + return None + if self._waiter.future.cancelled(): + self._waiter.cancelled = True + self._waiter = None + return None + self._lease = self._waiter.future.result() + self._waiter = None + return self._lease + + async def wait(self, timeout_s: float) -> Optional[LlamaAdmissionLease]: + lease = self.lease_nowait() + if lease is not None: + return lease + if self._waiter is None: + return None + waiter = self._waiter + try: + await asyncio.wait_for(asyncio.shield(waiter.future), timeout = timeout_s) + except asyncio.CancelledError: + if waiter.future.cancelled(): + waiter.cancelled = True + if self._waiter is waiter: + self._waiter = None + return None + raise + return self.lease_nowait() + + def cancel(self) -> None: + lease = self.lease_nowait() + if lease is not None: + lease.release() + self._lease = None + return + if self._queue is not None and self._waiter is not None: + self._queue.cancel(self._waiter) + self._waiter = None + + def snapshot_now(self) -> Optional[LlamaAdmissionSnapshot]: + if self._queue is None: + return self.snapshot + return self._queue.snapshot() + + +class LlamaAdmissionQueue: + def __init__(self, key: str): + self.key = key + self._lock = threading.Lock() + self._active = 0 + self._capacity = 1 + self._waiters: Deque[_Waiter] = deque() + + def reserve(self, *, capacity: int, config: LlamaAdmissionConfig) -> LlamaAdmissionReservation: + capacity = max(1, int(capacity or 1)) + if not config.enabled: + return LlamaAdmissionReservation( + queue = None, + lease = LlamaAdmissionLease(None), + snapshot = LlamaAdmissionSnapshot(self.key, capacity, 0, 0), + ) + + loop = asyncio.get_running_loop() + with self._lock: + self._capacity = capacity + self._prune_waiters_locked() + self._grant_waiters_locked() + if self._active < self._capacity and not self._waiters: + self._active += 1 + return LlamaAdmissionReservation( + queue = self, + lease = LlamaAdmissionLease(self), + snapshot = self._snapshot_locked(), + ) + if config.max_queue is not None and len(self._waiters) >= config.max_queue: + raise LlamaAdmissionQueueFull( + "llama-server generation queue is full", + snapshot = self._snapshot_locked(), + ) + waiter = _Waiter( + loop = loop, + future = loop.create_future(), + ) + self._waiters.append(waiter) + return LlamaAdmissionReservation( + queue = self, + waiter = waiter, + snapshot = self._snapshot_locked(), + ) + + def release(self) -> None: + with self._lock: + if self._active > 0: + self._active -= 1 + self._grant_waiters_locked() + + def cancel(self, waiter: _Waiter) -> None: + lease_to_release = None + with self._lock: + waiter.cancelled = True + try: + self._waiters.remove(waiter) + except ValueError: + pass + if waiter.granted_lease is not None: + lease_to_release = waiter.granted_lease + waiter.granted_lease = None + if not waiter.future.done(): + waiter.loop.call_soon_threadsafe(waiter.future.cancel) + if lease_to_release is not None: + lease_to_release.release() + + def snapshot(self) -> LlamaAdmissionSnapshot: + with self._lock: + self._prune_waiters_locked() + return self._snapshot_locked() + + def is_idle(self) -> bool: + with self._lock: + self._prune_waiters_locked() + return self._active == 0 and not self._waiters + + def _grant_waiters_locked(self) -> None: + self._prune_waiters_locked() + while self._waiters and self._active < self._capacity: + waiter = self._waiters.popleft() + if waiter.cancelled or waiter.future.done(): + continue + self._active += 1 + lease = LlamaAdmissionLease(self) + waiter.granted_lease = lease + waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease) + + def _deliver_lease(self, waiter: _Waiter, lease: LlamaAdmissionLease) -> None: + if waiter.cancelled or waiter.future.done(): + waiter.granted_lease = None + if not waiter.future.done(): + waiter.future.cancel() + lease.release() + return + try: + waiter.future.set_result(lease) + waiter.granted_lease = None + except asyncio.InvalidStateError: + waiter.granted_lease = None + lease.release() + + def _prune_waiters_locked(self) -> None: + self._waiters = deque( + waiter for waiter in self._waiters if not waiter.cancelled and not waiter.future.done() + ) + + def _snapshot_locked(self) -> LlamaAdmissionSnapshot: + return LlamaAdmissionSnapshot( + key = self.key, + capacity = self._capacity, + active = self._active, + queued = len(self._waiters), + ) + + +_QUEUES_LOCK = threading.Lock() +_QUEUES: dict[str, LlamaAdmissionQueue] = {} + + +def get_llama_admission_queue(key: str) -> LlamaAdmissionQueue: + with _QUEUES_LOCK: + queue = _QUEUES.get(key) + if queue is None: + queue = LlamaAdmissionQueue(key) + _QUEUES[key] = queue + # base_url carries a fresh ephemeral port on every model load, so + # each load registers a new key. Drop the now-idle queues from prior + # loads so the registry can't grow without bound on a long-running + # server. Queues with in-flight requests are kept until they drain. + for stale_key in [k for k in _QUEUES if k != key and _QUEUES[k].is_idle()]: + del _QUEUES[stale_key] + return queue + + +def reset_llama_admission_queues() -> None: + with _QUEUES_LOCK: + _QUEUES.clear() diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index c2e3adf815..b06c6eb5cb 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1541,6 +1541,7 @@ class LlamaCppBackend: self._context_length: Optional[int] = None self._effective_context_length: Optional[int] = None self._max_context_length: Optional[int] = None + self._effective_parallel_slots: int = 1 self._chat_template: Optional[str] = None self._chat_template_override: Optional[str] = None self._supports_reasoning: bool = False @@ -1726,6 +1727,15 @@ class LlamaCppBackend: """Return the effective context length the server is running at.""" return self._effective_context_length or self._context_length + @property + def effective_parallel_slots(self) -> int: + """Return the serving-slot count the active llama-server actually uses.""" + try: + slots = int(getattr(self, "_effective_parallel_slots", 1)) + except (TypeError, ValueError): + slots = 1 + return max(1, slots) + @property def max_context_length(self) -> Optional[int]: """Return the largest context that fits on this hardware at load time. @@ -1742,6 +1752,16 @@ class LlamaCppBackend: """Return the model's native context length from GGUF metadata.""" return self._context_length + def _commit_effective_parallel_slots(self, n_parallel: int) -> None: + try: + slots = int(n_parallel) + except (TypeError, ValueError): + slots = 1 + self._effective_parallel_slots = max(1, slots) + + def _reset_effective_parallel_slots(self) -> None: + self._effective_parallel_slots = 1 + @staticmethod def _read_rss_bytes(pid: int) -> Optional[int]: """Resident set size of ``pid`` in bytes, from /proc//status (Linux). @@ -7187,6 +7207,7 @@ class LlamaCppBackend: ) self._healthy = True + self._commit_effective_parallel_slots(n_parallel) # Commit caller intent only after _healthy=True so a failed start # can't poison the next inheritance check. None keeps prior, [] @@ -7724,6 +7745,7 @@ class LlamaCppBackend: self._context_length = None self._effective_context_length = None self._max_context_length = None + self._reset_effective_parallel_slots() self._chat_template = None self._chat_template_override = None self._supports_reasoning = False @@ -7779,6 +7801,7 @@ class LlamaCppBackend: # Stop the watchdog before a deliberate kill so a planned reload/unload # isn't seen as a crash; a real crash never routes through here. self._stop_mtp_crash_watchdog() + self._reset_effective_parallel_slots() if self._process is None: return try: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 826669d922..db350240ac 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -28,6 +28,16 @@ import re as _re from utils.models import extract_model_size_b as _extract_model_size_b from utils.api_errors import openai_error_body, anthropic_error_body +from core.inference.llama_admission import ( + LlamaAdmissionCancelled, + LlamaAdmissionConfig, + LlamaAdmissionLease, + LlamaAdmissionQueueFull, + LlamaAdmissionReservation, + LlamaAdmissionTimeout, + get_llama_admission_queue, + llama_admission_config_from_env, +) def _positive_int_or_none(value: Any) -> Optional[int]: @@ -1072,6 +1082,270 @@ _STREAM_DISCONNECT_POLL_TIMEOUT_S = 0.25 _OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S = 0.1 _OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S = 5.0 _OPENAI_PASSTHROUGH_SSE_KEEPALIVE = ": keep-alive\n\n" +_OPENAI_LLAMA_ADMISSION_POLL_S = 0.25 + + +def _openai_llama_admission_capacity(request: Optional[Request], llama_backend = None) -> int: + """Serving slots available for one local llama-server backend. + + The loaded backend is the source of truth because it may have reduced + ``--parallel`` at load time to keep the model on GPU. The app state is a + launch-intent fallback for tests and for the short window before a backend + reports its committed runtime slots. + """ + slots = _positive_int_or_none(getattr(llama_backend, "effective_parallel_slots", None)) + if slots is not None: + return slots + try: + slots = getattr(request.app.state, "llama_parallel_slots", None) + except Exception: + slots = None + return _positive_int_or_none(slots) or 1 + + +def _openai_llama_admission_reserve( + *, request: Optional[Request], llama_backend +) -> tuple[LlamaAdmissionReservation, LlamaAdmissionConfig]: + config = llama_admission_config_from_env() + capacity = _openai_llama_admission_capacity(request, llama_backend) + key = str(getattr(llama_backend, "base_url", "llama-server")) + reservation = get_llama_admission_queue(key).reserve( + capacity = capacity, + config = config, + ) + return reservation, config + + +def _openai_admission_request_path(request: Optional[Request]) -> Optional[str]: + try: + return str(request.url.path) if request is not None else None + except Exception: + return None + + +def _openai_admission_log( + event: str, + reservation: Optional[LlamaAdmissionReservation] = None, + *, + snapshot = None, + request: Optional[Request], + mode: str, + wait_started_at: Optional[float] = None, + completion_id: Optional[str] = None, + level: str = "debug", +) -> None: + if snapshot is None and reservation is not None: + snapshot = reservation.snapshot_now() + wait_ms = None + if wait_started_at is not None: + wait_ms = int(max(0.0, time.monotonic() - wait_started_at) * 1000) + log = getattr(logger, level, logger.debug) + log( + "openai admission %s: mode=%s path=%s completion_id=%s capacity=%s active=%s queued=%s wait_ms=%s", + event, + mode, + _openai_admission_request_path(request), + completion_id, + getattr(snapshot, "capacity", None), + getattr(snapshot, "active", None), + getattr(snapshot, "queued", None), + wait_ms, + ) + + +def _openai_admission_error_body(exc: Exception, *, status_code: int) -> dict: + snapshot = getattr(exc, "snapshot", None) + message = str(exc) + if snapshot is not None: + message = ( + f"{message} " + f"(active={snapshot.active}, queued={snapshot.queued}, capacity={snapshot.capacity})" + ) + return openai_error_body(message, status = status_code) + + +def _openai_admission_http_exception(exc: Exception, *, status_code: int) -> HTTPException: + return HTTPException( + status_code = status_code, + detail = _openai_admission_error_body(exc, status_code = status_code), + ) + + +def _openai_admission_timeout_error( + reservation: LlamaAdmissionReservation, +) -> LlamaAdmissionTimeout: + return LlamaAdmissionTimeout( + "Timed out waiting for an available local llama-server generation slot", + snapshot = reservation.snapshot_now(), + ) + + +def _openai_admission_cancelled_error( + reservation: LlamaAdmissionReservation, +) -> LlamaAdmissionCancelled: + return LlamaAdmissionCancelled( + "Client disconnected before an upstream llama-server generation slot was available", + snapshot = reservation.snapshot_now(), + ) + + +async def _raise_if_openai_admission_cancelled( + reservation: LlamaAdmissionReservation, *, request: Optional[Request], cancel_event +) -> None: + if reservation.is_cancelled: + raise _openai_admission_cancelled_error(reservation) + if await _preheader_cancelled(cancel_event, request): + reservation.cancel() + raise _openai_admission_cancelled_error(reservation) + + +async def _wait_for_openai_admission_non_streaming( + reservation: LlamaAdmissionReservation, + config: LlamaAdmissionConfig, + *, + request: Optional[Request], + cancel_event, +) -> LlamaAdmissionLease: + lease = reservation.lease_nowait() + if lease is not None: + try: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + lease.release() + raise + except LlamaAdmissionCancelled: + lease.release() + raise + return lease + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + deadline = None if config.queue_timeout_s is None else time.monotonic() + config.queue_timeout_s + try: + while True: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + lease = reservation.lease_nowait() + if lease is not None: + try: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + lease.release() + raise + except LlamaAdmissionCancelled: + lease.release() + raise + return lease + wait_s = _OPENAI_LLAMA_ADMISSION_POLL_S + if deadline is not None: + remaining_s = deadline - time.monotonic() + if remaining_s <= 0: + reservation.cancel() + raise _openai_admission_timeout_error(reservation) + wait_s = min(wait_s, max(remaining_s, 0.001)) + try: + lease = await reservation.wait(wait_s) + except asyncio.TimeoutError: + continue + if lease is not None: + return lease + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + reservation.cancel() + raise + + +async def _openai_admission_wait_stream_chunks( + reservation: LlamaAdmissionReservation, + config: LlamaAdmissionConfig, + *, + request: Optional[Request], + cancel_event, +): + lease = reservation.lease_nowait() + if lease is not None: + yield lease + return + + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + deadline = None if config.queue_timeout_s is None else time.monotonic() + config.queue_timeout_s + keepalive_interval_s = max(0.001, config.keepalive_interval_s) + next_keepalive_at = time.monotonic() + keepalive_interval_s + try: + while True: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + lease = reservation.lease_nowait() + if lease is not None: + yield lease + return + + now = time.monotonic() + wait_s = min(_OPENAI_LLAMA_ADMISSION_POLL_S, max(next_keepalive_at - now, 0.001)) + if deadline is not None: + remaining_s = deadline - now + if remaining_s <= 0: + reservation.cancel() + raise _openai_admission_timeout_error(reservation) + wait_s = min(wait_s, max(remaining_s, 0.001)) + try: + lease = await reservation.wait(wait_s) + except asyncio.TimeoutError: + lease = None + if lease is not None: + yield lease + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + now = time.monotonic() + if now >= next_keepalive_at: + next_keepalive_at = now + keepalive_interval_s + yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE + except asyncio.CancelledError: + reservation.cancel() + raise + + +async def _close_openai_admitted_stream_iterator(iterator, *, cancelled: bool) -> None: + if iterator is None: + return + if cancelled: + athrow = getattr(iterator, "athrow", None) + if athrow is not None: + try: + await athrow(asyncio.CancelledError()) + except (asyncio.CancelledError, StopAsyncIteration, RuntimeError): + return + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() def _openai_compat_stream_stall_timeout(): @@ -6649,6 +6923,24 @@ async def openai_chat_completions( bypass_permissions = bool(payload.bypass_permissions), ) + _tool_admission_mode = "chat_tool_stream" if payload.stream else "chat_tool_nonstream" + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = _tool_admission_mode, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + _tool_sentinel = object() _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) @@ -6657,6 +6949,8 @@ async def openai_chat_completions( async def gguf_tool_stream(): gen = None + next_task = None + stream_completed = False disconnect_watcher = asyncio.create_task( _await_disconnect_then_cancel(request, cancel_event) ) @@ -6694,7 +6988,14 @@ async def openai_chat_completions( api_monitor.finish(monitor_id, "cancelled") return - event = await asyncio.to_thread(next, gen, _tool_sentinel) + next_task = asyncio.create_task( + asyncio.to_thread(next, gen, _tool_sentinel) + ) + try: + event = await asyncio.shield(next_task) + finally: + if next_task.done(): + next_task = None if event is _tool_sentinel: break @@ -6793,6 +7094,7 @@ async def openai_chat_completions( api_monitor.finish( monitor_id, "cancelled" if cancel_event.is_set() else "completed" ) + stream_completed = True yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -6807,18 +7109,155 @@ async def openai_chat_completions( error_chunk = _openai_stream_error_chunk(e) yield _openai_stream_error_sse(error_chunk) finally: - await _stop_local_disconnect_cancel_watcher(disconnect_watcher) - if gen is not None: - try: - gen.close() - except (RuntimeError, ValueError): - pass - _tracker.__exit__(None, None, None) + try: + if not stream_completed: + cancel_event.set() + task_to_drain = next_task + next_task = None + while task_to_drain is not None and not task_to_drain.done(): + try: + await asyncio.shield(task_to_drain) + except asyncio.CancelledError: + cancel_event.set() + continue + except Exception: + break + if task_to_drain is not None and task_to_drain.done(): + try: + task_to_drain.exception() + except (asyncio.CancelledError, Exception): + pass + if gen is not None and not stream_completed: + try: + await asyncio.to_thread(gen.close) + except (RuntimeError, ValueError): + pass + except Exception: + logger.debug( + "Error closing GGUF tool stream generator during cleanup", + exc_info = True, + ) + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) + finally: + _tracker.__exit__(None, None, None) if payload.stream: + stream_lease = reservation.lease_nowait() + admission_wait_started_at = None + if stream_lease is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = _tool_admission_mode, + completion_id = completion_id, + level = "debug", + ) + + async def admitted_gguf_tool_stream(): + lease = stream_lease + stream_started = False + stream_cancelled = False + try: + if lease is None: + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ): + if isinstance(wait_item, str): + yield wait_item + continue + lease = wait_item + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + break + if lease is None: + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + iterator = gguf_tool_stream() + stream_started = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + stream_cancelled = True + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = stream_cancelled, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _openai_stream_error_sse( + _openai_admission_error_body(exc, status_code = 503) + ) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error) + finally: + if lease is not None: + lease.release() + if not stream_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + async def _gguf_tool_admission_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + if stream_lease is not None: + stream_lease.release() + reservation.cancel() + _tracker.__exit__(None, None, None) + return _SameTaskStreamingResponse( - gguf_tool_stream(), - unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), + admitted_gguf_tool_stream(), + unstarted_cleanup = _gguf_tool_admission_unstarted_cleanup, media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6864,10 +7303,61 @@ async def openai_chat_completions( except (RuntimeError, ValueError): pass + drain_task = None + + async def _drain_cancelled_gguf_tool_task(): + if drain_task is None: + return + while not drain_task.done(): + try: + await asyncio.shield(drain_task) + except asyncio.CancelledError: + cancel_event.set() + continue + except Exception: + break + if drain_task.done(): + try: + drain_task.exception() + except (asyncio.CancelledError, Exception): + pass + + admission_lease = None + admission_wait_started_at = None try: - full_text, completion_usage, completion_finish = await asyncio.to_thread( - _drain_gguf_tool_loop + if reservation.lease_nowait() is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = _tool_admission_mode, + completion_id = completion_id, + level = "debug", + ) + admission_lease = await _wait_for_openai_admission_non_streaming( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, ) + if admission_wait_started_at is not None: + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + drain_task = asyncio.create_task(asyncio.to_thread(_drain_gguf_tool_loop)) + full_text, completion_usage, completion_finish = await asyncio.shield(drain_task) reasoning_text, visible_text = _extract_responses_reasoning( full_text, parse_think_markers = _responses_should_parse_think_markers( @@ -6913,6 +7403,48 @@ async def openai_chat_completions( monitor_id, "cancelled" if cancel_event.is_set() else "completed" ) return _model_json_response(response) + except asyncio.CancelledError: + cancel_event.set() + await _drain_cancelled_gguf_tool_task() + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise _openai_admission_http_exception(exc, status_code = 503) + except LlamaAdmissionCancelled as exc: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) except Exception as e: logger.error(f"Error during GGUF tool completion: {e}", exc_info = True) api_monitor.fail(monitor_id, _friendly_error(e)) @@ -6933,6 +7465,8 @@ async def openai_chat_completions( ) raise HTTPException(status_code = 500, detail = safe_error_detail(e)) finally: + if admission_lease is not None: + admission_lease.release() _tracker.__exit__(None, None, None) # ── Standard GGUF path (no tools) ───────────────────── @@ -6967,6 +7501,23 @@ async def openai_chat_completions( _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _tracker.__exit__(None, None, None) + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_standard_stream", + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) async def gguf_stream_chunks(): disconnect_watcher = asyncio.create_task( @@ -7114,9 +7665,122 @@ async def openai_chat_completions( finally: _tracker.__exit__(None, None, None) + stream_lease = reservation.lease_nowait() + admission_wait_started_at = None + if stream_lease is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_standard_stream", + completion_id = completion_id, + level = "debug", + ) + + async def admitted_gguf_stream_chunks(): + lease = stream_lease + stream_started = False + stream_cancelled = False + try: + if lease is None: + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ): + if isinstance(wait_item, str): + yield wait_item + continue + lease = wait_item + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_standard_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + break + if lease is None: + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + iterator = gguf_stream_chunks() + stream_started = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + stream_cancelled = True + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = stream_cancelled, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_standard_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _openai_stream_error_sse( + _openai_admission_error_body(exc, status_code = 503) + ) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_standard_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error) + finally: + if lease is not None: + lease.release() + if not stream_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + async def _gguf_admission_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + if stream_lease is not None: + stream_lease.release() + reservation.cancel() + _tracker.__exit__(None, None, None) + return _SameTaskStreamingResponse( - gguf_stream_chunks(), - unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), + admitted_gguf_stream_chunks(), + unstarted_cleanup = _gguf_admission_unstarted_cleanup, media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -7125,6 +7789,101 @@ async def openai_chat_completions( }, ) else: + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_standard_nonstream", + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _tracker = _TrackedCancel(cancel_event, *_cancel_keys) + _tracker.__enter__() + admission_lease = None + admission_wait_started_at = None + try: + if reservation.lease_nowait() is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_standard_nonstream", + completion_id = completion_id, + level = "debug", + ) + admission_lease = await _wait_for_openai_admission_non_streaming( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ) + if admission_wait_started_at is not None: + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_standard_nonstream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_standard_nonstream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise _openai_admission_http_exception(exc, status_code = 503) + except LlamaAdmissionCancelled as exc: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_standard_nonstream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) + try: # ``n`` requests several independent completions; the single # decode slot yields one at a time, so loop sequentially. @@ -7268,6 +8027,10 @@ async def openai_chat_completions( ), ) raise HTTPException(status_code = 500, detail = safe_error_detail(e)) + finally: + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) # ── Standard Unsloth path ───────────────────────────────── # Decode image (from content parts OR legacy field) @@ -9399,6 +10162,52 @@ async def _responses_stream( ) body["stream_options"] = {"include_usage": True} target_url = f"{llama_backend.base_url}/v1/chat/completions" + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "responses_stream", + completion_id = resp_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + def _responses_admission_failed_sse(exc: Exception, *, status_code: int) -> str: + return ( + "event: response.failed\n" + "data: " + + json.dumps( + { + "type": "response.failed", + "response": { + "id": resp_id, + "object": "response", + "created_at": created_at, + "status": "failed", + "model": _llama_public_model_id(llama_backend, payload.model) + or payload.model, + "output": [], + "usage": { + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + }, + "error": { + "code": status_code, + "message": str(exc), + }, + }, + } + ) + + "\n\n" + ) async def event_generator(): # Clean public id for every response envelope. Prefer the loaded model's @@ -10208,14 +11017,111 @@ async def _responses_stream( api_monitor.finish(monitor_id) yield _sse("response.completed", completed_response) + async def admitted_event_generator(): + lease = reservation.lease_nowait() + admission_wait_started_at = None + stream_started = False + stream_cancelled = False + iterator = None + try: + if lease is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "responses_stream", + completion_id = resp_id, + level = "debug", + ) + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = None, + ): + if isinstance(wait_item, str): + yield wait_item + continue + lease = wait_item + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "responses_stream", + wait_started_at = admission_wait_started_at, + completion_id = resp_id, + level = "debug", + ) + break + if lease is None: + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = None, + ) + iterator = event_generator() + stream_started = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + stream_cancelled = True + api_monitor.finish(monitor_id, "cancelled") + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = stream_cancelled, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "responses_stream", + wait_started_at = admission_wait_started_at, + completion_id = resp_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _responses_admission_failed_sse(exc, status_code = 503) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "responses_stream", + wait_started_at = admission_wait_started_at, + completion_id = resp_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + finally: + if lease is not None: + lease.release() + if not stream_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + + async def _responses_admission_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + return _SameTaskStreamingResponse( - event_generator(), + admitted_event_generator(), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", "Connection": "close", "X-Accel-Buffering": "no", }, + unstarted_cleanup = _responses_admission_unstarted_cleanup, ) @@ -12059,7 +12965,202 @@ async def _openai_passthrough_stream( completion_id, monitor_id: Optional[str] = None, ): - """Streaming client-side pass-through for /v1/chat/completions. + _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _tracker = _TrackedCancel(cancel_event, *_cancel_keys) + _tracker.__enter__() + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _tracker.__exit__(None, None, None) + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_passthrough_stream", + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + lease = reservation.lease_nowait() + if lease is not None: + try: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + lease.release() + _tracker.__exit__(None, None, None) + raise + except LlamaAdmissionCancelled as exc: + lease.release() + _tracker.__exit__(None, None, None) + api_monitor.finish(monitor_id, "cancelled") + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) + return await _openai_passthrough_stream_admitted( + request, + cancel_event, + llama_backend, + payload, + model_name, + completion_id, + monitor_id = monitor_id, + admission_lease = lease, + tracker = _tracker, + ) + + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_passthrough_stream", + completion_id = completion_id, + level = "debug", + ) + + async def _queued_stream(): + admitted_started = False + admitted_body_owns_cleanup = False + admitted_response = None + admitted_body_cancelled = False + try: + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ): + if isinstance(wait_item, str): + yield wait_item + continue + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_passthrough_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + admitted_response = await _openai_passthrough_stream_admitted( + request, + cancel_event, + llama_backend, + payload, + model_name, + completion_id, + monitor_id = monitor_id, + admission_lease = wait_item, + tracker = _tracker, + ) + admitted_started = True + iterator = admitted_response.body_iterator + admitted_body_owns_cleanup = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + admitted_body_cancelled = True + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = admitted_body_cancelled, + ) + if not admitted_body_owns_cleanup: + cleanup = getattr(admitted_response, "_unstarted_cleanup", None) + if cleanup is not None: + await cleanup() + return + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_passthrough_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _openai_stream_error_sse(_openai_admission_error_body(exc, status_code = 503)) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_passthrough_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error) + finally: + if not admitted_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + async def _queued_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + return _SameTaskStreamingResponse( + _queued_stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, + unstarted_cleanup = _queued_unstarted_cleanup, + ) + + +async def _openai_passthrough_stream_admitted( + request, + cancel_event, + llama_backend, + payload, + model_name, + completion_id, + monitor_id: Optional[str] = None, + *, + admission_lease: LlamaAdmissionLease, + tracker, +): + """Streaming client-side pass-through after Studio granted an upstream slot. Forwards the client's OpenAI function-calling request to llama-server and relays the SSE stream back with minimal normalization (reasoning-only @@ -12073,9 +13174,7 @@ async def _openai_passthrough_stream( --reasoning-format auto``), so ``delta.content`` carries no raw markup and is deliberately not re-parsed locally, unlike the ``/completion`` paths. """ - _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) - _tracker = _TrackedCancel(cancel_event, *_cancel_keys) - _tracker.__enter__() + _tracker = tracker target_url = f"{llama_backend.base_url}/v1/chat/completions" upstream_headers = _openai_passthrough_upstream_headers(llama_backend = llama_backend) @@ -12166,7 +13265,10 @@ async def _openai_passthrough_stream( await _aclose_send_task(send_task) await _aclose_stream_resources(client = client) finally: - _tracker.__exit__(None, None, None) + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( iter(()), media_type = "text/event-stream", @@ -12712,7 +13814,10 @@ async def _openai_passthrough_stream( client = client, ) finally: - _tracker.__exit__(None, None, None) + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) async def _unstarted_cleanup() -> None: # Client disconnected before the body stream started, so _stream()'s @@ -12723,7 +13828,10 @@ async def _openai_passthrough_stream( await _aclose_send_task(send_task) await _aclose_stream_resources(resp = resp, client = client) finally: - _tracker.__exit__(None, None, None) + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( _stream(), @@ -12747,7 +13855,10 @@ async def _openai_passthrough_stream( await _aclose_send_task(send_task) await _aclose_stream_resources(resp = resp, client = client) finally: - _tracker.__exit__(None, None, None) + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) raise @@ -12759,6 +13870,106 @@ async def _openai_passthrough_non_streaming( *, request: Optional[Request] = None, cancel_event = None, +): + """Non-streaming pass-through guarded by local llama-server admission.""" + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_passthrough_nonstream", + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + lease = None + admission_wait_started_at = None + try: + if reservation.lease_nowait() is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + level = "debug", + ) + lease = await _wait_for_openai_admission_non_streaming( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ) + if admission_wait_started_at is not None: + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + wait_started_at = admission_wait_started_at, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + return await _openai_passthrough_non_streaming_upstream( + llama_backend, + payload, + model_name, + monitor_id = monitor_id, + request = request, + cancel_event = cancel_event, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + wait_started_at = admission_wait_started_at, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 503) + except LlamaAdmissionCancelled as exc: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + wait_started_at = admission_wait_started_at, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + raise + finally: + if lease is not None: + lease.release() + + +async def _openai_passthrough_non_streaming_upstream( + llama_backend, + payload, + model_name, + monitor_id: Optional[str] = None, + *, + request: Optional[Request] = None, + cancel_event = None, ): """Non-streaming client-side pass-through for /v1/chat/completions. diff --git a/studio/backend/tests/test_llama_admission.py b/studio/backend/tests/test_llama_admission.py new file mode 100644 index 0000000000..2f04e81926 --- /dev/null +++ b/studio/backend/tests/test_llama_admission.py @@ -0,0 +1,320 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import asyncio +import os +import sys +import threading + +import pytest + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference import llama_admission +from core.inference.llama_admission import ( + ADMISSION_CONTROL_ENV, + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S, + DEFAULT_ADMISSION_MAX_QUEUE, + DEFAULT_ADMISSION_QUEUE_TIMEOUT_S, + LlamaAdmissionConfig, + LlamaAdmissionQueueFull, + get_llama_admission_queue, + llama_admission_config_from_env, + reset_llama_admission_queues, +) + + +@pytest.fixture(autouse = True) +def _reset_queues(): + reset_llama_admission_queues() + yield + reset_llama_admission_queues() + + +def test_admission_config_defaults(monkeypatch): + for name in ( + ADMISSION_CONTROL_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ): + monkeypatch.delenv(name, raising = False) + + config = llama_admission_config_from_env() + + assert config.enabled is True + assert config.queue_timeout_s == DEFAULT_ADMISSION_QUEUE_TIMEOUT_S + assert config.keepalive_interval_s == DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S + assert config.max_queue == DEFAULT_ADMISSION_MAX_QUEUE + + +def test_admission_config_env_overrides(monkeypatch): + monkeypatch.setenv(ADMISSION_CONTROL_ENV, "off") + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0") + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.25") + monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "0") + + config = llama_admission_config_from_env() + + assert config.enabled is False + assert config.queue_timeout_s is None + assert config.keepalive_interval_s == 0.25 + assert config.max_queue is None + + +def test_admission_config_positive_queue_timeout_env(monkeypatch): + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "600") + + config = llama_admission_config_from_env() + + assert config.queue_timeout_s == 600.0 + + +def test_fifo_capacity_one_grants_next_waiter_on_release(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + third = queue.reserve(capacity = 1, config = config) + + first_lease = first.lease_nowait() + assert first_lease is not None + assert second.lease_nowait() is None + assert third.lease_nowait() is None + assert queue.snapshot().queued == 2 + + first_lease.release() + second_lease = await second.wait(0.1) + assert second_lease is not None + assert third.lease_nowait() is None + + second_lease.release() + third_lease = await third.wait(0.1) + assert third_lease is not None + third_lease.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_queue_full_rejects_excess_waiter(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig(max_queue = 1) + + first = queue.reserve(capacity = 1, config = config) + queued = queue.reserve(capacity = 1, config = config) + + assert first.lease_nowait() is not None + assert queued.lease_nowait() is None + with pytest.raises(LlamaAdmissionQueueFull): + queue.reserve(capacity = 1, config = config) + + asyncio.run(_run()) + + +def test_disabled_admission_bypasses_active_slot_limit(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig(enabled = False) + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + + assert first.lease_nowait() is not None + assert second.lease_nowait() is not None + assert queue.snapshot().active == 0 + assert queue.snapshot().queued == 0 + + asyncio.run(_run()) + + +def test_cancelling_promoted_waiter_releases_slot(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + + first_lease.release() + await asyncio.sleep(0) + second.cancel() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_cancelling_promoted_waiter_before_delivery_releases_slot(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + + first_lease.release() + second.cancel() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_external_waiter_future_cancel_invalidates_reservation(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + assert first_lease is not None + assert second._waiter is not None + + second._waiter.future.cancel() + + assert second.lease_nowait() is None + assert second.is_cancelled is True + assert await second.wait(0.01) is None + + first_lease.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_wait_returns_none_when_waiter_future_cancelled_during_wait(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + assert first_lease is not None + assert second._waiter is not None + + wait_task = asyncio.create_task(second.wait(1.0)) + await asyncio.sleep(0) + second._waiter.future.cancel() + + assert await asyncio.wait_for(wait_task, timeout = 0.1) is None + assert second.is_cancelled is True + + first_lease.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_capacity_increase_promotes_existing_waiter_fifo(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + + first_lease = first.lease_nowait() + assert first_lease is not None + assert second.lease_nowait() is None + assert queue.snapshot().active == 1 + assert queue.snapshot().queued == 1 + + third = queue.reserve(capacity = 2, config = config) + + second_lease = await second.wait(0.1) + assert second_lease is not None + assert third.lease_nowait() is None + + snapshot = queue.snapshot() + assert snapshot.capacity == 2 + assert snapshot.active == 2 + assert snapshot.queued == 1 + + first_lease.release() + third_lease = await third.wait(0.1) + assert third_lease is not None + + second_lease.release() + third_lease.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_lease_release_is_idempotent_under_concurrent_calls(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + reservation = queue.reserve(capacity = 1, config = config) + lease = reservation.lease_nowait() + assert lease is not None + + threads = [threading.Thread(target = lease.release) for _ in range(16)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_new_key_evicts_idle_prior_load_queues(): + # Each model load carries a fresh ephemeral port, so a new base_url key must + # not leave the drained queues from earlier loads accumulating forever. + get_llama_admission_queue("http://127.0.0.1:1001") + get_llama_admission_queue("http://127.0.0.1:1002") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1002"} + + get_llama_admission_queue("http://127.0.0.1:1003") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1003"} + + +def test_new_key_retains_in_flight_prior_load_queue(): + config = LlamaAdmissionConfig() + busy = get_llama_admission_queue("http://127.0.0.1:2001") + + async def _run(): + reservation = busy.reserve(capacity = 1, config = config) + lease = reservation.lease_nowait() + assert lease is not None + + # A new load must not drop a queue that still has an in-flight request. + get_llama_admission_queue("http://127.0.0.1:2002") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2001", "http://127.0.0.1:2002"} + + # Once it drains, the next load reclaims it. + lease.release() + get_llama_admission_queue("http://127.0.0.1:2003") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2003"} + + asyncio.run(_run()) diff --git a/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py b/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py new file mode 100644 index 0000000000..5525bc3ea9 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import os +import sys + +import pytest + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference import llama_cpp as llama_cpp_module +from core.inference.llama_cpp import LlamaCppBackend + + +@pytest.fixture +def backend(monkeypatch): + monkeypatch.setattr(LlamaCppBackend, "_kill_orphaned_servers", lambda self: 0) + monkeypatch.setattr(llama_cpp_module.atexit, "register", lambda *_args, **_kwargs: None) + return LlamaCppBackend() + + +def test_effective_parallel_slots_initial_value_is_one(backend): + assert backend.effective_parallel_slots == 1 + + +def test_effective_parallel_slots_commit_uses_final_positive_parallel(backend): + backend._commit_effective_parallel_slots(3) + + assert backend.effective_parallel_slots == 3 + + +@pytest.mark.parametrize("value", [None, 0, -2, "not-an-int"]) +def test_effective_parallel_slots_commit_invalid_value_falls_back_to_one(backend, value): + backend._commit_effective_parallel_slots(value) + + assert backend.effective_parallel_slots == 1 + + +def test_effective_parallel_slots_reset_returns_to_one(backend): + backend._commit_effective_parallel_slots(4) + + backend._reset_effective_parallel_slots() + + assert backend.effective_parallel_slots == 1 + + +def test_effective_parallel_slots_unload_resets_to_one(backend): + backend._commit_effective_parallel_slots(4) + + backend.unload_model() + + assert backend.effective_parallel_slots == 1 diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 10cbd2aa01..910818d7d8 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -30,6 +30,15 @@ from core.inference.anthropic_compat import ( anthropic_tool_choice_to_openai, ) from core.inference.api_monitor import ApiMonitor +from core.inference.llama_admission import ( + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + LlamaAdmissionCancelled, + LlamaAdmissionConfig, + get_llama_admission_queue, + reset_llama_admission_queues, +) from routes.inference import ( _aclose_stream_resources, _build_chat_request, @@ -50,13 +59,17 @@ from routes.inference import ( _monitor_openai_sse_event, _normalize_openai_passthrough_sse_line, _openai_compat_stream_stall_timeout, + _openai_llama_admission_capacity, _openai_messages_for_gguf_chat, _openai_passthrough_sse_line_terminal_state, _openai_passthrough_upstream_headers, _openai_passthrough_non_streaming, _openai_passthrough_stream, + _responses_stream, _openai_stream_error_sse, _openai_stream_usage_chunk, + _openai_admission_wait_stream_chunks, + _wait_for_openai_admission_non_streaming, _proxy_to_external_provider, _SameTaskStreamingResponse, _OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, @@ -68,6 +81,13 @@ from routes.inference import ( from state.tool_policy import reset_tool_policy, set_tool_policy +@pytest.fixture(autouse = True) +def _reset_admission_queues(): + reset_llama_admission_queues() + yield + reset_llama_admission_queues() + + def test_aclose_stream_resources_attempts_remaining_closes_after_cancel(): class Closeable: def __init__(self, *, cancel = False): @@ -1338,6 +1358,80 @@ class TestOpenAICompatibilityHelpers: assert headers["Authorization"] == "Bearer secret" assert headers["Connection"] == "close" + def test_openai_admission_capacity_prefers_backend_effective_slots(self): + request = SimpleNamespace( + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + ) + backend = SimpleNamespace(effective_parallel_slots = 3) + + assert _openai_llama_admission_capacity(request, backend) == 3 + + @pytest.mark.parametrize("backend_value", [None, 0, -1, "not-an-int"]) + def test_openai_admission_capacity_falls_back_to_app_state(self, backend_value): + request = SimpleNamespace( + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 2)) + ) + backend = SimpleNamespace(effective_parallel_slots = backend_value) + + assert _openai_llama_admission_capacity(request, backend) == 2 + + def test_openai_admission_capacity_falls_back_to_one_without_request(self): + assert _openai_llama_admission_capacity(None, SimpleNamespace()) == 1 + + def test_openai_admission_non_streaming_exits_invalidated_waiter(self): + async def _run(): + queue = get_llama_admission_queue("http://llama.invalidated.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + reservation = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()) + assert reservation._waiter is not None + + reservation._waiter.future.cancel() + + with pytest.raises(LlamaAdmissionCancelled): + await asyncio.wait_for( + _wait_for_openai_admission_non_streaming( + reservation, + LlamaAdmissionConfig(), + request = None, + cancel_event = None, + ), + timeout = 0.1, + ) + + blocker.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_openai_admission_stream_exits_invalidated_waiter(self): + async def _run(): + queue = get_llama_admission_queue("http://llama.invalidated.stream.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + reservation = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()) + assert reservation._waiter is not None + + reservation._waiter.future.cancel() + + chunks = _openai_admission_wait_stream_chunks( + reservation, + LlamaAdmissionConfig(), + request = None, + cancel_event = None, + ) + with pytest.raises(LlamaAdmissionCancelled): + await asyncio.wait_for(chunks.__anext__(), timeout = 0.1) + + blocker.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + def test_openai_compat_stream_stall_timeout_uses_default(self, monkeypatch): monkeypatch.delenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, raising = False) assert _openai_compat_stream_stall_timeout() == 120.0 @@ -1993,6 +2087,335 @@ class TestGgufVisionToolRouting: [entry] = result.monitor.snapshot() assert entry["reply"] == "visible" + def test_standard_gguf_stream_queued_request_sends_keepalive_before_generation( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + queue = get_llama_admission_queue("http://llama.standard.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + stream = True, + ) + response = await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_standard_gguf_stream_close_after_first_chunk_cleans_tracker(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + cancel_id = "standard-stream-close-cleanup" + + def _generate(**_kwargs): + yield "visible" + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + stream = True, + cancel_id = cancel_id, + ) + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + iterator = response.body_iterator + assert cancel_id in inf_mod._CANCEL_REGISTRY + await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + aclose = getattr(iterator, "aclose", None) + assert aclose is not None + await aclose() + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_standard_gguf_stream_task_cancel_after_first_chunk_finalizes_monitor( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + started = threading.Event() + released = threading.Event() + + def _generate(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + stream = True, + ) + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + iterator = response.body_iterator + assert await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + pending = asyncio.create_task(iterator.__anext__()) + assert await asyncio.to_thread(started.wait, 1.0) + + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + + assert released.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_gguf_tool_stream_queued_request_sends_keepalive_before_generation(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _generate(**_kwargs): + raise AssertionError("GGUF tool loop must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + queue = get_llama_admission_queue("http://llama.tool.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + stream = True, + ) + response = await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_gguf_tool_stream_task_cancel_after_first_chunk_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + started = threading.Event() + released = threading.Event() + + def _tools(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _tools, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + stream = True, + ) + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + iterator = response.body_iterator + assert await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + pending = asyncio.create_task(iterator.__anext__()) + assert await asyncio.to_thread(started.wait, 1.0) + + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + + assert released.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + assert get_llama_admission_queue("http://llama.tool.test").snapshot().active == 0 + + asyncio.run(_run()) + def test_global_enable_tools_does_not_preempt_response_format_passthrough(self, monkeypatch): import routes.inference as inf_mod @@ -2434,6 +2857,313 @@ class TestGgufVisionToolRouting: [entry] = result.monitor.snapshot() assert entry["reply"] == "visible" + def test_standard_gguf_non_streaming_admission_timeout_before_generation(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + queue = get_llama_admission_queue("http://llama.standard.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + ) + try: + with pytest.raises(HTTPException) as exc: + await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + assert exc.value.status_code == 503 + finally: + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_standard_gguf_non_streaming_cancel_id_stops_queued_request_before_generation( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start after cancel_id") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + queue = get_llama_admission_queue("http://llama.standard.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + cancel_id = "standard-nonstream-admission-cancel" + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + cancel_id = cancel_id, + ) + task = asyncio.create_task( + openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + ) + try: + for _ in range(50): + if cancel_id in inf_mod._CANCEL_REGISTRY: + break + await asyncio.sleep(0.01) + assert cancel_id in inf_mod._CANCEL_REGISTRY + assert inf_mod._cancel_by_cancel_id_or_stash(cancel_id) == 1 + with pytest.raises(HTTPException) as exc: + await asyncio.wait_for(task, timeout = 0.5) + assert exc.value.status_code == 499 + finally: + if not task.done(): + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + blocker.release() + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_standard_gguf_non_streaming_admission_task_cancel_cleans_tracker_and_slot( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + cancel_id = "standard-nonstream-task-cancel" + + async def fake_wait(*_args, **_kwargs): + raise asyncio.CancelledError() + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start after task cancel") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr( + inf_mod, + "_wait_for_openai_admission_non_streaming", + fake_wait, + ) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + cancel_id = cancel_id, + ) + with pytest.raises(asyncio.CancelledError): + await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_gguf_tool_non_streaming_admission_timeout_before_generation(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _generate(**_kwargs): + raise AssertionError("GGUF tool loop must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + queue = get_llama_admission_queue("http://llama.tool.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + ) + try: + with pytest.raises(HTTPException) as exc: + await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + assert exc.value.status_code == 503 + finally: + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_gguf_tool_non_streaming_cancel_drains_worker_before_releasing_slot(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + started = threading.Event() + released = threading.Event() + + def _tools(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _tools, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + ) + task = asyncio.create_task( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + assert await asyncio.to_thread(started.wait, 1.0) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout = 1.0) + + assert released.is_set() + assert get_llama_admission_queue("http://llama.tool.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_non_streaming_gguf_n_records_all_monitor_replies(self, monkeypatch): import routes.inference as inf_mod @@ -4086,6 +4816,270 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_stream_immediate_task_cancel_releases_admission_and_tracker( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + async def fake_cancel_check(*_args, **_kwargs): + raise asyncio.CancelledError() + + cancel_id = "passthrough-stream-immediate-task-cancel" + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_raise_if_openai_admission_cancelled", + fake_cancel_check, + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + backend = SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + with pytest.raises(asyncio.CancelledError): + await _openai_passthrough_stream( + self._Request(), + threading.Event(), + backend, + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert get_llama_admission_queue("http://llama.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_queued_cancel_before_inner_first_chunk_runs_cleanup( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + body_holder = {} + cleanup_called = threading.Event() + + async def fake_admitted(*_args, admission_lease, tracker, **_kwargs): + async def cleanup(): + admission_lease.release() + tracker.__exit__(None, None, None) + cleanup_called.set() + + class BlockingBody: + def __init__(self): + self.started = threading.Event() + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + self.started.set() + await asyncio.sleep(3600) + raise StopAsyncIteration + + async def aclose(self): + self.closed = True + await cleanup() + + body = BlockingBody() + body_holder["body"] = body + return _SameTaskStreamingResponse( + body, + media_type = "text/event-stream", + unstarted_cleanup = cleanup, + ) + + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_stream_admitted", + fake_admitted, + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + cancel_id = "queued-inner-unstarted-cleanup" + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + assert cancel_id in inf_mod._CANCEL_REGISTRY + + blocker.release() + pending = asyncio.create_task(iterator.__anext__()) + for _ in range(100): + if "body" in body_holder: + break + await asyncio.sleep(0.01) + body = body_holder["body"] + assert await asyncio.to_thread(body.started.wait, 1.0) + + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + assert body_holder["body"].closed + assert cleanup_called.is_set() + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert queue.snapshot().active == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_queued_cancel_after_inner_first_chunk_finalizes_monitor( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + async def fake_admitted( + *_args, + monitor_id = None, + admission_lease, + tracker, + **_kwargs, + ): + async def cleanup(): + admission_lease.release() + tracker.__exit__(None, None, None) + + async def body(): + try: + yield 'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n' + await asyncio.sleep(3600) + except asyncio.CancelledError: + inf_mod.api_monitor.finish(monitor_id, "cancelled") + raise + finally: + await cleanup() + + return _SameTaskStreamingResponse( + body(), + media_type = "text/event-stream", + unstarted_cleanup = cleanup, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_stream_admitted", + fake_admitted, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + cancel_id = "queued-inner-cancel-monitor" + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + + blocker.release() + first = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert "hello" in first + + pending = asyncio.create_task(iterator.__anext__()) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert queue.snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_passthrough_stream_synthesizes_missing_finish_reason(self, monkeypatch): async def _run(): result = await self._run_passthrough_stream( @@ -4189,6 +5183,292 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_stream_queued_request_sends_keepalive_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + url = SimpleNamespace(path = "/v1/chat/completions") + + async def is_disconnected(self): + return False + + async def fail_admitted(*_args, **_kwargs): + raise AssertionError("upstream must not start while request is queued") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_openai_passthrough_stream_admitted", fail_admitted) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_admission_timeout_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + url = SimpleNamespace(path = "/v1/chat/completions") + + async def is_disconnected(self): + return False + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start while request is queued") + + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.01") + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + try: + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + request = Request(), + cancel_event = threading.Event(), + ) + assert exc.value.status_code == 503 + finally: + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_admission_queue_full_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + url = SimpleNamespace(path = "/v1/chat/completions") + + async def is_disconnected(self): + return False + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start when admission queue is full") + + monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "1") + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve( + capacity = 1, + config = LlamaAdmissionConfig(max_queue = 1), + ).lease_nowait() + queued = queue.reserve(capacity = 1, config = LlamaAdmissionConfig(max_queue = 1)) + assert blocker is not None + assert queued.lease_nowait() is None + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + try: + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + request = Request(), + cancel_event = threading.Event(), + ) + assert exc.value.status_code == 429 + finally: + queued.cancel() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_immediate_cancel_stops_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start after client cancellation") + + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + + cancel_event = threading.Event() + cancel_event.set() + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + cancel_event = cancel_event, + ) + + assert exc.value.status_code == 499 + assert get_llama_admission_queue("http://llama.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_admission_task_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_wait(*_args, **_kwargs): + raise asyncio.CancelledError() + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start after admission task cancel") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_wait_for_openai_admission_non_streaming", + fake_wait, + ) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + + with pytest.raises(asyncio.CancelledError): + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + cancel_event = threading.Event(), + ) + + assert get_llama_admission_queue("http://llama.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_passthrough_non_streaming_cancel_finalizes_monitor(self, monkeypatch): async def _run(): import routes.inference as inf_mod @@ -5381,6 +6661,15 @@ class TestApiMonitorAudioInput: class TestResponsesChatTemplateKwargs: _messages = [ChatMessage(role = "user", content = "What is 100 - 67?")] + class _Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/responses") + method = "POST" + + async def is_disconnected(self): + return False + def test_enable_thinking_lifted_from_extra_body(self): payload = ResponsesRequest( model = "qwen-local", @@ -5413,6 +6702,113 @@ class TestResponsesChatTemplateKwargs: chat_req = _build_chat_request(payload, self._messages, stream = False) assert chat_req.enable_thinking is None + def test_responses_stream_queued_request_sends_keepalive_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fail_send(*_args, **_kwargs): + raise AssertionError("responses upstream must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + base_url = "http://llama.responses.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fail_send) + + queue = get_llama_admission_queue("http://llama.responses.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "qwen-local", + prompt = "hi", + ) + payload = ResponsesRequest(model = "qwen-local", input = "hi", stream = True) + + response = await _responses_stream( + payload, + [ChatMessage(role = "user", content = "hi")], + self._Request(), + monitor_id, + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_responses_stream_cancel_after_created_finalizes_monitor_and_slot(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fail_send(*_args, **_kwargs): + raise AssertionError("responses upstream must not start after created cancel") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + base_url = "http://llama.responses.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fail_send) + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "qwen-local", + prompt = "hi", + ) + payload = ResponsesRequest(model = "qwen-local", input = "hi", stream = True) + + response = await _responses_stream( + payload, + [ChatMessage(role = "user", content = "hi")], + self._Request(), + monitor_id, + ) + iterator = response.body_iterator + first = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert "event: response.created" in first + + with pytest.raises(asyncio.CancelledError): + await iterator.athrow(asyncio.CancelledError()) + + assert get_llama_admission_queue("http://llama.responses.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + # ===================================================================== # GGUF chat-template role alternation: coalesce orphaned user turns left diff --git a/tests/studio/test_stream_cancel_registration_timing.py b/tests/studio/test_stream_cancel_registration_timing.py index 7cfe9bf32c..a833006873 100644 --- a/tests/studio/test_stream_cancel_registration_timing.py +++ b/tests/studio/test_stream_cancel_registration_timing.py @@ -17,7 +17,7 @@ from pathlib import Path SOURCE_PATH = Path(__file__).resolve().parents[2] / "studio" / "backend" / "routes" / "inference.py" -SRC = SOURCE_PATH.read_text() +SRC = SOURCE_PATH.read_text(encoding = "utf-8") _TREE = ast.parse(SRC) @@ -166,16 +166,20 @@ def test_chat_completions_streams_avoid_starlette_task_group(): def test_openai_passthrough_stream_avoids_starlette_task_group(): - top = _async_function("_openai_passthrough_stream") + functions = [ + _async_function("_openai_passthrough_stream"), + _async_function("_openai_passthrough_stream_admitted"), + ] legacy_calls = [] same_task_calls = 0 - for sub in ast.walk(top): - if not (isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)): - continue - if sub.func.id == "StreamingResponse": - legacy_calls.append(sub.lineno) - if sub.func.id == "_SameTaskStreamingResponse": - same_task_calls += 1 + for fn in functions: + for sub in ast.walk(fn): + if not (isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)): + continue + if sub.func.id == "StreamingResponse": + legacy_calls.append(sub.lineno) + if sub.func.id == "_SameTaskStreamingResponse": + same_task_calls += 1 assert not legacy_calls, ( "OpenAI passthrough streams must use _SameTaskStreamingResponse, " "not Starlette's legacy task-group StreamingResponse. Lines: " @@ -197,7 +201,7 @@ def test_direct_llama_server_streams_install_disconnect_watcher(): "openai_completions", "_responses_stream", "_anthropic_passthrough_stream", - "_openai_passthrough_stream", + "_openai_passthrough_stream_admitted", } missing = [ name From 7bfa209623415e8ed71d739a80bcd12d3ccf74c3 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Fri, 10 Jul 2026 17:48:27 -0300 Subject: [PATCH 4/9] Studio: hint at Model auto-switch in the OpenAI "No model loaded" 400 (#7006) --- studio/backend/routes/inference.py | 29 +++++++--- .../backend/tests/test_openai_auto_switch.py | 57 +++++++++++++++++++ 2 files changed, 79 insertions(+), 7 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index db350240ac..ec5d309810 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -3344,6 +3344,21 @@ def _automatic_model_load_may_run() -> bool: return get_openai_auto_switch_enabled() or get_auto_unload_idle_seconds() > 0 +def _no_model_loaded_detail(base: str) -> str: + """Append a pointer to the opt-in auto-switch toggle to a "no model loaded" + error, but only when it's off. Auto-switch (default off) cold-loads a + requested downloaded GGUF, so an off toggle is the usual reason a request + naming a listed model still 400/503s; surface the fix. With it on the name + simply didn't resolve to a local GGUF, so the hint would mislead and is omitted.""" + from utils.openai_auto_switch_settings import get_openai_auto_switch_enabled + + if get_openai_auto_switch_enabled(): + return base + return base + ( + " Or enable Model auto-switch (Settings > API) to load a requested model automatically." + ) + + async def _maybe_auto_switch_model( requested_model: Optional[str], fastapi_request: Request, @@ -6451,7 +6466,7 @@ async def openai_chat_completions( if not backend.active_model_name: raise HTTPException( status_code = 400, - detail = "No model loaded. Call POST /inference/load first.", + detail = _no_model_loaded_detail("No model loaded. Call POST /inference/load first."), ) # Clean public id so the response never echoes a local path; the audio # branch below receives this sanitized label too. @@ -9184,7 +9199,7 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) if not isinstance(body, dict): # Re-read to re-raise a malformed-body error (post-503, pre-feature behavior); @@ -9400,7 +9415,7 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) if not isinstance(body, dict): # Re-read to re-raise a malformed-body error (post-503, pre-feature behavior); @@ -10140,7 +10155,7 @@ async def _responses_stream( # so the client sees a useful error instead of a dangling stream. raise HTTPException( status_code = 400, - detail = ( + detail = _no_model_loaded_detail( "Streaming /v1/responses requires a GGUF model loaded via " "llama-server. Use non-streaming /v1/responses, " "/v1/chat/completions, or load a GGUF model." @@ -11379,7 +11394,7 @@ async def anthropic_count_tokens( if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) # Same Anthropic → OpenAI translation as anthropic_messages: system is @@ -11452,7 +11467,7 @@ async def anthropic_messages( if not llama_backend.is_loaded and not _automatic_model_load_may_run(): raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) # max_tokens is a required field on the Anthropic Messages API; real Anthropic @@ -11503,7 +11518,7 @@ async def anthropic_messages( if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) # Advertised repo id after an auto-switch load, else a clean public id, never diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py index 7d2e2213b3..d02a2a4f7e 100644 --- a/studio/backend/tests/test_openai_auto_switch.py +++ b/studio/backend/tests/test_openai_auto_switch.py @@ -3037,3 +3037,60 @@ def test_acquire_swap_gate_is_cancellation_safe(): inference_route._auto_switch_process_lock.release() asyncio.run(asyncio.wait_for(main(), timeout = 5)) + + +def test_no_model_loaded_detail_appends_hint_only_when_off(monkeypatch): + # The "no model loaded" errors point at the opt-in auto-switch toggle so a + # request naming a listed-but-unloaded model is self-explanatory -- but only + # when it's off. With it on the name simply didn't resolve, so no hint. + base = "No GGUF model loaded. Load a GGUF model first." + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + off = inference_route._no_model_loaded_detail(base) + assert off.startswith(base) + assert "Model auto-switch" in off and "Settings > API" in off + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + assert inference_route._no_model_loaded_detail(base) == base + + +def _run_responses_stream_no_model(monkeypatch, *, enabled, active_model_name): + # Drive _responses_stream's GGUF-not-loaded guard: llama backend unloaded, + # inference backend maybe holding a non-GGUF model. Returns the 400 detail. + from fastapi import HTTPException + from models.inference import ResponsesRequest, ChatMessage + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: enabled) + monkeypatch.setattr( + inference_route, "get_llama_cpp_backend", lambda: _FakeBackend(loaded_id = None) + ) + monkeypatch.setattr( + inference_route, + "get_inference_backend", + lambda: type("_B", (), {"active_model_name": active_model_name})(), + ) + payload = ResponsesRequest(model = "unsloth/Qwen3.5-4B-GGUF", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route._responses_stream(payload, messages, None)) + assert exc.value.status_code == 400 + return exc.value.detail + + +def test_responses_stream_hint_matches_toggle_regardless_of_active_model(monkeypatch): + # Streaming /v1/responses shares the GGUF-only 400 with the other "no model + # loaded" sites, so the auto-switch hint attaches whenever the toggle is + # off -- including while a non-GGUF model is active, since auto-switch + # evicts it to load a resolved GGUF (_maybe_auto_switch_model's resolver + # branch has no active-model guard, unlike its reload-stash branch). Only + # the toggle being on suppresses it. + hinted = _run_responses_stream_no_model(monkeypatch, enabled = False, active_model_name = None) + assert "Model auto-switch" in hinted + + on = _run_responses_stream_no_model(monkeypatch, enabled = True, active_model_name = None) + assert "Model auto-switch" not in on + + non_gguf_loaded = _run_responses_stream_no_model( + monkeypatch, enabled = False, active_model_name = "unsloth/Llama-3.2-1B-Instruct" + ) + assert "Model auto-switch" in non_gguf_loaded From d105bd7b42ea8d4ecdb3e3364abb605b558c017e Mon Sep 17 00:00:00 2001 From: oobabooga Date: Fri, 10 Jul 2026 17:59:04 -0300 Subject: [PATCH 5/9] Studio: detect Windows Intel GPUs via the registry before WMI (#7064) --- .../tests/test_install_resolve_prebuilt.py | 234 ++++++++++++++++++ studio/install_llama_prebuilt.py | 104 ++++++-- 2 files changed, 318 insertions(+), 20 deletions(-) diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py index 090d2932ea..e97ca47717 100644 --- a/studio/backend/tests/test_install_resolve_prebuilt.py +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -483,3 +483,237 @@ def test_resolve_prebuilt_intel_host_routes_to_upstream(monkeypatch, capsys): seen, out = _run_resolve_capture_host(monkeypatch, capsys) assert seen["repo"] == UPSTREAM assert out["repo"] == UPSTREAM + + +# --------------------------------------------------------------------------- +# windows_intel_gpu_in_registry: the in-process Windows Intel probe. A fake +# winreg module stands in for the real registry so the walk runs anywhere. +# --------------------------------------------------------------------------- + + +class _FakeRegKey: + def __init__( + self, + subkeys = None, + values = None, + denied = False, + ): + self.subkeys = subkeys or {} + self.values = values or {} + self.denied = denied + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +class _FakeWinreg: + HKEY_LOCAL_MACHINE = object() + + def __init__(self, root_key): + self._root_key = root_key + + def OpenKey(self, parent, name): + if parent is self.HKEY_LOCAL_MACHINE: + # Pin the production constant: a typo'd class GUID must fail here, + # not silently return the fake tree. + if name != ilp._WINDOWS_DISPLAY_CLASS_KEY: + raise FileNotFoundError(name) + if self._root_key is None: + raise FileNotFoundError(name) + return self._root_key + key = parent.subkeys.get(name) + if key is None: + # Real winreg raises OSError, never KeyError, for a missing key. + raise FileNotFoundError(name) + if key.denied: + raise PermissionError(name) + return key + + def QueryInfoKey(self, key): + return (len(key.subkeys), len(key.values), 0) + + def EnumKey(self, key, index): + return list(key.subkeys)[index] + + def QueryValueEx(self, key, value_name): + if value_name not in key.values: + raise FileNotFoundError(value_name) + return (key.values[value_name], 1) + + +def _probe_with_display_class(monkeypatch, adapters): + # The helper lazily does `import winreg`; plant the fake in sys.modules the + # same way unsloth_cli/tests/test_start.py fakes it for _refresh_windows_path. + monkeypatch.setitem(sys.modules, "winreg", _FakeWinreg(_FakeRegKey(subkeys = adapters))) + return ilp.windows_intel_gpu_in_registry() + + +def test_windows_intel_registry_matches_vendor_id(monkeypatch): + assert ( + _probe_with_display_class( + monkeypatch, + { + "0000": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0&SUBSYS_12345678", + "DriverDesc": "Intel(R) Arc(TM) A770 Graphics", + } + ), + }, + ) + is True + ) + + +def test_windows_intel_registry_matches_driver_desc_without_device_id(monkeypatch): + assert ( + _probe_with_display_class( + monkeypatch, + { + "0000": _FakeRegKey(values = {"DriverDesc": "Intel(R) UHD Graphics 630"}), + }, + ) + is True + ) + + +def test_windows_intel_registry_ignores_non_intel_adapters(monkeypatch): + assert ( + _probe_with_display_class( + monkeypatch, + { + "0000": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684", + "DriverDesc": "NVIDIA GeForce RTX 4090", + } + ), + "0001": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_1002&DEV_744C", + "DriverDesc": "AMD Radeon RX 7900 XTX", + } + ), + }, + ) + is False + ) + + +def test_windows_intel_registry_skips_restricted_properties_subkey(monkeypatch): + # The real class key carries an ACL-restricted "Properties" subkey and can + # deny access to individual adapter keys; neither may abort the walk. + assert ( + _probe_with_display_class( + monkeypatch, + { + "Properties": _FakeRegKey(denied = True), + "0000": _FakeRegKey(denied = True), + "0001": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0", + } + ), + }, + ) + is True + ) + + +def test_windows_intel_registry_missing_class_key_is_false(monkeypatch): + monkeypatch.setitem(sys.modules, "winreg", _FakeWinreg(None)) + assert ilp.windows_intel_gpu_in_registry() is False + + +def _detect_windows_host( + monkeypatch, + winreg_fake, + powershell_stdout = "", +): + """Drive the real detect_host() as a GPU-less Windows host with a fake + registry, recording every run_capture invocation. Pins the wiring the + unit tests above cannot see: registry-first, CIM only on a registry miss.""" + monkeypatch.setitem(sys.modules, "winreg", winreg_fake) + monkeypatch.setattr(ilp.platform, "system", lambda: "Windows") + monkeypatch.setattr(ilp.platform, "machine", lambda: "AMD64") + for _env in ( + "CUDA_VISIBLE_DEVICES", + "HIP_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", + "HIP_PATH", + "ROCM_PATH", + ): + monkeypatch.delenv(_env, raising = False) + monkeypatch.setattr( + ilp.shutil, + "which", + lambda name: "powershell" if name in ("powershell", "pwsh") else None, + ) + captured = [] + + def _fake_run_capture(command, **kwargs): + captured.append(command[0]) + if command[0] == "powershell": + return SimpleNamespace(returncode = 0, stdout = powershell_stdout, stderr = "") + return SimpleNamespace(returncode = 1, stdout = "", stderr = "") + + monkeypatch.setattr(ilp, "run_capture", _fake_run_capture) + return ilp.detect_host(), captured + + +def test_detect_host_registry_intel_skips_cim_probe(monkeypatch): + winreg = _FakeWinreg( + _FakeRegKey( + subkeys = { + "0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0"}), + } + ) + ) + host, captured = _detect_windows_host(monkeypatch, winreg) + assert host.has_intel_gpu is True + assert "powershell" not in captured + + +def test_detect_host_cim_fallback_fires_on_registry_miss(monkeypatch): + winreg = _FakeWinreg( + _FakeRegKey( + subkeys = { + "0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684"}), + } + ) + ) + host, captured = _detect_windows_host( + monkeypatch, winreg, powershell_stdout = "Intel(R) Arc(TM) A770 Graphics" + ) + assert host.has_intel_gpu is True + assert "powershell" in captured + + +def test_windows_intel_registry_unexpected_error_is_false(monkeypatch): + # The probe is advisory: even a non-OSError bug in the walk must return + # False (deferring to the CIM fallback), never crash detect_host. + class _ExplodingWinreg: + HKEY_LOCAL_MACHINE = object() + + def OpenKey(self, parent, name): + raise TypeError(name) + + monkeypatch.setitem(sys.modules, "winreg", _ExplodingWinreg()) + assert ilp.windows_intel_gpu_in_registry() is False + + +def test_detect_host_cim_rescues_exploding_registry(monkeypatch): + class _ExplodingWinreg: + HKEY_LOCAL_MACHINE = object() + + def OpenKey(self, parent, name): + raise TypeError(name) + + host, captured = _detect_windows_host( + monkeypatch, _ExplodingWinreg(), powershell_stdout = "Intel(R) Arc(TM) A770 Graphics" + ) + assert host.has_intel_gpu is True + assert "powershell" in captured diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 40caebc040..ca1fe79efa 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -2762,6 +2762,64 @@ def _pick_rocm_gfx_target(out: str) -> str | None: return _tokens[0] +# Display-adapter device class: one NNNN subkey per installed display driver +# config, each carrying the driver's DriverDesc and PCI MatchingDeviceId. +_WINDOWS_DISPLAY_CLASS_KEY = ( + r"SYSTEM\CurrentControlSet\Control\Class\{4d36e968-e325-11ce-bfc1-08002be10318}" +) + + +def windows_intel_gpu_in_registry() -> bool: + """Whether the Windows registry lists an Intel display adapter. + + In-process Windows counterpart of the Linux DRM vendor-id check (0x8086), + with weaker semantics: the class key lists installed display-driver + configs, which can outlive removed hardware, where sysfs lists present + devices. A stale Intel entry at worst routes to the upstream Vulkan + prebuilt instead of the fork CPU bundle: inference still works (the + Vulkan build runs on CPU when no Vulkan device exists), at the cost of + fork-only extras such as the DiffusionGemma visual server. detect_host's + PowerShell + WMI probe can silently miss a real Intel GPU: a cold + powershell.exe start plus the first CIM query routinely exceeds the 15s + budget on hosts with slow AV scanning or a degraded WMI repository, and + the probe swallows the timeout (#4452, Arc A770 routed to the CPU + prebuilt). Reading the display-adapter class key needs no subprocess and + answers in microseconds. Matches the PCI vendor id in MatchingDeviceId + (ven_8086) or an Intel DriverDesc. + """ + try: + import winreg + except ImportError: + return False + try: + with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, _WINDOWS_DISPLAY_CLASS_KEY) as class_key: + for index in range(winreg.QueryInfoKey(class_key)[0]): + try: + name = winreg.EnumKey(class_key, index) + if not name.isdigit(): + # "Properties" is ACL-restricted and not an adapter. + continue + with winreg.OpenKey(class_key, name) as adapter_key: + for value_name, needle in ( + ("MatchingDeviceId", "ven_8086"), + ("DriverDesc", "intel"), + ): + try: + value, _ = winreg.QueryValueEx(adapter_key, value_name) + except OSError: + continue + if needle in str(value).lower(): + return True + except OSError: + continue + except Exception: + # Advisory probe: any unexpected failure must degrade to the CIM + # fallback, never crash the installer (mirrors detect_host's own + # swallow around the CIM probe). + return False + return False + + def detect_host() -> HostInfo: system = platform.system() machine = platform.machine().lower() @@ -2974,9 +3032,10 @@ def detect_host() -> HostInfo: # since the HIP SDK can be installed without an AMD GPU. # Detect an Intel GPU; gates the Vulkan prebuilt. Linux reads the DRM sysfs - # vendor id (0x8086); Windows queries the WMI video controller list. Only - # probed with no usable NVIDIA and no ROCm (matching the Vulkan branches), - # keeping the probe (notably the Windows powershell call) off that path. + # vendor id (0x8086); Windows reads the display-adapter registry class, + # then falls back to the WMI video controller list. Only probed with no + # usable NVIDIA and no ROCm (matching the Vulkan branches), keeping the + # probe (notably the Windows powershell call) off that path. has_intel_gpu = False if not has_usable_nvidia and not has_rocm: if is_linux: @@ -2989,23 +3048,28 @@ def detect_host() -> HostInfo: except OSError: continue elif is_windows: - _ps = shutil.which("powershell") or shutil.which("pwsh") - if _ps: - try: - _result = run_capture( - [ - _ps, - "-NoProfile", - "-Command", - "Get-CimInstance Win32_VideoController | " - "Select-Object -ExpandProperty Name", - ], - timeout = 15, - ) - if _result.returncode == 0 and "intel" in _result.stdout.lower(): - has_intel_gpu = True - except Exception: - pass + # Registry first (in-process; see windows_intel_gpu_in_registry). + # The CIM query stays as the fallback when the registry shows no + # Intel adapter. + has_intel_gpu = windows_intel_gpu_in_registry() + if not has_intel_gpu: + _ps = shutil.which("powershell") or shutil.which("pwsh") + if _ps: + try: + _result = run_capture( + [ + _ps, + "-NoProfile", + "-Command", + "Get-CimInstance Win32_VideoController | " + "Select-Object -ExpandProperty Name", + ], + timeout = 15, + ) + if _result.returncode == 0 and "intel" in _result.stdout.lower(): + has_intel_gpu = True + except Exception: + pass return HostInfo( system = system, From c3feac6160698c9c1990be89a1a2a74b01fa57aa Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 11 Jul 2026 05:08:07 -0700 Subject: [PATCH 6/9] Studio: route lfm2_moe (LFM2-8B-A1B) to transformers 5.3.0 (#7040) LFM2-8B-A1B and any other lfm2_moe checkpoint were missing from the transformers tier tables, so they fell through to the default 4.57.x sidecar, which does not register lfm2_moe and errors with "not supported yet in transformers==4.57.6". Only lfm2_vl was listed. Add Lfm2MoeForCausalLM / lfm2_moe to the 5.3.0 tier (lfm2_moe is registered in transformers 5.3.0). get_transformers_tier now returns 530 for LFM2-8B-A1B and the model loads and trains as expected. --- studio/backend/utils/transformers_version.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 2a63caa5b2..81d534a7c2 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -173,6 +173,7 @@ _TRANSFORMERS_530_ARCHITECTURES: set[str] = { "Qwen3MoeForCausalLM", "Qwen3NextForCausalLM", "Glm4MoeLiteForCausalLM", + "Lfm2MoeForCausalLM", "Lfm2VlForConditionalGeneration", } _TRANSFORMERS_530_MODEL_TYPES: set[str] = { @@ -183,6 +184,7 @@ _TRANSFORMERS_530_MODEL_TYPES: set[str] = { "qwen3_moe", "qwen3_next", "glm4_moe_lite", + "lfm2_moe", "lfm2_vl", } From 97161c89d6fd622ba5c3d9ccd7d62b2476e7466f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 11 Jul 2026 05:12:37 -0700 Subject: [PATCH 7/9] Studio: route models by CONFIG_MAPPING_NAMES instead of hardcoded tables (#7043) * Studio: route models by CONFIG_MAPPING_NAMES instead of hardcoded tables A model whose model_type is absent from an overlay's transformers cannot load there, so a new MoE arch not yet in the tier tables gets routed to default and fails (e.g. lfm2_moe, deepseek_v4). Add a static resolver that parses each overlay's CONFIG_MAPPING_NAMES straight from source (AST only, no import, no network, no trust_remote_code) and picks the lowest tier that ships the model_type. Runs after the existing checks and only ever upgrades default, so no existing routing changes and new archs no longer need a table edit. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio router: harden the CONFIG_MAPPING_NAMES resolver - Resolve the default tier map from the base install, skipping any .venv_t5_* sidecar on sys.path, so an in-process 5.x activation cannot make a 5.x-only model look loadable by 4.x. - Do not cache an overlay whose sidecar dir is absent, so a later call re-reads it once provisioned instead of serving a stale empty map. - Also collect model types added via CONFIG_MAPPING_NAMES.update({...}) and **{...} unpacking, not just the literal assignment (5.10 uses both). - Wrap the AST walk in the try/except so a malformed source can never crash tier resolution. - Feed the mapping fallback from _load_config_json so a config served from the hub cache during a transient outage still routes new architectures. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/utils/transformers_version.py | 132 +++++++++++++++++++ 1 file changed, 132 insertions(+) diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 81d534a7c2..a69673f081 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -28,7 +28,9 @@ Strategy: sys.path swap using the same directories pre-installed by setup.sh. """ +import ast import importlib +import importlib.util import json import structlog from loggers import get_logger @@ -872,6 +874,116 @@ def _cached_config_json(model_name: str, hf_token: str | None) -> dict | None: return _config_json_cache.get(_token_cache_key(model_name, hf_token)) +# --- Static tier from CONFIG_MAPPING_NAMES (AST only: no import/network/exec) --- +# A model_type absent from an overlay's mapping can't load there. Parse each sidecar's +# config map from source and pick the lowest tier that ships it, so a new arch routes +# correctly with no per-model table edit. Only ever upgrades default, never lowers. +_config_mapping_cache: dict[str, frozenset[str]] = {} + + +def _overlay_transformers_dir(tier: str) -> str | None: + """transformers source dir for a tier, located without importing it.""" + if tier != "default": + root = {"530": _VENV_T5_530_DIR, "550": _VENV_T5_550_DIR, "510": _VENV_T5_510_DIR}.get(tier) + src = os.path.join(root, "transformers") if root else None + return src if src and _safe_is_dir(Path(src)) else None + # default: the base 4.x transformers. find_spec resolves to a 5.x sidecar if one + # is already on sys.path, so skip any .venv_t5_* / llmcompressor overlay dir. + sidecars = tuple( + os.path.abspath(d) + os.sep + for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR, _VENV_LLMCOMPRESSOR_DIR) + ) + candidates = [] + try: + spec = importlib.util.find_spec("transformers") + if spec and spec.origin: + candidates.append(os.path.dirname(spec.origin)) + except Exception: + pass + candidates += [os.path.join(e, "transformers") for e in sys.path if e] + for c in candidates: + if _safe_is_dir(Path(c)) and not os.path.abspath(c).startswith(sidecars): + return c + return None + + +def _mapping_first_keys(value: ast.AST) -> set[str]: + """First keys of a dict literal, or of an OrderedDict(...)/dict(...)/.update(...) + built from 2-tuple lists and **{...} unpacking.""" + + def keys_of(node): + if isinstance(node, ast.Dict): + return list(node.keys) + if isinstance(node, (ast.List, ast.Tuple)): + return [ + el.elts[0] for el in node.elts if isinstance(el, (ast.Tuple, ast.List)) and el.elts + ] + return [] + + nodes = keys_of(value) + if isinstance(value, ast.Call): + for a in value.args: + nodes += keys_of(a) + for kw in value.keywords: # **{...} unpacking has kw.arg is None + if kw.arg is None: + nodes += keys_of(kw.value) + return {n.value for n in nodes if isinstance(n, ast.Constant) and isinstance(n.value, str)} + + +def _config_model_types(tier: str) -> frozenset[str]: + """model_type keys in a tier's CONFIG_MAPPING_NAMES (5.10 moved it to auto_mappings.py).""" + cached = _config_mapping_cache.get(tier) + if cached is not None: + return cached + tdir = _overlay_transformers_dir(tier) + if tdir is None: + return frozenset() # overlay not provisioned yet; do not cache so a later call re-reads + keys: set[str] = set() + for rel in ("models/auto/configuration_auto.py", "models/auto/auto_mappings.py"): + path = Path(tdir) / rel + if not _safe_is_file(path): + continue + try: + tree = ast.parse(path.read_text(encoding = "utf-8")) + for node in ast.walk(tree): + # direct binding, or a CONFIG_MAPPING_NAMES.update({...}) mutation + if isinstance(node, ast.Assign) and any( + isinstance(t, ast.Name) and t.id == "CONFIG_MAPPING_NAMES" for t in node.targets + ): + keys |= _mapping_first_keys(node.value) + elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Call): + fn = node.value.func + if ( + isinstance(fn, ast.Attribute) + and fn.attr == "update" + and isinstance(fn.value, ast.Name) + and fn.value.id == "CONFIG_MAPPING_NAMES" + ): + keys |= _mapping_first_keys(node.value) + except Exception: + continue + result = frozenset(keys) + _config_mapping_cache[tier] = result + return result + + +def _tier_from_config_mapping(cfg: dict) -> str | None: + """Lowest tier whose transformers ships cfg's model_type, or None if unknown.""" + model_type = cfg.get("model_type") + if not isinstance(model_type, str): + for key in _NESTED_CONFIG_KEYS: + sub = cfg.get(key) + if isinstance(sub, dict) and isinstance(sub.get("model_type"), str): + model_type = sub["model_type"] + break + if not isinstance(model_type, str): + return None + for tier in sorted(_TIER_RANK, key = _TIER_RANK.get): + if model_type in _config_model_types(tier): + return tier + return None + + # --- AutoConfig probe: general tier resolution for ambiguous models ---------- # When the cheap signals only say "needs some 5.x", parse config.json with the built-in # parser in each candidate sidecar (lowest first) instead of guessing. Generalizes beyond @@ -1212,6 +1324,14 @@ def get_transformers_tier( match, ) return tier + static = _tier_from_config_mapping(cfg) + if static is not None and static != "default": + logger.info( + "Transformers tier %s selected for %s (config mapping: model_type absent below)", + static, + model_name, + ) + return static local_tc = Path(model_name) / "tokenizer_config.json" if _safe_is_file(local_tc) and _check_tokenizer_config_needs_v5(model_name, hf_token): if not probe: @@ -1271,6 +1391,18 @@ def get_transformers_tier( return override logger.info("Transformers tier 530 selected for %s (config.json check)", model_name) return "530" + # _load_config_json (not the cache-only reader) so a config served from the hub + # cache during a transient outage still feeds the mapping resolver. + remote_cfg = _load_config_json(model_name, hf_token) + if remote_cfg is not None: + static = _tier_from_config_mapping(remote_cfg) + if static is not None and static != "default": + logger.info( + "Transformers tier %s selected for %s (config mapping: model_type absent below)", + static, + model_name, + ) + return static if _check_tokenizer_config_needs_v5(model_name, hf_token): if not probe: return "530" From 6412efd7d9ae98ea736110979cf6042b691a2249 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 11 Jul 2026 05:13:45 -0700 Subject: [PATCH 8/9] Studio: auto-detect completion masking markers, stop silent full-sequence training (#7054) * Auto-detect completion masking markers with template table fallback Studio's train_on_completions previously relied only on the hardcoded MODEL_TO_TEMPLATE_MAPPER / TEMPLATE_TO_RESPONSES_MAPPER tables and silently disabled masking when a model was not in the table, so unmapped models (LFM2-8B-A1B, DeepSeek, and others) trained on full sequences without telling the user. Several mapped templates (glm, mistral, llama, starling, zephyr, qwen3-thinking) also carried markers that mask every assistant token, which made every row drop in the post-masking filter. Both training callsites (CUDA trainer.py and MLX worker.py) now share utils.datasets.completion_masking.apply_completion_masking: - Try unsloth_zoo chat template auto-detection first; it raises loudly when the template cannot be parsed and never masks the EOS token. - gpt-oss models keep their manual markers so non-final assistant <|end|> tokens stay trained, matching current behavior. - If auto-detection raises, fall back to the template table exactly as before. - If the table also misses, emit an explicit user-visible warning that completion masking could not be applied and full-sequence training will occur, instead of a quiet log line. The >30 percent dropped-rows safety net in trainer.py now guards the auto path as well. Table consumers for inference and chat templates are unchanged. Validated against one representative tokenizer for every template in TEMPLATE_TO_RESPONSES_MAPPER plus the unmapped models: no template regresses; unit tests cover the four decision paths. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Restrict masking fallback to marker detection failures The auto branch wrapped the whole train_on_responses_only call, so a real failure while applying the masking (dataset map, tokenization) was treated as a detection miss and training silently proceeded on full sequences. Detect markers separately via get_chat_template_parts (test seam via detect_fn), then apply them with errors propagating, matching the manual path. Tokenizers with preset unsloth marker attrs skip detection and call bare so zoo reuses the stored parts. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fail the run when applying completion masking raises The helper already falls back internally on detection failures and returns applied=False on a double miss, so an exception reaching the callsites is a real failure applying the masking. Remove the callsite catches that downgraded it to full-sequence training; the run now fails visibly instead. Also use the explicit re-export alias form in utils/datasets/__init__.py for the two new names, satisfying the import-hoist source lint. * Import completion masking from its submodule The import-hoist source lint counts only real name loads, so package-level re-exports of the two new names cannot satisfy it. Import apply_completion_masking from utils.datasets.completion_masking directly at both callsites and leave utils/datasets/__init__.py untouched. * Completion masking: gpt-oss renames and MLX raw/alpaca parity Renamed or private gpt-oss checkpoints are name-detected as gpt-oss but miss the exact-name table; default them to the gpt-oss template markers instead of falling through to full-sequence training. Gate the MLX masking call on not raw_text_mode and format_type != alpaca, mirroring the CUDA path: raw/CPT text has no chat turns to mask and Alpaca-rendered text lacks the tokenizer's chat markers. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Define raw_text_mode outside the MLX feature-detect block With an older zoo lacking the append_eos config field, the masking gate referenced raw_text_mode before assignment. Hoist the assignment above the feature detection so both consumers see it. * Gate MLX masking on the formatter's resolved format format_type auto can resolve to alpaca or raw text; the masking skip checked only the requested value, so auto-detected Alpaca data got chat-template markers applied to rendered prompt text. Track the final_format returned by format_and_template_dataset and gate on it, matching the CUDA path. * Unwrap the mlx-lm TokenizerWrapper before marker checks The wrapper delegates plain reads to the wrapped HF tokenizer but hides underscore attrs, so preset unsloth markers were invisible and detection relied on the loader's call patch. Unwrap to the real tokenizer first, as the zoo MLX resolver does. * Tighten masking comments * gpt-oss: auto-detect markers first like every other template The quantized and BF16 gpt-oss checkpoints ship a chat template without the channel final header, so the pinned manual markers match nothing there and masking trained zero tokens. Auto-detection derives markers from whichever template the checkpoint ships and keeps the final terminator trained; the manual gpt-oss markers remain the detection failure fallback, including for renamed checkpoints. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/training/trainer.py | 168 ++++------ studio/backend/core/training/worker.py | 48 +-- .../backend/tests/test_completion_masking.py | 314 ++++++++++++++++++ .../utils/datasets/completion_masking.py | 144 ++++++++ 4 files changed, 556 insertions(+), 118 deletions(-) create mode 100644 studio/backend/tests/test_completion_masking.py create mode 100644 studio/backend/utils/datasets/completion_masking.py diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 958f8f4197..96d1b90b16 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -70,7 +70,7 @@ from core.inference.llama_cpp import _hf_offline_if_dns_dead from utils.models import is_vision_model, detect_audio_type from utils.models.model_config import _env_offline from utils.datasets import format_and_template_dataset -from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER +from utils.datasets.completion_masking import apply_completion_masking from utils.datasets.iterable import is_streaming_dataset as detect_streaming_dataset from utils.datasets.raw_text import prepare_raw_text_dataset, resolve_column_names from utils.paths import ( @@ -3455,8 +3455,6 @@ class UnslothTrainer: # ========== TRAIN ON RESPONSES ONLY ========== # Raw-text datasets always train on all tokens. - instruction_part = None - response_part = None is_cpt = training_args.get("is_cpt", False) train_on_responses_enabled = ( False @@ -3473,113 +3471,93 @@ class UnslothTrainer: # DeepSeek OCR handles this internally in its collator, so skip # Audio VLM handles label masking in its collator, so skip + # Markers auto-detected from the chat template first, manual table + # as fallback; gpt-oss stays on its manual markers. See + # apply_completion_masking. if ( train_on_responses_enabled and not self.is_audio_vlm and not self.is_audio and not (is_deepseek_ocr or dataset_final_format == "alpaca") ): - try: - logger.info("Configuring train on responses only...\n") + from unsloth.chat_templates import train_on_responses_only - # Template mapping for this model - model_name_lower = self.model_name.lower() + logger.info("Configuring train on responses only...\n") - if model_name_lower in MODEL_TO_TEMPLATE_MAPPER: - template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower] - logger.info(f"Detected template: {template_name}\n") + def _notify(level, message): + if level == "warning": + logger.warning(message) + else: + logger.info(f"{message}\n") - if template_name in TEMPLATE_TO_RESPONSES_MAPPER: - instruction_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name][ - "instruction" - ] - response_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["response"] + # No try/except: the helper handles detection failures and + # double misses itself, so an exception here is a real masking + # failure that must fail the run, not silently train on full + # sequences. + self.trainer, masking_applied = apply_completion_masking( + self.trainer, + self.model_name, + train_on_responses_only, + num_proc = config_args["dataset_num_proc"], + notify = _notify, + ) - logger.info(f"Instruction marker: {instruction_part[:50]}...\n") - logger.info(f"Response marker: {response_part[:50]}...\n") + if not masking_applied: + train_on_responses_enabled = False + + if masking_applied: + try: + # ── Safety net: check if all samples were filtered out ── + # train_on_responses_only masks non-response tokens with -100; a + # row becomes all -100 (Unsloth drops it) when the response + # template is not found in the formatted text. Usually a + # dataset/template mismatch (already-formatted data, or 'Train on + # completions' on data that doesn't match the model's chat + # template); only sometimes max_seq_length truncating the response + # away. Skip this len()-based check for streaming. + if detect_streaming_dataset(self.trainer.train_dataset): + logger.info("Skipping post-filter length check for streaming dataset\n") else: - logger.info( - f"No response mapping found for template: {template_name}\n" + filtered_len = len(self.trainer.train_dataset) + original_dataset_obj = ( + dataset["dataset"] if isinstance(dataset, dict) else dataset ) - train_on_responses_enabled = False - else: - logger.info(f"No template mapping found for model: {self.model_name}\n") - train_on_responses_enabled = False - - except Exception as e: - logger.warning(f"Could not configure train on responses: {e}") - train_on_responses_enabled = False - - # Apply train on responses only if we have valid parts - if ( - train_on_responses_enabled - and instruction_part - and response_part - and not self.is_audio_vlm - and not self.is_audio - and not (is_deepseek_ocr or dataset_final_format == "alpaca") - ): - try: - from unsloth.chat_templates import train_on_responses_only - - self.trainer = train_on_responses_only( - self.trainer, - instruction_part = instruction_part, - response_part = response_part, - num_proc = config_args["dataset_num_proc"], - ) - logger.info("Train on responses only configured successfully\n") - - # ── Safety net: check if all samples were filtered out ── - # train_on_responses_only masks non-response tokens with -100; - # a row becomes all -100 (and Unsloth drops it) when the response - # template is not found in the formatted text. That is usually a - # dataset/template mismatch (already-formatted data, or 'Train on - # completions' applied to data that doesn't match the model's chat - # template), and only sometimes max_seq_length truncating the - # response away. Skip this len()-based check for streaming. - if detect_streaming_dataset(self.trainer.train_dataset): - logger.info("Skipping post-filter length check for streaming dataset\n") - else: - filtered_len = len(self.trainer.train_dataset) - original_dataset_obj = ( - dataset["dataset"] if isinstance(dataset, dict) else dataset - ) - original_len = len(original_dataset_obj) - dropped = original_len - filtered_len - drop_pct = round(100 * dropped / original_len, 1) if original_len > 0 else 0 - - if filtered_len == 0 or drop_pct > 30: - max_seq = training_args.get("max_seq_length", 2048) - error_msg = ( - f"{dropped}/{original_len} samples ({drop_pct}%) were " - f"dropped after applying 'Train on completions': after " - f"masking, those rows had no trainable response tokens " - f"left. The usual cause is that this model's response " - f"template was not found in the formatted samples, so " - f"every token was masked out. That typically means the " - f"dataset is already formatted, or its structure does " - f"not match the model's chat template, so 'Train on " - f"completions' should be turned off for this dataset. " - f"Less commonly, a max_seq_length ({max_seq}) shorter " - f"than the prompt can truncate the response away; only " - f"raise it if your samples are actually longer than that." + original_len = len(original_dataset_obj) + dropped = original_len - filtered_len + drop_pct = ( + round(100 * dropped / original_len, 1) if original_len > 0 else 0 ) - logger.error(error_msg) - self._update_progress(error = error_msg, is_training = False) - return - if dropped > 0: - logger.info( - f"⚠️ {dropped}/{original_len} samples " - f"({drop_pct}%) were dropped (all labels " - f"masked). {filtered_len} samples remain.\n" - ) - logger.info(f"Post-filter dataset size: {filtered_len} samples\n") + if filtered_len == 0 or drop_pct > 30: + max_seq = training_args.get("max_seq_length", 2048) + error_msg = ( + f"{dropped}/{original_len} samples ({drop_pct}%) were " + f"dropped after applying 'Train on completions': after " + f"masking, those rows had no trainable response tokens " + f"left. The usual cause is that this model's response " + f"template was not found in the formatted samples, so " + f"every token was masked out. That typically means the " + f"dataset is already formatted, or its structure does " + f"not match the model's chat template, so 'Train on " + f"completions' should be turned off for this dataset. " + f"Less commonly, a max_seq_length ({max_seq}) shorter " + f"than the prompt can truncate the response away; only " + f"raise it if your samples are actually longer than that." + ) + logger.error(error_msg) + self._update_progress(error = error_msg, is_training = False) + return - except Exception as e: - logger.warning(f"Failed to apply train on responses only: {e}") - train_on_responses_enabled = False + if dropped > 0: + logger.info( + f"⚠️ {dropped}/{original_len} samples " + f"({drop_pct}%) were dropped (all labels " + f"masked). {filtered_len} samples remain.\n" + ) + logger.info(f"Post-filter dataset size: {filtered_len} samples\n") + + except Exception as e: + logger.warning(f"Post-masking dataset size check failed: {e}") else: if train_on_responses_enabled and is_deepseek_ocr: logger.info("Train on responses handled by DeepSeek OCR collator\n") diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 0ff4d517ed..5fb8fc2bb6 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1731,6 +1731,7 @@ def _run_mlx_training(event_queue, stop_queue, config): # sharegpt+images) and text (alpaca/sharegpt/chatml → "text" column). format_type = config.get("format_type", "") custom_format_mapping = config.get("custom_format_mapping") + dataset_final_format = "" try: from utils.datasets import format_and_template_dataset def _fmt_progress(status_message = "", **_kw): @@ -1796,6 +1797,7 @@ def _run_mlx_training(event_queue, stop_queue, config): ) if info.get("success", True): dataset = info.get("dataset", dataset) + dataset_final_format = str(info.get("final_format", "") or "").lower() if eval_dataset is not None: ev = format_and_template_dataset( eval_dataset, @@ -1894,6 +1896,9 @@ def _run_mlx_training(event_queue, stop_queue, config): eval_steps = eval_steps_val, ) + # Also gates the masking skip below, so defined outside the feature-detect block. + raw_text_mode = training_type == "Continued Pretraining" or format_type == "raw" + # Feature-detect optional fields so this PR works without the paired zoo bump. _supported_fields = getattr(MLXTrainingConfig, "__dataclass_fields__", {}) if "cast_norm_output_to_input_dtype" in _supported_fields: @@ -1907,7 +1912,6 @@ def _run_mlx_training(event_queue, stop_queue, config): if "max_grad_leaf_norm" in _supported_fields: mlx_config_kwargs["max_grad_leaf_norm"] = max_grad_leaf_norm if "append_eos" in _supported_fields: - raw_text_mode = training_type == "Continued Pretraining" or format_type == "raw" # Studio SFT formatting owns rendered examples; raw/CPT text still # needs MLX to append EOS like the CUDA raw-text path. mlx_config_kwargs["append_eos"] = bool(raw_text_mode) @@ -1928,29 +1932,27 @@ def _run_mlx_training(event_queue, stop_queue, config): _send("eval_configured") # ── 7. Apply train_on_responses_only if requested ── - if config.get("train_on_completions", False): + # Auto-detect markers from the chat template first, manual table as + # fallback. Mirror the CUDA skips: raw/CPT text has no chat turns and + # Alpaca-rendered text lacks the chat markers. Also check the resolved + # format, since format_type="auto" can land on alpaca or raw text. + if ( + config.get("train_on_completions", False) + and not raw_text_mode + and format_type != "alpaca" + and dataset_final_format not in ("alpaca", "raw_text") + ): _send("status", status_message = "Configuring response-only training...") - try: - from utils.datasets import ( - MODEL_TO_TEMPLATE_MAPPER, - TEMPLATE_TO_RESPONSES_MAPPER, - ) - - template_name = MODEL_TO_TEMPLATE_MAPPER.get(model_name.lower()) - markers = TEMPLATE_TO_RESPONSES_MAPPER.get(template_name) if template_name else None - if markers: - trainer = train_on_responses_only( - trainer, - instruction_part = markers["instruction"], - response_part = markers["response"], - ) - else: - _send( - "status", - status_message = f"train_on_completions skipped (no template for {model_name})", - ) - except Exception as e: - _send("status", status_message = f"train_on_completions failed: {e}") + # No catch: the helper handles detection failures and double misses, so + # an exception here is a real masking failure that must fail the run, + # not silently train on full sequences. + from utils.datasets.completion_masking import apply_completion_masking + trainer, _masking_applied = apply_completion_masking( + trainer, + model_name, + train_on_responses_only, + notify = lambda level, message: _send("status", status_message = message), + ) # ── 8. Setup wandb / tensorboard ── wandb_run = None diff --git a/studio/backend/tests/test_completion_masking.py b/studio/backend/tests/test_completion_masking.py new file mode 100644 index 0000000000..be0d8a69bd --- /dev/null +++ b/studio/backend/tests/test_completion_masking.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Completion-only masking policy: auto-detect first, manual table fallback. + +Covers utils.datasets.completion_masking.apply_completion_masking, shared by +the CUDA trainer (core/training/trainer.py) and the MLX worker +(core/training/worker.py): + - unmapped models use chat template auto-detection (previously masking was + silently disabled), + - gpt-oss goes auto-first too (its quantized checkpoints ship a template + the manual markers cannot match), + - an auto-detection failure falls back to the template table markers, + - a table miss after an auto failure warns and leaves the trainer unchanged. +""" + +from __future__ import annotations + +import pytest + +from utils.datasets.completion_masking import apply_completion_masking, lookup_manual_markers +from utils.datasets.model_mappings import TEMPLATE_TO_RESPONSES_MAPPER + + +class _Trainer: + """Sentinel trainer; train_fn wraps it in a new object when applied.""" + + +class _Recorder: + """Fake train_on_responses_only that records calls.""" + + def __init__(self): + self.calls = [] + + def __call__(self, trainer, **kwargs): + self.calls.append(kwargs) + wrapped = _Trainer() + wrapped.wrapped_from = trainer + return wrapped + + +def _detect_ok(processor): + return "", "" + + +def _detect_fail(processor): + raise ValueError( + "Unsloth: Could not reliably auto-detect response_part - " + "pass instruction_part and response_part." + ) + + +_AUTO = {"instruction_part": "", "response_part": ""} + + +class _Notes: + def __init__(self): + self.messages = [] + + def __call__(self, level, message): + self.messages.append((level, message)) + + def warnings(self): + return [m for level, m in self.messages if level == "warning"] + + +def test_unmapped_model_uses_auto_detection(): + # Unmapped model: the auto path applies masking (was silently disabled). + trainer = _Trainer() + train_fn = _Recorder() + notes = _Notes() + + result, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, notify = notes, detect_fn = _detect_ok + ) + + assert applied is True + assert result.wrapped_from is trainer + assert train_fn.calls == [dict(_AUTO)] # applied with the detected markers + assert notes.warnings() == [] + + +def test_mapped_model_prefers_auto_detection(): + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "unsloth/Qwen3-0.6B", train_fn, detect_fn = _detect_ok + ) + + assert applied is True + assert train_fn.calls == [dict(_AUTO)] + + +def test_gpt_oss_uses_auto_detection_first(): + # The quantized gpt-oss checkpoints ship a template without the + # <|channel|>final header, where the manual markers match nothing; auto + # derives markers from the template the checkpoint actually ships. + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "unsloth/gpt-oss-20b", train_fn, detect_fn = _detect_ok + ) + + assert applied is True + assert train_fn.calls == [dict(_AUTO)] + + +def test_gpt_oss_detection_failure_falls_back_to_manual_markers(): + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "unsloth/gpt-oss-20b", train_fn, detect_fn = _detect_fail + ) + + assert applied is True + expected = TEMPLATE_TO_RESPONSES_MAPPER["gpt-oss"] + assert train_fn.calls == [ + { + "instruction_part": expected["instruction"], + "response_part": expected["response"], + } + ] + + +def test_auto_failure_falls_back_to_template_table(): + trainer = _Trainer() + train_fn = _Recorder() + notes = _Notes() + + result, applied = apply_completion_masking( + trainer, "unsloth/Qwen3-0.6B", train_fn, notify = notes, detect_fn = _detect_fail + ) + + assert applied is True + assert result.wrapped_from is trainer + expected = TEMPLATE_TO_RESPONSES_MAPPER["qwen3"] + assert train_fn.calls == [ + { + "instruction_part": expected["instruction"], + "response_part": expected["response"], + }, + ] + assert any("falling back to the template table" in m for m in notes.warnings()) + + +def test_application_failure_propagates_not_fallback(): + # Detection succeeds; a failure while APPLYING the masking must propagate, + # never silently fall back to full-sequence training. + def train_fn(trainer, **kwargs): + raise RuntimeError("dataset map worker crashed") + + with pytest.raises(RuntimeError, match = "dataset map worker crashed"): + apply_completion_masking(_Trainer(), "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_ok) + + +def test_preset_tokenizer_markers_used_directly(): + # Preset unsloth marker attrs skip detection; zoo reuses them on a bare call. + class _Tok: + _unsloth_input_part = "" + _unsloth_output_part = "" + + trainer = _Trainer() + trainer.processing_class = _Tok() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_fail + ) + assert applied is True + assert train_fn.calls == [{}] # bare call, stored parts + + +def test_table_miss_warns_and_disables_without_crashing(): + trainer = _Trainer() + train_fn = _Recorder() + notes = _Notes() + + result, applied = apply_completion_masking( + trainer, "some-org/not-in-any-mapper", train_fn, notify = notes, detect_fn = _detect_fail + ) + + assert applied is False + assert result is trainer # unchanged: full sequence training + assert train_fn.calls == [] # detection failed; nothing applied + assert any("could not be applied" in m for m in notes.warnings()) + assert any("full sequences" in m for m in notes.warnings()) + + +def test_num_proc_forwarded_only_when_given(): + # CUDA path passes num_proc; the MLX path omits it. + train_fn = _Recorder() + apply_completion_masking( + _Trainer(), "unsloth/Qwen3-0.6B", train_fn, num_proc = 4, detect_fn = _detect_ok + ) + assert train_fn.calls == [dict(_AUTO, num_proc = 4)] + + train_fn = _Recorder() + apply_completion_masking( + _Trainer(), "unsloth/Qwen3-0.6B", train_fn, num_proc = 4, detect_fn = _detect_fail + ) + assert train_fn.calls[0]["num_proc"] == 4 + + train_fn = _Recorder() + apply_completion_masking(_Trainer(), "unsloth/Qwen3-0.6B", train_fn, detect_fn = _detect_ok) + assert train_fn.calls == [dict(_AUTO)] + + +def test_manual_fallback_failure_propagates_to_caller(): + # Errors while applying the manual fallback must propagate to the caller. + def train_fn(trainer, **kwargs): + raise RuntimeError("boom") + + with pytest.raises(RuntimeError, match = "boom"): + apply_completion_masking(_Trainer(), "unsloth/gpt-oss-20b", train_fn) + + +def test_notify_is_optional(): + train_fn = _Recorder() + _, applied = apply_completion_masking( + _Trainer(), "some-org/not-in-any-mapper", train_fn, detect_fn = _detect_fail + ) + assert applied is False + + +def test_lookup_manual_markers(): + template, instruction, response = lookup_manual_markers("unsloth/Qwen3-0.6B") + assert template == "qwen3" + assert instruction == TEMPLATE_TO_RESPONSES_MAPPER["qwen3"]["instruction"] + assert response == TEMPLATE_TO_RESPONSES_MAPPER["qwen3"]["response"] + + template, instruction, response = lookup_manual_markers("some-org/unknown") + assert (template, instruction, response) == (None, None, None) + + template, instruction, response = lookup_manual_markers(None) + assert (template, instruction, response) == (None, None, None) + + +def test_renamed_gpt_oss_gets_template_markers(): + # Name-detected as gpt-oss but not in the exact-name table: must use the + # gpt-oss markers, not fall through to full-sequence training. + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "some-org/gpt-oss-20b-sft", train_fn, detect_fn = _detect_fail + ) + assert applied is True + expected = TEMPLATE_TO_RESPONSES_MAPPER["gpt-oss"] + assert train_fn.calls == [ + { + "instruction_part": expected["instruction"], + "response_part": expected["response"], + } + ] + + +class _FakeTokenizerWrapper: + """mlx-lm TokenizerWrapper semantics: plain reads delegate to the wrapped + tokenizer, underscore attrs do not (so preset markers are hidden).""" + + def __init__(self, tokenizer): + object.__setattr__(self, "_tokenizer", tokenizer) + + def __getattr__(self, attr): + if attr.startswith("_"): + return object.__getattribute__(self, attr) + return getattr(object.__getattribute__(self, "_tokenizer"), attr) + + +_FakeTokenizerWrapper.__name__ = "TokenizerWrapper" + + +def test_mlx_tokenizer_wrapper_unwrapped_for_preset_markers(): + # Markers live on the inner HF tokenizer that the wrapper hides; the helper + # must unwrap so the preset bare-call path still fires on MLX. + class _Tok: + _unsloth_input_part = "" + _unsloth_output_part = "" + + trainer = _Trainer() + trainer.tokenizer = _FakeTokenizerWrapper(_Tok()) + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_fail + ) + assert applied is True + assert train_fn.calls == [{}] # bare call, stored parts + + +def test_mlx_tokenizer_wrapper_unwrapped_for_detection(): + # Detection must see the real tokenizer, not the wrapper, so it does not + # depend on the loader's __call__ patch. + class _Tok: + pass + + inner = _Tok() + trainer = _Trainer() + trainer.tokenizer = _FakeTokenizerWrapper(inner) + train_fn = _Recorder() + seen = [] + + def detect(processor): + seen.append(processor) + return "", "" + + _, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = detect + ) + assert applied is True + assert seen == [inner] diff --git a/studio/backend/utils/datasets/completion_masking.py b/studio/backend/utils/datasets/completion_masking.py new file mode 100644 index 0000000000..c7c4a474e3 --- /dev/null +++ b/studio/backend/utils/datasets/completion_masking.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Completion-only masking policy shared by the CUDA and MLX training paths. + +Decides how train_on_responses_only is applied for a model: chat template +auto-detection first, manual TEMPLATE_TO_RESPONSES_MAPPER markers as the +fallback. gpt-oss included: its quantized checkpoints ship a different +chat template, so only detection from the actual template is reliable. +""" + +from .model_mappings import ( + MODEL_TO_TEMPLATE_MAPPER, + TEMPLATE_TO_RESPONSES_MAPPER, + is_gpt_oss_model_name, +) + + +def lookup_manual_markers(model_name): + """Return (template_name, instruction_part, response_part) from the + manual template table, with None parts when the model or template is + not mapped.""" + template = MODEL_TO_TEMPLATE_MAPPER.get((model_name or "").lower()) + markers = TEMPLATE_TO_RESPONSES_MAPPER.get(template) if template else None + if markers: + return template, markers["instruction"], markers["response"] + return template, None, None + + +def apply_completion_masking( + trainer, + model_name, + train_fn, + num_proc = None, + notify = None, + detect_fn = None, +): + """Apply completion-only masking with auto-detection first and the manual + template table as fallback. + + Args: + trainer: The platform trainer (SFTTrainer or MLXTrainer). + model_name: Model repo id used for table lookup and the gpt-oss + renamed-checkpoint fallback. + train_fn: The platform train_on_responses_only callable. + num_proc: Forwarded to train_fn when not None (CUDA path only). + notify: Optional callback notify(level, message) with level "info" or + "warning" for user-visible progress and warnings. + detect_fn: Marker detector (tokenizer/processor) -> (instruction_part, + response_part). Defaults to unsloth_zoo's get_chat_template_parts, + which raises loudly when the template cannot be parsed. Test seam. + + Returns: + (trainer, applied): the possibly wrapped trainer and whether masking + was applied. When applied is False the trainer is unchanged and + training runs on full sequences. + + Only marker DETECTION failures trigger the table fallback. Exceptions + raised while applying the masking (dataset map, tokenization) propagate + to the caller in both the auto and manual paths, so a real failure stops + the run instead of silently changing the training objective. + """ + if notify is None: + notify = lambda level, message: None + kwargs = {} + if num_proc is not None: + kwargs["num_proc"] = num_proc + + template, instruction_part, response_part = lookup_manual_markers(model_name) + + # gpt-oss goes auto-first: quantized/BF16 checkpoints ship a channel-less + # template, so the manual markers match nothing (zero tokens trained). Auto + # derives markers from whichever template ships, and per the harmony format + # only the final terminator carries stop supervision. Renamed checkpoints + # miss the exact-name table, so give the fallback the gpt-oss markers. + if is_gpt_oss_model_name(model_name) and not (instruction_part and response_part): + markers = TEMPLATE_TO_RESPONSES_MAPPER.get("gpt-oss") + if markers: + template = "gpt-oss" + instruction_part = markers["instruction"] + response_part = markers["response"] + processor = getattr(trainer, "processing_class", None) or getattr(trainer, "tokenizer", None) + # mlx-lm TokenizerWrapper hides underscore attrs, so preset _unsloth_* + # markers are invisible through it. Unwrap to the real tokenizer (as + # zoo's MLX resolver does) before the preset check and detection. + if type(processor).__name__ == "TokenizerWrapper": + wrapped = getattr(processor, "_tokenizer", None) + if wrapped is not None: + processor = wrapped + inner = getattr(processor, "tokenizer", processor) + if hasattr(inner, "_unsloth_input_part") and hasattr(inner, "_unsloth_output_part"): + # Markers preset on the tokenizer; zoo reuses them on a bare call. + trainer = train_fn(trainer, **kwargs) + notify( + "info", + "Train on responses only configured via tokenizer preset markers", + ) + return trainer, True + auto_instruction = auto_response = None + try: + if detect_fn is None: + # Torch-backed import is fine: the MLX train_fn itself requires + # unsloth_zoo.dataset_utils, so a torch-free host cannot mask either way. + from unsloth_zoo.dataset_utils import get_chat_template_parts as detect_fn + auto_instruction, auto_response = detect_fn(processor) + except Exception as e: + notify( + "warning", + f"Auto-detection of instruction/response markers failed ({e}); " + f"falling back to the template table", + ) + if auto_instruction and auto_response: + trainer = train_fn( + trainer, + instruction_part = auto_instruction, + response_part = auto_response, + **kwargs, + ) + notify( + "info", + "Train on responses only configured via chat template auto-detection", + ) + return trainer, True + + if instruction_part and response_part: + trainer = train_fn( + trainer, + instruction_part = instruction_part, + response_part = response_part, + **kwargs, + ) + notify( + "info", + f"Train on responses only configured with template table markers ({template})", + ) + return trainer, True + + notify( + "warning", + f"'Train on completions' could not be applied for {model_name}: no " + f"auto-detected or mapped instruction/response markers. Training " + f"will run on full sequences (prompts included).", + ) + return trainer, False From 9fa6fd40e1a227a6c77b7e64d33dcfe3d5f617cd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 11 Jul 2026 08:15:43 -0700 Subject: [PATCH 9/9] scripts: refresh scan_packages allowlist baseline (#7078) New releases of huggingface-hub (1.23.0) and openai (2.45.0) shifted or added polling loops that the C2 polling/beaconing check flags, failing all three pip scan-packages shards (studio 1, hf-stack 1, extras 3 new CRITICAL findings) org-wide including on main. Regenerated with scan_packages.py --write-baseline per CI shard (same shard-to-requirements mapping and --with-deps as security-audit.yml) and merged. All entries were manually reviewed at the resolved versions: - huggingface-hub hf_api.py: create_repo 409-concurrency retry loop body changed in 1.23.0; refreshed evidence hash. The loop POSTs to the canonical Hub endpoint and retries only on a specific conflict error. Benign client retry. - openai beta/threads/runs/runs.py: create_and_poll run-status helper refactored in 2.45.0 (Assistants deprecation annotations); refreshed evidence hash. Documented polling helper against api.openai.com. - openai beta/responses/responses.py: new beta websocket client whose __aiter__ yields server events until the connection closes. New entry; standard event-stream iterator, not beaconing. - openai resources/responses/responses.py: evidence line number refreshed only, hash unchanged. The two dropped entries are the pre-refactor hashes of the same two loops above; they no longer occur at the resolved versions. Verified locally: all three shards exit 0 with 0 unsuppressed CRITICAL/HIGH (hf-stack 120, studio 151, extras 99 suppressed). --- scripts/scan_packages_baseline.json | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json index 3582517d31..27fa801a2a 100644 --- a/scripts/scan_packages_baseline.json +++ b/scripts/scan_packages_baseline.json @@ -159,8 +159,8 @@ "file": "huggingface_hub/hf_api.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L4613: while True: sha256:f764b6ca3118b23c7c0e670e77178c022a6905f825d7df6e528545fa10aae8f6", - "evidence_hash": "9c85d50c227285fa8dc69512999cbb082258cda4b299c7d0e0f69f5aff7accd4" + "evidence": "L4677: while True: sha256:04afb38843e4125d1476f3f04bdad0edf1f63f8d75ad49a713b13e4bc68612fb", + "evidence_hash": "18877a2502c862b46a5d7e33fa7c39ab4ef32da7e1b07f596fd455f4376770c6" }, { "package": "huggingface-hub", @@ -338,13 +338,21 @@ "evidence": "Env: L105: token = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\") | L150: environment_token = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\")\nNetwork: L415: http_client: httpx.Client | None = None, | L531: http_client: httpx.Client | None = None, | L649: http_client: httpx.AsyncClient | None = None, | L767: http_client: httpx.AsyncClient | None = None,", "evidence_hash": "92dbec8ccd79c1e0bc41e93cdd0bdbb091220616c6a1352873196e9dda6bd85c" }, + { + "package": "openai", + "file": "openai/resources/beta/responses/responses.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L3999: while True: sha256:df298b6eaf3416589b79f4ef283f8fb76e54d505bfda8840673f8e6419117e2e", + "evidence_hash": "10ce5cb5a7097fcff4042ddcfb4802edda60aa4b7b113c8b926a52ddb76f78c2" + }, { "package": "openai", "file": "openai/resources/beta/threads/runs/runs.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L1074: while True: sha256:ef6d59a4a10b73a5af491f10af2885b7a309fda9468eb0f9572d19558d3ceb9f", - "evidence_hash": "43c03b55fedcbc980e5e6649c3c4493729128d280cc868349ab9590908ea5f99" + "evidence": "L1053: while True: sha256:973bb1aeca2e17e022872dc343a1bf5d8fe33bfa59fe01e2f8fe875522db5bce", + "evidence_hash": "24626e4aa53047a515ead563b42c07c43f73a2c5b82978fa59f58ffc2859e19b" }, { "package": "openai", @@ -359,7 +367,7 @@ "file": "openai/resources/responses/responses.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L3803: while True: sha256:1ce0b5a388c747945cdfda1a71b77afdfd03ae840d7aa9fa62f02eb00aa5e29f", + "evidence": "L3950: while True: sha256:1ce0b5a388c747945cdfda1a71b77afdfd03ae840d7aa9fa62f02eb00aa5e29f", "evidence_hash": "6de300ebb5e6e17cb51c89cbcdf08515a44655182f0776f0908a9d1043ebbcd7" }, {