From 55726b91c831823436f86bd5bee40bf09f9e558f Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 09:28:47 -0700 Subject: [PATCH 1/3] 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 --- litellm/router.py | 35 ++- litellm/router_utils/cooldown_cache.py | 6 +- litellm/router_utils/cooldown_handlers.py | 20 +- litellm/router_utils/routing_read_batch.py | 15 +- tests/integration/_support/redis_process.py | 57 ++++ ...est_cooldown_reads_scoped_to_candidates.py | 256 ++++++++++++++++++ tests/unit/caching/test_redis_batch.py | 4 + .../test_request_redis_batch_pre_call.py | 184 ++++++++++++- 8 files changed, 547 insertions(+), 30 deletions(-) create mode 100644 tests/integration/routing/test_cooldown_reads_scoped_to_candidates.py diff --git a/litellm/router.py b/litellm/router.py index ac0933e158b..9c3b5ffbb3b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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, diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index dc21db6a601..3b93420f219 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -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] diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index c362c6899db..0ba62b2223b 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -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 = [] diff --git a/litellm/router_utils/routing_read_batch.py b/litellm/router_utils/routing_read_batch.py index 7f6267ec9b9..4b9a43b559d 100644 --- a/litellm/router_utils/routing_read_batch.py +++ b/litellm/router_utils/routing_read_batch.py @@ -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] diff --git a/tests/integration/_support/redis_process.py b/tests/integration/_support/redis_process.py index 1a325d73659..e14df6fe9dc 100644 --- a/tests/integration/_support/redis_process.py +++ b/tests/integration/_support/redis_process.py @@ -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" diff --git a/tests/integration/routing/test_cooldown_reads_scoped_to_candidates.py b/tests/integration/routing/test_cooldown_reads_scoped_to_candidates.py new file mode 100644 index 00000000000..6bcf62a7adb --- /dev/null +++ b/tests/integration/routing/test_cooldown_reads_scoped_to_candidates.py @@ -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 diff --git a/tests/unit/caching/test_redis_batch.py b/tests/unit/caching/test_redis_batch.py index cd270035416..66fd9b8199a 100644 --- a/tests/unit/caching/test_redis_batch.py +++ b/tests/unit/caching/test_redis_batch.py @@ -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 diff --git a/tests/unit/caching/test_request_redis_batch_pre_call.py b/tests/unit/caching/test_request_redis_batch_pre_call.py index b1da2a413b6..ac04bb434d2 100644 --- a/tests/unit/caching/test_request_redis_batch_pre_call.py +++ b/tests/unit/caching/test_request_redis_batch_pre_call.py @@ -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")) From 5559a4976f3cb7db03be90655ed6264a5e496fe0 Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 10:23:47 -0700 Subject: [PATCH 2/3] test(router): freeze the clock in the health-check cooldown error test The unit test stamped the unhealthy state with time.time(), so its result depended on when it ran. Freeze time at a fixed value, stamp the state with it, and assert the health check really dropped the deployment first. --- .../unit/caching/test_request_redis_batch_pre_call.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/unit/caching/test_request_redis_batch_pre_call.py b/tests/unit/caching/test_request_redis_batch_pre_call.py index ac04bb434d2..f72580f893b 100644 --- a/tests/unit/caching/test_request_redis_batch_pre_call.py +++ b/tests/unit/caching/test_request_redis_batch_pre_call.py @@ -6,7 +6,6 @@ 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 @@ -16,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 @@ -559,14 +559,19 @@ async def test_a_direct_cooldown_read_covers_only_the_groups_deployments_and_sti @pytest.mark.asyncio -async def test_a_cooling_deployment_the_health_check_already_dropped_still_shows_in_the_no_deployment_error() -> None: +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": time.time(), "reason": "probe failed"}} + {"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): From 67535847016bcd4652c388a59176f82ce59982ee Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:04:11 -0700 Subject: [PATCH 3/3] test(router): pin cooldown read scope under failures, streaming and a hung Redis Adds integration cells at two workers under simple-shuffle and usage-based-routing-v2: a cooled deployment stays skipped by fresh workers that learn it only from Redis, the no-deployments 429 still lists cooled deployments of other groups, streamed chat, messages and responses read only their group's keys, a complexity router reads only its entry and tier, and a burst across all three endpoints keeps answering while Redis hangs, lands one spend row per request, and reads only the group's keys once Redis is back. --- ...est_cooldown_reads_scoped_under_failure.py | 459 ++++++++++++++++++ 1 file changed, 459 insertions(+) create mode 100644 tests/integration/routing/test_cooldown_reads_scoped_under_failure.py diff --git a/tests/integration/routing/test_cooldown_reads_scoped_under_failure.py b/tests/integration/routing/test_cooldown_reads_scoped_under_failure.py new file mode 100644 index 00000000000..0ee81f54b8e --- /dev/null +++ b/tests/integration/routing/test_cooldown_reads_scoped_under_failure.py @@ -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)}"