mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-10 08:28:25 +00:00
perf(router): read cooldown state only for the request's candidate deployments
Routing used to read the cooldown key of every deployment on the router on each request, so the Redis MGET and the in-memory scan grew with the size of the whole model list. The cooldown reads on the routing path now take the ids of the request's own deployments, and the proxy's routing prefetch arms the cooldown keys of the route's team-aware candidates instead of every id. The no-deployments error paths keep reading every deployment, so the cooldown_list in the 429 message is unchanged. Resolves LIT-9861
This commit is contained in:
yucheng-berri
parent
8d2bf426f5
commit
55726b91c8
8 changed files with 547 additions and 30 deletions
|
|
@ -2164,7 +2164,12 @@ class Router:
|
|||
)
|
||||
deployments: Final = self.get_model_list(model_name=model)
|
||||
if deployments:
|
||||
RoutingPrefetch.arm(self, usage_selector, deployments)
|
||||
RoutingPrefetch.arm(
|
||||
self,
|
||||
usage_selector,
|
||||
deployments,
|
||||
self.get_candidate_model_ids_for_route(model=model, team_id=get_request_team_id(request_kwargs)),
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # a prefetch is an optimisation, never a reason to fail the request
|
||||
verbose_router_logger.debug(
|
||||
"routing read prefetch not armed for %s: %s", _without_line_breaks(model), _without_line_breaks(e)
|
||||
|
|
@ -8245,17 +8250,15 @@ class Router:
|
|||
prev_excluded: Final = set(meta.get("_failover_excluded_ids") or [])
|
||||
excluded: Final = prev_excluded | {failed_id}
|
||||
|
||||
all_ids: Final = {
|
||||
(d.get("model_info") or {}).get("id")
|
||||
for d in all_deployments
|
||||
if (d.get("model_info") or {}).get("id") is not None
|
||||
}
|
||||
all_ids: Final = self._deployment_ids(all_deployments)
|
||||
# Only consider deployments that are currently healthy (not in cooldown).
|
||||
# Using all_ids here would cause a wasteful run_async_fallback invocation
|
||||
# that fails with RouterRateLimitError whenever the "remaining" entries
|
||||
# are all in cooldown — the inner async_get_healthy_deployments call
|
||||
# would find an empty list and raise immediately.
|
||||
cooldown_ids = set(await async_get_cooldown_deployments(litellm_router_instance=self, parent_otel_span=None))
|
||||
cooldown_ids = set(
|
||||
await async_get_cooldown_deployments(litellm_router_instance=self, parent_otel_span=None, model_ids=all_ids)
|
||||
)
|
||||
remaining: Final = (all_ids - cooldown_ids) - excluded
|
||||
if not remaining:
|
||||
return None
|
||||
|
|
@ -9712,7 +9715,9 @@ class Router:
|
|||
pass
|
||||
|
||||
unhealthy_deployments: Final = await async_get_cooldown_deployments(
|
||||
litellm_router_instance=self, parent_otel_span=parent_otel_span
|
||||
litellm_router_instance=self,
|
||||
parent_otel_span=parent_otel_span,
|
||||
model_ids=self._deployment_ids(_all_deployments),
|
||||
)
|
||||
# Convert to set for O(1) lookup instead of O(n)
|
||||
unhealthy_deployments_set: Final = set(unhealthy_deployments)
|
||||
|
|
@ -14643,13 +14648,17 @@ class Router:
|
|||
health_check_probe=health_check_probe,
|
||||
)
|
||||
|
||||
request_deployment_ids: Final = self._deployment_ids(request_deployments)
|
||||
routing_read_batch: Final = RoutingReadBatch.active()
|
||||
cooldown_deployments: Final = (
|
||||
await async_get_cooldown_deployments(litellm_router_instance=self, parent_otel_span=parent_otel_span)
|
||||
await async_get_cooldown_deployments(
|
||||
litellm_router_instance=self, parent_otel_span=parent_otel_span, model_ids=request_deployment_ids
|
||||
)
|
||||
if routing_read_batch is None
|
||||
else await routing_read_batch.async_get_cooldown_deployments(
|
||||
litellm_router_instance=self,
|
||||
healthy_deployments=healthy_deployments,
|
||||
model_ids=request_deployment_ids,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
)
|
||||
|
|
@ -15918,7 +15927,9 @@ class Router:
|
|||
)
|
||||
|
||||
cooldown_deployments: Final = get_cooldown_deployments(
|
||||
litellm_router_instance=self, parent_otel_span=parent_otel_span
|
||||
litellm_router_instance=self,
|
||||
parent_otel_span=parent_otel_span,
|
||||
model_ids=self._deployment_ids(request_deployments),
|
||||
)
|
||||
_pre_cooldown_deployments: Final = healthy_deployments
|
||||
healthy_deployments = self._filter_cooldown_deployments(
|
||||
|
|
@ -16166,7 +16177,9 @@ class Router:
|
|||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
cooldown_deployments: Final = get_cooldown_deployments(
|
||||
litellm_router_instance=self, parent_otel_span=parent_otel_span
|
||||
litellm_router_instance=self,
|
||||
parent_otel_span=parent_otel_span,
|
||||
model_ids=self._deployment_ids(request_deployments),
|
||||
)
|
||||
pass_through_deployments = self._filter_cooldown_deployments(
|
||||
healthy_deployments=pass_through_deployments,
|
||||
|
|
|
|||
|
|
@ -160,7 +160,7 @@ class CooldownCache:
|
|||
return cooldown_cache_value
|
||||
|
||||
async def async_get_active_cooldowns(
|
||||
self, model_ids: list[str], parent_otel_span: Span | None
|
||||
self, model_ids: Sequence[str], parent_otel_span: Span | None
|
||||
) -> list[tuple[str, CooldownCacheValue]]:
|
||||
# Generate the keys for the deployments
|
||||
keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids]
|
||||
|
|
@ -172,7 +172,7 @@ class CooldownCache:
|
|||
return self.active_cooldowns_from_results(model_ids, results)
|
||||
|
||||
def active_cooldowns_from_results(
|
||||
self, model_ids: list[str], results: Sequence[object] | None
|
||||
self, model_ids: Sequence[str], results: Sequence[object] | None
|
||||
) -> list[tuple[str, CooldownCacheValue]]:
|
||||
"""The cooldowns still active in a `cooldown_store` batch read of `get_cooldown_cache_key(model_id)` per id."""
|
||||
active_cooldowns: Final[list[tuple[str, CooldownCacheValue]]] = []
|
||||
|
|
@ -191,7 +191,7 @@ class CooldownCache:
|
|||
return active_cooldowns
|
||||
|
||||
def get_active_cooldowns(
|
||||
self, model_ids: list[str], parent_otel_span: Span | None
|
||||
self, model_ids: Sequence[str], parent_otel_span: Span | None
|
||||
) -> list[tuple[str, CooldownCacheValue]]:
|
||||
# Generate the keys for the deployments
|
||||
keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids]
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ Router cooldown handlers
|
|||
import asyncio
|
||||
import math
|
||||
from collections.abc import Collection, Mapping, Sequence
|
||||
from collections.abc import Set as AbstractSet
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
|
@ -574,16 +575,21 @@ def set_cooldown_deployments(
|
|||
_set_cooldown_deployments = set_cooldown_deployments
|
||||
|
||||
|
||||
def _cooldown_read_ids(litellm_router_instance: LitellmRouter, model_ids: AbstractSet[str] | None) -> Sequence[str]:
|
||||
"""The deployments a cooldown lookup reads, every deployment on the router when no ids are given"""
|
||||
return litellm_router_instance.get_model_ids() if model_ids is None else tuple(model_ids)
|
||||
|
||||
|
||||
async def async_get_cooldown_deployments(
|
||||
litellm_router_instance: LitellmRouter,
|
||||
parent_otel_span: Span | None,
|
||||
model_ids: AbstractSet[str] | None = None,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Async implementation of 'get_cooldown_deployments'
|
||||
"""
|
||||
model_ids: Final = litellm_router_instance.get_model_ids()
|
||||
cooldown_models: Final = await litellm_router_instance.cooldown_cache.async_get_active_cooldowns(
|
||||
model_ids=model_ids,
|
||||
model_ids=_cooldown_read_ids(litellm_router_instance, model_ids),
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
|
||||
|
|
@ -622,7 +628,11 @@ async def async_get_cooldown_deployments_with_debug_info(
|
|||
_async_get_cooldown_deployments_with_debug_info = async_get_cooldown_deployments_with_debug_info
|
||||
|
||||
|
||||
def get_cooldown_deployments(litellm_router_instance: LitellmRouter, parent_otel_span: Span | None) -> list[str]:
|
||||
def get_cooldown_deployments(
|
||||
litellm_router_instance: LitellmRouter,
|
||||
parent_otel_span: Span | None,
|
||||
model_ids: AbstractSet[str] | None = None,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Get the list of models being cooled down for this minute
|
||||
"""
|
||||
|
|
@ -631,10 +641,8 @@ def get_cooldown_deployments(litellm_router_instance: LitellmRouter, parent_otel
|
|||
# ----------------------
|
||||
# Return cooldown models
|
||||
# ----------------------
|
||||
model_ids: Final = litellm_router_instance.get_model_ids()
|
||||
|
||||
cooldown_models: Final = litellm_router_instance.cooldown_cache.get_active_cooldowns(
|
||||
model_ids=model_ids, parent_otel_span=parent_otel_span
|
||||
model_ids=_cooldown_read_ids(litellm_router_instance, model_ids), parent_otel_span=parent_otel_span
|
||||
)
|
||||
|
||||
cached_value_deployment_ids = []
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ the usage slice to the strategy, so selection does not read again.
|
|||
import asyncio
|
||||
import itertools
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from collections.abc import Set as AbstractSet
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass
|
||||
|
|
@ -97,14 +98,13 @@ class RoutingPrefetch:
|
|||
litellm_router_instance: "Router",
|
||||
usage_selector: LowestTPMLoggingHandler_v2 | None,
|
||||
deployments: list,
|
||||
cooldown_model_ids: AbstractSet[str],
|
||||
) -> None:
|
||||
request: Final = active_request_redis_batches()
|
||||
redis_cache: Final = litellm_router_instance.cache.redis_cache
|
||||
if request is None or redis_cache is None or _PREFETCH_SLOT in request.prefetched:
|
||||
return
|
||||
cooldown_keys: Final = tuple(
|
||||
CooldownCache.get_cooldown_cache_key(model_id) for model_id in litellm_router_instance.get_model_ids()
|
||||
)
|
||||
cooldown_keys: Final = tuple(CooldownCache.get_cooldown_cache_key(model_id) for model_id in cooldown_model_ids)
|
||||
usage_keys: Final = (
|
||||
() if usage_selector is None else tuple(itertools.chain(*usage_selector.usage_counter_keys(deployments)))
|
||||
)
|
||||
|
|
@ -188,14 +188,15 @@ class RoutingReadBatch:
|
|||
self,
|
||||
litellm_router_instance: "Router",
|
||||
healthy_deployments: list,
|
||||
model_ids: AbstractSet[str],
|
||||
parent_otel_span: "Span | None",
|
||||
) -> list[str]:
|
||||
"""
|
||||
`_async_get_cooldown_deployments`, with the strategy's tpm/rpm counters for
|
||||
`_async_get_cooldown_deployments` for `model_ids`, with the strategy's tpm/rpm counters for
|
||||
`healthy_deployments` fetched in the same MGET and kept as `prefetched_usage`.
|
||||
"""
|
||||
model_ids: Final = litellm_router_instance.get_model_ids()
|
||||
cooldown_keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids]
|
||||
read_ids: Final = tuple(model_ids)
|
||||
cooldown_keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in read_ids]
|
||||
selector: Final = self.usage_selector
|
||||
usage_keys: Final = (
|
||||
() if selector is None else tuple(itertools.chain(*selector.usage_counter_keys(healthy_deployments)))
|
||||
|
|
@ -217,7 +218,7 @@ class RoutingReadBatch:
|
|||
)
|
||||
|
||||
cooldown_models: Final = litellm_router_instance.cooldown_cache.active_cooldowns_from_results(
|
||||
model_ids, cooldown_results
|
||||
read_ids, cooldown_results
|
||||
)
|
||||
verbose_router_logger.debug("retrieve cooldown models: %s", cooldown_models)
|
||||
return [model_id for model_id, _ in cooldown_models]
|
||||
|
|
|
|||
|
|
@ -2,19 +2,24 @@ import os
|
|||
import shutil
|
||||
import signal
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from typing import Final, TextIO
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
from redis import Redis
|
||||
from redis.exceptions import ConnectionError as RedisConnectionError
|
||||
|
||||
from tests.integration._support.ports import allocate_port
|
||||
|
||||
_MONITOR_RECORD: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
@dataclass
|
||||
class OwnedRedis:
|
||||
|
|
@ -186,3 +191,55 @@ def owned_redis(
|
|||
finally:
|
||||
if server.process is not None:
|
||||
server.stop()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RedisCommands:
|
||||
queue: SimpleQueue[str]
|
||||
|
||||
def drain(self) -> tuple[tuple[str, ...], ...]:
|
||||
"""Every command MONITOR saw since the last drain, split on spaces into its name and arguments"""
|
||||
lines: Final = tuple(self.queue.get_nowait() for _ in range(self.queue.qsize()))
|
||||
return tuple(tuple(line.split(" ")) for line in lines)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def redis_commands(host: str, port: int) -> Iterator[RedisCommands]:
|
||||
"""Records every command the Redis server at `host:port` executes while the block runs"""
|
||||
commands: Final = RedisCommands(SimpleQueue())
|
||||
started: Final = threading.Event()
|
||||
armed: Final = threading.Event()
|
||||
ready_marker: Final = f"monitor-ready-{uuid.uuid4().hex}"
|
||||
stop_marker: Final = f"monitor-stop-{uuid.uuid4().hex}"
|
||||
|
||||
def capture() -> None:
|
||||
# No socket timeout: a timeout raised inside `listen()` ends the generator, so capture would stop at the
|
||||
# first quiet second. The stop marker sent on exit is what wakes and ends this read.
|
||||
with Redis(host=host, port=port, decode_responses=True) as client:
|
||||
with client.monitor() as monitor:
|
||||
started.set()
|
||||
for raw in monitor.listen():
|
||||
record: Final = _MONITOR_RECORD.validate_python(raw)
|
||||
command: Final = record.get("command")
|
||||
if not isinstance(command, str):
|
||||
continue
|
||||
if stop_marker in command:
|
||||
return
|
||||
if ready_marker in command:
|
||||
armed.set()
|
||||
continue
|
||||
commands.queue.put(command)
|
||||
|
||||
thread: Final = threading.Thread(target=capture, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
assert started.wait(timeout=5), "Redis MONITOR did not start"
|
||||
with Redis(host=host, port=port, socket_timeout=1, decode_responses=True) as client:
|
||||
client.set(ready_marker, "ready", ex=1)
|
||||
assert armed.wait(timeout=5), "Redis MONITOR did not capture its readiness command"
|
||||
yield commands
|
||||
finally:
|
||||
with Redis(host=host, port=port, socket_timeout=1, decode_responses=True) as client:
|
||||
client.set(stop_marker, "stop", ex=1)
|
||||
thread.join(timeout=5)
|
||||
assert not thread.is_alive(), "Redis MONITOR did not reach its stop command"
|
||||
|
|
|
|||
|
|
@ -0,0 +1,256 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.redis_process import OwnedRedis, owned_redis, redis_commands
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
pytestmark: Final = [pytest.mark.timeout(300)]
|
||||
|
||||
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
MODEL_INFO_ENTRIES: Final = TypeAdapter(tuple[dict[str, JsonValue], ...])
|
||||
OPENAI_MODEL: Final = "gpt-4o-mini"
|
||||
MASTER_KEY: Final = "sk-integration-cooldown-read-scope"
|
||||
API_KEY: Final = "synthetic-cooldown-read-scope-key"
|
||||
COOLDOWN_KEY: Final = re.compile(r"^deployment:(.+):cooldown$")
|
||||
REDIS_READS: Final = frozenset({"GET", "MGET"})
|
||||
ENDPOINTS: Final = ("/v1/chat/completions", "/v1/messages", "/v1/responses")
|
||||
OTHER_GROUP_SIZE: Final = 4
|
||||
CHAT_RESPONSE: Final = json.dumps(
|
||||
{
|
||||
"id": "chatcmpl_cooldown_read_scope",
|
||||
"object": "chat.completion",
|
||||
"created": 1700000000,
|
||||
"model": OPENAI_MODEL,
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "scoped"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 7, "completion_tokens": 4, "total_tokens": 11},
|
||||
}
|
||||
).encode()
|
||||
RESPONSES_RESPONSE: Final = json.dumps(
|
||||
{
|
||||
"id": "resp_cooldown_read_scope",
|
||||
"object": "response",
|
||||
"created_at": 1700000000,
|
||||
"status": "completed",
|
||||
"model": OPENAI_MODEL,
|
||||
"output": [
|
||||
{
|
||||
"id": "msg_cooldown_read_scope",
|
||||
"type": "message",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "scoped", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 7, "output_tokens": 4, "total_tokens": 11},
|
||||
}
|
||||
).encode()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
gateway: Gateway
|
||||
redis: OwnedRedis
|
||||
wire: Wire
|
||||
groups: Mapping[str, tuple[str, frozenset[str]]]
|
||||
wildcard_prefix: str
|
||||
wildcard_ids: frozenset[str]
|
||||
other_ids: frozenset[str]
|
||||
|
||||
@property
|
||||
def configured_ids(self) -> frozenset[str]:
|
||||
return frozenset().union(*(ids for _, ids in self.groups.values()), self.wildcard_ids, self.other_ids)
|
||||
|
||||
|
||||
def _reply(request: Request) -> Reply:
|
||||
if request.target == "/v1/models":
|
||||
return Reply(body=json.dumps({"object": "list", "data": [{"id": OPENAI_MODEL, "object": "model"}]}).encode())
|
||||
return Reply(body=RESPONSES_RESPONSE if request.target == "/v1/responses" else CHAT_RESPONSE)
|
||||
|
||||
|
||||
def _deployments(model_name: str, model: str, api_base: str, ids: frozenset[str]) -> tuple[dict[str, JsonValue], ...]:
|
||||
return tuple(
|
||||
{
|
||||
"model_name": model_name,
|
||||
"litellm_params": {"model": model, "api_base": api_base, "api_key": API_KEY},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
for deployment_id in sorted(ids)
|
||||
)
|
||||
|
||||
|
||||
def _ids(prefix: str, count: int) -> frozenset[str]:
|
||||
return frozenset(f"{prefix}-{uuid.uuid4().hex[:8]}" for _ in range(count))
|
||||
|
||||
|
||||
def _registered_ids(gateway: Gateway) -> frozenset[str]:
|
||||
entries: Final = MODEL_INFO_ENTRIES.validate_python(gateway.get("/model/info")["data"])
|
||||
return frozenset(string_value(object_value(entry["model_info"])["id"]) for entry in entries)
|
||||
|
||||
|
||||
def _payload(endpoint: str, model: str) -> dict[str, JsonValue]:
|
||||
marker: Final = f"cooldown-scope-{uuid.uuid4().hex}"
|
||||
if endpoint == "/v1/responses":
|
||||
return {"model": model, "input": marker, "max_output_tokens": 16, "store": False}
|
||||
return {"model": model, "messages": [{"role": "user", "content": marker}], "max_tokens": 16}
|
||||
|
||||
|
||||
def _headers(endpoint: str) -> Mapping[str, str]:
|
||||
return {"anthropic-version": "2023-06-01"} if endpoint == "/v1/messages" else {}
|
||||
|
||||
|
||||
def _cooldown_reads(commands: tuple[tuple[str, ...], ...]) -> tuple[frozenset[str], ...]:
|
||||
"""The deployment ids whose cooldown key each Redis read carried, one entry per read that carried any"""
|
||||
reads: Final = tuple(command for command in commands if command and command[0].upper() in REDIS_READS)
|
||||
matched: Final = tuple(
|
||||
frozenset(match.group(1) for argument in command[1:] if (match := COOLDOWN_KEY.match(argument)))
|
||||
for command in reads
|
||||
)
|
||||
return tuple(ids for ids in matched if ids)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=("simple-shuffle", "usage-based-routing-v2"))
|
||||
def rig(request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
strategy: Final = str(request.param)
|
||||
tmp_path: Final = tmp_path_factory.mktemp(f"cooldown-scope-{strategy}")
|
||||
with owned_redis(tmp_path) as cache, wire_server(_reply) as wire:
|
||||
run_id: Final = uuid.uuid4().hex[:8]
|
||||
api_base: Final = f"{wire.url}/v1"
|
||||
groups: Final = MappingProxyType(
|
||||
{
|
||||
endpoint: (f"scoped-{index}-{run_id}", _ids(f"scoped-{index}", 2))
|
||||
for index, endpoint in enumerate(ENDPOINTS)
|
||||
}
|
||||
)
|
||||
wildcard_prefix: Final = f"wild-{run_id}"
|
||||
wildcard_ids: Final = _ids("wildcard", 1)
|
||||
other_ids: Final = _ids("other", OTHER_GROUP_SIZE)
|
||||
model_list: Final = (
|
||||
*(
|
||||
deployment
|
||||
for name, ids in groups.values()
|
||||
for deployment in _deployments(name, f"openai/{OPENAI_MODEL}", api_base, ids)
|
||||
),
|
||||
*_deployments(f"{wildcard_prefix}/*", "openai/*", api_base, wildcard_ids),
|
||||
*_deployments(f"other-{run_id}", f"openai/{OPENAI_MODEL}", api_base, other_ids),
|
||||
)
|
||||
configuration: Final = JSON_OBJECT.validate_python(
|
||||
yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
)
|
||||
config_path: Final = tmp_path / "cooldown-scope.yaml"
|
||||
config_path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
**configuration,
|
||||
"general_settings": {
|
||||
**JSON_OBJECT.validate_python(configuration["general_settings"]),
|
||||
"store_model_in_db": True,
|
||||
},
|
||||
"model_list": list(model_list),
|
||||
"router_settings": {
|
||||
"routing_strategy": strategy,
|
||||
"redis_host": cache.host,
|
||||
"redis_port": cache.port,
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
with httpx.Client(base_url=wire.url, timeout=15, trust_env=False) as bootstrap_client:
|
||||
bootstrap: Final = Gateway(bootstrap_client, MASTER_KEY, wire.url)
|
||||
with owned_proxy(
|
||||
bootstrap,
|
||||
tmp_path,
|
||||
{
|
||||
"REDIS_HOST": cache.host,
|
||||
"REDIS_PORT": str(cache.port),
|
||||
"STORE_MODEL_IN_DB": "True",
|
||||
"DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS": "0.001",
|
||||
},
|
||||
config=config_path,
|
||||
) as gateway:
|
||||
built: Final = Rig(gateway, cache, wire, groups, wildcard_prefix, wildcard_ids, other_ids)
|
||||
eventually(
|
||||
lambda: _registered_ids(gateway),
|
||||
lambda registered: built.configured_ids <= registered,
|
||||
seconds=30,
|
||||
)
|
||||
yield built
|
||||
|
||||
|
||||
def _send(rig: Rig, endpoint: str, model: str, key: str | None = None) -> tuple[httpx.Response, frozenset[str]]:
|
||||
"""Sends one request and returns it with the ids of every cooldown key the proxy read from Redis for it"""
|
||||
rig.wire.drain()
|
||||
with redis_commands(rig.redis.host, rig.redis.port) as commands:
|
||||
response: Final = rig.gateway.request(
|
||||
"POST", endpoint, _payload(endpoint, model), key=key, headers=_headers(endpoint)
|
||||
)
|
||||
captured: Final = commands.drain()
|
||||
reads: Final = _cooldown_reads(captured)
|
||||
assert len(reads) == 1, f"cooldown keys should arrive in one Redis read, saw {reads} in {captured}"
|
||||
assert len(rig.wire.drain()) == 1
|
||||
return response, reads[0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("endpoint", ENDPOINTS, ids=("chat-completions", "messages", "responses"))
|
||||
def test_request_reads_cooldown_keys_of_its_model_group_only(rig: Rig, endpoint: str) -> None:
|
||||
model_name, group_ids = rig.groups[endpoint]
|
||||
|
||||
response, read_ids = _send(rig, endpoint, model_name)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["x-litellm-model-id"] in group_ids
|
||||
assert read_ids == group_ids, f"read {sorted(read_ids - group_ids)} outside group {sorted(group_ids)}"
|
||||
|
||||
|
||||
def test_wildcard_route_reads_cooldown_key_of_matching_deployment_only(rig: Rig) -> None:
|
||||
response, read_ids = _send(rig, "/v1/chat/completions", f"{rig.wildcard_prefix}/{OPENAI_MODEL}")
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["x-litellm-model-id"] in rig.wildcard_ids
|
||||
assert read_ids == rig.wildcard_ids, f"read {sorted(read_ids - rig.wildcard_ids)} outside the wildcard route"
|
||||
|
||||
|
||||
def test_team_route_reads_cooldown_key_of_team_deployment_only(rig: Rig) -> None:
|
||||
with rig.gateway.scenario() as scenario:
|
||||
team_id: Final = scenario.team()
|
||||
public_name: Final = f"team-scoped-{uuid.uuid4().hex[:8]}"
|
||||
team_model_id: Final = _team_model(scenario, public_name, team_id, rig.wire.url)
|
||||
team_key: Final = scenario.key(team_id=team_id)
|
||||
eventually(
|
||||
lambda: _registered_ids(rig.gateway),
|
||||
lambda registered: team_model_id in registered,
|
||||
seconds=30,
|
||||
)
|
||||
|
||||
response, read_ids = _send(rig, "/v1/chat/completions", public_name, key=team_key)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["x-litellm-model-id"] == team_model_id
|
||||
assert read_ids == {team_model_id}, f"read {sorted(read_ids - {team_model_id})} outside the team route"
|
||||
|
||||
|
||||
def _team_model(scenario: Scenario, public_name: str, team_id: str, upstream_url: str) -> str:
|
||||
created: Final = scenario.gateway.post(
|
||||
"/model/new",
|
||||
{
|
||||
"model_name": public_name,
|
||||
"litellm_params": {"model": f"openai/{OPENAI_MODEL}", "api_key": API_KEY, "api_base": f"{upstream_url}/v1"},
|
||||
"model_info": {"team_id": team_id},
|
||||
},
|
||||
)
|
||||
model_id: Final = string_value(object_value(created["model_info"])["id"])
|
||||
scenario.cleanups.callback(scenario.delete_model, model_id)
|
||||
return model_id
|
||||
|
|
@ -106,6 +106,10 @@ class FakeRedisCache(RedisCache):
|
|||
self.alone.append(("MGET", tuple(key_list)))
|
||||
return {key: self.store.get(key) for key in key_list}
|
||||
|
||||
def batch_get_cache(self, key_list: Sequence[str], **kwargs: object) -> dict[str, Any]: # pyright: ignore[reportIncompatibleMethodOverride] # records the direct sync read
|
||||
self.alone.append(("MGET", tuple(key_list)))
|
||||
return {key: self.store.get(key) for key in key_list}
|
||||
|
||||
async def async_increment(self, key: str, value: float, ttl: int | None = None, **kwargs: object) -> float: # pyright: ignore[reportIncompatibleMethodOverride] # records the direct write
|
||||
self.alone.append(("INCRBYFLOAT", key, value))
|
||||
self.store[key] = float(self.store.get(key, 0.0)) + value
|
||||
|
|
|
|||
|
|
@ -6,6 +6,8 @@ from __future__ import annotations
|
|||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from datetime import datetime, timezone
|
||||
from itertools import chain
|
||||
from typing import Any, Final
|
||||
|
|
@ -36,6 +38,7 @@ from litellm.router_utils.routing_read_batch import (
|
|||
RoutingPrefetch,
|
||||
_routing_read_target, # pyright: ignore[reportPrivateUsage] # the family rule under test
|
||||
)
|
||||
from litellm.types.router import RouterRateLimitError
|
||||
|
||||
from .test_redis_batch import FakeClient, FakeRedisCache, replies
|
||||
|
||||
|
|
@ -274,9 +277,9 @@ async def test_without_a_request_scope_descriptor_groups_run_the_script_directly
|
|||
assert client.pipelines == []
|
||||
|
||||
|
||||
def _deployment(deployment_id: str) -> dict:
|
||||
def _deployment(deployment_id: str, model_name: str = _MODEL_GROUP) -> dict:
|
||||
return {
|
||||
"model_name": _MODEL_GROUP,
|
||||
"model_name": model_name,
|
||||
"litellm_params": {"model": "anthropic/claude-x", "api_key": "test", "mock_response": "pong"},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
|
|
@ -390,7 +393,7 @@ async def test_a_prefetch_with_incomplete_usage_keys_releases_cooldown_reservati
|
|||
)
|
||||
|
||||
with request_redis_batch_scope():
|
||||
RoutingPrefetch.arm(router, router.lowesttpm_logger_v2, router.model_list[:1])
|
||||
RoutingPrefetch.arm(router, router.lowesttpm_logger_v2, router.model_list[:1], frozenset({"dep-a", "dep-b"}))
|
||||
deployment: Final = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
|
|
@ -403,6 +406,181 @@ async def test_a_prefetch_with_incomplete_usage_keys_releases_cooldown_reservati
|
|||
assert len(fallback_cooldown_mgets) == 1
|
||||
|
||||
|
||||
_GROUP_IDS: Final = frozenset({"dep-a", "dep-b"})
|
||||
_OTHER_GROUP_IDS: Final = ("other-a", "other-b", "other-c")
|
||||
_TEAM_ID: Final = "team-1"
|
||||
_TEAM_DEPLOYMENT_ID: Final = "team-dep"
|
||||
_KNOWN_IDS: Final = (*sorted(_GROUP_IDS), *_OTHER_GROUP_IDS, _TEAM_DEPLOYMENT_ID)
|
||||
_ACTIVE_COOLDOWN: Final = {
|
||||
"exception_received": "429",
|
||||
"status_code": "429",
|
||||
"timestamp": _FAR_FUTURE,
|
||||
"cooldown_time": 60,
|
||||
}
|
||||
|
||||
|
||||
def _pass_through_deployment(deployment_id: str) -> dict:
|
||||
deployment: Final = _deployment(deployment_id)
|
||||
return {**deployment, "litellm_params": {**deployment["litellm_params"], "use_in_pass_through": True}}
|
||||
|
||||
|
||||
def _router_beside_other_groups(redis_cache: FakeRedisCache, routing_strategy: str = "simple-shuffle") -> Router:
|
||||
team_deployment: Final = {
|
||||
"model_name": f"{_MODEL_GROUP}_{_TEAM_ID}",
|
||||
"litellm_params": {"model": "anthropic/claude-x", "api_key": "test", "mock_response": "pong"},
|
||||
"model_info": {"id": _TEAM_DEPLOYMENT_ID, "team_id": _TEAM_ID, "team_public_model_name": _MODEL_GROUP},
|
||||
}
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
*(_pass_through_deployment(deployment_id) for deployment_id in sorted(_GROUP_IDS)),
|
||||
*(_deployment(deployment_id, model_name="other") for deployment_id in _OTHER_GROUP_IDS),
|
||||
team_deployment,
|
||||
],
|
||||
routing_strategy=routing_strategy,
|
||||
)
|
||||
router.update_redis_cache(cache=redis_cache)
|
||||
return router
|
||||
|
||||
|
||||
def _cooldown_ids_read(client: FakeClient, redis_cache: FakeRedisCache) -> tuple[frozenset[str], ...]:
|
||||
"""The deployment ids whose cooldown keys each Redis read carried, pipelined or direct, for reads carrying any"""
|
||||
pipelined: Final = tuple(
|
||||
command[1:] for command in chain.from_iterable(p.commands for p in client.pipelines) if command[0] == "MGET"
|
||||
)
|
||||
direct: Final = tuple(keys for command, keys in redis_cache.alone if command == "MGET")
|
||||
per_read: Final = tuple(
|
||||
frozenset(model_id for model_id in _KNOWN_IDS if CooldownCache.get_cooldown_cache_key(model_id) in keys)
|
||||
for keys in (*pipelined, *direct)
|
||||
)
|
||||
return tuple(ids for ids in per_read if ids)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("routing_strategy", ["simple-shuffle", "usage-based-routing-v2"])
|
||||
async def test_the_armed_routing_read_carries_only_the_requested_groups_cooldown_keys(routing_strategy: str) -> None:
|
||||
client: Final = FakeClient(_lua_ok_replies)
|
||||
redis_cache: Final = FakeRedisCache(client)
|
||||
router: Final = _router_beside_other_groups(redis_cache, routing_strategy)
|
||||
|
||||
with request_redis_batch_scope():
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
deployment: Final = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
|
||||
assert deployment["model_info"]["id"] in _GROUP_IDS
|
||||
assert _cooldown_ids_read(client, redis_cache) == (_GROUP_IDS,)
|
||||
assert redis_cache.alone == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_team_routes_armed_read_carries_the_team_deployments_cooldown_keys() -> None:
|
||||
client: Final = FakeClient(_lua_ok_replies)
|
||||
redis_cache: Final = FakeRedisCache(client)
|
||||
router: Final = _router_beside_other_groups(redis_cache)
|
||||
request_kwargs: Final = {"metadata": {"user_api_key_team_id": _TEAM_ID}}
|
||||
|
||||
with request_redis_batch_scope():
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, request_kwargs)
|
||||
deployment: Final = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs=request_kwargs
|
||||
)
|
||||
|
||||
assert deployment["model_info"]["id"] == _TEAM_DEPLOYMENT_ID
|
||||
assert _cooldown_ids_read(client, redis_cache) == (frozenset({_TEAM_DEPLOYMENT_ID}),)
|
||||
assert redis_cache.alone == []
|
||||
|
||||
|
||||
async def _picked(router: Router) -> frozenset[str]:
|
||||
deployment: Final = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
return frozenset({deployment["model_info"]["id"]})
|
||||
|
||||
|
||||
async def _picked_sync(router: Router) -> frozenset[str]:
|
||||
deployment: Final = router.get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
return frozenset({deployment["model_info"]["id"]})
|
||||
|
||||
|
||||
async def _picked_for_pass_through(router: Router) -> frozenset[str]:
|
||||
deployment: Final = router.get_available_deployment_for_pass_through(model=_MODEL_GROUP, request_kwargs={})
|
||||
return frozenset({deployment["model_info"]["id"]})
|
||||
|
||||
|
||||
async def _healthy(router: Router) -> frozenset[str]:
|
||||
healthy, _ = await router._async_get_healthy_deployments( # pyright: ignore[reportPrivateUsage] # retry and scheduler read
|
||||
model=_MODEL_GROUP, parent_otel_span=None
|
||||
)
|
||||
return frozenset(deployment["model_info"]["id"] for deployment in healthy)
|
||||
|
||||
|
||||
async def _weighted_failover_targets(router: Router) -> frozenset[str]:
|
||||
failed: Final = Exception("dep-a failed")
|
||||
failed.failed_deployment_id = "dep-a" # pyright: ignore[reportAttributeAccessIssue] # stamped like the router stamps it
|
||||
result: Final = await router._maybe_run_weighted_failover( # pyright: ignore[reportPrivateUsage] # same-group retry
|
||||
exception=failed,
|
||||
original_model_group=_MODEL_GROUP,
|
||||
all_deployments=[_deployment(deployment_id) for deployment_id in sorted(_GROUP_IDS)],
|
||||
args=(),
|
||||
kwargs={"metadata": {}},
|
||||
input_kwargs={},
|
||||
)
|
||||
return frozenset() if result is None else frozenset({"retried"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("read", "expected"),
|
||||
[
|
||||
(_picked, frozenset({"dep-a"})),
|
||||
(_picked_sync, frozenset({"dep-a"})),
|
||||
(_picked_for_pass_through, frozenset({"dep-a"})),
|
||||
(_healthy, frozenset({"dep-a"})),
|
||||
(_weighted_failover_targets, frozenset()),
|
||||
],
|
||||
ids=["async-pick", "sync-pick", "pass-through-pick", "async-healthy", "weighted-failover"],
|
||||
)
|
||||
async def test_a_direct_cooldown_read_covers_only_the_groups_deployments_and_still_sees_their_cooldowns(
|
||||
read: Callable[[Router], Awaitable[frozenset[str]]], expected: frozenset[str]
|
||||
) -> None:
|
||||
client: Final = FakeClient(_lua_ok_replies)
|
||||
redis_cache: Final = FakeRedisCache(client)
|
||||
router: Final = _router_beside_other_groups(redis_cache)
|
||||
for model_id in ("dep-b", *_OTHER_GROUP_IDS):
|
||||
redis_cache.store[CooldownCache.get_cooldown_cache_key(model_id)] = _ACTIVE_COOLDOWN
|
||||
|
||||
observed: Final = await read(router)
|
||||
|
||||
assert observed == expected
|
||||
assert _cooldown_ids_read(client, redis_cache) == (_GROUP_IDS,)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_cooling_deployment_the_health_check_already_dropped_still_shows_in_the_no_deployment_error() -> None:
|
||||
router: Final = Router(
|
||||
model_list=[_deployment(deployment_id) for deployment_id in sorted(_GROUP_IDS)],
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
router.health_state_cache.set_deployment_health_states(
|
||||
{"dep-a": {"is_healthy": False, "timestamp": time.time(), "reason": "probe failed"}}
|
||||
)
|
||||
memory: Final = router.cooldown_cache.cooldown_store.in_memory_cache
|
||||
assert memory is not None
|
||||
for model_id in sorted(_GROUP_IDS):
|
||||
memory.set_cache(CooldownCache.get_cooldown_cache_key(model_id), _ACTIVE_COOLDOWN)
|
||||
|
||||
with pytest.raises(RouterRateLimitError) as raised:
|
||||
await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
|
||||
assert raised.value.cooldown_deployments is not None
|
||||
assert frozenset(info["model_id"] for info in raised.value.cooldown_deployments) == _GROUP_IDS
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_prefetch_falls_back_to_the_shared_read():
|
||||
client = FakeClient(_lua_ok_replies, fail=ConnectionError("redis down"))
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue