Merge remote-tracking branch 'origin/main' into ig_merge

This commit is contained in:
Daniel Han 2026-07-11 15:15:54 +00:00
commit 57f08ebfb2
21 changed files with 4722 additions and 242 deletions

View file

@ -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

View file

@ -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

View file

@ -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", {

View file

@ -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"
},
{

View file

@ -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()

View file

@ -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/<pid>/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:

View file

@ -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")

View file

@ -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

File diff suppressed because it is too large Load diff

View file

@ -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 "<INS>", "<RES>"
def _detect_fail(processor):
raise ValueError(
"Unsloth: Could not reliably auto-detect response_part - "
"pass instruction_part and response_part."
)
_AUTO = {"instruction_part": "<INS>", "response_part": "<RES>"}
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 = "<I>"
_unsloth_output_part = "<O>"
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 = "<I>"
_unsloth_output_part = "<O>"
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 "<INS>", "<RES>"
_, applied = apply_completion_masking(
trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = detect
)
assert applied is True
assert seen == [inner]

View file

@ -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

View file

@ -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())

View file

@ -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

View file

@ -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

File diff suppressed because it is too large Load diff

View file

@ -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

View file

@ -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
@ -173,6 +175,7 @@ _TRANSFORMERS_530_ARCHITECTURES: set[str] = {
"Qwen3MoeForCausalLM",
"Qwen3NextForCausalLM",
"Glm4MoeLiteForCausalLM",
"Lfm2MoeForCausalLM",
"Lfm2VlForConditionalGeneration",
}
_TRANSFORMERS_530_MODEL_TYPES: set[str] = {
@ -183,6 +186,7 @@ _TRANSFORMERS_530_MODEL_TYPES: set[str] = {
"qwen3_moe",
"qwen3_next",
"glm4_moe_lite",
"lfm2_moe",
"lfm2_vl",
}
@ -870,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
@ -1210,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:
@ -1269,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"

View file

@ -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,

View file

@ -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("<s>"), "<s>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("<s>"), "Hello world") == "Hello world"

View file

@ -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

View file

@ -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