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"))