From e136e91273bad186baf9bf89beab3677cfdef0cd Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 14 Oct 2026 21:10:09 -0700 Subject: [PATCH] fix(router): keep the prompt caching affinity pin alive as long as the prefix's cache_control ttl (#46084) * fix(router): keep the prompt caching affinity pin alive as long as the prefix's cache_control ttl A prefix marked cache_control ttl 1h lost its pinned deployment after the fixed 300 s, so the next request went to another deployment and paid a cold cache. The pin TTL now follows the longest ephemeral breakpoint in the cacheable prefix, tools included: 3600 s for a 1h breakpoint and 300 s otherwise. * fix(router): give prompt cache pins their own tier and honour any declared ttl Pins now live in a DualCache of their own instead of the router's 200-slot in-memory tier, so an hour-long pin no longer evicts cooldowns and rate windows. The pin lifetime follows any s|m|h cache_control ttl, floored at 300 s, not only the literal "1h". * fix(router): hash the pin ttl off the event loop and flush prompt cache pins with the router cache Compute the cache_control ttl scan through offload_token_count beside the prefix hash, clear the dedicated pin cache in Router.flush_cache, and point the order-fallback and priority-group pin tests at the pin cache they read * fix(router): pin 1h prefixes only on 1h-capable models, cap the pin at 1h, and give pins 10k in-memory slots * fix(router): match the 1h pin gate to what Bedrock forwards Bedrock deployments now decide the 1h pin with is_claude_4_5_on_bedrock, the same check that decides whether the 1h ttl reaches Bedrock, so an application inference profile ARN keeps a 5 minute pin whatever its base_model prices * test(router): give the anthropic_messages pin fixture its provider * fix(router): let each pin write set its own lifetime and require a positive 1h price In a group mixing 1h and 5m deployments the in-memory pin kept the lifetime of its first write, so a 5m deployment could hold an hour-long pin. Each write now clears the in-memory entry first, matching Redis. The Bedrock gate uses supports_bedrock_prompt_cache_ttl after #46058, and a zero 1h price no longer counts as 1h support. * test(router): integration cells for the one-hour prompt cache pin lifetime Cover every surface the pin lifetime touches: chat and /v1/messages through both SDKs, the responses bridge, response-cache hits, each provider family's 1h gate, malformed and duplicated ttl values, rewrites of one prefix, mixed and cross-group routing, concurrency, every Redis attach path, the in-memory slot pressure without Redis reads, a Redis outage and a killed worker * fix(caching): bound the in-memory expiry heap and count prompt cache pins in the memory debug routes * test(integration): cover prompt cache pins in the memory debug routes and the bounded expiration heap * fix(caching): serialize in-memory cache writes so a heap rebuild cannot lose another thread's expiry entry * fix(caching): keep the in-memory cache bounded when entries have no usable expiry * test(caching): cover the response cache bound when cached entries have no usable ttl --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/caching/in_memory_cache.py | 56 +- litellm/constants.py | 1 + litellm/proxy/common_utils/debug_utils.py | 48 +- litellm/router.py | 9 +- .../prompt_caching_deployment_check.py | 14 +- litellm/router_utils/prompt_caching_cache.py | 57 +- .../test_in_memory_cache_heap_bound.py | 361 ++++++++++ .../routing/_prompt_cache_pin_support.py | 281 ++++++++ .../test_prompt_cache_pin_memory_routes.py | 204 ++++++ .../routing/test_prompt_cache_pin_ttl.py | 646 ++++++++++++++++++ .../test_prompt_cache_pin_ttl_chaos.py | 248 +++++++ tests/unit/caching/test_in_memory_cache.py | 64 ++ .../proxy/common_utils/test_debug_utils.py | 54 +- .../test_router_routing_groups.py | 2 +- .../test_prompt_caching_deployment_check.py | 110 ++- .../router_utils/test_prompt_caching_cache.py | 270 +++++++- tests/unit/test_router/test_router.py | 125 +++- tests/unit/test_router_order_fallback.py | 2 +- 18 files changed, 2503 insertions(+), 49 deletions(-) create mode 100644 tests/integration/caching/test_in_memory_cache_heap_bound.py create mode 100644 tests/integration/routing/_prompt_cache_pin_support.py create mode 100644 tests/integration/routing/test_prompt_cache_pin_memory_routes.py create mode 100644 tests/integration/routing/test_prompt_cache_pin_ttl.py create mode 100644 tests/integration/routing/test_prompt_cache_pin_ttl_chaos.py diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 56c9147e066..cb6426f64ff 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -51,6 +51,7 @@ class InMemoryCache(BaseCache): self.ttl_dict: dict = {} self.expiration_heap: list[tuple[float, str]] = [] self._increment_lock = threading.Lock() + self._write_lock = threading.Lock() self._clock = clock if clock is not None else lambda: time.time() def check_value_size(self, value: Any): @@ -96,6 +97,9 @@ class InMemoryCache(BaseCache): """ return key in self.ttl_dict and self._clock() > self.ttl_dict[key] + def _cache_keys(self) -> tuple[str, ...]: + return tuple(self.cache_dict) # pyright: ignore[reportUnknownVariableType] # cache_dict is untyped legacy state + def _remove_key(self, key: str) -> None: """ Remove a key from both cache_dict and ttl_dict @@ -133,13 +137,23 @@ class InMemoryCache(BaseCache): # Case 3: Entry is valid and not expired break + if len(self.expiration_heap) > 2 * self.max_size_in_memory: + self.expiration_heap = [entry for entry in self.expiration_heap if self.ttl_dict.get(entry[1]) == entry[0]] + heapq.heapify(self.expiration_heap) + # Step 2: Evict if cache is still full - while len(self.cache_dict) >= self.max_size_in_memory: + while self.expiration_heap and len(self.cache_dict) >= self.max_size_in_memory: expiration_time, key = heapq.heappop(self.expiration_heap) # Skip if key was removed or updated if self.ttl_dict.get(key) == expiration_time: self._remove_key(key) + if not self.expiration_heap: + untracked_keys: Final = self._cache_keys() + overflow: Final = len(untracked_keys) - self.max_size_in_memory + 1 + for untracked_key in untracked_keys[: max(overflow, 0)]: + self._remove_key(untracked_key) + # de-reference the removed item # https://www.geeksforgeeks.org/diagnosing-and-fixing-memory-leaks-in-python/ # One of the most common causes of memory leaks in Python is the retention of objects that are no longer being used. @@ -150,7 +164,7 @@ class InMemoryCache(BaseCache): Check if ttl is set for a key """ ttl_time: Final = self.ttl_dict.get(key) - if ttl_time is None or float(ttl_time) < self._clock(): # if ttl is not set, allow override + if ttl_time is None or ttl_time < self._clock(): # if ttl is not set, allow override return True else: return False @@ -160,21 +174,22 @@ class InMemoryCache(BaseCache): if self.max_size_in_memory == 0: return # Don't cache anything if max size is 0 - # Always prune expired/outdated heap roots before inserting. - # This keeps expiration_heap bounded even when the live cache stays - # below max_size_in_memory and keys are reinserted after TTL expiry. - self.evict_cache() - if not self.check_value_size(value): - return + with self._write_lock: + # Always prune expired/outdated heap roots before inserting. + # This keeps expiration_heap bounded even when the live cache stays + # below max_size_in_memory and keys are reinserted after TTL expiry. + self.evict_cache() + if not self.check_value_size(value): + return - self.cache_dict[key] = value - if self.allow_ttl_override(key): # if ttl is not set, set it to default ttl - if "ttl" in kwargs and kwargs["ttl"] is not None: - self.ttl_dict[key] = self._clock() + float(kwargs["ttl"]) - heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key)) - else: - self.ttl_dict[key] = self._clock() + self.default_ttl - heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key)) + self.cache_dict[key] = value + if self.allow_ttl_override(key): # if ttl is not set, set it to default ttl + if "ttl" in kwargs and kwargs["ttl"] is not None: + self.ttl_dict[key] = self._clock() + float(kwargs["ttl"]) + heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key)) + else: + self.ttl_dict[key] = self._clock() + self.default_ttl + heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key)) async def async_set_cache(self, key, value, **kwargs): self.set_cache(key=key, value=value, **kwargs) @@ -258,14 +273,15 @@ class InMemoryCache(BaseCache): return results def flush_cache(self): - self.cache_dict.clear() - self.ttl_dict.clear() - self.expiration_heap.clear() + with self._write_lock: + self.cache_dict.clear() + self.ttl_dict.clear() + self.expiration_heap.clear() async def disconnect(self): pass - def delete_cache(self, key): + def delete_cache(self, key: str) -> None: self._remove_key(key) async def async_get_ttl(self, key: str) -> int | None: diff --git a/litellm/constants.py b/litellm/constants.py index e78e3e36364..c72e7ab777b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -478,6 +478,7 @@ MINIMUM_PROMPT_CACHE_TOKEN_COUNT: Final = ( else DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT ) PROMPT_CACHE_LOOKBACK_POSITIONS: Final = 20 +PROMPT_CACHE_PIN_CACHE_MAX_SIZE: Final = 10_000 DEFAULT_TRIM_RATIO: Final = get_env_float( "DEFAULT_TRIM_RATIO", 0.75 ) # default ratio of tokens to trim from the end of a prompt diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 5b77761d6dc..9c2cc768ae8 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -8,7 +8,7 @@ import sys import tracemalloc from collections import Counter from collections.abc import Mapping, Sequence -from typing import Annotated, Any, Final, NamedTuple, Protocol, TypedDict +from typing import TYPE_CHECKING, Annotated, Any, Final, NamedTuple, Protocol, TypedDict from fastapi import APIRouter, Depends, HTTPException, Query from typing_extensions import ReadOnly @@ -22,6 +22,9 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.bug_report_config import build_proxy_environment_report from litellm.proxy.common_utils.resource_ownership import is_proxy_admin +if TYPE_CHECKING: + from litellm.router import Router + router: Final = APIRouter() @@ -127,6 +130,13 @@ if os.environ.get("LITELLM_PROFILE", "false").lower() == "true": return {"top_50_memory_usage": result} +def _prompt_caching_pins(llm_router: "Router | None") -> tuple[Mapping[str, object], Mapping[str, float]]: + if llm_router is None: + return {}, {} + pin_cache: Final = llm_router.prompt_caching_pin_cache.in_memory_cache + return pin_cache.cache_dict, pin_cache.ttl_dict + + @router.get("/memory-usage-in-mem-cache", include_in_schema=False) async def memory_usage_in_mem_cache( _: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -151,6 +161,9 @@ async def memory_usage_in_mem_cache( llm_router.cache.in_memory_cache.ttl_dict ) + pin_entries, pin_expiries = _prompt_caching_pins(llm_router) + num_items_in_prompt_caching_pin_cache: Final = len(pin_entries) + len(pin_expiries) + num_items_in_user_api_key_cache: Final = ( len(user_api_key_cache.in_memory_cache.cache_dict) + len(user_api_key_cache.in_memory_cache.ttl_dict) @@ -165,6 +178,7 @@ async def memory_usage_in_mem_cache( return { "num_items_in_user_api_key_cache": num_items_in_user_api_key_cache, "num_items_in_llm_router_cache": num_items_in_llm_router_cache, + "num_items_in_prompt_caching_pin_cache": num_items_in_prompt_caching_pin_cache, "num_items_in_proxy_logging_obj_cache": num_items_in_proxy_logging_obj_cache, } @@ -186,20 +200,17 @@ async def memory_usage_in_mem_cache_items( user_api_key_cache, ) - if llm_router is None: - llm_router_in_memory_cache_dict = {} - llm_router_in_memory_ttl_dict = {} - else: - llm_router_in_memory_cache_dict = llm_router.cache.in_memory_cache.cache_dict - llm_router_in_memory_ttl_dict = llm_router.cache.in_memory_cache.ttl_dict + pin_entries, pin_expiries = _prompt_caching_pins(llm_router) return { "user_api_key_cache": user_api_key_cache.in_memory_cache.cache_dict, "user_api_key_ttl": user_api_key_cache.in_memory_cache.ttl_dict, "user_key_object_cache": user_api_key_cache.key_object_cache.in_memory_cache.cache_dict, "user_key_object_ttl": user_api_key_cache.key_object_cache.in_memory_cache.ttl_dict, - "llm_router_cache": llm_router_in_memory_cache_dict, - "llm_router_ttl": llm_router_in_memory_ttl_dict, + "llm_router_cache": {} if llm_router is None else llm_router.cache.in_memory_cache.cache_dict, + "llm_router_ttl": {} if llm_router is None else llm_router.cache.in_memory_cache.ttl_dict, + "prompt_caching_pin_cache": pin_entries, + "prompt_caching_pin_ttl": pin_expiries, "proxy_logging_obj_cache": proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict, "proxy_logging_obj_ttl": proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.ttl_dict, } @@ -363,12 +374,18 @@ async def get_memory_summary( # Router cache if llm_router is not None: router_cache_items: Final = len(llm_router.cache.in_memory_cache.cache_dict) - total_cache_items += router_cache_items + pin_cache_items: Final = len(_prompt_caching_pins(llm_router)[0]) + total_cache_items += router_cache_items + pin_cache_items caches["llm_responses"] = { "count": router_cache_items, "count_readable": f"{router_cache_items:,}", "what_it_stores": "LLM responses for identical requests", } + caches["prompt_caching_pins"] = { + "count": pin_cache_items, + "count_readable": f"{pin_cache_items:,}", + "what_it_stores": "Deployment pins for cached prompt prefixes", + } # Proxy logging cache logging_cache_items: Final = len(proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict) @@ -482,7 +499,7 @@ def _get_uncollectable_objects_info() -> Mapping[str, object]: def _get_cache_memory_stats( - user_api_key_cache, llm_router, proxy_logging_obj, redis_usage_cache + user_api_key_cache, llm_router: "Router | None", proxy_logging_obj, redis_usage_cache ) -> Mapping[str, object]: """Calculate memory usage for all caches.""" cache_stats: Final[dict[str, object]] = {} @@ -513,6 +530,15 @@ def _get_cache_memory_stats( "ttl_dict_size_bytes": router_ttl_size, "total_size_mb": round((router_cache_size + router_ttl_size) / (1024 * 1024), 2), } + pin_entries, pin_expiries = _prompt_caching_pins(llm_router) + pin_cache_size: Final = sys.getsizeof(pin_entries) + pin_ttl_size: Final = sys.getsizeof(pin_expiries) + cache_stats["prompt_caching_pin_cache"] = { + "num_items": len(pin_entries), + "cache_dict_size_bytes": pin_cache_size, + "ttl_dict_size_bytes": pin_ttl_size, + "total_size_mb": round((pin_cache_size + pin_ttl_size) / (1024 * 1024), 2), + } # Proxy logging cache logging_cache_size = sys.getsizeof(proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict) diff --git a/litellm/router.py b/litellm/router.py index 0696012c2a5..fb8dbe40eb8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -69,6 +69,7 @@ from litellm.constants import ( DEFAULT_MAX_LRU_CACHE_SIZE, INTERNAL_CALL_ORIGIN_METADATA_KEY, OUTPUT_TOKEN_CEILING_PARAMS, + PROMPT_CACHE_PIN_CACHE_MAX_SIZE, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY, RUNTIME_UPDATABLE_ROUTER_SETTINGS, @@ -1204,6 +1205,10 @@ class Router: redis_cache=redis_cache, in_memory_cache=InMemoryCache(), ) + self.prompt_caching_pin_cache: DualCache = DualCache( + redis_cache=redis_cache, + in_memory_cache=InMemoryCache(max_size_in_memory=PROMPT_CACHE_PIN_CACHE_MAX_SIZE), + ) ### SCHEDULER ### self.scheduler = Scheduler(polling_interval=polling_interval, redis_cache=redis_cache) @@ -1533,6 +1538,7 @@ class Router: """ self.cache.attach_redis_cache(cache) self._claude_code_session_router_cache.attach_redis_cache(cache) + self.prompt_caching_pin_cache.attach_redis_cache(cache) _update_redis_cache = update_redis_cache @@ -2613,7 +2619,7 @@ class Router: continue if pre_call_check == "prompt_caching": _callback = PromptCachingDeploymentCheck( - cache=self.cache, is_priority_group=self._is_priority_routing_group + cache=self.prompt_caching_pin_cache, is_priority_group=self._is_priority_routing_group ) elif pre_call_check == "router_budget_limiting": if self.get_router_deployment_budget_limiter() is not None: @@ -15460,6 +15466,7 @@ class Router: session_in_memory_cache: Final = self._claude_code_session_router_cache.in_memory_cache if session_in_memory_cache is not None: session_in_memory_cache.flush_cache() + self.prompt_caching_pin_cache.in_memory_cache.flush_cache() def reset(self): ## clean up on close diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 9f1558dedae..091472c1ea3 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -16,6 +16,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( ) from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.llms.bedrock.common_utils import supports_bedrock_prompt_cache_ttl from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes, StandardLoggingPayload from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt @@ -48,6 +49,13 @@ def _get_min_token_count_for_deployments(healthy_deployments: list[dict]) -> int ) +def _honors_one_hour_cache_ttl(standard_logging_object: StandardLoggingPayload) -> bool: + if standard_logging_object["custom_llm_provider"] == "bedrock": + return supports_bedrock_prompt_cache_ttl(standard_logging_object["model"]) + model_map_value: Final = standard_logging_object["model_map_information"]["model_map_value"] + return model_map_value is not None and (model_map_value.get("cache_creation_input_token_cost_above_1hr") or 0) > 0 + + class PromptCachingDeploymentCheck(CustomLogger): def __init__(self, cache: DualCache, is_priority_group: Callable[[str], bool] | None = None): super().__init__() @@ -146,18 +154,20 @@ class PromptCachingDeploymentCheck(CustomLogger): ) return + prompt_messages: Final = cast(list[AllMessageValues], messages) ## PROMPT CACHING - cache model id, if prompt caching valid prompt + provider if await offload_token_count(is_prompt_caching_valid_prompt)( model=model, - messages=cast(list[AllMessageValues], messages), + messages=prompt_messages, ): cache: Final = PromptCachingCache( cache=self.cache, ) await cache.async_add_model_id( model_id=model_id, - messages=messages, + messages=prompt_messages, tools=None, # [TODO]: add tools once standard_logging_object supports it + honors_one_hour_ttl=_honors_one_hour_cache_ttl(standard_logging_object), ) return diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index e6a8dc88086..07607f9943c 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -6,7 +6,7 @@ import hashlib import json from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass -from itertools import accumulate +from itertools import accumulate, chain from typing import TYPE_CHECKING, Any, Final, cast from pydantic import JsonValue, TypeAdapter @@ -37,6 +37,7 @@ class PromptCachingCacheValue(TypedDict): PROMPT_CACHE_PIN_TTL_SECONDS: Final = 300 +PROMPT_CACHE_ONE_HOUR_PIN_TTL_SECONDS: Final = 3600 _PROMPT_CACHE_PINS_TARGET: Final = "prompt_cache_pins" _TOOL_RUN_BLOCK_TYPES: Final = frozenset({"tool_use", "tool_result"}) _PREFIX_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...]) @@ -117,6 +118,20 @@ def _lookback_keys(positions: tuple[PrefixPosition, ...]) -> tuple[str, ...]: return tuple(entry.cache_key for entry in reversed(positions) if entry.position > oldest_probed_position) +def _breakpoint_pin_ttl_seconds(carrier: JsonValue | Mapping[str, JsonValue]) -> int | None: + if not isinstance(carrier, Mapping): + return None + cache_control: Final = carrier.get("cache_control") + if not isinstance(cache_control, dict) or cache_control.get("type") != "ephemeral": + return None + return PROMPT_CACHE_ONE_HOUR_PIN_TTL_SECONDS if cache_control.get("ttl") == "1h" else PROMPT_CACHE_PIN_TTL_SECONDS + + +def _breakpoint_carriers(message: Mapping[str, JsonValue]) -> tuple[JsonValue | Mapping[str, JsonValue], ...]: + content: Final = message.get("content") + return (message, *content) if isinstance(content, list) else (message,) + + def _pinned_value(value: JsonValue) -> PromptCachingCacheValue | None: if not isinstance(value, dict): return None @@ -263,6 +278,27 @@ class PromptCachingCache: return () return await offload_token_count(PromptCachingCache.prefix_positions)(messages, tools) + @staticmethod + def pin_ttl_seconds( + messages: Sequence[AllMessageValues] | None, + tools: Sequence[ChatCompletionToolParam] | None, + ) -> int: + if not messages: + return PROMPT_CACHE_PIN_TTL_SECONDS + prefix: Final = _PREFIX_ADAPTER.validate_python( + to_jsonable_python( + PromptCachingCache.extract_cacheable_prefix(list(messages)), serialize_unknown=True, bytes_mode="base64" + ) + ) + tool_carriers: Final = _TOOLS_ADAPTER.validate_python( + to_jsonable_python(tools or (), serialize_unknown=True, bytes_mode="base64") + ) + carriers: Final = chain(tool_carriers, *map(_breakpoint_carriers, prefix)) + return max( + (ttl for ttl in map(_breakpoint_pin_ttl_seconds, carriers) if ttl is not None), + default=PROMPT_CACHE_PIN_TTL_SECONDS, + ) + @staticmethod def get_prompt_caching_cache_key( messages: list[AllMessageValues] | None, @@ -276,28 +312,43 @@ class PromptCachingCache: model_id: str, messages: list[AllMessageValues] | None, tools: Sequence[ChatCompletionToolParam] | None, + honors_one_hour_ttl: bool = False, ) -> None: cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools) if cache_key is None: return - self.cache.set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=PROMPT_CACHE_PIN_TTL_SECONDS) + self.cache.in_memory_cache.delete_cache(cache_key) + self.cache.set_cache( + cache_key, + PromptCachingCacheValue(model_id=model_id), + ttl=PromptCachingCache.pin_ttl_seconds(messages, tools) + if honors_one_hour_ttl + else PROMPT_CACHE_PIN_TTL_SECONDS, + ) async def async_add_model_id( self, model_id: str, messages: list[AllMessageValues] | None, tools: Sequence[ChatCompletionToolParam] | None, + honors_one_hour_ttl: bool = False, ) -> None: positions: Final = await PromptCachingCache.async_prefix_positions(messages, tools) if not positions: return + ttl: Final = ( + await offload_token_count(PromptCachingCache.pin_ttl_seconds)(messages, tools) + if honors_one_hour_ttl + else PROMPT_CACHE_PIN_TTL_SECONDS + ) + self.cache.in_memory_cache.delete_cache(positions[-1].cache_key) with service_target(_PROMPT_CACHE_PINS_TARGET): await self.cache.async_set_cache( positions[-1].cache_key, PromptCachingCacheValue(model_id=model_id), - ttl=PROMPT_CACHE_PIN_TTL_SECONDS, + ttl=ttl, ) async def async_get_model_id( diff --git a/tests/integration/caching/test_in_memory_cache_heap_bound.py b/tests/integration/caching/test_in_memory_cache_heap_bound.py new file mode 100644 index 00000000000..7cbfaa7f29a --- /dev/null +++ b/tests/integration/caching/test_in_memory_cache_heap_bound.py @@ -0,0 +1,361 @@ +import itertools +import json +import uuid +from collections.abc import Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +import yaml +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.openai_wire import answering_model_discovery, chat_reply, responses_reply +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import anthropic_peer, marker_of +from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 300) + +_FIRST: Final = "heap-bound-first" +_SECOND: Final = "heap-bound-second" +_SLOTS: Final = 8 +_CYCLES: Final = 6 * _SLOTS +_FILLERS: Final = 2 * _SLOTS +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def _rig_config() -> dict[str, JsonValue]: + return _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + + +def _config(directory: Path, api_base: str) -> Path: + config: Final = _rig_config() + deployments: Final = [ + { + "model_name": name, + "litellm_params": {"model": "anthropic/claude-opus-5-5", "api_base": api_base, "api_key": "synthetic"}, + } + for name in (_FIRST, _SECOND) + ] + merged: Final = { + **config, + "model_list": deployments, + "general_settings": { + **object_value(config["general_settings"]), + "user_api_key_cache_ttl": 3600, + "user_api_key_cache_max_size": _SLOTS, + }, + "litellm_settings": { + **object_value(config["litellm_settings"]), + "cache": False, + "enable_redis_auth_cache": False, + }, + } + path: Final = directory / f"heap-bound-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(merged)) + return path + + +def _chat(gateway: Gateway, key: str, model: str) -> tuple[int, str]: + marker: Final = uuid.uuid4().hex + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": f"Reply. marker-{marker}"}]}, + key=key, + ) + return response.status_code, marker + + +def _cycle(gateway: Gateway, key: str, turn: int) -> tuple[tuple[int, int], str]: + allowed, denied = (_SECOND, _FIRST) if turn % 2 == 0 else (_FIRST, _SECOND) + gateway.post("/key/update", {"key": key, "models": [allowed]}) + allowed_status, marker = _chat(gateway, key, allowed) + denied_status, _ = _chat(gateway, key, denied) + return (allowed_status, denied_status), marker + + +def _fillers(scenario: Scenario, count: int) -> Sequence[str]: + keys: Final = tuple(string_value(scenario.gateway.post("/key/generate", {})["key"]) for _ in range(count)) + scenario.cleanups.callback(scenario.gateway.post, "/key/delete", {"keys": list(keys)}) + return keys + + +def _reached(peer: Wire) -> list[str]: + return sorted(marker_of(request) for request in peer.drain()) + + +def test_key_updates_past_twice_the_key_cache_slots_take_effect_on_the_next_request( + gateway: Gateway, tmp_path: Path +) -> None: + with wire_server(anthropic_peer) as peer: + config: Final = _config(tmp_path, peer.url) + with ( + owned_proxy_process(gateway, tmp_path, {}, config=config, workers=1) as owned, + owned.gateway.scenario() as scenario, + ): + anchor: Final = scenario.key(models=[_FIRST]) + anchor_status, anchor_marker = _chat(owned.gateway, anchor, _FIRST) + churned: Final = scenario.key(models=[_FIRST]) + cycles: Final = tuple(_cycle(owned.gateway, churned, turn) for turn in range(_CYCLES)) + assert [statuses for statuses, _ in cycles] == [(200, 403)] * _CYCLES + filled: Final = tuple( + owned.gateway.request("GET", "/v1/models", key=key).status_code for key in _fillers(scenario, _FILLERS) + ) + assert filled == (200,) * _FILLERS + after_anchor, after_anchor_marker = _chat(owned.gateway, anchor, _FIRST) + last_allowed: Final = _SECOND if (_CYCLES - 1) % 2 == 0 else _FIRST + after_churned, after_churned_marker = _chat(owned.gateway, churned, last_allowed) + assert (anchor_status, after_anchor, after_churned) == (200, 200, 200) + assert _reached(peer) == sorted( + (anchor_marker, *(marker for _, marker in cycles), after_anchor_marker, after_churned_marker) + ) + + +_RESPONSE_SLOTS: Final = DEFAULT_MAX_SIZE_IN_MEMORY +_OVERFLOW: Final = 20 +_BURST_PER_ENDPOINT: Final = 6 +_NO_STORE: Final[Mapping[str, JsonValue]] = MappingProxyType({"no-store": True}) +_FRESH_TTL: Final[Mapping[str, JsonValue]] = MappingProxyType({"ttl": 600}) +_EMBEDDING: Final = json.dumps( + { + "object": "list", + "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}], + "model": "text-embedding-3-large", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + } +).encode() + + +class _PeerBody(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + model: str = "" + stream: bool = False + + +def _response_cache_peer(request: Request) -> Reply: + body: Final = _PeerBody.model_validate_json(request.body or b"{}") + identity: Final = f"{marker_of(request)}-{uuid.uuid4().hex[:8]}" + match request.target.split("?", 1)[0]: + case "/v1/messages": + return anthropic_peer(request) + case "/v1/chat/completions": + return chat_reply(f"chatcmpl-{identity}", body.model, "scripted", stream=body.stream) + case "/v1/responses": + return responses_reply(f"resp_{identity}", body.model, "scripted", stream=body.stream) + case "/v1/embeddings": + return Reply(body=_EMBEDDING) + case _: + return Reply(status=404, body=b'{"error": {"message": "unscripted route"}}') + + +_NAN_TTL: Final[Mapping[str, JsonValue]] = MappingProxyType({"ttl": "NaN"}) +_UNUSABLE_TTLS: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"lowercase-nan": "nan", "word": "abc", "empty-string": "", "list": [1], "object": {"x": 1}} +) + + +def _prompt_fields(path: str, text: str) -> dict[str, JsonValue]: + match path: + case "/v1/chat/completions": + return {"messages": [{"role": "user", "content": text}]} + case "/v1/messages": + return {"max_tokens": 16, "messages": [{"role": "user", "content": text}]} + case _: + return {"input": text} + + +@dataclass(frozen=True, slots=True) +class _Endpoint: + path: str + model: str + streamed: bool + + def body(self, marker: str, cache: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return { + "model": self.model, + **_prompt_fields(self.path, f"Reply. marker-{marker}"), + **({"stream": True} if self.streamed else {}), + "cache": dict(cache), + } + + +_CHAT: Final = _Endpoint("/v1/chat/completions", "rc-openai", streamed=False) +_ENDPOINTS: Final = MappingProxyType( + { + "chat": _CHAT, + "chat-stream": _Endpoint("/v1/chat/completions", "rc-openai", streamed=True), + "messages": _Endpoint("/v1/messages", "rc-anthropic", streamed=False), + "messages-stream": _Endpoint("/v1/messages", "rc-anthropic", streamed=True), + "responses": _Endpoint("/v1/responses", "rc-openai", streamed=False), + "responses-stream": _Endpoint("/v1/responses", "rc-openai", streamed=True), + "embeddings": _Endpoint("/v1/embeddings", "rc-embedding", streamed=False), + } +) + + +def _posted_markers(peer: Wire) -> frozenset[str]: + return frozenset(marker_of(request) for request in peer.drain() if request.method == "POST") + + +@dataclass(frozen=True, slots=True) +class _CacheRig: + gateway: Gateway + peer: Wire + + def send(self, endpoint: _Endpoint, marker: str, cache: Mapping[str, JsonValue]) -> int: + return self.gateway.request("POST", endpoint.path, endpoint.body(marker, cache)).status_code + + def served(self, endpoint: _Endpoint, marker: str) -> bool: + self.peer.drain() + response: Final = self.gateway.request("POST", endpoint.path, endpoint.body(marker, _NO_STORE)) + assert response.status_code == 200, response.text + return marker not in _posted_markers(self.peer) + + +def _response_cache_config(directory: Path, peer_url: str) -> Path: + config: Final = _rig_config() + deployments: Final = [ + { + "model_name": "rc-openai", + "litellm_params": {"model": "openai/gpt-5.6", "api_base": f"{peer_url}/v1", "api_key": "synthetic"}, + }, + { + "model_name": "rc-anthropic", + "litellm_params": {"model": "anthropic/claude-haiku-5-5", "api_base": peer_url, "api_key": "synthetic"}, + }, + { + "model_name": "rc-embedding", + "litellm_params": { + "model": "openai/text-embedding-3-large", + "api_base": f"{peer_url}/v1", + "api_key": "synthetic", + }, + }, + ] + merged: Final = { + **config, + "model_list": deployments, + "litellm_settings": { + **object_value(config["litellm_settings"]), + "cache": True, + "cache_params": {"type": "local"}, + "enable_redis_auth_cache": False, + }, + } + path: Final = directory / f"response-cache-bound-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(merged)) + return path + + +@pytest.fixture(scope="module") +def response_cache_rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_CacheRig]: + directory: Final = tmp_path_factory.mktemp("response-cache-bound") + with ( + wire_server(answering_model_discovery(_response_cache_peer)) as peer, + owned_proxy_process( + session_gateway, directory, {}, config=_response_cache_config(directory, peer.url), workers=1 + ) as owned, + ): + yield _CacheRig(owned.gateway, peer) + + +def _flood(rig: _CacheRig, endpoint: _Endpoint, cache: Mapping[str, JsonValue], count: int) -> tuple[str, ...]: + markers: Final = tuple(uuid.uuid4().hex for _ in range(count)) + statuses: Final = tuple(rig.send(endpoint, marker, cache) for marker in markers) + assert statuses == (200,) * count + return markers + + +def _served_indexes(rig: _CacheRig, endpoint: _Endpoint, markers: Sequence[str]) -> tuple[int, ...]: + return tuple(index for index, marker in enumerate(markers) if rig.served(endpoint, marker)) + + +def _assert_newest_suffix(served: Sequence[int], flooded: int) -> None: + summary: Final = f"{len(served)} of {flooded} served, first {served[:3]}, last {served[-3:]}" + assert served, summary + assert tuple(served) == tuple(range(served[0], flooded)), summary + + +def _assert_next_response_cached(rig: _CacheRig, endpoint: _Endpoint) -> None: + fresh: Final = uuid.uuid4().hex + assert rig.send(endpoint, fresh, _FRESH_TTL) == 200 + assert eventually(lambda: rig.served(endpoint, fresh), lambda hit: hit, return_last_on_timeout=True), ( + "a response with a usable ttl was not cached after the flood" + ) + + +def _assert_flood_keeps_the_newest_responses( + rig: _CacheRig, endpoint: _Endpoint, cache: Mapping[str, JsonValue] +) -> None: + markers: Final = _flood(rig, endpoint, cache, _RESPONSE_SLOTS + _OVERFLOW) + eventually(lambda: rig.served(endpoint, markers[-1]), lambda hit: hit, return_last_on_timeout=True) + served: Final = _served_indexes(rig, endpoint, markers) + _assert_newest_suffix(served, len(markers)) + assert _RESPONSE_SLOTS - 1 <= len(served) <= _RESPONSE_SLOTS, f"{len(served)} served" + _assert_next_response_cached(rig, endpoint) + + +@pytest.mark.parametrize("endpoint", tuple(_ENDPOINTS.values()), ids=tuple(_ENDPOINTS)) +def test_nan_ttl_flood_past_the_response_cache_size_keeps_the_newest_responses_and_caches_the_next_one( + response_cache_rig: _CacheRig, endpoint: _Endpoint +) -> None: + _assert_flood_keeps_the_newest_responses(response_cache_rig, endpoint, _NAN_TTL) + + +@pytest.mark.parametrize("ttl", tuple(_UNUSABLE_TTLS.values()), ids=tuple(_UNUSABLE_TTLS)) +def test_unusable_ttl_shapes_flooded_past_the_response_cache_size_keep_the_newest_responses( + response_cache_rig: _CacheRig, ttl: JsonValue +) -> None: + _assert_flood_keeps_the_newest_responses(response_cache_rig, _CHAT, {"ttl": ttl}) + + +def _post(gateway: Gateway, endpoint: _Endpoint, marker: str) -> int: + with httpx.Client(base_url=gateway.client.base_url, timeout=60, trust_env=False) as client: + return client.post( + endpoint.path, + json=endpoint.body(marker, _NAN_TTL), + headers={"Authorization": f"Bearer {gateway.key}"}, + ).status_code + + +def _liveliness(gateway: Gateway) -> int: + with httpx.Client(base_url=gateway.client.base_url, timeout=60, trust_env=False) as client: + return client.get("/health/liveliness").status_code + + +def test_concurrent_burst_into_a_response_cache_full_of_nan_ttls_caches_every_burst_response( + response_cache_rig: _CacheRig, +) -> None: + rig: Final = response_cache_rig + fill: Final = _flood(rig, _CHAT, _NAN_TTL, _RESPONSE_SLOTS) + eventually(lambda: rig.served(_CHAT, fill[-1]), lambda hit: hit) + burst: Final = tuple( + (endpoint, uuid.uuid4().hex) + for endpoint, _ in itertools.product(_ENDPOINTS.values(), range(_BURST_PER_ENDPOINT)) + ) + with ThreadPoolExecutor(2 * len(burst)) as pool: + answers: Final = tuple(pool.submit(_post, rig.gateway, endpoint, marker) for endpoint, marker in burst) + probes: Final = tuple(pool.submit(_liveliness, rig.gateway) for _ in burst) + statuses: Final = tuple(answer.result() for answer in answers) + liveliness: Final = tuple(probe.result() for probe in probes) + assert statuses == (200,) * len(burst) + assert liveliness == (200,) * len(burst) + missing: Final = eventually( + lambda: tuple(marker for endpoint, marker in burst if not rig.served(endpoint, marker)), + lambda left: not left, + return_last_on_timeout=True, + ) + assert not missing, f"{len(missing)} of {len(burst)} burst responses were never served from the cache" + served_fill: Final = _served_indexes(rig, _CHAT, fill) + _assert_newest_suffix(served_fill, len(fill)) + assert _RESPONSE_SLOTS - len(burst) <= len(served_fill) + len(burst) <= _RESPONSE_SLOTS, ( + f"{len(served_fill)} fill and {len(burst)} burst responses served" + ) + _assert_next_response_cached(rig, _CHAT) diff --git a/tests/integration/routing/_prompt_cache_pin_support.py b/tests/integration/routing/_prompt_cache_pin_support.py new file mode 100644 index 00000000000..426c2053265 --- /dev/null +++ b/tests/integration/routing/_prompt_cache_pin_support.py @@ -0,0 +1,281 @@ +import json +import time +import uuid +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import yaml +from integration._support.client import Gateway, eventually, held, object_value +from integration._support.database import read_rows +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire +from integration.providers._cache_control_marks_support import anthropic_peer, bedrock_peer, marker_of +from pydantic import JsonValue, TypeAdapter +from redis import Redis + +FIVE_MINUTES: Final = 300 +ONE_HOUR: Final = 3600 +PROVIDER_KEY: Final = "synthetic-provider-key" +BEDROCK_TOKEN: Final = "synthetic-bedrock-bearer" +LONG_TEXT: Final = "Answer from the reference below. " + "lorem ipsum " * 2500 +ONE_HOUR_MARK: Final[dict[str, JsonValue]] = {"type": "ephemeral", "ttl": "1h"} +FIVE_MINUTE_MARK: Final[dict[str, JsonValue]] = {"type": "ephemeral", "ttl": "5m"} +DEFAULT_MARK: Final[dict[str, JsonValue]] = {"type": "ephemeral"} +OVERLOAD_TOKEN: Final = "scripted-overload-on-this-turn" +BEDROCK_ONE_HOUR_MODEL: Final = "us.anthropic.claude-sonnet-5-5" +BEDROCK_FIVE_MINUTE_MODEL: Final = "anthropic.claude-sonnet-4-20250514-v1:0" +INFERENCE_PROFILE_ARN: Final = "arn:aws:bedrock:us-east-1:000000000000:application-inference-profile/pinprofile" +_PIN_VALUE: Final = TypeAdapter(dict[str, str]) +_JSON: Final = TypeAdapter(JsonValue) +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +@dataclass(frozen=True, slots=True) +class Pin: + model_id: str + ttl: int + + +def lives_one_hour(pin: Pin) -> bool: + return FIVE_MINUTES < pin.ttl <= ONE_HOUR + + +def lives_five_minutes(pin: Pin) -> bool: + return 0 < pin.ttl <= FIVE_MINUTES + + +def new_marker() -> str: + return uuid.uuid4().hex + + +def marked(text: str, mark: JsonValue | None) -> dict[str, JsonValue]: + block: Final[dict[str, JsonValue]] = {"type": "text", "text": text} + return block if mark is None else {**block, "cache_control": mark} + + +def reference(session: str) -> str: + return f"{LONG_TEXT} session {session}" + + +def question(marker: str) -> str: + return f"Reply in one word. marker-{marker}" + + +def system_turn(session: str, mark: JsonValue | None) -> dict[str, JsonValue]: + return {"role": "system", "content": [marked(reference(session), mark)]} + + +def ask_turn(marker: str, mark: JsonValue | None = None) -> dict[str, JsonValue]: + if mark is None: + return {"role": "user", "content": question(marker)} + return {"role": "user", "content": [marked(question(marker), mark)]} + + +def follow_up(first_turn: Sequence[JsonValue], marker: str) -> list[JsonValue]: + return [*first_turn, {"role": "assistant", "content": "sunny"}, ask_turn(marker)] + + +def chat_body(model: str, messages: Sequence[JsonValue], **fields: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "messages": list(messages), "max_tokens": 16, **fields} + + +def read_pins(redis: Redis) -> Mapping[str, Pin]: + def pin_of(key: bytes) -> tuple[str, Pin] | None: + value: Final = redis.get(key) + ttl: Final = redis.ttl(key) + if value is None or ttl < 0: + return None + return key.decode(), Pin(model_id=_PIN_VALUE.validate_json(value)["model_id"], ttl=ttl) + + found: Final = (pin_of(key) for key in redis.scan_iter(match="*:prompt_caching", count=1000)) + return MappingProxyType(dict(entry for entry in found if entry is not None)) + + +def pins_since(redis: Redis, before: frozenset[str]) -> Mapping[str, Pin]: + return MappingProxyType({key: pin for key, pin in read_pins(redis).items() if key not in before}) + + +def one_new_pin(redis: Redis, before: frozenset[str]) -> tuple[str, Pin]: + pins: Final = eventually(lambda: pins_since(redis, before), lambda found: len(found) >= 1, seconds=30) + assert len(pins) == 1, pins + return next(iter(pins.items())) + + +def pin_at(redis: Redis, key: str) -> Pin | None: + return read_pins(redis).get(key) + + +def spend_rows(response_id: str) -> int: + return len(read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (response_id,))) + + +def spend_row_lands(response_id: str) -> None: + eventually(lambda: spend_rows(response_id), lambda count: count >= 1, seconds=70) + assert spend_rows(response_id) == 1 + + +def end_user_row_lands(user: str) -> None: + eventually( + lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE end_user=%s', (user,)), + lambda rows: len(rows) >= 1, + seconds=70, + ) + assert len(read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE end_user=%s', (user,))) == 1 + + +def no_pin_written(redis: Redis, before: frozenset[str], response_id: str) -> None: + spend_row_lands(response_id) + held(lambda: pins_since(redis, before), lambda found: not found, holding=3, seconds=20) + + +def _carries(value: JsonValue, ttl: str) -> bool: + if isinstance(value, list): + return any(_carries(item, ttl) for item in value) + if not isinstance(value, dict): + return False + marks: Final = (value.get("cache_control"), value.get("cachePoint")) + return any(isinstance(mark, dict) and mark.get("ttl") == ttl for mark in marks) or any( + _carries(item, ttl) for item in value.values() + ) + + +def wire_carries_one_hour(request: Request) -> bool: + return _carries(_JSON.validate_json(request.body), "1h") + + +def sent_with(wires: Sequence[Wire], marker: str) -> tuple[Request, ...]: + return tuple(request for wire in wires for request in wire.drain() if marker_of(request) == marker) + + +def _overloaded() -> Reply: + return Reply( + status=529, + body=json.dumps({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}).encode(), + ) + + +def claude_peer(request: Request) -> Reply: + if OVERLOAD_TOKEN.encode() in request.body: + return _overloaded() + return anthropic_peer(request) + + +def slow_claude_peer(request: Request) -> Reply: + time.sleep(1) + return claude_peer(request) + + +def vertex_peer(request: Request) -> Reply: + if request.target == "/_oauth/token": + return Reply( + body=json.dumps( + {"access_token": "synthetic-vertex-token", "token_type": "Bearer", "expires_in": 3600} + ).encode() + ) + return claude_peer(request) + + +def bedrock_runtime_peer(request: Request) -> Reply: + if request.target.endswith("/converse"): + return bedrock_peer(request) + return claude_peer(request) + + +def openai_peer(request: Request) -> Reply: + body: Final = _JSON_OBJECT.validate_json(request.body) + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{marker_of(request)}", + "object": "chat.completion", + "created": 1760000000, + "model": str(body.get("model")), + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "sunny"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 12, "completion_tokens": 1, "total_tokens": 13}, + } + ).encode() + ) + + +def anthropic_deployment(name: str, identity: str, api_base: str, **fields: JsonValue) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": { + "model": "anthropic/claude-opus-5-5", + "api_base": api_base, + "api_key": PROVIDER_KEY, + **fields, + }, + "model_info": {"id": identity}, + } + + +def bedrock_deployment( + name: str, identity: str, model: str, api_base: str, model_info: Mapping[str, JsonValue] = MappingProxyType({}) +) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": { + "model": model, + "api_key": BEDROCK_TOKEN, + "aws_region_name": "us-east-1", + "api_base": api_base, + "aws_bedrock_runtime_endpoint": api_base, + }, + "model_info": {"id": identity, **model_info}, + } + + +def deployment(name: str, identity: str, **litellm_params: JsonValue) -> dict[str, JsonValue]: + return {"model_name": name, "litellm_params": dict(litellm_params), "model_info": {"id": identity}} + + +def vertex_deployment(name: str, identity: str, api_base: str) -> dict[str, JsonValue]: + return deployment( + name, + identity, + model="vertex_ai/claude-opus-5-5", + api_base=api_base, + vertex_project="scripted-project", + vertex_location="us-east5", + vertex_credentials=service_account_json("scripted-project", api_base), + ) + + +def pin_config( + directory: Path, + model_list: Sequence[JsonValue], + *, + litellm_settings: Mapping[str, JsonValue] = MappingProxyType({}), + router_settings: Mapping[str, JsonValue] = MappingProxyType({}), + general_settings: Mapping[str, JsonValue] = MappingProxyType({}), +) -> Path: + config: Final = _JSON_OBJECT.validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + merged: Final = { + **config, + "model_list": list(model_list), + "general_settings": {**object_value(config["general_settings"]), **general_settings}, + "litellm_settings": {**object_value(config["litellm_settings"]), **litellm_settings}, + "router_settings": { + **object_value(config["router_settings"]), + "num_retries": 0, + "optional_pre_call_checks": ["prompt_caching"], + **router_settings, + }, + } + path: Final = directory / f"prompt-cache-pins-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(merged)) + return path + + +def post_chat(gateway: Gateway, body: Mapping[str, JsonValue]) -> tuple[str, str]: + response: Final = gateway.request("POST", "/v1/chat/completions", body) + assert response.status_code == 200, response.text + return str(_JSON_OBJECT.validate_json(response.content)["id"]), response.headers["x-litellm-model-id"] diff --git a/tests/integration/routing/test_prompt_cache_pin_memory_routes.py b/tests/integration/routing/test_prompt_cache_pin_memory_routes.py new file mode 100644 index 00000000000..9ca064fba59 --- /dev/null +++ b/tests/integration/routing/test_prompt_cache_pin_memory_routes.py @@ -0,0 +1,204 @@ +import time +from collections.abc import Callable, Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually, gateway_from_environment, object_value +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.redis_process import owned_redis +from integration._support.wire import wire_server +from integration.routing._prompt_cache_pin_support import ( + FIVE_MINUTE_MARK, + FIVE_MINUTES, + ONE_HOUR, + ONE_HOUR_MARK, + Pin, + anthropic_deployment, + ask_turn, + chat_body, + claude_peer, + lives_one_hour, + new_marker, + pin_config, + pins_since, + post_chat, + read_pins, + system_turn, +) +from pydantic import JsonValue +from redis import Redis + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 300) + +_GROUP: Final = "pin-memory" +_LEFT: Final = "pin-memory-left" +_RIGHT: Final = "pin-memory-right" +_ITEMS: Final = "/memory-usage-in-mem-cache-items" +_MASTER_KEY_ONLY: Final = ("/memory-usage-in-mem-cache", _ITEMS) +_POLLED_ROUTES: Final = (*_MASTER_KEY_ONLY, "/debug/memory/summary") +_NO_RESPONSE_CACHE: Final[Mapping[str, JsonValue]] = {"cache": False, "enable_redis_auth_cache": False} +_BURST: Final = 24 + + +@dataclass(frozen=True, slots=True) +class _Rig: + owned: OwnedProxy + redis: Redis + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("prompt-cache-pin-memory-routes") + with ( + gateway_from_environment() as gateway, + owned_redis(directory) as cache, + wire_server(claude_peer) as left, + wire_server(claude_peer) as right, + ): + config: Final = pin_config( + directory, + [anthropic_deployment(_GROUP, _LEFT, left.url), anthropic_deployment(_GROUP, _RIGHT, right.url)], + litellm_settings=_NO_RESPONSE_CACHE, + ) + with ( + owned_proxy_process( + gateway, directory, {"REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port)}, config=config, workers=1 + ) as owned, + Redis(host=cache.host, port=cache.port, socket_timeout=2) as redis, + ): + yield _Rig(owned, redis) + + +def _field(body: JsonValue, *path: str) -> JsonValue: + if not path: + return body + return _field(object_value(body)[path[0]], *path[1:]) + + +def _read(gateway: Gateway, path: str) -> dict[str, JsonValue]: + response: Final = gateway.request("GET", path) + assert response.status_code == 200, (path, response.status_code, response.text) + return object_value(response.json()) + + +def _write_pins(rig: _Rig, marks: Sequence[JsonValue]) -> Mapping[str, Pin]: + before: Final = frozenset(read_pins(rig.redis)) + for mark in marks: + post_chat(rig.owned.gateway, chat_body(_GROUP, [system_turn(new_marker(), mark), ask_turn(new_marker())])) + pins: Final = eventually(lambda: pins_since(rig.redis, before), lambda found: len(found) >= len(marks), seconds=30) + assert len(pins) == len(marks), pins + return pins + + +def test_items_route_lists_every_pin_with_its_own_lifetime_and_keeps_it_out_of_the_router_cache(rig: _Rig) -> None: + pins: Final = _write_pins(rig, (ONE_HOUR_MARK, FIVE_MINUTE_MARK)) + body: Final = _read(rig.owned.gateway, _ITEMS) + read_at: Final = time.time() + entries: Final = object_value(body["prompt_caching_pin_cache"]) + expiries: Final = object_value(body["prompt_caching_pin_ttl"]) + router_entries: Final = object_value(body["llm_router_cache"]) + assert sorted(lives_one_hour(pin) for pin in pins.values()) == [False, True], pins + for key, pin in pins.items(): + assert entries[key] == {"model_id": pin.model_id}, (key, entries.get(key), pin) + remaining: Final = float(str(expiries[key])) - read_at + assert 0 < remaining <= ONE_HOUR, (key, remaining) + assert (remaining > FIVE_MINUTES) is lives_one_hour(pin), (key, remaining, pin) + assert key not in router_entries, key + + +_COUNT_ROUTES: Final[Mapping[str, tuple[str, Callable[[dict[str, JsonValue]], JsonValue], int]]] = { + "mem-cache": ("/memory-usage-in-mem-cache", lambda body: body["num_items_in_prompt_caching_pin_cache"], 2), + "summary": ( + "/debug/memory/summary", + lambda body: _field(body, "caches", "breakdown", "prompt_caching_pins", "count"), + 1, + ), + "details": ( + "/debug/memory/details?include_process_info=false", + lambda body: _field(body, "cache_memory", "prompt_caching_pin_cache", "num_items"), + 1, + ), +} + + +@pytest.mark.parametrize("route", sorted(_COUNT_ROUTES)) +def test_count_routes_count_exactly_the_pins_the_pin_cache_holds(rig: _Rig, route: str) -> None: + path, count_of, per_pin = _COUNT_ROUTES[route] + pins: Final = _write_pins(rig, (ONE_HOUR_MARK, ONE_HOUR_MARK, ONE_HOUR_MARK)) + counted: Final = _read(rig.owned.gateway, path) + held_pins: Final = object_value(_read(rig.owned.gateway, _ITEMS)["prompt_caching_pin_cache"]) + assert frozenset(pins) <= frozenset(held_pins), frozenset(pins) - frozenset(held_pins) + assert count_of(counted) == per_pin * len(held_pins), (count_of(counted), len(held_pins)) + + +def test_summary_total_adds_the_pins_to_every_other_cache(rig: _Rig) -> None: + _write_pins(rig, (ONE_HOUR_MARK,)) + caches: Final = object_value(_read(rig.owned.gateway, "/debug/memory/summary")["caches"]) + breakdown: Final = object_value(caches["breakdown"]) + assert "prompt_caching_pins" in breakdown, breakdown + assert caches["total_items"] == sum(int(str(_field(entry, "count"))) for entry in breakdown.values()), caches + + +def _poll_routes(gateway: Gateway, rounds: int) -> tuple[int, ...]: + return tuple(gateway.request("GET", path).status_code for _ in range(rounds) for path in _POLLED_ROUTES) + + +def test_memory_routes_answer_while_a_burst_writes_pins(rig: _Rig) -> None: + before: Final = frozenset(read_pins(rig.redis)) + bodies: Final = tuple( + chat_body(_GROUP, [system_turn(new_marker(), ONE_HOUR_MARK), ask_turn(new_marker())]) for _ in range(_BURST) + ) + with ThreadPoolExecutor(max_workers=_BURST + 4) as pool: + polls: Final = tuple(pool.submit(_poll_routes, rig.owned.gateway, 8) for _ in range(4)) + answers: Final = tuple(pool.map(lambda body: post_chat(rig.owned.gateway, body), bodies)) + statuses: Final = tuple(status for poll in polls for status in poll.result()) + assert frozenset(statuses) == {200}, statuses + pins: Final = eventually(lambda: pins_since(rig.redis, before), lambda found: len(found) >= _BURST, seconds=60) + assert len(pins) == _BURST, pins + assert sorted(pin.model_id for pin in pins.values()) == sorted(served for _, served in answers) + held_pins: Final = object_value(_read(rig.owned.gateway, _ITEMS)["prompt_caching_pin_cache"]) + assert {key: held_pins.get(key) for key in pins} == {key: {"model_id": pin.model_id} for key, pin in pins.items()} + + +def test_master_key_only_memory_routes_refuse_a_virtual_key_and_show_it_no_pin(rig: _Rig) -> None: + pins: Final = _write_pins(rig, (ONE_HOUR_MARK,)) + with rig.owned.gateway.scenario() as scenario: + virtual_key: Final = scenario.key() + refused: Final = tuple(rig.owned.gateway.request("GET", path, key=virtual_key) for path in _MASTER_KEY_ONLY) + assert [response.status_code for response in refused] == [401, 401], [response.text for response in refused] + assert not any(key in response.text for key in pins for response in refused) + + +def test_memory_routes_report_zero_pins_on_a_proxy_with_no_deployments(tmp_path: Path) -> None: + config: Final = pin_config( + tmp_path, [], litellm_settings=_NO_RESPONSE_CACHE, general_settings={"store_model_in_db": False} + ) + with gateway_from_environment() as gateway, owned_proxy_process(gateway, tmp_path, {}, config=config) as owned: + counts: Final = _read(owned.gateway, "/memory-usage-in-mem-cache") + items: Final = _read(owned.gateway, _ITEMS) + summary: Final = _read(owned.gateway, "/debug/memory/summary") + details: Final = _read(owned.gateway, "/debug/memory/details?include_process_info=false") + assert counts["num_items_in_prompt_caching_pin_cache"] == 0, counts + assert (items["prompt_caching_pin_cache"], items["prompt_caching_pin_ttl"]) == ({}, {}), items + assert _field(summary, "caches", "breakdown", "prompt_caching_pins", "count") == 0, summary + assert _field(details, "cache_memory", "prompt_caching_pin_cache", "num_items") == 0, details + + +def test_memory_routes_report_zero_pins_on_a_proxy_that_loaded_no_config(tmp_path: Path) -> None: + unloaded: Final = {"CONFIG_FILE_PATH": str(tmp_path / "absent.yaml"), "STORE_MODEL_IN_DB": "False"} + with gateway_from_environment() as gateway, owned_proxy_process(gateway, tmp_path, unloaded) as owned: + counts: Final = _read(owned.gateway, "/memory-usage-in-mem-cache") + items: Final = _read(owned.gateway, _ITEMS) + breakdown: Final = object_value(_field(_read(owned.gateway, "/debug/memory/summary"), "caches", "breakdown")) + details: Final = object_value( + _read(owned.gateway, "/debug/memory/details?include_process_info=false")["cache_memory"] + ) + assert (counts["num_items_in_prompt_caching_pin_cache"], counts["num_items_in_llm_router_cache"]) == (0, 0), counts + assert [items[name] for name in ("prompt_caching_pin_cache", "prompt_caching_pin_ttl", "llm_router_cache")] == [ + {} + ] * 3, items + assert "prompt_caching_pins" not in breakdown and "llm_responses" not in breakdown, breakdown + assert "prompt_caching_pin_cache" not in details and "llm_router_cache" not in details, details diff --git a/tests/integration/routing/test_prompt_cache_pin_ttl.py b/tests/integration/routing/test_prompt_cache_pin_ttl.py new file mode 100644 index 00000000000..d4328cd0e31 --- /dev/null +++ b/tests/integration/routing/test_prompt_cache_pin_ttl.py @@ -0,0 +1,646 @@ +import asyncio +import json +from collections.abc import Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import anthropic +import openai +import pytest +from integration._support.client import eventually, gateway_from_environment, held +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.redis_process import owned_redis +from integration._support.wire import Wire, wire_server +from integration.routing._prompt_cache_pin_support import ( + BEDROCK_FIVE_MINUTE_MODEL, + BEDROCK_ONE_HOUR_MODEL, + DEFAULT_MARK, + FIVE_MINUTE_MARK, + INFERENCE_PROFILE_ARN, + ONE_HOUR_MARK, + OVERLOAD_TOKEN, + PROVIDER_KEY, + Pin, + anthropic_deployment, + ask_turn, + bedrock_deployment, + bedrock_runtime_peer, + chat_body, + claude_peer, + deployment, + end_user_row_lands, + follow_up, + lives_five_minutes, + lives_one_hour, + marked, + marker_of, + new_marker, + no_pin_written, + one_new_pin, + openai_peer, + pin_at, + pin_config, + pins_since, + post_chat, + question, + read_pins, + reference, + sent_with, + slow_claude_peer, + spend_row_lands, + spend_rows, + system_turn, + vertex_deployment, + vertex_peer, + wire_carries_one_hour, +) +from pydantic import JsonValue +from redis import Redis + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 300) + +_CLAUDE: Final = "pin-claude" +_SLOW: Final = "pin-slow" +_MIXED: Final = "pin-mixed" +_OPENAI: Final = "pin-openai" +_LEFT: Final = "pin-claude-left" +_RIGHT: Final = "pin-claude-right" +_PRIMARY: Final = "pin-mixed-primary" +_SECONDARY: Final = "pin-mixed-secondary" +_MALFORMED_TTLS: Final[tuple[JsonValue, ...]] = (3600, "2h", "1H", "", None, ["1h"], "1h" * 2500) + + +@dataclass(frozen=True, slots=True) +class _Family: + group: str + identity: str + one_hour: bool + wire_one_hour: bool | None + + +_FAMILIES: Final = ( + _Family("pin-bedrock-converse-1h", "pin-bedrock-converse-1h", one_hour=True, wire_one_hour=True), + _Family("pin-bedrock-converse-5m", "pin-bedrock-converse-5m", one_hour=False, wire_one_hour=False), + _Family("pin-bedrock-invoke-1h", "pin-bedrock-invoke-1h", one_hour=True, wire_one_hour=True), + _Family("pin-bedrock-arn-base", "pin-bedrock-arn-base", one_hour=False, wire_one_hour=False), + _Family("pin-bedrock-arn", "pin-bedrock-arn", one_hour=False, wire_one_hour=False), + _Family("pin-azure-ai", "pin-azure-ai", one_hour=True, wire_one_hour=True), + _Family("pin-vertex", "pin-vertex", one_hour=True, wire_one_hour=True), + _Family(_OPENAI, "pin-openai", one_hour=False, wire_one_hour=None), + _Family("pin-unmapped", "pin-unmapped", one_hour=False, wire_one_hour=True), + _Family("pin-zero-one-hour-price", "pin-zero-one-hour-price", one_hour=False, wire_one_hour=None), +) + + +@dataclass(frozen=True, slots=True) +class _Rig: + owned: OwnedProxy + redis: Redis + left: Wire + right: Wire + misc: Wire + bedrock: Wire + openai: Wire + primary: Wire + slow: Wire + + def wires(self) -> tuple[Wire, ...]: + return (self.left, self.right, self.misc, self.bedrock, self.openai, self.primary, self.slow) + + def base_url(self) -> str: + return str(self.owned.gateway.client.base_url).rstrip("/") + + def served_by(self, marker: str) -> str: + identities: Final = {_LEFT: self.left, _RIGHT: self.right} + found: Final = [identity for identity, wire in identities.items() if sent_with((wire,), marker)] + assert len(found) == 1, found + return found[0] + + +def _second_in_order(entry: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + params: Final = entry["litellm_params"] + assert isinstance(params, dict) + return {**entry, "litellm_params": {**params, "order": 2}} + + +def _model_list(rig_wires: Mapping[str, str]) -> list[JsonValue]: + bedrock: Final = rig_wires["bedrock"] + misc: Final = rig_wires["misc"] + return [ + anthropic_deployment(_CLAUDE, _LEFT, rig_wires["left"]), + anthropic_deployment(_CLAUDE, _RIGHT, rig_wires["right"]), + anthropic_deployment(_SLOW, "pin-slow-left", rig_wires["slow"]), + anthropic_deployment(_SLOW, "pin-slow-right", rig_wires["slow"]), + anthropic_deployment(_MIXED, _PRIMARY, rig_wires["primary"], order=1), + _second_in_order( + bedrock_deployment(_MIXED, _SECONDARY, f"bedrock/converse/{BEDROCK_FIVE_MINUTE_MODEL}", bedrock) + ), + bedrock_deployment( + "pin-bedrock-converse-1h", "pin-bedrock-converse-1h", f"bedrock/converse/{BEDROCK_ONE_HOUR_MODEL}", bedrock + ), + bedrock_deployment( + "pin-bedrock-converse-5m", + "pin-bedrock-converse-5m", + f"bedrock/converse/{BEDROCK_FIVE_MINUTE_MODEL}", + bedrock, + ), + bedrock_deployment( + "pin-bedrock-invoke-1h", "pin-bedrock-invoke-1h", f"bedrock/invoke/{BEDROCK_ONE_HOUR_MODEL}", bedrock + ), + bedrock_deployment( + "pin-bedrock-arn-base", + "pin-bedrock-arn-base", + f"bedrock/converse/{INFERENCE_PROFILE_ARN}", + bedrock, + model_info={"base_model": f"bedrock/{BEDROCK_ONE_HOUR_MODEL}"}, + ), + bedrock_deployment("pin-bedrock-arn", "pin-bedrock-arn", f"bedrock/converse/{INFERENCE_PROFILE_ARN}", bedrock), + deployment( + "pin-azure-ai", "pin-azure-ai", model="azure_ai/claude-opus-5-5", api_base=misc, api_key=PROVIDER_KEY + ), + vertex_deployment("pin-vertex", "pin-vertex", misc), + deployment( + _OPENAI, "pin-openai", model="openai/gpt-5.6", api_base=f"{rig_wires['openai']}/v1", api_key=PROVIDER_KEY + ), + deployment( + "pin-unmapped", + "pin-unmapped", + model="anthropic/claude-pin-model-the-cost-map-lacks", + api_base=misc, + api_key=PROVIDER_KEY, + ), + deployment( + "pin-zero-one-hour-price", + "pin-zero-one-hour-price", + model="sarvam/sarvam-m", + api_base=f"{rig_wires['openai']}/v1", + api_key=PROVIDER_KEY, + ), + ] + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("prompt-cache-pin-ttl") + with ( + gateway_from_environment() as gateway, + owned_redis(directory) as cache, + wire_server(claude_peer) as left, + wire_server(claude_peer) as right, + wire_server(vertex_peer) as misc, + wire_server(bedrock_runtime_peer) as bedrock, + wire_server(openai_peer) as openai_wire, + wire_server(claude_peer) as primary, + wire_server(slow_claude_peer) as slow, + ): + urls: Final = { + "left": left.url, + "right": right.url, + "misc": misc.url, + "bedrock": bedrock.url, + "openai": openai_wire.url, + "primary": primary.url, + "slow": slow.url, + } + config: Final = pin_config(directory, _model_list(urls)) + with ( + owned_proxy_process( + gateway, + directory, + {"REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port)}, + config=config, + workers=2, + ) as owned, + Redis(host=cache.host, port=cache.port, socket_timeout=2) as redis, + ): + yield _Rig(owned, redis, left, right, misc, bedrock, openai_wire, primary, slow) + + +def _before(rig: _Rig) -> frozenset[str]: + for wire in rig.wires(): + wire.drain() + return frozenset(read_pins(rig.redis)) + + +def _first_turn(session: str, marker: str, mark: JsonValue) -> list[JsonValue]: + return [system_turn(session, mark), ask_turn(marker)] + + +def _breakpoint_shapes(session: str, marker: str) -> Mapping[str, list[JsonValue]]: + return { + "system-1h": _first_turn(session, marker, ONE_HOUR_MARK), + "system-5m": _first_turn(session, marker, FIVE_MINUTE_MARK), + "system-default-ttl": _first_turn(session, marker, DEFAULT_MARK), + "system-1h-then-turn-5m": [ + system_turn(session, ONE_HOUR_MARK), + {"role": "user", "content": "Earlier question."}, + {"role": "assistant", "content": "sunny"}, + ask_turn(marker, FIVE_MINUTE_MARK), + ], + "earlier-user-turn-1h": [ + system_turn(session, None), + {"role": "user", "content": [marked("Earlier question.", ONE_HOUR_MARK)]}, + {"role": "assistant", "content": "sunny"}, + ask_turn(marker), + ], + "message-level-1h-on-string-content": [ + {"role": "system", "content": reference(session), "cache_control": ONE_HOUR_MARK}, + ask_turn(marker), + ], + } + + +_ONE_HOUR_SHAPES: Final = frozenset( + {"system-1h", "system-1h-then-turn-5m", "earlier-user-turn-1h", "message-level-1h-on-string-content"} +) + + +@pytest.mark.parametrize("shape", sorted(_breakpoint_shapes("s", "m"))) +def test_chat_pin_lifetime_follows_the_longest_breakpoint_ttl_in_the_cacheable_prefix(rig: _Rig, shape: str) -> None: + before: Final = _before(rig) + marker: Final = new_marker() + response_id, served = post_chat( + rig.owned.gateway, chat_body(_CLAUDE, _breakpoint_shapes(new_marker(), marker)[shape]) + ) + _, pin = one_new_pin(rig.redis, before) + assert served == rig.served_by(marker) == pin.model_id, (served, pin) + assert (lives_one_hour if shape in _ONE_HOUR_SHAPES else lives_five_minutes)(pin), (shape, pin) + spend_row_lands(response_id) + + +def test_chat_one_hour_mark_on_a_tool_definition_alone_writes_no_pin(rig: _Rig) -> None: + before: Final = _before(rig) + marker: Final = new_marker() + tool: Final[dict[str, JsonValue]] = { + "type": "function", + "function": {"name": "lookup_weather", "parameters": {"type": "object"}}, + "cache_control": ONE_HOUR_MARK, + } + response_id, served = post_chat( + rig.owned.gateway, chat_body(_CLAUDE, _first_turn(new_marker(), marker, None), tools=[tool]) + ) + assert served == rig.served_by(marker), served + no_pin_written(rig.redis, before, response_id) + + +def _openai_client(rig: _Rig) -> openai.OpenAI: + return openai.OpenAI(base_url=f"{rig.base_url()}/v1", api_key=rig.owned.gateway.key, max_retries=0) + + +def _openai_chat(rig: _Rig, messages: list[JsonValue], *, asynchronous: bool, stream: bool) -> str: + body: Final = chat_body(_CLAUDE, messages) + if asynchronous: + return asyncio.run(_openai_chat_async(rig, body, stream=stream)) + with _openai_client(rig) as client: + if not stream: + return client.chat.completions.create(**body).id + chunks: Final = tuple(client.chat.completions.create(**body, stream=True)) + return chunks[0].id + + +async def _openai_chat_async(rig: _Rig, body: Mapping[str, JsonValue], *, stream: bool) -> str: + async with openai.AsyncOpenAI( + base_url=f"{rig.base_url()}/v1", api_key=rig.owned.gateway.key, max_retries=0 + ) as client: + if not stream: + return (await client.chat.completions.create(**body)).id + chunks: Final = [chunk async for chunk in await client.chat.completions.create(**body, stream=True)] + return chunks[0].id + + +def _messages_body(session: str, marker: str) -> dict[str, JsonValue]: + return { + "model": _CLAUDE, + "max_tokens": 16, + "messages": [ + {"role": "user", "content": [marked(reference(session), ONE_HOUR_MARK), marked(question(marker), None)]} + ], + } + + +def _anthropic_messages(rig: _Rig, body: Mapping[str, JsonValue], *, asynchronous: bool, stream: bool) -> str: + if asynchronous: + return asyncio.run(_anthropic_messages_async(rig, body, stream=stream)) + with anthropic.Anthropic(base_url=rig.base_url(), api_key=rig.owned.gateway.key, max_retries=0) as client: + if not stream: + return client.messages.create(**body).id + with client.messages.stream(**body) as events: + return events.get_final_message().id + + +async def _anthropic_messages_async(rig: _Rig, body: Mapping[str, JsonValue], *, stream: bool) -> str: + async with anthropic.AsyncAnthropic( + base_url=rig.base_url(), api_key=rig.owned.gateway.key, max_retries=0 + ) as client: + if not stream: + return (await client.messages.create(**body)).id + async with client.messages.stream(**body) as events: + return (await events.get_final_message()).id + + +_CLIENT_MODES: Final = ("sync", "sync-stream", "async", "async-stream") + + +@pytest.mark.parametrize("mode", _CLIENT_MODES) +def test_openai_sdk_chat_pins_a_one_hour_breakpoint_for_an_hour(rig: _Rig, mode: str) -> None: + before: Final = _before(rig) + marker: Final = new_marker() + response_id: Final = _openai_chat( + rig, + _first_turn(new_marker(), marker, ONE_HOUR_MARK), + asynchronous=mode.startswith("async"), + stream=mode.endswith("stream"), + ) + _, pin = one_new_pin(rig.redis, before) + assert pin.model_id == rig.served_by(marker), pin + assert lives_one_hour(pin), pin + spend_row_lands(response_id) + + +@pytest.mark.parametrize("mode", _CLIENT_MODES) +def test_anthropic_sdk_messages_pins_a_one_hour_breakpoint_for_an_hour(rig: _Rig, mode: str) -> None: + before: Final = _before(rig) + marker: Final = new_marker() + response_id: Final = _anthropic_messages( + rig, + _messages_body(new_marker(), marker), + asynchronous=mode.startswith("async"), + stream=mode.endswith("stream"), + ) + assert response_id == f"msg_{marker}", response_id + _, pin = one_new_pin(rig.redis, before) + assert pin.model_id == rig.served_by(marker), pin + assert lives_one_hour(pin), pin + + +def _responses_id(rig: _Rig, body: Mapping[str, JsonValue], *, stream: bool) -> str: + with _openai_client(rig) as client: + if not stream: + return client.responses.create(**body).id + events: Final = tuple(client.responses.create(**body, stream=True)) + return next(event.response.id for event in events if event.type == "response.completed") + + +@pytest.mark.parametrize("stream", [False, True], ids=["non-stream", "stream"]) +def test_responses_api_one_hour_breakpoint_answers_and_writes_no_pin(rig: _Rig, stream: bool) -> None: + before: Final = _before(rig) + marker: Final = new_marker() + user: Final = f"pin-responses-{marker}" + body: Final[dict[str, JsonValue]] = { + "model": _CLAUDE, + "max_output_tokens": 16, + "user": user, + "input": [ + { + "role": "system", + "content": [{"type": "input_text", "text": reference(new_marker()), "cache_control": ONE_HOUR_MARK}], + }, + {"role": "user", "content": question(marker)}, + ], + } + assert _responses_id(rig, body, stream=stream).startswith("resp_") + assert rig.served_by(marker) in {_LEFT, _RIGHT} + end_user_row_lands(user) + held(lambda: pins_since(rig.redis, before), lambda found: not found, holding=3, seconds=20) + + +def _ask_surface(rig: _Rig, surface: str, session: str, marker: str) -> str: + if surface == "messages": + return _anthropic_messages(rig, _messages_body(session, marker), asynchronous=False, stream=False) + return _openai_chat( + rig, _first_turn(session, marker, ONE_HOUR_MARK), asynchronous=False, stream=surface == "chat-stream" + ) + + +@pytest.mark.parametrize("surface", ["chat", "chat-stream", "messages"]) +def test_response_cache_hit_rewrites_the_pin_with_the_breakpoint_lifetime(rig: _Rig, surface: str) -> None: + before: Final = _before(rig) + marker: Final = new_marker() + session: Final = new_marker() + first: Final = _ask_surface(rig, surface, session, marker) + key, pin = one_new_pin(rig.redis, before) + assert lives_one_hour(pin), pin + assert len(sent_with(rig.wires(), marker)) == 1 + spend_row_lands(first) + rig.redis.delete(key) + second: Final = _ask_surface(rig, surface, session, marker) + assert sent_with(rig.wires(), marker) == () + assert second == first + rewritten: Final = eventually(lambda: pin_at(rig.redis, key), lambda found: found is not None, seconds=30) + assert rewritten is not None and rewritten.model_id in {_LEFT, _RIGHT}, rewritten + assert lives_one_hour(rewritten), rewritten + + +@pytest.mark.parametrize("family", _FAMILIES, ids=[family.group for family in _FAMILIES]) +def test_provider_family_pin_lifetime_agrees_with_the_ttl_the_wire_carries(rig: _Rig, family: _Family) -> None: + before: Final = _before(rig) + marker: Final = new_marker() + response_id, served = post_chat( + rig.owned.gateway, chat_body(family.group, _first_turn(new_marker(), marker, ONE_HOUR_MARK)) + ) + _, pin = one_new_pin(rig.redis, before) + sent: Final = sent_with(rig.wires(), marker) + assert served == pin.model_id == family.identity, (served, pin) + assert lives_one_hour(pin) is family.one_hour, pin + if family.wire_one_hour is not None: + assert [wire_carries_one_hour(request) for request in sent] == [family.wire_one_hour], sent + spend_row_lands(response_id) + + +def test_messages_through_bedrock_invoke_pins_an_hour_and_sends_the_hour(rig: _Rig) -> None: + before: Final = _before(rig) + marker: Final = new_marker() + response: Final = rig.owned.gateway.request( + "POST", "/v1/messages", {**_messages_body(new_marker(), marker), "model": "pin-bedrock-invoke-1h"} + ) + assert response.status_code == 200, response.text + _, pin = one_new_pin(rig.redis, before) + assert pin.model_id == "pin-bedrock-invoke-1h", pin + assert lives_one_hour(pin), pin + assert [wire_carries_one_hour(request) for request in sent_with((rig.bedrock,), marker)] == [True] + + +@pytest.mark.parametrize("ttl", _MALFORMED_TTLS, ids=["int", "2h", "upper-1H", "empty", "null", "list", "5kb"]) +def test_malformed_ttl_answers_and_pins_five_minutes_without_a_wire_ttl(rig: _Rig, ttl: JsonValue) -> None: + before: Final = _before(rig) + marker: Final = new_marker() + response_id, _ = post_chat( + rig.owned.gateway, + chat_body("pin-bedrock-converse-1h", _first_turn(new_marker(), marker, {"type": "ephemeral", "ttl": ttl})), + ) + _, pin = one_new_pin(rig.redis, before) + assert lives_five_minutes(pin), pin + assert [wire_carries_one_hour(request) for request in sent_with((rig.bedrock,), marker)] == [False] + spend_row_lands(response_id) + + +@pytest.mark.parametrize("order", [("5m", "1h"), ("1h", "5m")], ids=["5m-then-1h", "1h-then-5m"]) +def test_duplicated_ttl_key_takes_the_last_value_for_pin_and_wire(rig: _Rig, order: tuple[str, str]) -> None: + before: Final = _before(rig) + marker: Final = new_marker() + placeholder: Final = new_marker() + body: Final = json.dumps( + chat_body( + "pin-bedrock-converse-1h", _first_turn(new_marker(), marker, {"type": "ephemeral", "ttl": placeholder}) + ) + ).replace(f'"ttl": "{placeholder}"', f'"ttl": "{order[0]}", "ttl": "{order[1]}"') + response: Final = rig.owned.gateway.client.post( + "/v1/chat/completions", + content=body.encode(), + headers={"Authorization": f"Bearer {rig.owned.gateway.key}", "Content-Type": "application/json"}, + ) + assert response.status_code == 200, response.text + _, pin = one_new_pin(rig.redis, before) + assert lives_one_hour(pin) is (order[1] == "1h"), pin + assert [wire_carries_one_hour(request) for request in sent_with((rig.bedrock,), marker)] == [order[1] == "1h"] + + +@pytest.mark.parametrize( + "mark", ["ephemeral", {"type": "persistent", "ttl": "1h"}], ids=["string", "non-ephemeral-type-1h"] +) +def test_non_ephemeral_cache_control_never_pins_an_hour(rig: _Rig, mark: JsonValue) -> None: + before: Final = _before(rig) + marker: Final = new_marker() + response: Final = rig.owned.gateway.request( + "POST", "/v1/chat/completions", chat_body(_CLAUDE, _first_turn(new_marker(), marker, mark)) + ) + assert response.status_code == 200, response.text + no_pin_written(rig.redis, before, str(response.json()["id"])) + + +def test_failed_first_turn_reports_the_error_and_writes_no_pin(rig: _Rig) -> None: + before: Final = _before(rig) + response: Final = rig.owned.gateway.request( + "POST", + "/v1/chat/completions", + chat_body(_CLAUDE, _first_turn(new_marker(), f"{new_marker()} {OVERLOAD_TOKEN}", ONE_HOUR_MARK)), + ) + assert response.status_code == 500, response.text + assert "overloaded_error" in response.text, response.text + healthy, _ = post_chat(rig.owned.gateway, chat_body(_CLAUDE, _first_turn(new_marker(), new_marker(), None))) + spend_row_lands(healthy) + assert pins_since(rig.redis, before) == {} + + +def test_one_hour_breakpoint_after_a_five_minute_one_reports_the_provider_400_and_writes_no_pin(rig: _Rig) -> None: + before: Final = _before(rig) + marker: Final = new_marker() + messages: Final = [ + system_turn(new_marker(), FIVE_MINUTE_MARK), + ask_turn(marker, ONE_HOUR_MARK), + ] + response: Final = rig.owned.gateway.request("POST", "/v1/chat/completions", chat_body(_CLAUDE, messages)) + assert response.status_code == 400, response.text + assert "must not come after" in response.text, response.text + healthy, _ = post_chat(rig.owned.gateway, chat_body(_CLAUDE, _first_turn(new_marker(), new_marker(), None))) + spend_row_lands(healthy) + assert pins_since(rig.redis, before) == {} + + +@pytest.mark.parametrize( + "marks", [(FIVE_MINUTE_MARK, ONE_HOUR_MARK), (ONE_HOUR_MARK, FIVE_MINUTE_MARK)], ids=["5m-then-1h", "1h-then-5m"] +) +def test_rewriting_one_prefix_moves_its_pin_to_the_latest_lifetime( + rig: _Rig, marks: tuple[Mapping[str, JsonValue], Mapping[str, JsonValue]] +) -> None: + before: Final = _before(rig) + session: Final = new_marker() + first_id, _ = post_chat(rig.owned.gateway, chat_body(_CLAUDE, _first_turn(session, new_marker(), dict(marks[0])))) + key, first = one_new_pin(rig.redis, before) + spend_row_lands(first_id) + second_id, served = post_chat( + rig.owned.gateway, chat_body(_CLAUDE, _first_turn(session, new_marker(), dict(marks[1]))) + ) + spend_row_lands(second_id) + second: Final = pin_at(rig.redis, key) + assert second is not None and second.model_id == first.model_id == served, (first, second) + assert lives_one_hour(first) is (marks[0] == ONE_HOUR_MARK), first + assert lives_one_hour(second) is (marks[1] == ONE_HOUR_MARK), second + assert frozenset(pins_since(rig.redis, before)) == {key} + + +def test_mixed_group_pin_takes_the_lifetime_of_the_deployment_that_served_last(rig: _Rig) -> None: + before: Final = _before(rig) + session: Final = new_marker() + first_turn: Final = _first_turn(session, new_marker(), ONE_HOUR_MARK) + first_id, first_served = post_chat(rig.owned.gateway, chat_body(_MIXED, first_turn)) + key, first = one_new_pin(rig.redis, before) + assert first_served == first.model_id == _PRIMARY, first + assert lives_one_hour(first), first + spend_row_lands(first_id) + marker: Final = new_marker() + second_id, second_served = post_chat( + rig.owned.gateway, chat_body(_MIXED, follow_up(first_turn, f"{marker} {OVERLOAD_TOKEN}")) + ) + spend_row_lands(second_id) + assert second_served == _SECONDARY + assert [wire_carries_one_hour(request) for request in sent_with((rig.bedrock,), marker)] == [False] + second: Final = pin_at(rig.redis, key) + assert second is not None and second.model_id == _SECONDARY, second + assert lives_five_minutes(second), second + + +def test_same_prefix_moving_from_a_five_minute_group_to_a_one_hour_group_repins_for_an_hour(rig: _Rig) -> None: + before: Final = _before(rig) + session: Final = new_marker() + first_id, _ = post_chat(rig.owned.gateway, chat_body(_OPENAI, _first_turn(session, new_marker(), ONE_HOUR_MARK))) + key, first = one_new_pin(rig.redis, before) + assert first.model_id == "pin-openai" and lives_five_minutes(first), first + spend_row_lands(first_id) + second_id, served = post_chat( + rig.owned.gateway, chat_body(_CLAUDE, _first_turn(session, new_marker(), ONE_HOUR_MARK)) + ) + spend_row_lands(second_id) + second: Final = pin_at(rig.redis, key) + assert second is not None and second.model_id == served, (second, served) + assert lives_one_hour(second), second + + +def test_follow_ups_stick_to_the_pinned_deployment_and_keep_an_hour(rig: _Rig) -> None: + before: Final = _before(rig) + first_turn: Final = _first_turn(new_marker(), new_marker(), ONE_HOUR_MARK) + first_id, first_served = post_chat(rig.owned.gateway, chat_body(_CLAUDE, first_turn)) + key, first = one_new_pin(rig.redis, before) + spend_row_lands(first_id) + follow_ups: Final = tuple( + post_chat(rig.owned.gateway, chat_body(_CLAUDE, follow_up(first_turn, new_marker()))) for _ in range(5) + ) + assert [served for _, served in follow_ups] == [first_served] * 5 + for response_id, _ in follow_ups: + spend_row_lands(response_id) + pin: Final = pin_at(rig.redis, key) + assert pin is not None and pin.model_id == first.model_id == first_served, pin + assert lives_one_hour(pin), pin + assert frozenset(pins_since(rig.redis, before)) == {key} + + +def test_concurrent_first_turns_on_one_prefix_write_one_hour_long_pin(rig: _Rig) -> None: + before: Final = _before(rig) + session: Final = new_marker() + bodies: Final = tuple(chat_body(_CLAUDE, _first_turn(session, new_marker(), ONE_HOUR_MARK)) for _ in range(8)) + with ThreadPoolExecutor(max_workers=8) as pool: + answers: Final = tuple(pool.map(lambda body: post_chat(rig.owned.gateway, body), bodies)) + for response_id, _ in answers: + spend_row_lands(response_id) + pins: Final = pins_since(rig.redis, before) + assert len(pins) == 1, pins + pin: Final[Pin] = next(iter(pins.values())) + assert pin.model_id in {served for _, served in answers}, pin + assert lives_one_hour(pin), pin + + +def test_slow_upstream_burst_writes_every_one_hour_pin(rig: _Rig) -> None: + before: Final = _before(rig) + markers: Final = tuple(new_marker() for _ in range(24)) + bodies: Final = tuple(chat_body(_SLOW, _first_turn(new_marker(), marker, ONE_HOUR_MARK)) for marker in markers) + with ThreadPoolExecutor(max_workers=24) as pool: + answers: Final = tuple(pool.map(lambda body: post_chat(rig.owned.gateway, body), bodies)) + for response_id, _ in answers: + spend_row_lands(response_id) + pins: Final = pins_since(rig.redis, before) + assert len(pins) == 24, pins + assert sorted(pin.model_id for pin in pins.values()) == sorted(served for _, served in answers) + assert all(lives_one_hour(pin) for pin in pins.values()), pins + assert sorted(marker_of(request) for request in rig.slow.drain()) == sorted(markers) diff --git a/tests/integration/routing/test_prompt_cache_pin_ttl_chaos.py b/tests/integration/routing/test_prompt_cache_pin_ttl_chaos.py new file mode 100644 index 00000000000..4cdda116d23 --- /dev/null +++ b/tests/integration/routing/test_prompt_cache_pin_ttl_chaos.py @@ -0,0 +1,248 @@ +from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.redis_process import OwnedRedis, owned_redis +from integration._support.wire import wire_server +from integration.providers._cache_control_marks_support import anthropic_peer +from integration.routing._prompt_cache_pin_support import ( + FIVE_MINUTE_MARK, + ONE_HOUR_MARK, + OVERLOAD_TOKEN, + Pin, + anthropic_deployment, + ask_turn, + chat_body, + claude_peer, + follow_up, + lives_five_minutes, + lives_one_hour, + marked, + new_marker, + one_new_pin, + pin_at, + pin_config, + pins_since, + post_chat, + read_pins, + spend_row_lands, + system_turn, +) +from pydantic import JsonValue +from redis import Redis + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 600) + +_GROUP: Final = "pin-chaos" +_LEFT: Final = "pin-chaos-left" +_RIGHT: Final = "pin-chaos-right" +_REDIS_ENVIRONMENT: Final = ("REDIS_HOST", "REDIS_PORT", "REDIS_PASSWORD", "REDIS_URL", "REDIS_SSL") +_NO_RESPONSE_CACHE: Final[Mapping[str, JsonValue]] = {"cache": False, "enable_redis_auth_cache": False} +_AFTER_GROUP: Final = "pin-chaos-after" +_AFTER_LEFT: Final = "pin-chaos-after-left" +_AFTER_RIGHT: Final = "pin-chaos-after-right" +_AFTER_IDS: Final = frozenset({_AFTER_LEFT, _AFTER_RIGHT}) +_PIN_SLOTS: Final = 10_000 +_FILLERS: Final = _PIN_SLOTS - 1 +_FILLER_BATCH: Final = 250 + + +def _attach_path(path: str, cache: OwnedRedis) -> tuple[Mapping[str, str], Mapping[str, Mapping[str, JsonValue]]]: + target: Final[dict[str, JsonValue]] = {"host": cache.host, "port": cache.port} + if path == "router-settings": + return {}, {"router_settings": {"redis_host": cache.host, "redis_port": cache.port}} + if path == "coordination-redis": + return {}, {"general_settings": {"coordination_redis": target}} + return {"REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port)}, {} + + +@pytest.mark.parametrize("path", ["router-settings", "coordination-redis", "env-only"]) +def test_one_hour_pin_lands_in_the_redis_of_every_attach_path(gateway: Gateway, tmp_path: Path, path: str) -> None: + with ( + owned_redis(tmp_path) as cache, + wire_server(claude_peer) as left, + wire_server(claude_peer) as right, + Redis(host=cache.host, port=cache.port, socket_timeout=2) as redis, + ): + environment, settings = _attach_path(path, cache) + config: Final = pin_config( + tmp_path, + [anthropic_deployment(_GROUP, _LEFT, left.url), anthropic_deployment(_GROUP, _RIGHT, right.url)], + litellm_settings=_NO_RESPONSE_CACHE, + router_settings=settings.get("router_settings", {}), + general_settings=settings.get("general_settings", {}), + ) + with owned_proxy_process( + gateway, tmp_path, environment, config=config, remove_environment=_REDIS_ENVIRONMENT, workers=2 + ) as owned: + before: Final = frozenset(read_pins(redis)) + response_id, served = post_chat( + owned.gateway, + chat_body(_GROUP, [system_turn(new_marker(), ONE_HOUR_MARK), ask_turn(new_marker())]), + ) + _, pin = one_new_pin(redis, before) + assert pin.model_id == served, (pin, served) + assert lives_one_hour(pin), pin + spend_row_lands(response_id) + + +def _short_turn(session: str, mark: JsonValue, ask: str) -> list[JsonValue]: + return [{"role": "system", "content": [marked(f"Reference {session}", mark)]}, {"role": "user", "content": ask}] + + +def _served(gateway: Gateway, session: str, mark: JsonValue, *, overload: bool = False) -> str: + ask: Final = f"Reply in one word. marker-{new_marker()}{f' {OVERLOAD_TOKEN}' if overload else ''}" + return post_chat(gateway, chat_body(_GROUP, _short_turn(session, mark, ask)))[1] + + +def _pinned_to_right( + gateway: Gateway, redis: Redis, session: str, mark: JsonValue, *, overload: bool, lives: Callable[[Pin], bool] +) -> None: + before: Final = frozenset(read_pins(redis)) + assert _served(gateway, session, mark, overload=overload) == _RIGHT + pin: Final = eventually( + lambda: tuple(pins_since(redis, before).values()), + lambda found: len(found) == 1 and lives(found[0]), + seconds=30, + )[0] + assert pin.model_id == _RIGHT, pin + + +def _pin_keys(redis: Redis) -> tuple[bytes, ...]: + return tuple(redis.scan_iter(match="*:prompt_caching", count=1000)) + + +def _fill_batch(gateway: Gateway, redis: Redis, size: int) -> None: + before: Final = frozenset(_pin_keys(redis)) + with ThreadPoolExecutor(max_workers=8) as pool: + served: Final = frozenset(pool.map(lambda _: _served(gateway, new_marker(), ONE_HOUR_MARK), range(size))) + assert served == {_LEFT}, served + eventually(lambda: len(frozenset(_pin_keys(redis)) - before), lambda count: count >= size, seconds=120) + + +def test_in_memory_pins_keep_the_latest_lifetime_when_the_slots_fill(gateway: Gateway, tmp_path: Path) -> None: + upgraded: Final = new_marker() + downgraded: Final = new_marker() + with ( + owned_redis(tmp_path) as cache, + wire_server(claude_peer) as left, + wire_server(anthropic_peer) as right, + Redis(host=cache.host, port=cache.port, socket_timeout=5) as redis, + ): + config: Final = pin_config( + tmp_path, + [ + anthropic_deployment(_GROUP, _LEFT, left.url, order=1), + anthropic_deployment(_GROUP, _RIGHT, right.url, order=2), + ], + litellm_settings=_NO_RESPONSE_CACHE, + router_settings={"redis_host": cache.host, "redis_port": cache.port}, + general_settings={"disable_spend_logs": True}, + ) + with owned_proxy_process( + gateway, + tmp_path, + {"MINIMUM_PROMPT_CACHE_TOKEN_COUNT": "1"}, + config=config, + remove_environment=_REDIS_ENVIRONMENT, + workers=1, + ) as owned: + _pinned_to_right(owned.gateway, redis, upgraded, FIVE_MINUTE_MARK, overload=True, lives=lives_five_minutes) + upgraded_key: Final = next(iter(read_pins(redis))) + assert _served(owned.gateway, upgraded, ONE_HOUR_MARK) == _RIGHT + eventually( + lambda: pin_at(redis, upgraded_key), + lambda found: found is not None and found.model_id == _RIGHT and lives_one_hour(found), + seconds=30, + ) + settled: Final = frozenset(read_pins(redis)) + _pinned_to_right(owned.gateway, redis, downgraded, ONE_HOUR_MARK, overload=True, lives=lives_one_hour) + downgraded_key: Final = next(iter(pins_since(redis, settled))) + assert _served(owned.gateway, downgraded, FIVE_MINUTE_MARK) == _RIGHT + eventually( + lambda: pin_at(redis, downgraded_key), + lambda found: found is not None and found.model_id == _RIGHT and lives_five_minutes(found), + seconds=30, + ) + for start in range(0, _FILLERS, _FILLER_BATCH): + _fill_batch(owned.gateway, redis, min(_FILLER_BATCH, _FILLERS - start)) + written: Final = _pin_keys(redis) + assert redis.delete(*written) == len(written) + assert _served(owned.gateway, upgraded, ONE_HOUR_MARK) == _RIGHT + assert _served(owned.gateway, downgraded, FIVE_MINUTE_MARK) == _LEFT + + +def _first_turns(gateway: Gateway, count: int, group: str = _GROUP) -> tuple[tuple[str, str], ...]: + bodies: Final = tuple( + chat_body(group, [system_turn(new_marker(), ONE_HOUR_MARK), ask_turn(new_marker())]) for _ in range(count) + ) + with ThreadPoolExecutor(max_workers=count) as pool: + return tuple(pool.map(lambda body: post_chat(gateway, body), bodies)) + + +def _probe_pin(gateway: Gateway, redis: Redis) -> Mapping[str, Pin]: + before: Final = frozenset(read_pins(redis)) + _first_turns(gateway, 1) + return eventually( + lambda: pins_since(redis, before), lambda found: len(found) >= 1, seconds=10, return_last_on_timeout=True + ) + + +def _after_group_pins(redis: Redis) -> Mapping[str, Pin]: + return MappingProxyType({key: pin for key, pin in read_pins(redis).items() if pin.model_id in _AFTER_IDS}) + + +def test_pins_resume_after_a_redis_outage_and_survive_a_killed_worker(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_redis(tmp_path) as cache, + wire_server(claude_peer) as left, + wire_server(claude_peer) as right, + Redis(host=cache.host, port=cache.port, socket_timeout=2) as redis, + ): + config: Final = pin_config( + tmp_path, + [ + anthropic_deployment(_GROUP, _LEFT, left.url), + anthropic_deployment(_GROUP, _RIGHT, right.url), + anthropic_deployment(_AFTER_GROUP, _AFTER_LEFT, left.url), + anthropic_deployment(_AFTER_GROUP, _AFTER_RIGHT, right.url), + ], + ) + with owned_proxy_process( + gateway, + tmp_path, + {"REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port), "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT": "1"}, + config=config, + workers=2, + ) as owned: + healthy: Final = _first_turns(owned.gateway, 8) + cache.stop() + during: Final = _first_turns(owned.gateway, 16) + cache.start() + eventually(lambda: _probe_pin(owned.gateway, redis), lambda found: len(found) >= 1, seconds=90) + recovered: Final = _first_turns(owned.gateway, 8, _AFTER_GROUP) + for response_id, _ in (*healthy, *during, *recovered): + spend_row_lands(response_id) + resumed: Final = eventually(lambda: _after_group_pins(redis), lambda found: len(found) >= 8, seconds=30) + assert len(resumed) == 8, resumed + assert all(lives_one_hour(pin) for pin in resumed.values()), resumed + assert sorted(pin.model_id for pin in resumed.values()) == sorted(served for _, served in recovered) + turn: Final = [system_turn(new_marker(), ONE_HOUR_MARK), ask_turn(new_marker())] + first_id, first_served = post_chat(owned.gateway, chat_body(_AFTER_GROUP, turn)) + grown: Final = eventually(lambda: _after_group_pins(redis), lambda found: len(found) >= 9, seconds=30) + assert len(grown) == 9, grown + ((key, pin),) = tuple((key, pin) for key, pin in grown.items() if key not in resumed) + assert pin.model_id == first_served and lives_one_hour(pin), pin + spend_row_lands(first_id) + owned.kill_worker(sorted(owned.worker_pids())[0]) + follow_ups: Final = tuple( + post_chat(owned.gateway, chat_body(_AFTER_GROUP, follow_up(turn, new_marker()))) for _ in range(6) + ) + assert [served for _, served in follow_ups] == [first_served] * 6 + after: Final = pin_at(redis, key) + assert after is not None and after.model_id == first_served and lives_one_hour(after), after diff --git a/tests/unit/caching/test_in_memory_cache.py b/tests/unit/caching/test_in_memory_cache.py index 40ad4f0c6f0..214f4ecf301 100644 --- a/tests/unit/caching/test_in_memory_cache.py +++ b/tests/unit/caching/test_in_memory_cache.py @@ -1,8 +1,12 @@ import asyncio import json +import random +import sys import threading import time +from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor +from typing import Final from unittest.mock import MagicMock, patch import httpx @@ -232,6 +236,66 @@ def test_in_memory_cache_heap_size_staus_bounded(): assert len(in_memory_cache.expiration_heap) == 1 +def test_in_memory_cache_heap_stays_bounded_when_a_live_key_is_rewritten_with_a_new_ttl() -> None: + in_memory_cache = InMemoryCache(max_size_in_memory=10) + in_memory_cache.set_cache(key="expires_first", value="pin", ttl=300) + + for i in range(1_000): + in_memory_cache.delete_cache("hot_key") + in_memory_cache.set_cache(key="hot_key", value=f"value_{i}", ttl=3600) + + assert in_memory_cache.get_cache("hot_key") == "value_999" + assert len(in_memory_cache.expiration_heap) <= 2 * in_memory_cache.max_size_in_memory + 1 + + +@pytest.fixture +def fine_grained_thread_switching() -> Iterator[None]: + previous_interval: Final = sys.getswitchinterval() + sys.setswitchinterval(1e-6) + yield + sys.setswitchinterval(previous_interval) + + +def _keys_until(deadline: float, rng: random.Random) -> Iterator[str]: + while time.monotonic() < deadline: + yield f"hot_{rng.randrange(5)}" if rng.random() < 0.7 else f"cold_{rng.randrange(40)}" + + +def _rewrite_keys_until(cache: InMemoryCache, deadline: float, seed: int) -> None: + rng: Final = random.Random(seed) + for key in _keys_until(deadline, rng): + cache.delete_cache(key) + cache.set_cache(key, "pin", ttl=rng.choice((300, 3600))) + + +@pytest.mark.usefixtures("fine_grained_thread_switching") +def test_in_memory_cache_survives_threads_rewriting_keys_in_a_full_cache() -> None: + cache: Final = InMemoryCache(max_size_in_memory=5) + deadline: Final = time.monotonic() + 0.5 + thread_count: Final = 8 + + with ThreadPoolExecutor(max_workers=thread_count) as executor: + futures: Final = tuple( + executor.submit(_rewrite_keys_until, cache, deadline, seed) for seed in range(thread_count) + ) + + assert [future.exception() for future in futures] == [None] * thread_count + keys: Final = (*(f"hot_{i}" for i in range(5)), *(f"cold_{i}" for i in range(40))) + assert sum(cache.get_cache(key) is not None for key in keys) <= cache.max_size_in_memory + + +def test_in_memory_cache_stays_bounded_when_entries_have_no_usable_expiry() -> None: + cache: Final = InMemoryCache(max_size_in_memory=10) + for i in range(30): + cache.set_cache(key=f"nan_{i}", value="response", ttl="NaN") + cache.set_cache(key="fresh", value="response", ttl=60) + + keys: Final = (*(f"nan_{i}" for i in range(30)), "fresh") + retained: Final = tuple(key for key in keys if cache.get_cache(key) is not None) + assert len(retained) <= cache.max_size_in_memory + assert cache.get_cache("fresh") == "response" + + def test_in_memory_cache_prunes_expired_heap_entries_below_capacity(): """ Re-inserting expired keys below capacity should not grow expiration_heap diff --git a/tests/unit/proxy/common_utils/test_debug_utils.py b/tests/unit/proxy/common_utils/test_debug_utils.py index a64c94c2991..456be4e8c00 100644 --- a/tests/unit/proxy/common_utils/test_debug_utils.py +++ b/tests/unit/proxy/common_utils/test_debug_utils.py @@ -1,7 +1,7 @@ import json import os import socket -from collections.abc import Iterator, Mapping +from collections.abc import Callable, Iterator, Mapping from dataclasses import asdict from pathlib import Path @@ -20,6 +20,8 @@ from litellm.proxy.common_utils.debug_utils import ( get_memory_summary, ) from litellm.proxy.common_utils.debug_utils import router as debug_router +from litellm.router import Router +from litellm.router_utils.prompt_caching_cache import PromptCachingCache PAGE_SIZE = 4096 STATM_SIZE_PAGES = 100_000 @@ -146,3 +148,53 @@ def test_debug_report_returns_what_the_bug_report_link_carries_and_nothing_from_ "model_list[*].provider = [azure]", ] assert not any(hostile in response.text for hostile in HOSTILE_STRINGS), response.text + + +PINNED_PREFIXES = 3 + + +@pytest.fixture +def router_with_pins(monkeypatch: pytest.MonkeyPatch) -> None: + router = Router(model_list=[{"model_name": "claude", "litellm_params": {"model": "anthropic/claude-sonnet-4-5"}}]) + pins = PromptCachingCache(cache=router.prompt_caching_pin_cache) + for prefix in range(PINNED_PREFIXES): + pins.add_model_id( + model_id=f"deployment-{prefix}", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": f"prefix {prefix} " * 300, "cache_control": {"type": "ephemeral"}} + ], + } + ], + tools=None, + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + + +@pytest.mark.parametrize( + ("path", "pin_count", "expected"), + [ + ( + "/memory-usage-in-mem-cache", + lambda body: body["num_items_in_prompt_caching_pin_cache"], + 2 * PINNED_PREFIXES, + ), + ("/memory-usage-in-mem-cache-items", lambda body: len(body["prompt_caching_pin_cache"]), PINNED_PREFIXES), + ("/debug/memory/summary", lambda body: body["caches"]["breakdown"]["prompt_caching_pins"]["count"], PINNED_PREFIXES), + ( + "/debug/memory/details?include_process_info=false", + lambda body: body["cache_memory"]["prompt_caching_pin_cache"]["num_items"], + PINNED_PREFIXES, + ), + ], +) +@pytest.mark.usefixtures("router_with_pins") +def test_memory_routes_count_prompt_caching_pins( + path: str, pin_count: Callable[[dict[str, object]], int], expected: int +) -> None: + response = _debug_client(UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)).get(path) + + assert response.status_code == 200 + assert pin_count(response.json()) == expected diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index e18c9948e39..7e4b470f1c1 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -2034,7 +2034,7 @@ async def test_priority_group_ignores_cached_backup_affinity() -> None: "role": "user", "content": [{"type": "text", "text": "word " * 5000, "cache_control": {"type": "ephemeral"}}], }] - cache: Final = PromptCachingCache(cache=router.cache) + cache: Final = PromptCachingCache(cache=router.prompt_caching_pin_cache) await cache.async_add_model_id(model_id="deploy-3", messages=messages, tools=None) assert await cache.async_get_model_id(messages=messages, tools=None) == {"model_id": "deploy-3"} response: Final = await router.acompletion(model="priority-group", messages=messages) diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index c6ab259991e..018fa808f70 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -9,6 +9,7 @@ import pytest import litellm from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, PROMPT_CACHE_LOOKBACK_POSITIONS from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.integrations.custom_logger import CustomLogger @@ -16,7 +17,7 @@ from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck, _get_min_token_count_for_deployments, ) -from litellm.router_utils.prompt_caching_cache import PromptCachingCache +from litellm.router_utils.prompt_caching_cache import PROMPT_CACHE_PIN_TTL_SECONDS, PromptCachingCache from litellm.types.llms.openai import AllMessageValues from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt, token_counter @@ -314,9 +315,11 @@ async def test_replayed_redacted_thinking_block_still_records_and_pins(): kwargs={ "standard_logging_object": { "call_type": "anthropic_messages", + "custom_llm_provider": "anthropic", "model": model, "messages": messages, "model_id": "dep-2", + "model_map_information": {"model_map_key": model, "model_map_value": None}, } }, response_obj=None, @@ -431,7 +434,7 @@ async def test_repeated_auto_cached_prefix_pins_to_one_deployment(monkeypatch, l served_by = first._hidden_params["model_id"] affinity_key = PromptCachingCache.get_prompt_caching_cache_key(_affinity_messages(messages), None) - assert await _eventually(lambda: router.cache.get_cache(key=affinity_key)) is not None + assert await _eventually(lambda: router.prompt_caching_pin_cache.get_cache(key=affinity_key)) is not None subsequent = [ (await router.acompletion(model=MODEL_GROUP_ALIAS, messages=messages, mock_response="ok"))._hidden_params[ @@ -661,9 +664,11 @@ async def test_async_log_success_event_counts_the_prompt_off_the_event_loop(): ) standard_logging_object = { "call_type": "acompletion", + "custom_llm_provider": "anthropic", "model": "anthropic/claude-fable-5", "messages": messages, "model_id": "dep-1", + "model_map_information": {"model_map_key": "anthropic/claude-fable-5", "model_map_value": None}, } _, took, lags = await timed_with_loop_lags( @@ -681,6 +686,103 @@ async def test_async_log_success_event_counts_the_prompt_off_the_event_loop(): assert_loop_stayed_free(took, lags) +class _ManualClock: + def __init__(self) -> None: + self.now = 1_700_000_000.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now = self.now + seconds + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("custom_llm_provider", "model", "model_map_value", "pinned_past_five_minutes"), + [ + pytest.param( + "anthropic", + "claude-sonnet-4-5", + {"cache_creation_input_token_cost_above_1hr": 6e-06}, + True, + id="one-hour-cache-writes-priced", + ), + pytest.param( + "anthropic", + "claude-sonnet-4-5", + {"cache_creation_input_token_cost_above_1hr": None}, + False, + id="no-one-hour-cache-writes", + ), + pytest.param("anthropic", "claude-sonnet-4-5", None, False, id="model-not-in-cost-map"), + pytest.param( + "anthropic", + "claude-sonnet-4-5", + {"cache_creation_input_token_cost_above_1hr": 0.0}, + False, + id="one-hour-cache-writes-priced-at-zero", + ), + pytest.param( + "bedrock", + "apac.anthropic.claude-haiku-4-5-20251001-v1:0", + {"cache_creation_input_token_cost_above_1hr": None}, + True, + id="bedrock-forwards-1h-from-the-base-model-entry", + ), + pytest.param( + "bedrock", + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", + {"cache_creation_input_token_cost_above_1hr": 6.6e-06}, + False, + id="bedrock-drops-1h-for-a-profile-arn-whatever-its-base-model", + ), + ], +) +async def test_one_hour_pin_follows_the_deployments_one_hour_cache_support( + custom_llm_provider: str, + model: str, + model_map_value: dict[str, float | None] | None, + pinned_past_five_minutes: bool, +): + clock: Final = _ManualClock() + cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock)) + check: Final = PromptCachingDeploymentCheck(cache=cache) + messages: Final = cast( + list[AllMessageValues], + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "word " * 6000, "cache_control": {"type": "ephemeral", "ttl": "1h"}} + ], + } + ], + ) + + await check.async_log_success_event( + kwargs={ + "standard_logging_object": { + "call_type": "acompletion", + "custom_llm_provider": custom_llm_provider, + "model": model, + "messages": messages, + "model_id": "dep-1", + "model_map_information": {"model_map_key": model, "model_map_value": model_map_value}, + } + }, + response_obj=None, + start_time=None, + end_time=None, + ) + pins: Final = PromptCachingCache(cache=cache) + assert await pins.async_get_model_id(messages=messages, tools=None) == {"model_id": "dep-1"} + + clock.advance(PROMPT_CACHE_PIN_TTL_SECONDS + 1) + + assert (await pins.async_get_model_id(messages=messages, tools=None) is not None) is pinned_past_five_minutes + + LONG_PROMPT = "word " * 3000 ONE_PIXEL_PNG = ( "data:image/png;base64," @@ -969,7 +1071,7 @@ async def test_claude_code_style_session_stays_on_one_deployment_across_turns(lo response = await router.acompletion(model=MODEL_GROUP_ALIAS, messages=request, mock_response="ok") served.append(response._hidden_params["model_id"]) pin_key = PromptCachingCache.get_prompt_caching_cache_key(request, None) - assert await _eventually(functools.partial(router.cache.get_cache, key=pin_key)) is not None + assert await _eventually(functools.partial(router.prompt_caching_pin_cache.get_cache, key=pin_key)) is not None history = [*history, {"role": "user", "content": [_text(text)]}, {"role": "assistant", "content": "ok"}] assert served == [served[0]] * len(user_turns) @@ -1054,7 +1156,7 @@ async def test_router_with_prompt_caching(anthropic_messages): initial_model_id = response._hidden_params["model_id"] cache = PromptCachingCache( - cache=router.cache, + cache=router.prompt_caching_pin_cache, ) cached_model_id = await _eventually(lambda: cache.get_model_id(messages=anthropic_messages, tools=None)) diff --git a/tests/unit/router_utils/test_prompt_caching_cache.py b/tests/unit/router_utils/test_prompt_caching_cache.py index 3a1132354c4..b81107b47fa 100644 --- a/tests/unit/router_utils/test_prompt_caching_cache.py +++ b/tests/unit/router_utils/test_prompt_caching_cache.py @@ -3,7 +3,273 @@ from typing import Final import pytest -from litellm.router_utils.prompt_caching_cache import PromptCachingCache +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.router_utils.prompt_caching_cache import ( + PROMPT_CACHE_ONE_HOUR_PIN_TTL_SECONDS, + PROMPT_CACHE_PIN_TTL_SECONDS, + PromptCachingCache, +) +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam + + +class _ManualClock: + def __init__(self) -> None: + self.now = 1_700_000_000.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now = self.now + seconds + + +def _pin_cache(clock: _ManualClock) -> PromptCachingCache: + return PromptCachingCache(cache=DualCache(in_memory_cache=InMemoryCache(clock=clock))) + + +def _system_block_messages( + cache_control: dict[str, str], document: str = "a complex legal agreement" +) -> list[AllMessageValues]: + return [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": f"Here is the full text of {document}" * 200, + "cache_control": cache_control, + } + ], + }, + {"role": "user", "content": "Reply with the single word ok."}, + ] + + +def _message_level_messages(cache_control: dict[str, str]) -> list[AllMessageValues]: + return [ + { + "role": "system", + "content": "Here is the full text of a complex legal agreement" * 200, + "cache_control": cache_control, + }, + {"role": "user", "content": "Reply with the single word ok."}, + ] + + +def _tool(cache_control: dict[str, str] | None) -> ChatCompletionToolParam: + tool: ChatCompletionToolParam = { + "type": "function", + "function": {"name": "lookup_clause", "parameters": {"type": "object", "properties": {}}}, + } + return {**tool, "cache_control": cache_control} if cache_control else tool + + +@pytest.mark.parametrize( + ("messages", "tools", "expected_ttl"), + [ + pytest.param( + _system_block_messages({"type": "ephemeral", "ttl": "1h"}), + None, + PROMPT_CACHE_ONE_HOUR_PIN_TTL_SECONDS, + id="content-block-1h", + ), + pytest.param( + _message_level_messages({"type": "ephemeral", "ttl": "1h"}), + None, + PROMPT_CACHE_ONE_HOUR_PIN_TTL_SECONDS, + id="message-level-1h", + ), + pytest.param( + _system_block_messages({"type": "ephemeral", "ttl": "5m"}), + None, + PROMPT_CACHE_PIN_TTL_SECONDS, + id="explicit-5m", + ), + pytest.param( + _system_block_messages({"type": "ephemeral"}), + None, + PROMPT_CACHE_PIN_TTL_SECONDS, + id="default-ttl", + ), + pytest.param( + [ + *_system_block_messages({"type": "ephemeral", "ttl": "1h"}), + {"role": "assistant", "content": "ok"}, + {"role": "user", "content": "Again.", "cache_control": {"type": "ephemeral", "ttl": "5m"}}, + ], + None, + PROMPT_CACHE_ONE_HOUR_PIN_TTL_SECONDS, + id="1h-system-then-5m-turn", + ), + pytest.param( + _system_block_messages({"type": "ephemeral"}), + [_tool({"type": "ephemeral", "ttl": "1h"})], + PROMPT_CACHE_ONE_HOUR_PIN_TTL_SECONDS, + id="1h-on-tools-only", + ), + pytest.param( + _system_block_messages({"type": "ephemeral"}), + [_tool(None)], + PROMPT_CACHE_PIN_TTL_SECONDS, + id="unmarked-tools", + ), + pytest.param( + [ + *_system_block_messages({"type": "ephemeral", "ttl": "5m"}), + {"role": "assistant", "content": "ok"}, + {"role": "user", "content": "Again.", "cache_control": {"type": "something_else", "ttl": "1h"}}, + ], + None, + PROMPT_CACHE_PIN_TTL_SECONDS, + id="non-ephemeral-1h-ignored", + ), + pytest.param( + _system_block_messages({"type": "ephemeral", "ttl": "3600s"}), + None, + PROMPT_CACHE_PIN_TTL_SECONDS, + id="seconds-form-keeps-the-floor", + ), + pytest.param( + _system_block_messages({"type": "ephemeral", "ttl": "2h"}), + None, + PROMPT_CACHE_PIN_TTL_SECONDS, + id="longer-than-1h-keeps-the-floor", + ), + pytest.param( + _system_block_messages({"type": "ephemeral", "ttl": "99999999h"}), + None, + PROMPT_CACHE_PIN_TTL_SECONDS, + id="huge-ttl-keeps-the-floor", + ), + pytest.param( + _system_block_messages({"type": "ephemeral", "ttl": "9" * 400 + "h"}), + None, + PROMPT_CACHE_PIN_TTL_SECONDS, + id="overflowing-ttl-keeps-the-floor", + ), + pytest.param( + _system_block_messages({"type": "ephemeral", "ttl": "90s"}), + None, + PROMPT_CACHE_PIN_TTL_SECONDS, + id="shorter-than-5m-keeps-the-floor", + ), + pytest.param( + _system_block_messages({"type": "ephemeral", "ttl": "soon"}), + None, + PROMPT_CACHE_PIN_TTL_SECONDS, + id="unparseable-ttl-keeps-the-floor", + ), + ], +) +def test_pin_ttl_follows_the_longest_breakpoint_ttl( + messages: list[AllMessageValues], + tools: list[ChatCompletionToolParam] | None, + expected_ttl: int, +) -> None: + assert PromptCachingCache.pin_ttl_seconds(messages, tools) == expected_ttl + + +@pytest.mark.asyncio +async def test_one_hour_prefix_pin_outlives_the_five_minute_window() -> None: + clock: Final = _ManualClock() + cache: Final = _pin_cache(clock) + one_hour: Final = _system_block_messages({"type": "ephemeral", "ttl": "1h"}, document="a lease") + five_minutes: Final = _system_block_messages({"type": "ephemeral"}, document="a will") + + await cache.async_add_model_id(model_id="deployment-1h", messages=one_hour, tools=None, honors_one_hour_ttl=True) + await cache.async_add_model_id( + model_id="deployment-5m", messages=five_minutes, tools=None, honors_one_hour_ttl=True + ) + + clock.advance(PROMPT_CACHE_PIN_TTL_SECONDS + 1) + assert await cache.async_get_model_id(messages=five_minutes, tools=None) is None + pinned_after_five_minutes: Final = await cache.async_get_model_id(messages=one_hour, tools=None) + assert pinned_after_five_minutes is not None + assert pinned_after_five_minutes["model_id"] == "deployment-1h" + + clock.advance(PROMPT_CACHE_ONE_HOUR_PIN_TTL_SECONDS - PROMPT_CACHE_PIN_TTL_SECONDS) + assert await cache.async_get_model_id(messages=one_hour, tools=None) is None + + +@pytest.mark.asyncio +async def test_one_hour_prefix_on_a_deployment_without_one_hour_caching_keeps_the_five_minute_pin() -> None: + clock: Final = _ManualClock() + cache: Final = _pin_cache(clock) + one_hour: Final = _system_block_messages({"type": "ephemeral", "ttl": "1h"}) + + await cache.async_add_model_id( + model_id="deployment-5m-only", messages=one_hour, tools=None, honors_one_hour_ttl=False + ) + + clock.advance(PROMPT_CACHE_PIN_TTL_SECONDS + 1) + assert await cache.async_get_model_id(messages=one_hour, tools=None) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("first_write", "second_write", "pinned_past_five_minutes"), + [ + pytest.param( + ("deployment-1h", True), + ("deployment-5m-only", False), + None, + id="five-minute-deployment-after-a-one-hour-pin", + ), + pytest.param( + ("deployment-5m-only", False), + ("deployment-1h", True), + {"model_id": "deployment-1h"}, + id="one-hour-deployment-after-a-five-minute-pin", + ), + ], +) +async def test_latest_write_sets_the_pin_lifetime_in_a_mixed_group( + first_write: tuple[str, bool], + second_write: tuple[str, bool], + pinned_past_five_minutes: dict[str, str] | None, +) -> None: + clock: Final = _ManualClock() + cache: Final = _pin_cache(clock) + one_hour: Final = _system_block_messages({"type": "ephemeral", "ttl": "1h"}) + + await cache.async_add_model_id( + model_id=first_write[0], messages=one_hour, tools=None, honors_one_hour_ttl=first_write[1] + ) + await cache.async_add_model_id( + model_id=second_write[0], messages=one_hour, tools=None, honors_one_hour_ttl=second_write[1] + ) + + clock.advance(PROMPT_CACHE_PIN_TTL_SECONDS + 1) + assert await cache.async_get_model_id(messages=one_hour, tools=None) == pinned_past_five_minutes + + +def test_sync_five_minute_deployment_after_a_one_hour_pin_keeps_the_five_minute_pin() -> None: + clock: Final = _ManualClock() + cache: Final = _pin_cache(clock) + one_hour: Final = _message_level_messages({"type": "ephemeral", "ttl": "1h"}) + + cache.add_model_id(model_id="deployment-1h", messages=one_hour, tools=None, honors_one_hour_ttl=True) + cache.add_model_id(model_id="deployment-5m-only", messages=one_hour, tools=None, honors_one_hour_ttl=False) + + clock.advance(PROMPT_CACHE_PIN_TTL_SECONDS + 1) + assert cache.get_model_id(messages=one_hour, tools=None) is None + + +def test_sync_one_hour_message_level_pin_outlives_the_five_minute_window() -> None: + clock: Final = _ManualClock() + cache: Final = _pin_cache(clock) + one_hour: Final = _message_level_messages({"type": "ephemeral", "ttl": "1h"}) + + cache.add_model_id(model_id="deployment-1h", messages=one_hour, tools=None, honors_one_hour_ttl=True) + + clock.advance(PROMPT_CACHE_PIN_TTL_SECONDS + 1) + pinned: Final = cache.get_model_id(messages=one_hour, tools=None) + assert pinned is not None + assert pinned["model_id"] == "deployment-1h" + + clock.advance(PROMPT_CACHE_ONE_HOUR_PIN_TTL_SECONDS - PROMPT_CACHE_PIN_TTL_SECONDS) + assert cache.get_model_id(messages=one_hour, tools=None) is None async def _first_non_none(read): @@ -167,7 +433,7 @@ async def test_same_cacheable_prefix_routes_to_same_deployment() -> None: enable_prompt_caching=None, request_kwargs={}, ) - cache: Final = PromptCachingCache(cache=router.cache) + cache: Final = PromptCachingCache(cache=router.prompt_caching_pin_cache) cached_2: Final = await asyncio.wait_for( _first_non_none(lambda: cache.async_get_model_id(affinity_messages, None)), timeout=10, diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index aa89a6ad488..a7c11cc55f3 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -28,7 +28,7 @@ from litellm import APIConnectionError, Router from litellm._logging import verbose_logger, verbose_router_logger from litellm.caching import RedisCache, RedisClusterCache from litellm.caching.caching import DualCache -from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache from litellm.caching.redis_cache import _redis_circuit_breaker_guard from litellm.exceptions import GuardrailRaisedException, MidStreamFallbackError, ModifyResponseException from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper @@ -69,9 +69,10 @@ from litellm.router_utils.fallback_event_handlers import ( DISABLE_FALLBACKS_METADATA_KEY, MID_STREAM_FALLBACK_CONTROLS_KEY, ) +from litellm.router_utils.prompt_caching_cache import PromptCachingCache from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute from litellm.scheduler import FlowItem -from litellm.types.llms.openai import ChatCompletionRequest +from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest from litellm.types.router import ( CustomRoutingStrategyBase, Deployment, @@ -8108,7 +8109,6 @@ async def test_anthropic_messages_call_type_is_cached(): from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( PromptCachingDeploymentCheck, ) - from litellm.router_utils.prompt_caching_cache import PromptCachingCache from litellm.types.utils import ( CallTypes, StandardLoggingHiddenParams, @@ -8135,6 +8135,7 @@ async def test_anthropic_messages_call_type_is_cached(): model_map_key="gpt-3.5-turbo", model_map_value=None ), model="gpt-3.5-turbo", + custom_llm_provider="anthropic", model_id="model-123", model_group="openai-gpt", api_base="https://api.openai.com", @@ -26230,6 +26231,124 @@ async def test_batch_completion_fastest_response_streams_mocked_chunks() -> None assert chunks[-1].choices[0].finish_reason == "stop" +@pytest.mark.asyncio +async def test_prompt_caching_pins_do_not_evict_router_state(): + router: Final = Router( + model_list=[ + { + "model_name": "claude", + "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test"}, + "model_info": {"id": "dep-a"}, + }, + { + "model_name": "claude", + "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test"}, + "model_info": {"id": "dep-b"}, + }, + ], + optional_pre_call_checks=["prompt_caching"], + ) + router.cache.set_cache("deployment:dep-a:cooldown", True, ttl=5) + pins: Final = PromptCachingCache(cache=router.prompt_caching_pin_cache) + for index in range(DEFAULT_MAX_SIZE_IN_MEMORY): + await pins.async_add_model_id( + model_id="dep-a", + messages=[ + { + "role": "system", + "content": [ + {"type": "text", "text": f"policy {index}", "cache_control": {"type": "ephemeral", "ttl": "1h"}} + ], + }, + {"role": "user", "content": "ok"}, + ], + tools=None, + ) + + assert router.cache.get_cache("deployment:dep-a:cooldown") is True + + +@pytest.mark.asyncio +async def test_one_hour_pins_past_the_routers_default_cache_size_keep_fresh_five_minute_pins(): + router: Final = Router( + model_list=[ + { + "model_name": "claude", + "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test"}, + "model_info": {"id": "dep-a"}, + }, + { + "model_name": "claude", + "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test"}, + "model_info": {"id": "dep-b"}, + }, + ], + optional_pre_call_checks=["prompt_caching"], + ) + pins: Final = PromptCachingCache(cache=router.prompt_caching_pin_cache) + + def policy(index: int, ttl: str) -> list[AllMessageValues]: + return [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": f"policy {ttl} {index}", + "cache_control": {"type": "ephemeral", "ttl": ttl}, + } + ], + }, + {"role": "user", "content": "ok"}, + ] + + for index in range(DEFAULT_MAX_SIZE_IN_MEMORY): + await pins.async_add_model_id( + model_id="dep-a", messages=policy(index, "1h"), tools=None, honors_one_hour_ttl=True + ) + for index in range(3): + await pins.async_add_model_id( + model_id="dep-b", messages=policy(index, "5m"), tools=None, honors_one_hour_ttl=True + ) + + assert [await pins.async_get_model_id(messages=policy(index, "5m"), tools=None) for index in range(3)] == [ + {"model_id": "dep-b"} + ] * 3 + + +@pytest.mark.asyncio +async def test_flush_cache_drops_prompt_caching_pins(): + router: Final = Router( + model_list=[ + { + "model_name": "claude", + "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test"}, + "model_info": {"id": "dep-a"}, + }, + { + "model_name": "claude", + "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test"}, + "model_info": {"id": "dep-b"}, + }, + ], + optional_pre_call_checks=["prompt_caching"], + ) + messages: Final = [ + { + "role": "system", + "content": [{"type": "text", "text": "policy", "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + }, + {"role": "user", "content": "ok"}, + ] + pins: Final = PromptCachingCache(cache=router.prompt_caching_pin_cache) + await pins.async_add_model_id(model_id="dep-a", messages=messages, tools=None) + assert await pins.async_get_model_id(messages=messages, tools=None) == {"model_id": "dep-a"} + + router.flush_cache() + + assert await pins.async_get_model_id(messages=messages, tools=None) is None + + class _FirstPullStream: """A stream whose first pull yields `first`, raises it, or ends (first=None).""" diff --git a/tests/unit/test_router_order_fallback.py b/tests/unit/test_router_order_fallback.py index 37859761858..47e0be28635 100644 --- a/tests/unit/test_router_order_fallback.py +++ b/tests/unit/test_router_order_fallback.py @@ -502,7 +502,7 @@ async def test_router_order_fallback_ignores_prompt_cache_pin_on_target_order(): num_retries=0, optional_pre_call_checks=["prompt_caching"], ) - await PromptCachingCache(cache=router.cache).async_add_model_id( + await PromptCachingCache(cache=router.prompt_caching_pin_cache).async_add_model_id( model_id="1", messages=messages, tools=None,