mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
feat(proxy): per-worker admission control that rejects excess requests with 503 (#39352)
* feat(proxy): reject excess per-worker requests with 503 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): drop redundant suppressions in admission middleware Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): exempt the /metrics/ redirect target from admission control Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: allowlist live Granian saturation benchmark Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): regenerate dashboard API types Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): normalize root_path for admission exemptions, validate settings, inject state Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover prometheus metric factory, lifespan scope, and prefix lookalike paths Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): queue behind pending waiters, cache admission settings parsing, log invalid limits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): simplify invalid admission settings handling Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
bd10977a9a
commit
aec083cdac
9 changed files with 940 additions and 1 deletions
5
.github/ci-coverage-allowlist.yml
vendored
5
.github/ci-coverage-allowlist.yml
vendored
|
|
@ -79,6 +79,11 @@ test_paths:
|
|||
- tests/load_tests/test_otel_load_test.py
|
||||
- tests/load_tests/test_vertex_embeddings_load_test.py
|
||||
- tests/load_tests/test_vertex_load_tests.py
|
||||
- reason: >-
|
||||
Env-gated saturation benchmark requires a live proxy and provider credentials, so it is run
|
||||
locally rather than in pull-request jobs
|
||||
paths:
|
||||
- tests/load_tests/test_granian_admission_saturation.py
|
||||
- reason: >-
|
||||
A local-only agent rig: test_a2a_completion_bridge.py needs a LangGraph server on
|
||||
localhost:2024 and test_a2a.py drives a live A2A endpoint, so neither can run in a
|
||||
|
|
|
|||
|
|
@ -2404,6 +2404,15 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
"""
|
||||
|
||||
completion_model: str | None = Field(None, description="proxy level default model for all chat completion calls")
|
||||
max_in_flight_requests_per_worker: int | None = Field(
|
||||
None, gt=0, description="maximum concurrent requests handled by each worker"
|
||||
)
|
||||
max_queued_requests_per_worker: int | None = Field(
|
||||
None, ge=0, description="maximum requests waiting for a worker slot"
|
||||
)
|
||||
admission_queue_timeout_seconds: float = Field(
|
||||
1.0, gt=0, description="maximum time a request waits for a worker slot"
|
||||
)
|
||||
plugins: list[PluginConfig] | None = Field(
|
||||
None, description="external services registered as embeddable UI plugins"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -50,6 +50,9 @@ from litellm.proxy.health_check import (
|
|||
perform_health_check,
|
||||
run_with_timeout,
|
||||
)
|
||||
from litellm.proxy.middleware.admission_control_middleware import (
|
||||
get_admission_control_stats,
|
||||
)
|
||||
from litellm.proxy.middleware.in_flight_requests_middleware import (
|
||||
get_in_flight_requests,
|
||||
)
|
||||
|
|
@ -63,6 +66,13 @@ from litellm.secret_managers.main import get_secret_bool
|
|||
#### Health ENDPOINTS ####
|
||||
|
||||
|
||||
class _HealthBacklogResponse(TypedDict):
|
||||
in_flight_requests: ReadOnly[int]
|
||||
admitted_requests: ReadOnly[int]
|
||||
queued_requests: ReadOnly[int]
|
||||
rejected_requests: ReadOnly[int]
|
||||
|
||||
|
||||
def _reject_os_environ_references(params: dict) -> None:
|
||||
"""
|
||||
Validate that the provided params do not contain any ``os.environ/``
|
||||
|
|
@ -1759,7 +1769,14 @@ async def health_backlog():
|
|||
for the event loop to get to them, adding latency before LiteLLM even starts
|
||||
its own timer.
|
||||
"""
|
||||
return {"in_flight_requests": get_in_flight_requests()}
|
||||
stats: Final = get_admission_control_stats()
|
||||
response: Final[_HealthBacklogResponse] = {
|
||||
"in_flight_requests": get_in_flight_requests(),
|
||||
"admitted_requests": stats.admitted,
|
||||
"queued_requests": stats.queued,
|
||||
"rejected_requests": stats.rejected_total,
|
||||
}
|
||||
return response
|
||||
|
||||
|
||||
@router.get(
|
||||
|
|
|
|||
315
litellm/proxy/middleware/admission_control_middleware.py
Normal file
315
litellm/proxy/middleware/admission_control_middleware.py
Normal file
|
|
@ -0,0 +1,315 @@
|
|||
import asyncio
|
||||
import os
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from functools import lru_cache
|
||||
from typing import Annotated, Final, Protocol, TypeAlias, runtime_checkable
|
||||
|
||||
from pydantic import Field, TypeAdapter, ValidationError
|
||||
from starlette.responses import JSONResponse
|
||||
from starlette.types import ASGIApp, Receive, Scope, Send
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
_EXEMPT_PATHS: Final[frozenset[str]] = frozenset(
|
||||
{
|
||||
"/health/liveliness",
|
||||
"/health/liveness",
|
||||
"/health/readiness",
|
||||
"/health/readiness/details",
|
||||
"/health/backlog",
|
||||
"/health/drain",
|
||||
"/metrics",
|
||||
"/metrics/",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AdmissionControlSettings:
|
||||
max_in_flight_requests: int
|
||||
max_queued_requests: int
|
||||
queue_timeout_seconds: float
|
||||
|
||||
|
||||
AdmissionControlSettingsGetter: TypeAlias = Callable[[], AdmissionControlSettings | None] # mutable-ok: Callable params
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AdmissionControlStats:
|
||||
admitted: int
|
||||
queued: int
|
||||
rejected_total: int
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _Gauge(Protocol):
|
||||
def inc(self, amount: float = 1) -> None: ...
|
||||
|
||||
def dec(self, amount: float = 1) -> None: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _CounterChild(Protocol):
|
||||
def inc(self, amount: float = 1) -> None: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _Counter(Protocol):
|
||||
def labels(self, reason: str) -> _CounterChild: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AdmissionControlMetrics:
|
||||
admitted_gauge: _Gauge
|
||||
queued_gauge: _Gauge
|
||||
rejected_counter: _Counter
|
||||
|
||||
|
||||
AdmissionControlMetricsFactory: TypeAlias = Callable[[], AdmissionControlMetrics | None] # mutable-ok: Callable params
|
||||
|
||||
|
||||
class AdmissionControlState:
|
||||
"""Per-process admission counters and the in-flight semaphore shared by one worker's requests."""
|
||||
|
||||
def __init__(self, metrics_factory: AdmissionControlMetricsFactory) -> None:
|
||||
self._metrics_factory = metrics_factory
|
||||
self._metrics: AdmissionControlMetrics | None = None
|
||||
self._metrics_init_attempted = False
|
||||
self._admitted = 0
|
||||
self._queued = 0
|
||||
self._rejected_total = 0
|
||||
self._semaphore: asyncio.Semaphore | None = None
|
||||
self._semaphore_loop: asyncio.AbstractEventLoop | None = None
|
||||
|
||||
def get_stats(self) -> AdmissionControlStats:
|
||||
return AdmissionControlStats(
|
||||
admitted=self._admitted,
|
||||
queued=self._queued,
|
||||
rejected_total=self._rejected_total,
|
||||
)
|
||||
|
||||
def get_semaphore(self, max_in_flight_requests: int) -> asyncio.Semaphore:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
if self._semaphore_loop is not loop:
|
||||
self._semaphore = asyncio.Semaphore(max_in_flight_requests)
|
||||
self._semaphore_loop = loop
|
||||
semaphore: Final = self._semaphore
|
||||
if semaphore is None:
|
||||
raise RuntimeError("Admission control semaphore was not initialized")
|
||||
return semaphore
|
||||
|
||||
def record_admission(self) -> None:
|
||||
self._admitted += 1
|
||||
metrics: Final = self._get_metrics()
|
||||
if metrics is not None:
|
||||
metrics.admitted_gauge.inc()
|
||||
|
||||
def record_release(self) -> None:
|
||||
self._admitted -= 1
|
||||
metrics: Final = self._get_metrics()
|
||||
if metrics is not None:
|
||||
metrics.admitted_gauge.dec()
|
||||
|
||||
def record_queue(self) -> None:
|
||||
self._queued += 1
|
||||
metrics: Final = self._get_metrics()
|
||||
if metrics is not None:
|
||||
metrics.queued_gauge.inc()
|
||||
|
||||
def record_dequeue(self) -> None:
|
||||
self._queued -= 1
|
||||
metrics: Final = self._get_metrics()
|
||||
if metrics is not None:
|
||||
metrics.queued_gauge.dec()
|
||||
|
||||
def record_rejection(self, reason: str) -> None:
|
||||
self._rejected_total += 1
|
||||
metrics: Final = self._get_metrics()
|
||||
if metrics is not None:
|
||||
metrics.rejected_counter.labels(reason=reason).inc()
|
||||
|
||||
def _get_metrics(self) -> AdmissionControlMetrics | None:
|
||||
if not self._metrics_init_attempted:
|
||||
self._metrics_init_attempted = True
|
||||
self._metrics = self._metrics_factory()
|
||||
return self._metrics
|
||||
|
||||
|
||||
class AdmissionControlMiddleware:
|
||||
def __init__(
|
||||
self,
|
||||
app: ASGIApp,
|
||||
get_settings: AdmissionControlSettingsGetter,
|
||||
state: AdmissionControlState,
|
||||
) -> None:
|
||||
self.app = app
|
||||
self.get_settings = get_settings
|
||||
self.state = state
|
||||
|
||||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
if scope["type"] != "http":
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
settings: Final = self.get_settings()
|
||||
if settings is None or _get_route_path(scope) in _EXEMPT_PATHS:
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
state: Final = self.state
|
||||
semaphore: Final = state.get_semaphore(settings.max_in_flight_requests)
|
||||
if not semaphore.locked():
|
||||
await semaphore.acquire()
|
||||
state.record_admission()
|
||||
elif state.get_stats().queued >= settings.max_queued_requests:
|
||||
state.record_rejection("queue_full")
|
||||
await _overloaded_response(state)(scope, receive, send)
|
||||
return
|
||||
else:
|
||||
state.record_queue()
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
semaphore.acquire(),
|
||||
timeout=settings.queue_timeout_seconds,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
state.record_dequeue()
|
||||
state.record_rejection("queue_timeout")
|
||||
await _overloaded_response(state)(scope, receive, send)
|
||||
return
|
||||
except asyncio.CancelledError:
|
||||
state.record_dequeue()
|
||||
raise
|
||||
state.record_dequeue()
|
||||
state.record_admission()
|
||||
|
||||
try:
|
||||
await self.app(scope, receive, send)
|
||||
finally:
|
||||
semaphore.release()
|
||||
state.record_release()
|
||||
|
||||
|
||||
def _get_route_path(scope: Scope) -> str:
|
||||
"""Strip the ASGI root_path (SERVER_ROOT_PATH) the same way Starlette does before route matching."""
|
||||
path: Final[str] = scope["path"]
|
||||
root_path: Final[str] = scope.get("root_path", "")
|
||||
if not root_path or not path.startswith(root_path):
|
||||
return path
|
||||
if path == root_path:
|
||||
return ""
|
||||
if path[len(root_path)] == "/":
|
||||
return path[len(root_path) :]
|
||||
return path
|
||||
|
||||
|
||||
def _create_gauge(gauge_type: Callable[..., object], name: str, description: str) -> _Gauge:
|
||||
metric: Final = (
|
||||
gauge_type(name, description, multiprocess_mode="livesum")
|
||||
if "PROMETHEUS_MULTIPROC_DIR" in os.environ
|
||||
else gauge_type(name, description)
|
||||
)
|
||||
if not isinstance(metric, _Gauge):
|
||||
raise TypeError("Admission gauge has an unexpected type")
|
||||
return metric
|
||||
|
||||
|
||||
def create_prometheus_admission_metrics() -> AdmissionControlMetrics | None:
|
||||
try:
|
||||
from prometheus_client import Counter, Gauge
|
||||
|
||||
return AdmissionControlMetrics(
|
||||
admitted_gauge=_create_gauge(
|
||||
Gauge,
|
||||
"litellm_admission_admitted_requests",
|
||||
"Number of requests admitted by this worker",
|
||||
),
|
||||
queued_gauge=_create_gauge(
|
||||
Gauge,
|
||||
"litellm_admission_queued_requests",
|
||||
"Number of requests queued by this worker",
|
||||
),
|
||||
rejected_counter=Counter( # mutable-ok: Prometheus requires runtime Counter construction
|
||||
"litellm_admission_rejected_requests_total",
|
||||
"Number of requests rejected by this worker",
|
||||
labelnames=("reason",),
|
||||
),
|
||||
)
|
||||
except (ImportError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
admission_control_state: Final = AdmissionControlState(create_prometheus_admission_metrics)
|
||||
|
||||
|
||||
def get_admission_control_stats() -> AdmissionControlStats:
|
||||
return admission_control_state.get_stats()
|
||||
|
||||
|
||||
_PositiveInt: TypeAlias = Annotated[int, Field(gt=0)]
|
||||
_NonNegativeInt: TypeAlias = Annotated[int, Field(ge=0)]
|
||||
_PositiveFloat: TypeAlias = Annotated[float, Field(gt=0)]
|
||||
_AdmissionControlRaw: TypeAlias = int | float | str | None
|
||||
|
||||
|
||||
def _hashable(value: object) -> _AdmissionControlRaw:
|
||||
return value if value is None or isinstance(value, (int, float, str)) else repr(value)
|
||||
|
||||
|
||||
_POSITIVE_INT_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(_PositiveInt)
|
||||
_NON_NEGATIVE_INT_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(_NonNegativeInt)
|
||||
_POSITIVE_FLOAT_ADAPTER: Final[TypeAdapter[float]] = TypeAdapter(_PositiveFloat)
|
||||
|
||||
|
||||
@lru_cache(maxsize=16)
|
||||
def _parse_admission_control_settings(
|
||||
max_in_flight_raw: _AdmissionControlRaw,
|
||||
max_queued_raw: _AdmissionControlRaw,
|
||||
queue_timeout_raw: _AdmissionControlRaw,
|
||||
) -> AdmissionControlSettings | None:
|
||||
try:
|
||||
max_in_flight: Final = _POSITIVE_INT_ADAPTER.validate_python(max_in_flight_raw)
|
||||
max_queued: Final = (
|
||||
max_in_flight if max_queued_raw is None else _NON_NEGATIVE_INT_ADAPTER.validate_python(max_queued_raw)
|
||||
)
|
||||
queue_timeout: Final = _POSITIVE_FLOAT_ADAPTER.validate_python(queue_timeout_raw)
|
||||
except ValidationError as exc:
|
||||
verbose_proxy_logger.error(
|
||||
"Ignoring invalid admission control settings, per-worker admission control is disabled: %s",
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
return AdmissionControlSettings(
|
||||
max_in_flight_requests=max_in_flight,
|
||||
max_queued_requests=max_queued,
|
||||
queue_timeout_seconds=queue_timeout,
|
||||
)
|
||||
|
||||
|
||||
def get_admission_control_settings(settings: Mapping[str, object]) -> AdmissionControlSettings | None:
|
||||
max_in_flight_raw: Final = settings.get("max_in_flight_requests_per_worker")
|
||||
if max_in_flight_raw is None:
|
||||
return None
|
||||
return _parse_admission_control_settings(
|
||||
_hashable(max_in_flight_raw),
|
||||
_hashable(settings.get("max_queued_requests_per_worker")),
|
||||
_hashable(settings.get("admission_queue_timeout_seconds", 1.0)),
|
||||
)
|
||||
|
||||
|
||||
def _overloaded_response(state: AdmissionControlState) -> JSONResponse:
|
||||
stats: Final = state.get_stats()
|
||||
return JSONResponse(
|
||||
status_code=503,
|
||||
headers={"retry-after": "1"}, # mutable-ok: Starlette expects a plain headers mapping
|
||||
content={ # mutable-ok: Starlette serializes a plain response mapping
|
||||
"error": { # mutable-ok: nested response mapping
|
||||
"message": (
|
||||
f"Worker at capacity: {stats.admitted} in-flight, {stats.queued} queued requests. Retry later."
|
||||
),
|
||||
"type": "overloaded_error",
|
||||
"code": "503",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
|
@ -583,6 +583,11 @@ try:
|
|||
except ImportError:
|
||||
build_billing_metrics_recorder = None
|
||||
shutdown_billing_metrics_recorder = None
|
||||
from litellm.proxy.middleware.admission_control_middleware import (
|
||||
AdmissionControlMiddleware,
|
||||
admission_control_state,
|
||||
get_admission_control_settings,
|
||||
)
|
||||
from litellm.proxy.middleware.in_flight_requests_middleware import (
|
||||
InFlightRequestsMiddleware,
|
||||
)
|
||||
|
|
@ -16502,6 +16507,9 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro
|
|||
{
|
||||
"max_parallel_requests": "Integer",
|
||||
"global_max_parallel_requests": "Integer",
|
||||
"max_in_flight_requests_per_worker": "Integer",
|
||||
"max_queued_requests_per_worker": "Integer",
|
||||
"admission_queue_timeout_seconds": "Float",
|
||||
"max_request_size_mb": "Integer",
|
||||
"max_batch_file_size_mb": "Integer",
|
||||
"max_file_size_mb": "Integer",
|
||||
|
|
@ -18177,6 +18185,11 @@ app.add_middleware(
|
|||
get_max_request_size_mb=lambda: general_settings.get("max_request_size_mb"),
|
||||
is_request_size_limit_enabled=lambda: premium_user is True,
|
||||
)
|
||||
app.add_middleware(
|
||||
AdmissionControlMiddleware,
|
||||
get_settings=lambda: get_admission_control_settings(general_settings),
|
||||
state=admission_control_state,
|
||||
)
|
||||
|
||||
|
||||
async def _stream_mcp_asgi_response(handle_fn, scope: dict, receive) -> "StreamingResponse":
|
||||
|
|
|
|||
150
tests/load_tests/test_granian_admission_saturation.py
Normal file
150
tests/load_tests/test_granian_admission_saturation.py
Normal file
|
|
@ -0,0 +1,150 @@
|
|||
import asyncio
|
||||
import os
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
os.environ.get("LITELLM_RUN_SATURATION_BENCHMARK") != "1",
|
||||
reason="set LITELLM_RUN_SATURATION_BENCHMARK=1 to run the saturation benchmark",
|
||||
)
|
||||
|
||||
|
||||
def _free_port() -> int:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
|
||||
listener.bind(("127.0.0.1", 0))
|
||||
return int(listener.getsockname()[1])
|
||||
|
||||
|
||||
def _percentile(values: list[float], percentile: float) -> float:
|
||||
return sorted(values)[min(int(len(values) * percentile), len(values) - 1)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_granian_admission_control_saturation(tmp_path: Path) -> None:
|
||||
fake_port: Final = _free_port()
|
||||
proxy_port: Final = _free_port()
|
||||
fake_script: Final = Path(__file__).parents[1] / "_fake_openai_endpoint_server.py"
|
||||
config_path: Final = tmp_path / "saturation_config.yaml"
|
||||
config_path.write_text(
|
||||
f"""model_list:
|
||||
- model_name: slow-endpoint
|
||||
litellm_params:
|
||||
model: openai/slow-endpoint
|
||||
api_base: http://127.0.0.1:{fake_port}/v1
|
||||
general_settings:
|
||||
master_key: sk-saturation
|
||||
max_in_flight_requests_per_worker: 8
|
||||
max_queued_requests_per_worker: 8
|
||||
admission_queue_timeout_seconds: 0.5
|
||||
"""
|
||||
)
|
||||
fake_process: Final = subprocess.Popen(
|
||||
[sys.executable, str(fake_script), "--host", "127.0.0.1", "--port", str(fake_port)],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
try:
|
||||
proxy_process: Final = subprocess.Popen(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"litellm.proxy.proxy_cli",
|
||||
"--config",
|
||||
str(config_path),
|
||||
"--run_granian",
|
||||
"--num_workers",
|
||||
"1",
|
||||
"--port",
|
||||
str(proxy_port),
|
||||
],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
try:
|
||||
async with httpx.AsyncClient(base_url=f"http://127.0.0.1:{proxy_port}") as client:
|
||||
deadline: Final = time.monotonic() + 60
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
response: Final = await client.get("/health/liveliness", timeout=2)
|
||||
if response.status_code == 200:
|
||||
break
|
||||
except httpx.HTTPError:
|
||||
pass
|
||||
await asyncio.sleep(0.25)
|
||||
else:
|
||||
raise AssertionError("Granian proxy did not become healthy")
|
||||
|
||||
liveness_latencies: Final[list[float]] = []
|
||||
stop_sampling: Final = asyncio.Event()
|
||||
|
||||
async def sample_liveness() -> None:
|
||||
while not stop_sampling.is_set():
|
||||
start: Final = time.perf_counter()
|
||||
try:
|
||||
response = await client.get("/health/liveliness", timeout=2)
|
||||
response.raise_for_status()
|
||||
liveness_latencies.append(time.perf_counter() - start)
|
||||
except httpx.HTTPError:
|
||||
pass
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
async def send_completion() -> tuple[int, float, bool]:
|
||||
start: Final = time.perf_counter()
|
||||
response = await client.post(
|
||||
"/chat/completions",
|
||||
headers={"Authorization": "Bearer sk-saturation"},
|
||||
json={
|
||||
"model": "slow-endpoint",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
},
|
||||
timeout=10,
|
||||
)
|
||||
return response.status_code, time.perf_counter() - start, "retry-after" in response.headers
|
||||
|
||||
sampler: Final = asyncio.create_task(sample_liveness())
|
||||
results: Final = await asyncio.gather(*(send_completion() for _ in range(200)))
|
||||
stop_sampling.set()
|
||||
await sampler
|
||||
|
||||
statuses: Final = [result[0] for result in results]
|
||||
latencies: Final = [result[1] for result in results]
|
||||
rejected: Final = [result for result in results if result[0] == 503]
|
||||
assert set(statuses) <= {200, 503}
|
||||
assert rejected
|
||||
assert all(result[2] for result in rejected)
|
||||
assert _percentile(latencies, 0.99) < 5
|
||||
assert liveness_latencies
|
||||
assert _percentile(liveness_latencies, 0.95) < 0.5
|
||||
|
||||
duration: Final = max(latencies)
|
||||
print(
|
||||
"\nmetric value\n"
|
||||
f"rps {len(results) / duration:.2f}\n"
|
||||
f"200 count {statuses.count(200)}\n"
|
||||
f"503 count {statuses.count(503)}\n"
|
||||
f"p50 {_percentile(latencies, 0.50):.3f}s\n"
|
||||
f"p95 {_percentile(latencies, 0.95):.3f}s\n"
|
||||
f"p99 {_percentile(latencies, 0.99):.3f}s\n"
|
||||
f"liveness p95 {_percentile(liveness_latencies, 0.95):.3f}s"
|
||||
)
|
||||
finally:
|
||||
proxy_process.terminate()
|
||||
try:
|
||||
proxy_process.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
proxy_process.kill()
|
||||
proxy_process.wait()
|
||||
finally:
|
||||
fake_process.terminate()
|
||||
try:
|
||||
fake_process.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
fake_process.kill()
|
||||
fake_process.wait()
|
||||
|
|
@ -1238,6 +1238,18 @@ def test_health_liveness_endpoint(proxy_client):
|
|||
print(f"\n/health/liveness response time: {duration_ms:.2f}ms")
|
||||
|
||||
|
||||
def test_health_backlog_includes_admission_control_stats(proxy_client):
|
||||
response = proxy_client.get("/health/backlog")
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert set(response.json()) == {
|
||||
"in_flight_requests",
|
||||
"admitted_requests",
|
||||
"queued_requests",
|
||||
"rejected_requests",
|
||||
}
|
||||
|
||||
|
||||
def test_health_readiness(proxy_client):
|
||||
"""
|
||||
Test /health/readiness endpoint.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,402 @@
|
|||
import asyncio
|
||||
import json
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
||||
|
||||
from litellm.proxy.middleware.admission_control_middleware import (
|
||||
AdmissionControlMetrics,
|
||||
AdmissionControlMiddleware,
|
||||
AdmissionControlSettings,
|
||||
AdmissionControlState,
|
||||
AdmissionControlStats,
|
||||
_parse_admission_control_settings,
|
||||
create_prometheus_admission_metrics,
|
||||
get_admission_control_settings,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def state() -> AdmissionControlState:
|
||||
return AdmissionControlState(lambda: None)
|
||||
|
||||
|
||||
async def _call(
|
||||
middleware: AdmissionControlMiddleware,
|
||||
path: str = "/",
|
||||
root_path: str = "",
|
||||
) -> tuple[Message, ...]:
|
||||
messages: Final[list[Message]] = []
|
||||
|
||||
async def receive() -> Message:
|
||||
return {"type": "http.request", "body": b"", "more_body": False}
|
||||
|
||||
async def send(message: Message) -> None:
|
||||
messages.append(message)
|
||||
|
||||
scope: Final[Scope] = {
|
||||
"type": "http",
|
||||
"path": path,
|
||||
"root_path": root_path,
|
||||
"method": "GET",
|
||||
"headers": [],
|
||||
}
|
||||
await middleware(scope, receive, send)
|
||||
return tuple(messages)
|
||||
|
||||
|
||||
def _handler_with_release(
|
||||
started: asyncio.Event,
|
||||
release: asyncio.Event,
|
||||
) -> ASGIApp:
|
||||
async def handler(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
started.set()
|
||||
await release.wait()
|
||||
await send({"type": "http.response.start", "status": 200, "headers": []})
|
||||
await send({"type": "http.response.body", "body": b"ok", "more_body": False})
|
||||
|
||||
return handler
|
||||
|
||||
|
||||
def test_is_not_base_http_middleware() -> None:
|
||||
assert not issubclass(AdmissionControlMiddleware, BaseHTTPMiddleware)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_capacity_rejects_excess_and_releases_queued_request(state: AdmissionControlState) -> None:
|
||||
started: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
middleware: Final = AdmissionControlMiddleware(
|
||||
_handler_with_release(started, release),
|
||||
lambda: AdmissionControlSettings(1, 1, 1.0),
|
||||
state,
|
||||
)
|
||||
|
||||
first: Final = asyncio.create_task(_call(middleware))
|
||||
await started.wait()
|
||||
second: Final = asyncio.create_task(_call(middleware))
|
||||
await asyncio.sleep(0)
|
||||
assert state.get_stats().queued == 1
|
||||
|
||||
third: Final = await _call(middleware)
|
||||
assert third[0]["status"] == 503
|
||||
headers: Final = dict(third[0]["headers"])
|
||||
assert headers[b"retry-after"] == b"1"
|
||||
assert headers[b"content-type"] == b"application/json"
|
||||
assert json.loads(third[1]["body"])["error"] == {
|
||||
"message": "Worker at capacity: 1 in-flight, 1 queued requests. Retry later.",
|
||||
"type": "overloaded_error",
|
||||
"code": "503",
|
||||
}
|
||||
assert state.get_stats().rejected_total == 1
|
||||
|
||||
release.set()
|
||||
assert (await first)[0]["status"] == 200
|
||||
assert (await second)[0]["status"] == 200
|
||||
assert state.get_stats() == AdmissionControlStats(0, 0, 1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_waiter_is_not_skipped_after_admission_is_released(state: AdmissionControlState) -> None:
|
||||
started: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
third_trigger: Final = asyncio.Event()
|
||||
middleware: Final = AdmissionControlMiddleware(
|
||||
_handler_with_release(started, release),
|
||||
lambda: AdmissionControlSettings(1, 2, 1.0),
|
||||
state,
|
||||
)
|
||||
|
||||
first: Final = asyncio.create_task(_call(middleware))
|
||||
await started.wait()
|
||||
second: Final = asyncio.create_task(_call(middleware))
|
||||
await asyncio.sleep(0)
|
||||
|
||||
async def call_third() -> tuple[Message, ...]:
|
||||
await third_trigger.wait()
|
||||
return await _call(middleware)
|
||||
|
||||
third: Final = asyncio.create_task(call_third())
|
||||
await asyncio.sleep(0)
|
||||
release.set()
|
||||
third_trigger.set()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert state.get_stats().queued == 2
|
||||
await asyncio.gather(first, second, third)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_queue_timeout_rejects_and_decrements_queue(state: AdmissionControlState) -> None:
|
||||
started: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
middleware: Final = AdmissionControlMiddleware(
|
||||
_handler_with_release(started, release),
|
||||
lambda: AdmissionControlSettings(1, 1, 0.05),
|
||||
state,
|
||||
)
|
||||
|
||||
first: Final = asyncio.create_task(_call(middleware))
|
||||
await started.wait()
|
||||
start_time: Final = asyncio.get_running_loop().time()
|
||||
second: Final = await _call(middleware)
|
||||
elapsed: Final = asyncio.get_running_loop().time() - start_time
|
||||
|
||||
assert second[0]["status"] == 503
|
||||
assert elapsed < 0.5
|
||||
assert state.get_stats().queued == 0
|
||||
assert state.get_stats().rejected_total == 1
|
||||
release.set()
|
||||
await first
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("root_path", "probe_path"),
|
||||
(
|
||||
("", "/health/liveliness"),
|
||||
("/proxy", "/proxy/health/liveliness"),
|
||||
("/proxy", "/proxy/metrics"),
|
||||
),
|
||||
)
|
||||
async def test_exempt_path_passes_through_when_saturated(
|
||||
state: AdmissionControlState,
|
||||
root_path: str,
|
||||
probe_path: str,
|
||||
) -> None:
|
||||
started: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
|
||||
async def handler(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
if scope["path"] == "/":
|
||||
started.set()
|
||||
await release.wait()
|
||||
await send({"type": "http.response.start", "status": 200, "headers": []})
|
||||
await send({"type": "http.response.body", "body": b"ok", "more_body": False})
|
||||
|
||||
middleware: Final = AdmissionControlMiddleware(handler, lambda: AdmissionControlSettings(1, 0, 1.0), state)
|
||||
|
||||
first: Final = asyncio.create_task(_call(middleware))
|
||||
await started.wait()
|
||||
health: Final = await _call(middleware, probe_path, root_path)
|
||||
assert health[0]["status"] == 200
|
||||
blocked: Final = await _call(middleware, "/proxy/v1/chat/completions", root_path)
|
||||
assert blocked[0]["status"] == 503
|
||||
lookalike: Final = await _call(middleware, "/proxyhealth/liveliness", "/proxy")
|
||||
assert lookalike[0]["status"] == 503
|
||||
release.set()
|
||||
await first
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_http_scope_passes_through_when_saturated(state: AdmissionControlState) -> None:
|
||||
seen: Final[list[str]] = []
|
||||
|
||||
async def handler(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
seen.append(scope["type"])
|
||||
|
||||
middleware: Final = AdmissionControlMiddleware(handler, lambda: AdmissionControlSettings(1, 0, 1.0), state)
|
||||
state.record_admission()
|
||||
|
||||
async def receive() -> Message:
|
||||
return {"type": "lifespan.startup"}
|
||||
|
||||
async def send(message: Message) -> None:
|
||||
return None
|
||||
|
||||
await middleware({"type": "lifespan"}, receive, send)
|
||||
assert seen == ["lifespan"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_none_settings_does_not_limit_concurrency() -> None:
|
||||
active: Final = [0]
|
||||
peak: Final = [0]
|
||||
all_started: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
|
||||
async def handler(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
active[0] += 1
|
||||
peak[0] = max(peak[0], active[0])
|
||||
if active[0] == 3:
|
||||
all_started.set()
|
||||
await release.wait()
|
||||
active[0] -= 1
|
||||
await send({"type": "http.response.start", "status": 200, "headers": []})
|
||||
await send({"type": "http.response.body", "body": b"ok", "more_body": False})
|
||||
|
||||
middleware: Final = AdmissionControlMiddleware(handler, lambda: None, AdmissionControlState(lambda: None))
|
||||
requests: Final = tuple(asyncio.create_task(_call(middleware)) for _ in range(3))
|
||||
await all_started.wait()
|
||||
assert peak[0] == 3
|
||||
release.set()
|
||||
results: Final = await asyncio.gather(*requests)
|
||||
assert tuple(result[0]["status"] for result in results) == (200, 200, 200)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelling_queued_request_does_not_leak_counter(state: AdmissionControlState) -> None:
|
||||
started: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
middleware: Final = AdmissionControlMiddleware(
|
||||
_handler_with_release(started, release),
|
||||
lambda: AdmissionControlSettings(1, 1, 1.0),
|
||||
state,
|
||||
)
|
||||
|
||||
first: Final = asyncio.create_task(_call(middleware))
|
||||
await started.wait()
|
||||
queued: Final = asyncio.create_task(_call(middleware))
|
||||
await asyncio.sleep(0)
|
||||
queued.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await queued
|
||||
assert state.get_stats().queued == 0
|
||||
release.set()
|
||||
await first
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_response_holds_admission_until_final_body(state: AdmissionControlState) -> None:
|
||||
first_chunk_sent: Final = asyncio.Event()
|
||||
finish_stream: Final = asyncio.Event()
|
||||
|
||||
async def handler(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
await send({"type": "http.response.start", "status": 200, "headers": []})
|
||||
await send({"type": "http.response.body", "body": b"first", "more_body": True})
|
||||
first_chunk_sent.set()
|
||||
await finish_stream.wait()
|
||||
await send({"type": "http.response.body", "body": b"last", "more_body": False})
|
||||
|
||||
middleware: Final = AdmissionControlMiddleware(
|
||||
handler,
|
||||
lambda: AdmissionControlSettings(1, 1, 1.0),
|
||||
state,
|
||||
)
|
||||
first: Final = asyncio.create_task(_call(middleware))
|
||||
await first_chunk_sent.wait()
|
||||
second: Final = asyncio.create_task(_call(middleware))
|
||||
await asyncio.sleep(0)
|
||||
assert not second.done()
|
||||
assert state.get_stats().queued == 1
|
||||
finish_stream.set()
|
||||
assert (await first)[0]["status"] == 200
|
||||
assert (await second)[0]["status"] == 200
|
||||
assert state.get_stats().admitted == 0
|
||||
assert state.get_stats().queued == 0
|
||||
|
||||
|
||||
class _FakeGauge:
|
||||
def __init__(self) -> None:
|
||||
self.value = 0.0
|
||||
|
||||
def inc(self, amount: float = 1) -> None:
|
||||
self.value += amount
|
||||
|
||||
def dec(self, amount: float = 1) -> None:
|
||||
self.value -= amount
|
||||
|
||||
|
||||
class _FakeCounter:
|
||||
def __init__(self) -> None:
|
||||
self.by_reason: Final[dict[str, _FakeGauge]] = {}
|
||||
|
||||
def labels(self, reason: str) -> _FakeGauge:
|
||||
return self.by_reason.setdefault(reason, _FakeGauge())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metrics_track_admitted_queued_and_rejected() -> None:
|
||||
admitted: Final = _FakeGauge()
|
||||
queued: Final = _FakeGauge()
|
||||
rejected: Final = _FakeCounter()
|
||||
state: Final = AdmissionControlState(
|
||||
lambda: AdmissionControlMetrics(admitted_gauge=admitted, queued_gauge=queued, rejected_counter=rejected)
|
||||
)
|
||||
started: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
middleware: Final = AdmissionControlMiddleware(
|
||||
_handler_with_release(started, release),
|
||||
lambda: AdmissionControlSettings(1, 1, 0.05),
|
||||
state,
|
||||
)
|
||||
|
||||
first: Final = asyncio.create_task(_call(middleware))
|
||||
await started.wait()
|
||||
second: Final = asyncio.create_task(_call(middleware))
|
||||
await asyncio.sleep(0)
|
||||
assert (admitted.value, queued.value) == (1.0, 1.0)
|
||||
await _call(middleware)
|
||||
assert rejected.by_reason["queue_full"].value == 1.0
|
||||
await second
|
||||
assert rejected.by_reason["queue_timeout"].value == 1.0
|
||||
release.set()
|
||||
await first
|
||||
assert (admitted.value, queued.value) == (0.0, 0.0)
|
||||
|
||||
|
||||
def test_create_prometheus_admission_metrics_registers_named_metrics() -> None:
|
||||
from prometheus_client import REGISTRY
|
||||
|
||||
metrics: Final = create_prometheus_admission_metrics()
|
||||
if metrics is not None:
|
||||
metrics.admitted_gauge.inc()
|
||||
metrics.queued_gauge.inc()
|
||||
metrics.rejected_counter.labels(reason="queue_full").inc()
|
||||
assert REGISTRY.get_sample_value("litellm_admission_admitted_requests") == 1.0
|
||||
assert REGISTRY.get_sample_value("litellm_admission_queued_requests") == 1.0
|
||||
assert REGISTRY.get_sample_value("litellm_admission_rejected_requests_total", {"reason": "queue_full"}) is not None
|
||||
assert create_prometheus_admission_metrics() is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("settings", "expected"),
|
||||
(
|
||||
({}, None),
|
||||
({"max_in_flight_requests_per_worker": None}, None),
|
||||
({"max_in_flight_requests_per_worker": 0}, None),
|
||||
({"max_in_flight_requests_per_worker": "many"}, None),
|
||||
({"max_in_flight_requests_per_worker": 3, "max_queued_requests_per_worker": -1}, None),
|
||||
({"max_in_flight_requests_per_worker": 3, "admission_queue_timeout_seconds": 0}, None),
|
||||
({"max_in_flight_requests_per_worker": 3, "admission_queue_timeout_seconds": -0.5}, None),
|
||||
(
|
||||
{"max_in_flight_requests_per_worker": 3, "max_queued_requests_per_worker": 0},
|
||||
AdmissionControlSettings(3, 0, 1.0),
|
||||
),
|
||||
(
|
||||
{"max_in_flight_requests_per_worker": 3},
|
||||
AdmissionControlSettings(3, 3, 1.0),
|
||||
),
|
||||
(
|
||||
{
|
||||
"max_in_flight_requests_per_worker": 3,
|
||||
"max_queued_requests_per_worker": 5,
|
||||
"admission_queue_timeout_seconds": 0.25,
|
||||
},
|
||||
AdmissionControlSettings(3, 5, 0.25),
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_get_admission_control_settings(
|
||||
settings: dict[str, object],
|
||||
expected: AdmissionControlSettings | None,
|
||||
) -> None:
|
||||
assert get_admission_control_settings(settings) == expected
|
||||
|
||||
|
||||
def test_invalid_admission_control_settings_logs_once(caplog: pytest.LogCaptureFixture) -> None:
|
||||
_parse_admission_control_settings.cache_clear()
|
||||
caplog.set_level("ERROR")
|
||||
settings: Final = {"max_in_flight_requests_per_worker": [1]}
|
||||
|
||||
assert get_admission_control_settings(settings) is None
|
||||
assert get_admission_control_settings(settings) is None
|
||||
|
||||
messages: Final = tuple(
|
||||
record.message
|
||||
for record in caplog.records
|
||||
if record.message.startswith("Ignoring invalid admission control settings")
|
||||
)
|
||||
assert len(messages) == 1
|
||||
16
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
16
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -25501,6 +25501,12 @@ export interface components {
|
|||
* @description Documents all the fields supported by `general_settings` in config.yaml
|
||||
*/
|
||||
ConfigGeneralSettings: {
|
||||
/**
|
||||
* Admission Queue Timeout Seconds
|
||||
* @description maximum time a request waits for a worker slot
|
||||
* @default 1
|
||||
*/
|
||||
admission_queue_timeout_seconds: number;
|
||||
/**
|
||||
* Alert To Webhook Url
|
||||
* @description Mapping of alert type to webhook url. e.g. `alert_to_webhook_url: {'budget_alerts': 'https://nothooks.slack.com/services/T00000000/B00000000/XXXXXXXXXXXXXXXXXXXXXXXX'}`
|
||||
|
|
@ -25709,11 +25715,21 @@ export interface components {
|
|||
* @description max file size in MB for /v1/files uploads, for any purpose, if a file is larger than this size it will be rejected before being forwarded to the provider
|
||||
*/
|
||||
max_file_size_mb?: number | null;
|
||||
/**
|
||||
* Max In Flight Requests Per Worker
|
||||
* @description maximum concurrent requests handled by each worker
|
||||
*/
|
||||
max_in_flight_requests_per_worker?: number | null;
|
||||
/**
|
||||
* Max Parallel Requests
|
||||
* @description maximum parallel requests for each api key
|
||||
*/
|
||||
max_parallel_requests?: number | null;
|
||||
/**
|
||||
* Max Queued Requests Per Worker
|
||||
* @description maximum requests waiting for a worker slot
|
||||
*/
|
||||
max_queued_requests_per_worker?: number | null;
|
||||
/**
|
||||
* Max Request Size Mb
|
||||
* @description max request size in MB, if a request is larger than this size it will be rejected
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue