Merge remote-tracking branch 'origin/main' into ig_merge
This commit is contained in:
commit
57f08ebfb2
21 changed files with 4722 additions and 242 deletions
44
.github/workflows/studio-inference-smoke.yml
vendored
44
.github/workflows/studio-inference-smoke.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
44
.github/workflows/studio-mac-inference-smoke.yml
vendored
44
.github/workflows/studio-mac-inference-smoke.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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", {
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
},
|
||||
{
|
||||
|
|
|
|||
368
studio/backend/core/inference/llama_admission.py
Normal file
368
studio/backend/core/inference/llama_admission.py
Normal 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()
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
314
studio/backend/tests/test_completion_masking.py
Normal file
314
studio/backend/tests/test_completion_masking.py
Normal 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]
|
||||
|
|
@ -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
|
||||
|
|
|
|||
320
studio/backend/tests/test_llama_admission.py
Normal file
320
studio/backend/tests/test_llama_admission.py
Normal 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())
|
||||
|
|
@ -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
|
||||
|
|
@ -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
144
studio/backend/utils/datasets/completion_masking.py
Normal file
144
studio/backend/utils/datasets/completion_masking.py
Normal 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
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
46
tests/python/test_remove_special_tokens_no_bos.py
Normal file
46
tests/python/test_remove_special_tokens_no_bos.py
Normal 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"
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue