diff --git a/litellm/router.py b/litellm/router.py index 233e7581a40..7785d06707f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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, 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 a1b1e327802..aa4ad085af1 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 @@ -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 = [] 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/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)}" 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..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,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"))