mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-08 08:08:37 +00:00
Merge 6753584701 into 24b8889965
This commit is contained in:
yucheng-berri
GitHub
commit
e30bfefb41
9 changed files with 1011 additions and 30 deletions
|
|
@ -2566,7 +2566,12 @@ class Router(DeprecatedRouterMembers):
|
|||
)
|
||||
deployments: Final = self._core_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)
|
||||
|
|
@ -8829,17 +8834,15 @@ class Router(DeprecatedRouterMembers):
|
|||
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
|
||||
|
|
@ -10387,7 +10390,9 @@ class Router(DeprecatedRouterMembers):
|
|||
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)
|
||||
|
|
@ -15467,13 +15472,17 @@ class Router(DeprecatedRouterMembers):
|
|||
)
|
||||
)
|
||||
|
||||
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,
|
||||
)
|
||||
)
|
||||
|
|
@ -16754,7 +16763,9 @@ class Router(DeprecatedRouterMembers):
|
|||
)
|
||||
|
||||
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(
|
||||
|
|
@ -17009,7 +17020,9 @@ class Router(DeprecatedRouterMembers):
|
|||
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
|
||||
|
|
@ -641,16 +642,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,
|
||||
)
|
||||
|
||||
|
|
@ -689,7 +695,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
|
||||
"""
|
||||
|
|
@ -698,10 +708,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
|
||||
|
|
@ -0,0 +1,459 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import json
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from concurrent.futures import ThreadPoolExecutor, wait
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually, held, object_value, string_value
|
||||
from integration._support.process import OwnedProxy, owned_proxy_process
|
||||
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
|
||||
from redis import Redis
|
||||
|
||||
pytestmark: Final = [pytest.mark.timeout(600)]
|
||||
|
||||
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-failure"
|
||||
GOOD_KEY: Final = "synthetic-cooldown-good-key"
|
||||
BAD_KEY: Final = "synthetic-cooldown-revoked-key"
|
||||
COOLDOWN_KEY: Final = re.compile(r"^deployment:(.+):cooldown$")
|
||||
COOLDOWN_LIST: Final = re.compile(r"cooldown_list=\[([^\]]*)\]")
|
||||
REDIS_READS: Final = frozenset({"GET", "MGET"})
|
||||
ENDPOINTS: Final = ("/v1/chat/completions", "/v1/messages", "/v1/responses")
|
||||
ENDPOINT_IDS: Final = ("chat-completions", "messages", "responses")
|
||||
WORKERS: Final = 2
|
||||
BURST_ROUNDS: Final = 4
|
||||
REDIS_SOCKET_TIMEOUT_SECONDS: Final = 1.0
|
||||
UNAUTHORIZED: Final = json.dumps(
|
||||
{"error": {"message": "Incorrect API key provided", "type": "invalid_request_error", "code": "invalid_api_key"}}
|
||||
).encode()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Group:
|
||||
name: str
|
||||
ids: frozenset[str]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
owned: OwnedProxy
|
||||
gateway: Gateway
|
||||
redis: OwnedRedis
|
||||
wire: Wire
|
||||
flaky: Group
|
||||
flaky_bad_id: str
|
||||
neighbour: Group
|
||||
dead: Group
|
||||
streamed: Group
|
||||
burst: Group
|
||||
tier: Group
|
||||
smart: Group
|
||||
bystanders: Group
|
||||
|
||||
@property
|
||||
def configured_ids(self) -> frozenset[str]:
|
||||
groups: Final = (
|
||||
self.flaky,
|
||||
self.neighbour,
|
||||
self.dead,
|
||||
self.streamed,
|
||||
self.burst,
|
||||
self.tier,
|
||||
self.smart,
|
||||
self.bystanders,
|
||||
)
|
||||
return frozenset().union(*(group.ids for group in groups))
|
||||
|
||||
|
||||
def _streams(request: Request) -> bool:
|
||||
body: Final = JSON_OBJECT.validate_python(json.loads(request.body or b"{}"))
|
||||
return body.get("stream") is True
|
||||
|
||||
|
||||
def _sse(*events: Mapping[str, JsonValue], done: bool) -> tuple[bytes, ...]:
|
||||
frames: Final = tuple(b"data: " + json.dumps(event).encode() + b"\n\n" for event in events)
|
||||
return (*frames[:-1], frames[-1] + (b"data: [DONE]\n\n" if done else b""))
|
||||
|
||||
|
||||
def _chat_reply(stream: bool) -> Reply:
|
||||
identity: Final = f"chatcmpl-{uuid.uuid4().hex}"
|
||||
usage: Final = {"prompt_tokens": 7, "completion_tokens": 4, "total_tokens": 11}
|
||||
if not stream:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion",
|
||||
"created": 1700000000,
|
||||
"model": OPENAI_MODEL,
|
||||
"choices": [
|
||||
{"index": 0, "message": {"role": "assistant", "content": "scoped"}, "finish_reason": "stop"}
|
||||
],
|
||||
"usage": usage,
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
chunk: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1700000000, "model": OPENAI_MODEL}
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=_sse(
|
||||
{**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "scoped"}}]},
|
||||
{**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
|
||||
{**chunk, "choices": [], "usage": usage},
|
||||
done=True,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _responses_reply(stream: bool) -> Reply:
|
||||
identity: Final = f"resp_{uuid.uuid4().hex}"
|
||||
completed: Final = {
|
||||
"id": identity,
|
||||
"object": "response",
|
||||
"created_at": 1700000000,
|
||||
"status": "completed",
|
||||
"model": OPENAI_MODEL,
|
||||
"output": [
|
||||
{
|
||||
"id": f"msg_{identity}",
|
||||
"type": "message",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "scoped", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"parallel_tool_calls": True,
|
||||
"tool_choice": "auto",
|
||||
"tools": [],
|
||||
"usage": {
|
||||
"input_tokens": 7,
|
||||
"input_tokens_details": {"cached_tokens": 0},
|
||||
"output_tokens": 4,
|
||||
"output_tokens_details": {"reasoning_tokens": 0},
|
||||
"total_tokens": 11,
|
||||
},
|
||||
}
|
||||
if not stream:
|
||||
return Reply(body=json.dumps(completed).encode())
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=_sse(
|
||||
{
|
||||
"type": "response.created",
|
||||
"sequence_number": 0,
|
||||
"response": {**completed, "status": "in_progress", "output": [], "usage": None},
|
||||
},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 1,
|
||||
"item_id": f"msg_{identity}",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "scoped",
|
||||
},
|
||||
{"type": "response.completed", "sequence_number": 2, "response": completed},
|
||||
done=False,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _reply(request: Request) -> Reply:
|
||||
if request.target == "/v1/models":
|
||||
return Reply(body=json.dumps({"object": "list", "data": [{"id": OPENAI_MODEL, "object": "model"}]}).encode())
|
||||
if BAD_KEY in request.headers.get("authorization", ""):
|
||||
return Reply(status=401, body=UNAUTHORIZED)
|
||||
if request.target == "/v1/responses":
|
||||
return _responses_reply(_streams(request))
|
||||
return _chat_reply(_streams(request))
|
||||
|
||||
|
||||
def _group(name: str, prefix: str, count: int) -> Group:
|
||||
return Group(name, frozenset(f"{prefix}-{uuid.uuid4().hex[:8]}" for _ in range(count)))
|
||||
|
||||
|
||||
def _deployments(
|
||||
group: Group, api_base: str, bad_ids: frozenset[str] = frozenset()
|
||||
) -> tuple[dict[str, JsonValue], ...]:
|
||||
return tuple(
|
||||
{
|
||||
"model_name": group.name,
|
||||
"litellm_params": {
|
||||
"model": f"openai/{OPENAI_MODEL}",
|
||||
"api_base": api_base,
|
||||
"api_key": BAD_KEY if deployment_id in bad_ids else GOOD_KEY,
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
for deployment_id in sorted(group.ids)
|
||||
)
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
@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-failure-{strategy}")
|
||||
run_id: Final = uuid.uuid4().hex[:8]
|
||||
flaky: Final = _group(f"flaky-{run_id}", "flaky", 2)
|
||||
flaky_bad_id: Final = min(flaky.ids)
|
||||
neighbour: Final = _group(f"neighbour-{run_id}", "neighbour", 1)
|
||||
dead: Final = _group(f"dead-{run_id}", "dead", 1)
|
||||
streamed: Final = _group(f"streamed-{run_id}", "streamed", 2)
|
||||
burst: Final = _group(f"burst-{run_id}", "burst", 2)
|
||||
tier: Final = _group(f"tier-{run_id}", "tier", 2)
|
||||
smart: Final = _group(f"smart-{run_id}", "smart", 1)
|
||||
bystanders: Final = _group(f"bystanders-{run_id}", "bystander", 4)
|
||||
with owned_redis(tmp_path) as cache, wire_server(_reply) as wire:
|
||||
api_base: Final = f"{wire.url}/v1"
|
||||
model_list: Final = (
|
||||
*_deployments(flaky, api_base, frozenset({flaky_bad_id})),
|
||||
*_deployments(neighbour, api_base, neighbour.ids),
|
||||
*_deployments(dead, api_base, dead.ids),
|
||||
*_deployments(streamed, api_base),
|
||||
*_deployments(burst, api_base),
|
||||
*_deployments(tier, api_base),
|
||||
{
|
||||
"model_name": smart.name,
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_default_model": tier.name,
|
||||
"complexity_router_config": {
|
||||
"tiers": {level: tier.name for level in ("SIMPLE", "MEDIUM", "COMPLEX", "REASONING")}
|
||||
},
|
||||
},
|
||||
"model_info": {"id": min(smart.ids)},
|
||||
},
|
||||
*_deployments(bystanders, api_base),
|
||||
)
|
||||
configuration: Final = JSON_OBJECT.validate_python(
|
||||
yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
)
|
||||
config_path: Final = tmp_path / "cooldown-failure.yaml"
|
||||
config_path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
**configuration,
|
||||
"model_list": list(model_list),
|
||||
"router_settings": {
|
||||
"routing_strategy": strategy,
|
||||
"redis_host": cache.host,
|
||||
"redis_port": cache.port,
|
||||
"num_retries": 2,
|
||||
"cooldown_time": 60,
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
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_process(
|
||||
bootstrap,
|
||||
tmp_path,
|
||||
{
|
||||
"REDIS_HOST": cache.host,
|
||||
"REDIS_PORT": str(cache.port),
|
||||
"REDIS_SOCKET_TIMEOUT": str(REDIS_SOCKET_TIMEOUT_SECONDS),
|
||||
"DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS": "0.001",
|
||||
},
|
||||
config=config_path,
|
||||
workers=WORKERS,
|
||||
client_timeout=120,
|
||||
) as owned:
|
||||
# A fresh connection per request lets every worker answer, and survives a killed worker
|
||||
with httpx.Client(
|
||||
base_url=str(owned.gateway.client.base_url),
|
||||
timeout=120,
|
||||
trust_env=False,
|
||||
limits=httpx.Limits(max_keepalive_connections=0),
|
||||
) as unpooled:
|
||||
gateway: Final = dataclasses.replace(owned.gateway, client=unpooled)
|
||||
built: Final = Rig(
|
||||
owned,
|
||||
gateway,
|
||||
cache,
|
||||
wire,
|
||||
flaky,
|
||||
flaky_bad_id,
|
||||
neighbour,
|
||||
dead,
|
||||
streamed,
|
||||
burst,
|
||||
tier,
|
||||
smart,
|
||||
bystanders,
|
||||
)
|
||||
eventually(
|
||||
lambda: _registered_ids(gateway),
|
||||
lambda registered: built.configured_ids <= registered,
|
||||
seconds=30,
|
||||
)
|
||||
yield built
|
||||
|
||||
|
||||
def _payload(endpoint: str, model: str, stream: bool = False) -> dict[str, JsonValue]:
|
||||
marker: Final = f"cooldown-failure-{uuid.uuid4().hex}"
|
||||
if endpoint == "/v1/responses":
|
||||
return {"model": model, "input": marker, "max_output_tokens": 16, "store": False, "stream": stream}
|
||||
return {"model": model, "messages": [{"role": "user", "content": marker}], "max_tokens": 16, "stream": stream}
|
||||
|
||||
|
||||
def _headers(endpoint: str) -> Mapping[str, str]:
|
||||
return {"anthropic-version": "2023-06-01"} if endpoint == "/v1/messages" else {}
|
||||
|
||||
|
||||
def _post(rig: Rig, endpoint: str, model: str, stream: bool = False) -> httpx.Response:
|
||||
return rig.gateway.request("POST", endpoint, _payload(endpoint, model, stream), headers=_headers(endpoint))
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
def _sent(
|
||||
rig: Rig, endpoint: str, model: str, stream: bool = False
|
||||
) -> tuple[httpx.Response, tuple[frozenset[str], ...]]:
|
||||
"""Sends one request and returns it with the deployment ids of every cooldown read Redis served for it"""
|
||||
with redis_commands(rig.redis.host, rig.redis.port) as commands:
|
||||
response: Final = _post(rig, endpoint, model, stream)
|
||||
return response, _cooldown_reads(commands.drain())
|
||||
|
||||
|
||||
def _read_ids(reads: tuple[frozenset[str], ...]) -> frozenset[str]:
|
||||
return frozenset().union(*reads)
|
||||
|
||||
|
||||
def _bad_key_requests(requests: tuple[Request, ...]) -> int:
|
||||
return sum(BAD_KEY in request.headers.get("authorization", "") for request in requests)
|
||||
|
||||
|
||||
def _is_cooling(rig: Rig, deployment_id: str) -> bool:
|
||||
with Redis(host=rig.redis.host, port=rig.redis.port, socket_timeout=1) as client:
|
||||
return bool(client.exists(f"deployment:{deployment_id}:cooldown"))
|
||||
|
||||
|
||||
def _cool_flaky_deployment(rig: Rig) -> None:
|
||||
"""Sends to the flaky group until its revoked deployment has failed once and is cooling in Redis"""
|
||||
|
||||
def attempt() -> bool:
|
||||
response: Final = _post(rig, "/v1/chat/completions", rig.flaky.name)
|
||||
assert response.status_code == 200, response.text
|
||||
return _is_cooling(rig, rig.flaky_bad_id)
|
||||
|
||||
eventually(attempt, lambda cooling: cooling, seconds=60)
|
||||
|
||||
|
||||
def test_a_cooled_deployment_stays_skipped_by_workers_that_never_saw_it_fail(rig: Rig) -> None:
|
||||
_cool_flaky_deployment(rig)
|
||||
for pid in sorted(rig.owned.worker_pids()):
|
||||
rig.owned.kill_worker(pid)
|
||||
rig.wire.drain()
|
||||
|
||||
sent: Final = tuple(_sent(rig, "/v1/chat/completions", rig.flaky.name) for _ in range(8))
|
||||
|
||||
assert tuple(response.status_code for response, _ in sent) == (200,) * len(sent), [r.text for r, _ in sent]
|
||||
assert {response.headers["x-litellm-model-id"] for response, _ in sent} <= rig.flaky.ids - {rig.flaky_bad_id}
|
||||
assert _bad_key_requests(rig.wire.drain()) == 0
|
||||
# A worker that already holds the cooling key in memory reads only the keys it lacks
|
||||
for _, reads in sent:
|
||||
assert reads, "every request should read cooldown state from Redis"
|
||||
assert _read_ids(reads) <= rig.flaky.ids, f"read {sorted(_read_ids(reads) - rig.flaky.ids)} outside the group"
|
||||
|
||||
|
||||
def test_the_no_deployments_error_still_lists_cooled_deployments_outside_the_group(rig: Rig) -> None:
|
||||
neighbour_id: Final = min(rig.neighbour.ids)
|
||||
dead_id: Final = min(rig.dead.ids)
|
||||
for group, deployment_id in ((rig.neighbour, neighbour_id), (rig.dead, dead_id)):
|
||||
failed = _post(rig, "/v1/chat/completions", group.name)
|
||||
assert failed.status_code == 401, failed.text
|
||||
eventually(partial(_is_cooling, rig, deployment_id), lambda cooling: cooling, seconds=30)
|
||||
|
||||
refused: Final = _post(rig, "/v1/chat/completions", rig.dead.name)
|
||||
|
||||
assert refused.status_code == 429, refused.text
|
||||
assert "No deployments available for selected model" in refused.text
|
||||
listed: Final = COOLDOWN_LIST.search(refused.text)
|
||||
assert listed is not None, refused.text
|
||||
listed_ids: Final = frozenset(name.strip().strip("'\"") for name in listed.group(1).split(","))
|
||||
assert {neighbour_id, dead_id} <= listed_ids, listed_ids
|
||||
|
||||
|
||||
@pytest.mark.parametrize("endpoint", ENDPOINTS, ids=ENDPOINT_IDS)
|
||||
def test_a_streamed_request_reads_cooldown_keys_of_its_model_group_only(rig: Rig, endpoint: str) -> None:
|
||||
response, reads = _sent(rig, endpoint, rig.streamed.name, stream=True)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream")
|
||||
assert "scoped" in response.text
|
||||
assert response.headers["x-litellm-model-id"] in rig.streamed.ids
|
||||
assert reads, "the request should read cooldown state from Redis"
|
||||
assert _read_ids(reads) <= rig.streamed.ids, f"read {sorted(_read_ids(reads) - rig.streamed.ids)} outside the group"
|
||||
|
||||
|
||||
def test_a_complexity_router_request_reads_cooldown_keys_of_its_entry_and_tier_only(rig: Rig) -> None:
|
||||
response, reads = _sent(rig, "/v1/chat/completions", rig.smart.name)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["x-litellm-model-id"] in rig.tier.ids
|
||||
allowed: Final = rig.smart.ids | rig.tier.ids
|
||||
assert 1 <= len(reads) <= 2, reads
|
||||
assert _read_ids(reads) <= allowed, f"read {sorted(_read_ids(reads) - allowed)} outside the route"
|
||||
|
||||
|
||||
def test_requests_keep_answering_while_redis_hangs_and_each_lands_once(rig: Rig) -> None:
|
||||
calls: Final = tuple((endpoint, stream) for endpoint in ENDPOINTS for stream in (False, True)) * BURST_ROUNDS
|
||||
|
||||
with rig.redis.paused(), ThreadPoolExecutor(max_workers=8) as pool:
|
||||
futures: Final = tuple(pool.submit(_post, rig, endpoint, rig.burst.name, stream) for endpoint, stream in calls)
|
||||
done, pending = wait(futures, timeout=120)
|
||||
assert not pending, f"{len(pending)} of {len(calls)} requests still waiting on a hung Redis after 120s"
|
||||
statuses: Final = tuple(future.result().status_code for future in futures)
|
||||
assert statuses == (200,) * len(calls), [
|
||||
future.result().text for future in done if future.result().status_code != 200
|
||||
]
|
||||
|
||||
# The breaker that opened while Redis hung skips Redis until it closes, so the first reads may come from memory
|
||||
recoveries: Final[list[tuple[httpx.Response, tuple[frozenset[str], ...]]]] = []
|
||||
|
||||
def recover() -> tuple[httpx.Response, tuple[frozenset[str], ...]]:
|
||||
recoveries.append(_sent(rig, "/v1/chat/completions", rig.burst.name))
|
||||
return recoveries[-1]
|
||||
|
||||
recovered, reads = eventually(recover, lambda sent: bool(sent[1]), seconds=60)
|
||||
|
||||
assert all(response.status_code == 200 for response, _ in recoveries), [r.text for r, _ in recoveries]
|
||||
landed: Final = len(calls) + len(recoveries)
|
||||
rows: Final = held(
|
||||
lambda: rig.gateway.rows(
|
||||
'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE model_group = %s', (rig.burst.name,)
|
||||
),
|
||||
lambda found: len(found) == landed,
|
||||
holding=3,
|
||||
seconds=90,
|
||||
)
|
||||
assert len({string_value(row["request_id"]) for row in rows}) == landed
|
||||
assert recovered.headers["x-litellm-model-id"] in rig.burst.ids
|
||||
assert reads == (rig.burst.ids,), f"read {reads}, expected one read of {sorted(rig.burst.ids)}"
|
||||
|
|
@ -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,7 @@ from __future__ import annotations
|
|||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable
|
||||
from datetime import datetime, timezone
|
||||
from itertools import chain
|
||||
from typing import Any, Final
|
||||
|
|
@ -14,6 +15,7 @@ from unittest.mock import AsyncMock, MagicMock
|
|||
import pytest
|
||||
|
||||
import litellm.caching.dual_cache as dual_cache_module
|
||||
import litellm.router_utils.health_state_cache as health_state_cache_module
|
||||
from litellm import Router
|
||||
from litellm._internal_context import current_service_target
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
|
|
@ -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,186 @@ 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(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
clock: Final = 1_000_000.0
|
||||
monkeypatch.setattr(health_state_cache_module.time, "time", lambda: clock)
|
||||
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": clock, "reason": "probe failed"}}
|
||||
)
|
||||
assert router.health_state_cache.get_unhealthy_deployment_ids() == {"dep-a"}
|
||||
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