From 755efd48a2671c0465b82ec35be31bf18389044a Mon Sep 17 00:00:00 2001 From: Ondra Zahradnik Date: Thu, 15 Oct 2026 15:52:56 +0200 Subject: [PATCH] fix(vertex_ai): reuse known context caches instead of listing on every request Every cache-marked request listed all cachedContents in the location to find its cache by display name. Vertex meters each list page against a per-minute, per-location quota, so under load the lookups themselves return 429 RESOURCE_EXHAUSTED even though generation quota is available. The process now remembers each cache it created or found, keyed by endpoint scope and cache key, until shortly before the expireTime Vertex reported, and concurrent async requests for the same prefix share one lookup-or-create. Listed caches about to expire are skipped, so a request no longer reaches generateContent with a cache that expires in flight. Co-Authored-By: Claude Opus 5.5 --- .../vertex_ai_context_caching.py | 285 ++++++++++++++---- .../vertex_ai/context_caching/conftest.py | 12 + .../test_known_cached_contents.py | 238 +++++++++++++++ 3 files changed, 472 insertions(+), 63 deletions(-) create mode 100644 tests/unit/llms/vertex_ai/context_caching/conftest.py create mode 100644 tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 06b8f3a3fbb..65116b407be 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -1,4 +1,11 @@ -from collections.abc import Sequence +import asyncio +import hashlib +import re +import threading +import time +import weakref +from collections.abc import Awaitable, Callable, Sequence +from datetime import datetime, timezone from typing import Final, Literal import httpx @@ -30,11 +37,19 @@ local_cache_obj: Final = Cache(type=LiteLLMCacheType.LOCAL) # only used for cal MAX_PAGINATION_PAGES: Final = 100 # Reasonable upper bound for pagination +# A remembered cache stops being reused this long before Vertex expires it, so a request +# resolved just before expiry does not reach generateContent with an expired cache. +KNOWN_CACHE_EXPIRY_MARGIN_SECONDS: Final = 10.0 +MAX_KNOWN_CACHED_CONTENTS: Final = 10_000 + +_RFC3339_FRACTION: Final = re.compile(r"\.(\d{1,6})\d*") + @with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) class _CachedContentItem(TypedDict): displayName: ReadOnly[NotRequired[object]] name: ReadOnly[NotRequired[str | None]] + expireTime: ReadOnly[NotRequired[str | None]] @with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) @@ -51,6 +66,7 @@ class _CachedContentUsageMetadata(TypedDict): @with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) class _CreatedCachedContent(TypedDict): name: ReadOnly[NotRequired[str | None]] + expireTime: ReadOnly[NotRequired[str | None]] usageMetadata: ReadOnly[NotRequired[_CachedContentUsageMetadata | None]] @@ -67,6 +83,118 @@ def _record_cache_creation_tokens(logging_obj: Logging, created: _CreatedCachedC logging_obj.vertex_ai_cache_creation_input_tokens = max(tokens, 0) # rebind-ok: set provider count +def _parse_expire_time(expire_time: object) -> float | None: + """Parse a Vertex RFC 3339 ``expireTime`` (nanosecond precision, ``Z`` suffix) to epoch seconds.""" + if not isinstance(expire_time, str) or not expire_time: + return None + normalized: Final = _RFC3339_FRACTION.sub(lambda m: "." + m.group(1), expire_time.replace("Z", "+00:00"), count=1) + try: + parsed: Final = datetime.fromisoformat(normalized) + except ValueError: + return None + if parsed.tzinfo is None: + return parsed.replace(tzinfo=timezone.utc).timestamp() + return parsed.timestamp() + + +class _KnownCachedContents: + """Process-local record of the cachedContents this process created or found. + + Vertex meters every cachedContents list page against a per-minute, per-location quota, + so listing on every cache-marked request exhausts it under load. A cache is only + remembered with the ``expireTime`` Vertex reported, and is never reused within + ``KNOWN_CACHE_EXPIRY_MARGIN_SECONDS`` of it. + """ + + def __init__(self) -> None: + self._entries: dict[tuple[str, str], tuple[str, float]] = {} # mutable-ok: private, guarded by _lock + self._lock: Final = threading.Lock() + self._inflight: weakref.WeakKeyDictionary[ # mutable-ok: private, guarded by _lock + asyncio.AbstractEventLoop, dict[tuple[str, str], asyncio.Future[str | None]] + ] = weakref.WeakKeyDictionary() + + @staticmethod + def scope( + base_url: str, + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], + api_key: str | None, + ) -> str: + """Where a cache name is valid: the Vertex URL names project and location, a Gemini key names the project.""" + if custom_llm_provider == "gemini": + return f"{base_url}#{hashlib.sha256((api_key or '').encode()).hexdigest()}" + return base_url + + @staticmethod + def is_usable(expire_time: object) -> bool: + expires_at: Final = _parse_expire_time(expire_time) + return expires_at is None or expires_at - KNOWN_CACHE_EXPIRY_MARGIN_SECONDS > time.time() + + def get(self, scope: str, cache_key: str) -> str | None: + with self._lock: + entry: Final = self._entries.get((scope, cache_key)) + if entry is None: + return None + name, reuse_until = entry + if reuse_until > time.time(): + return name + del self._entries[(scope, cache_key)] + return None + + def remember(self, scope: str, cache_key: str, name: str | None, expire_time: object) -> None: + expires_at: Final = _parse_expire_time(expire_time) + if not name or expires_at is None: + return + reuse_until: Final = expires_at - KNOWN_CACHE_EXPIRY_MARGIN_SECONDS + now: Final = time.time() + if reuse_until <= now: + return + with self._lock: + if len(self._entries) >= MAX_KNOWN_CACHED_CONTENTS: + self._entries = {k: v for k, v in self._entries.items() if v[1] > now} + while len(self._entries) >= MAX_KNOWN_CACHED_CONTENTS: + del self._entries[next(iter(self._entries))] + self._entries[(scope, cache_key)] = (name, reuse_until) + + async def single_flight( + self, scope: str, cache_key: str, resolve: Callable[[], Awaitable[str | None]] + ) -> str | None: + """Run ``resolve`` once per (scope, cache_key) at a time within this event loop. + + Concurrent requests for the same prefix wait for the first one's lookup-or-create + instead of each listing (and creating a duplicate cache). If that first attempt + fails, each waiter resolves on its own. + """ + loop: Final = asyncio.get_running_loop() + own: Final[asyncio.Future[str | None]] = loop.create_future() + with self._lock: + inflight: Final = self._inflight.setdefault(loop, {}) + leader: Final = inflight.setdefault((scope, cache_key), own) + if leader is not own: + name: Final = await asyncio.shield(leader) + if name: + return name + return await resolve() + try: + resolved: Final = await resolve() + except BaseException: + own.set_result(None) + raise + else: + own.set_result(resolved) + return resolved + finally: + with self._lock: + inflight.pop((scope, cache_key), None) + + def clear(self) -> None: + with self._lock: + self._entries.clear() + self._inflight.clear() + + +known_cached_contents: Final = _KnownCachedContents() + + class ContextCachingEndpoints(VertexBase): """ Covers context caching endpoints for Vertex AI + Google AI Studio @@ -159,6 +287,10 @@ class ContextCachingEndpoints(VertexBase): vertex_auth_header=vertex_auth_header, model=model, ) + scope: Final = _KnownCachedContents.scope(base_url, custom_llm_provider, api_key) + known_name: Final = known_cached_contents.get(scope, cache_key) + if known_name is not None: + return known_name page_token: str | None = None @@ -201,8 +333,14 @@ class ContextCachingEndpoints(VertexBase): # Check current page for matching cache_key for cached_item in raw_response["cachedContents"]: display_name = cached_item.get("displayName") - if display_name is not None and display_name == cache_key: - return cached_item.get("name") + if display_name is None or display_name != cache_key: + continue + expire_time = cached_item.get("expireTime") + if not _KnownCachedContents.is_usable(expire_time): + continue + name = cached_item.get("name") + known_cached_contents.remember(scope, cache_key, name, expire_time) + return name # Check if there are more pages page_token = raw_response.get("nextPageToken") @@ -246,6 +384,10 @@ class ContextCachingEndpoints(VertexBase): vertex_auth_header=vertex_auth_header, model=model, ) + scope: Final = _KnownCachedContents.scope(base_url, custom_llm_provider, api_key) + known_name: Final = known_cached_contents.get(scope, cache_key) + if known_name is not None: + return known_name page_token: str | None = None @@ -288,8 +430,14 @@ class ContextCachingEndpoints(VertexBase): # Check current page for matching cache_key for cached_item in raw_response["cachedContents"]: display_name = cached_item.get("displayName") - if display_name is not None and display_name == cache_key: - return cached_item.get("name") + if display_name is None or display_name != cache_key: + continue + expire_time = cached_item.get("expireTime") + if not _KnownCachedContents.is_usable(expire_time): + continue + name = cached_item.get("name") + known_cached_contents.remember(scope, cache_key, name, expire_time) + return name # Check if there are more pages page_token = raw_response.get("nextPageToken") @@ -451,6 +599,12 @@ class ContextCachingEndpoints(VertexBase): raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root _record_cache_creation_tokens(logging_obj, raw_response_cached) + known_cached_contents.remember( + _KnownCachedContents.scope(url, custom_llm_provider, api_key), + generated_cache_key, + raw_response_cached.get("name"), + raw_response_cached.get("expireTime"), + ) return ( non_cached_messages, optional_params, @@ -548,68 +702,73 @@ class ContextCachingEndpoints(VertexBase): generated_cache_key: Final = local_cache_obj.get_cache_key( messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model ) - google_cache_name: Final = await self.async_check_cache( - cache_key=generated_cache_key, - client=client, - headers=headers, - api_key=api_key, - api_base=api_base, - logging_obj=logging_obj, - custom_llm_provider=custom_llm_provider, - vertex_project=vertex_project, - vertex_location=vertex_location, - vertex_auth_header=vertex_auth_header, - model=model, - ) + scope: Final = _KnownCachedContents.scope(url, custom_llm_provider, api_key) - if google_cache_name: - return non_cached_messages, optional_params, google_cache_name - - ## TRANSFORM REQUEST - cached_content_request_body: Final = transform_openai_messages_to_gemini_context_caching( - model=model, - messages=cached_messages, - cache_key=generated_cache_key, - custom_llm_provider=custom_llm_provider, - vertex_project=vertex_project, - vertex_location=vertex_location, - ) - - cached_content_request_body["tools"] = tools - if tool_choice is not None: - cached_content_request_body["toolConfig"] = tool_choice - - ## LOGGING - logging_obj.pre_call( - input=messages, - api_key="", - additional_args={ - "complete_input_dict": cached_content_request_body, - "api_base": url, - "headers": headers, - }, - ) - - try: - response: Final = await client.post( - url=url, + async def lookup_or_create() -> str | None: + google_cache_name: Final = await self.async_check_cache( + cache_key=generated_cache_key, + client=client, headers=headers, - json=cached_content_request_body, + api_key=api_key, + api_base=api_base, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_auth_header=vertex_auth_header, + model=model, ) - response.raise_for_status() - except httpx.HTTPStatusError as err: - error_code: Final = err.response.status_code - raise VertexAIError(status_code=error_code, message=err.response.text) - except httpx.TimeoutException: - raise VertexAIError(status_code=408, message="Timeout error occurred.") - raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root - _record_cache_creation_tokens(logging_obj, raw_response_cached) - return ( - non_cached_messages, - optional_params, - raw_response_cached.get("name"), - ) + if google_cache_name: + return google_cache_name + + ## TRANSFORM REQUEST + cached_content_request_body: Final = transform_openai_messages_to_gemini_context_caching( + model=model, + messages=cached_messages, + cache_key=generated_cache_key, + custom_llm_provider=custom_llm_provider, + vertex_project=vertex_project, + vertex_location=vertex_location, + ) + + cached_content_request_body["tools"] = tools + if tool_choice is not None: + cached_content_request_body["toolConfig"] = tool_choice + + ## LOGGING + logging_obj.pre_call( + input=messages, + api_key="", + additional_args={ + "complete_input_dict": cached_content_request_body, + "api_base": url, + "headers": headers, + }, + ) + + try: + response: Final = await client.post( + url=url, + headers=headers, + json=cached_content_request_body, + ) + response.raise_for_status() + except httpx.HTTPStatusError as err: + error_code: Final = err.response.status_code + raise VertexAIError(status_code=error_code, message=err.response.text) + except httpx.TimeoutException: + raise VertexAIError(status_code=408, message="Timeout error occurred.") + + raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root + _record_cache_creation_tokens(logging_obj, raw_response_cached) + known_cached_contents.remember( + scope, generated_cache_key, raw_response_cached.get("name"), raw_response_cached.get("expireTime") + ) + return raw_response_cached.get("name") + + cache_name: Final = await known_cached_contents.single_flight(scope, generated_cache_key, lookup_or_create) + return non_cached_messages, optional_params, cache_name def get_cache(self): pass diff --git a/tests/unit/llms/vertex_ai/context_caching/conftest.py b/tests/unit/llms/vertex_ai/context_caching/conftest.py new file mode 100644 index 00000000000..75bcb5cde8a --- /dev/null +++ b/tests/unit/llms/vertex_ai/context_caching/conftest.py @@ -0,0 +1,12 @@ +import pytest + +from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( + known_cached_contents, +) + + +@pytest.fixture(autouse=True) +def _forget_known_cached_contents(): + known_cached_contents.clear() + yield + known_cached_contents.clear() diff --git a/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py b/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py new file mode 100644 index 00000000000..089ed9cb69b --- /dev/null +++ b/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py @@ -0,0 +1,238 @@ +import asyncio +from datetime import datetime, timedelta, timezone +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.vertex_ai.common_utils import VertexAIError +from litellm.llms.vertex_ai.context_caching import vertex_ai_context_caching +from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( + ContextCachingEndpoints, + _parse_expire_time, + known_cached_contents, +) + +_MODULE: Final = "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching" + + +@pytest.fixture(autouse=True) +def _cacheable_prompt_and_fixed_url(): + with ( + patch(f"{_MODULE}.is_prompt_caching_valid_prompt", return_value=True), + patch.object( + ContextCachingEndpoints, + "_get_token_and_url_context_caching", + return_value=("token", "https://vertex.test/v1/projects/p/locations/us/cachedContents"), + ), + ): + yield + + +def _expire_time(seconds_from_now: float) -> str: + expires: Final = datetime.now(timezone.utc) + timedelta(seconds=seconds_from_now) + return expires.strftime("%Y-%m-%dT%H:%M:%S.%f") + "123Z" + + +def _messages(prefix: str = "cacheable prefix ") -> list: + return [ + {"role": "system", "content": [{"type": "text", "text": prefix * 50, "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": "question"}, + ] + + +def _response(body: dict) -> MagicMock: + return MagicMock(json=lambda: body) + + +def _created(name: str, expire_time: str | None) -> dict: + body: Final = {"name": name, "usageMetadata": {"totalTokenCount": 2048}} + return body if expire_time is None else {**body, "expireTime": expire_time} + + +def _sync_call(client: HTTPHandler, api_key: str = "k", custom_llm_provider: str = "vertex_ai", messages=None): + return ContextCachingEndpoints().check_and_create_cache( + messages=messages or _messages(), + optional_params={}, + api_key=api_key, + api_base=None, + model="gemini-3.1-pro", + client=client, + timeout=30.0, + logging_obj=MagicMock(model_call_details={}), + custom_llm_provider=custom_llm_provider, + vertex_project="p", + vertex_location="us", + vertex_auth_header=None, + ) + + +async def _async_call(client: AsyncHTTPHandler): + return await ContextCachingEndpoints().async_check_and_create_cache( + messages=_messages(), + optional_params={}, + api_key="k", + api_base=None, + model="gemini-3.1-pro", + client=client, + timeout=30.0, + logging_obj=MagicMock(model_call_details={}), + custom_llm_provider="vertex_ai", + vertex_project="p", + vertex_location="us", + vertex_auth_header=None, + ) + + +def test_created_cache_is_reused_without_listing(): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.return_value = _response(_created("cachedContents/1", _expire_time(60))) + + assert _sync_call(client)[2] == "cachedContents/1" + assert _sync_call(client)[2] == "cachedContents/1" + + assert client.get.call_count == 1 + assert client.post.call_count == 1 + + +def test_listed_cache_is_reused_without_listing_again(): + client: Final = MagicMock(spec=HTTPHandler) + key: Final = vertex_ai_context_caching.local_cache_obj.get_cache_key( + messages=_messages()[:1], tools=None, tool_choice=None, model="gemini-3.1-pro" + ) + client.get.return_value = _response( + {"cachedContents": [{"name": "cachedContents/listed", "displayName": key, "expireTime": _expire_time(60)}]} + ) + + assert _sync_call(client)[2] == "cachedContents/listed" + assert _sync_call(client)[2] == "cachedContents/listed" + + assert client.get.call_count == 1 + client.post.assert_not_called() + + +def test_listed_cache_about_to_expire_is_not_used(): + client: Final = MagicMock(spec=HTTPHandler) + key: Final = vertex_ai_context_caching.local_cache_obj.get_cache_key( + messages=_messages()[:1], tools=None, tool_choice=None, model="gemini-3.1-pro" + ) + client.get.return_value = _response( + {"cachedContents": [{"name": "cachedContents/stale", "displayName": key, "expireTime": _expire_time(3)}]} + ) + client.post.return_value = _response(_created("cachedContents/fresh", _expire_time(60))) + + assert _sync_call(client)[2] == "cachedContents/fresh" + + +def test_cache_without_expire_time_is_not_remembered(): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.return_value = _response(_created("cachedContents/1", None)) + + _sync_call(client) + _sync_call(client) + + assert client.get.call_count == 2 + + +def test_remembered_cache_is_listed_again_once_its_reuse_window_closes(): + client: Final = MagicMock(spec=HTTPHandler) + key: Final = vertex_ai_context_caching.local_cache_obj.get_cache_key( + messages=_messages()[:1], tools=None, tool_choice=None, model="gemini-3.1-pro" + ) + client.get.side_effect = [ + _response({}), + _response( + {"cachedContents": [{"name": "cachedContents/listed", "displayName": key, "expireTime": _expire_time(600)}]} + ), + ] + client.post.return_value = _response(_created("cachedContents/1", _expire_time(60))) + assert _sync_call(client)[2] == "cachedContents/1" + + with patch(f"{_MODULE}.time.time", return_value=datetime.now(timezone.utc).timestamp() + 55): + assert _sync_call(client)[2] == "cachedContents/listed" + + +def test_gemini_caches_are_not_shared_across_api_keys(): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.return_value = _response(_created("cachedContents/1", _expire_time(60))) + + _sync_call(client, api_key="key-a", custom_llm_provider="gemini") + _sync_call(client, api_key="key-b", custom_llm_provider="gemini") + + assert client.post.call_count == 2 + + +def test_different_prefixes_do_not_share_a_cache(): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _response({}) + client.post.side_effect = [ + _response(_created("cachedContents/a", _expire_time(60))), + _response(_created("cachedContents/b", _expire_time(60))), + ] + + assert _sync_call(client, messages=_messages("first "))[2] == "cachedContents/a" + assert _sync_call(client, messages=_messages("second "))[2] == "cachedContents/b" + + +@pytest.mark.asyncio +async def test_concurrent_requests_for_one_prefix_list_and_create_once(): + client: Final = MagicMock(spec=AsyncHTTPHandler) + + async def slow_list(**_): + await asyncio.sleep(0.01) + return _response({}) + + client.get = AsyncMock(side_effect=slow_list) + client.post = AsyncMock(return_value=_response(_created("cachedContents/1", _expire_time(60)))) + + results: Final = await asyncio.gather(*(_async_call(client) for _ in range(5))) + + assert {result[2] for result in results} == {"cachedContents/1"} + assert client.get.await_count == 1 + assert client.post.await_count == 1 + + +@pytest.mark.asyncio +async def test_waiters_resolve_on_their_own_when_the_first_request_fails(): + client: Final = MagicMock(spec=AsyncHTTPHandler) + + async def slow_list(**_): + await asyncio.sleep(0.01) + return _response({}) + + client.get = AsyncMock(side_effect=slow_list) + client.post = AsyncMock( + side_effect=[ + VertexAIError(status_code=500, message="create failed"), + _response(_created("cachedContents/2", _expire_time(60))), + ] + ) + + results: Final = await asyncio.gather(_async_call(client), _async_call(client), return_exceptions=True) + + assert isinstance(results[0], VertexAIError) + assert results[1][2] == "cachedContents/2" + + +def test_parse_expire_time_accepts_nanosecond_precision(): + assert _parse_expire_time("2026-10-14T18:06:21.123456789Z") == pytest.approx( + datetime(2026, 10, 14, 18, 6, 21, 123456, tzinfo=timezone.utc).timestamp() + ) + assert ( + _parse_expire_time("2026-10-14T18:06:21Z") == datetime(2026, 10, 14, 18, 6, 21, tzinfo=timezone.utc).timestamp() + ) + assert _parse_expire_time("not a time") is None + assert _parse_expire_time(None) is None + + +def test_registry_is_bounded(): + with patch(f"{_MODULE}.MAX_KNOWN_CACHED_CONTENTS", 3): + for index in range(5): + known_cached_contents.remember("scope", f"key-{index}", f"cachedContents/{index}", _expire_time(60)) + + assert known_cached_contents.get("scope", "key-0") is None + assert known_cached_contents.get("scope", "key-4") == "cachedContents/4"