From c8706c8b036a1cb15e405ff4802062950ff35584 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 23:06:52 -0700 Subject: [PATCH] fix(vertex_ai): forget a remembered context cache that generateContent rejects A cache deleted out of band stayed remembered until its expireTime, so every request failed with 400 (Vertex) or 403 (Gemini) until then. Drop the name on a 400, 403, or 404 so the next request lists and recreates it. Move the registry into its own module so the Gemini handler can import it without the circular import through the context caching module's Cache(). --- .../context_caching/known_cached_contents.py | 139 +++++++++++++++++ .../vertex_ai_context_caching.py | 145 +----------------- .../vertex_and_google_ai_studio_gemini.py | 24 ++- .../vertex_ai/context_caching/conftest.py | 4 +- .../test_known_cached_contents.py | 77 ++++++++-- 5 files changed, 229 insertions(+), 160 deletions(-) create mode 100644 litellm/llms/vertex_ai/context_caching/known_cached_contents.py diff --git a/litellm/llms/vertex_ai/context_caching/known_cached_contents.py b/litellm/llms/vertex_ai/context_caching/known_cached_contents.py new file mode 100644 index 00000000000..daef9076349 --- /dev/null +++ b/litellm/llms/vertex_ai/context_caching/known_cached_contents.py @@ -0,0 +1,139 @@ +import asyncio +import re +import threading +import time +import weakref +from collections.abc import Awaitable, Callable, Mapping +from datetime import datetime, timezone +from typing import Final, Literal + +# 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 +CACHE_REJECTED_STATUS_CODES: Final = frozenset({400, 403, 404}) + +_RFC3339_FRACTION: Final = re.compile(r"\.(\d{1,6})\d*") + + +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[tuple[str, str], str], tuple[str, float]] = {} # mutable-ok: guarded by _lock + self._lock: Final = threading.Lock() + self._inflight: weakref.WeakKeyDictionary[ # mutable-ok: private, guarded by _lock + asyncio.AbstractEventLoop, dict[tuple[tuple[str, str], str], asyncio.Future[str | None]] + ] = weakref.WeakKeyDictionary() + + @staticmethod + def scope( + base_url: str, + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], + headers: Mapping[str, str], + ) -> tuple[str, str]: + """Where a cache name is valid: the Vertex URL names project and location, a Gemini key names the project. + + The Gemini key is read from the headers actually sent, so an ``extra_headers`` override scopes the + cache to the caller's project. It is kept as-is in this in-memory scope and never logged or persisted. + """ + if custom_llm_provider != "gemini": + return (base_url, "") + sent_keys: Final = [value for name, value in headers.items() if name.lower() == "x-goog-api-key"] + return (base_url, sent_keys[-1] if sent_keys else "") + + @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: tuple[str, 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: tuple[str, 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) + + def forget_rejected(self, cached_content: str | None, status_code: int) -> None: + if cached_content is None or status_code not in CACHE_REJECTED_STATUS_CODES: + return + with self._lock: + self._entries = {key: entry for key, entry in self._entries.items() if entry[0] != cached_content} + + async def single_flight( + self, scope: tuple[str, 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() 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 c423b0ccaf3..326872a367f 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,10 +1,4 @@ -import asyncio -import re -import threading -import time -import weakref -from collections.abc import Awaitable, Callable, Mapping, Sequence -from datetime import datetime, timezone +from collections.abc import Sequence from typing import Final, Literal import httpx @@ -26,6 +20,7 @@ from litellm.utils import is_prompt_caching_valid_prompt from ..common_utils import VertexAIError, get_vertex_base_url from ..vertex_llm_base import VertexBase +from .known_cached_contents import KnownCachedContents, known_cached_contents from .transformation import ( cached_messages_end_on_supported_turn, separate_cached_messages, @@ -36,13 +31,6 @@ 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): @@ -82,123 +70,6 @@ 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[tuple[str, str], str], tuple[str, float]] = {} # mutable-ok: guarded by _lock - self._lock: Final = threading.Lock() - self._inflight: weakref.WeakKeyDictionary[ # mutable-ok: private, guarded by _lock - asyncio.AbstractEventLoop, dict[tuple[tuple[str, str], str], asyncio.Future[str | None]] - ] = weakref.WeakKeyDictionary() - - @staticmethod - def scope( - base_url: str, - custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], - headers: Mapping[str, str], - ) -> tuple[str, str]: - """Where a cache name is valid: the Vertex URL names project and location, a Gemini key names the project. - - The Gemini key is read from the headers actually sent, so an ``extra_headers`` override scopes the - cache to the caller's project. It is kept as-is in this in-memory scope and never logged or persisted. - """ - if custom_llm_provider != "gemini": - return (base_url, "") - sent_keys: Final = [value for name, value in headers.items() if name.lower() == "x-goog-api-key"] - return (base_url, sent_keys[-1] if sent_keys else "") - - @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: tuple[str, 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: tuple[str, 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: tuple[str, 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 @@ -291,7 +162,7 @@ class ContextCachingEndpoints(VertexBase): vertex_auth_header=vertex_auth_header, model=model, ) - scope: Final = _KnownCachedContents.scope(base_url, custom_llm_provider, headers) + scope: Final = KnownCachedContents.scope(base_url, custom_llm_provider, headers) known_name: Final = known_cached_contents.get(scope, cache_key) if known_name is not None: return known_name @@ -340,7 +211,7 @@ class ContextCachingEndpoints(VertexBase): if display_name is None or display_name != cache_key: continue expire_time = cached_item.get("expireTime") - if not _KnownCachedContents.is_usable(expire_time): + if not KnownCachedContents.is_usable(expire_time): continue name = cached_item.get("name") known_cached_contents.remember(scope, cache_key, name, expire_time) @@ -388,7 +259,7 @@ class ContextCachingEndpoints(VertexBase): vertex_auth_header=vertex_auth_header, model=model, ) - scope: Final = _KnownCachedContents.scope(base_url, custom_llm_provider, headers) + scope: Final = KnownCachedContents.scope(base_url, custom_llm_provider, headers) known_name: Final = known_cached_contents.get(scope, cache_key) if known_name is not None: return known_name @@ -437,7 +308,7 @@ class ContextCachingEndpoints(VertexBase): if display_name is None or display_name != cache_key: continue expire_time = cached_item.get("expireTime") - if not _KnownCachedContents.is_usable(expire_time): + if not KnownCachedContents.is_usable(expire_time): continue name = cached_item.get("name") known_cached_contents.remember(scope, cache_key, name, expire_time) @@ -604,7 +475,7 @@ 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, headers), + KnownCachedContents.scope(url, custom_llm_provider, headers), generated_cache_key, raw_response_cached.get("name"), raw_response_cached.get("expireTime"), @@ -706,7 +577,7 @@ class ContextCachingEndpoints(VertexBase): generated_cache_key: Final = local_cache_obj.get_cache_key( messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model ) - scope: Final = _KnownCachedContents.scope(url, custom_llm_provider, headers) + scope: Final = KnownCachedContents.scope(url, custom_llm_provider, headers) async def lookup_or_create() -> str | None: google_cache_name: Final = await self.async_check_cache( diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 257fafbb6b8..f30e8c69324 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -96,6 +96,7 @@ from ..common_utils import ( build_vertex_schema, supports_response_json_schema, ) +from ..context_caching.known_cached_contents import known_cached_contents from ..vertex_llm_base import VertexBase from .grounding_requests import calculate_grounding_requests from .long_decode import ( @@ -2831,6 +2832,7 @@ async def make_call( ) response.raise_for_status() except httpx.HTTPStatusError as e: + known_cached_contents.forget_rejected(request_body.get("cachedContent"), e.response.status_code) exception_string: Final = str(await e.response.aread()) raise VertexAIError( status_code=e.response.status_code, @@ -2891,14 +2893,18 @@ def make_sync_call( if client is None: client = HTTPHandler() # Create a new client if none provided - response: Final = client.post( - url_for(api_base, request_body), - headers=headers, - data=data, - stream=True, - logging_obj=logging_obj, - timeout=timeout, - ) + try: + response: Final = client.post( + url_for(api_base, request_body), + headers=headers, + data=data, + stream=True, + logging_obj=logging_obj, + timeout=timeout, + ) + except httpx.HTTPStatusError as e: + known_cached_contents.forget_rejected(request_body.get("cachedContent"), e.response.status_code) + raise if response.status_code != 200 and response.status_code != 201: raise VertexAIError( @@ -3175,6 +3181,7 @@ class VertexLLM(VertexBase): ) response.raise_for_status() except httpx.HTTPStatusError as err: + known_cached_contents.forget_rejected(body.get("cachedContent"), err.response.status_code) raise VertexAIError( status_code=err.response.status_code, message=err.response.text, @@ -3385,6 +3392,7 @@ class VertexLLM(VertexBase): ) response.raise_for_status() except httpx.HTTPStatusError as err: + known_cached_contents.forget_rejected(body.get("cachedContent"), err.response.status_code) raise VertexAIError( status_code=err.response.status_code, message=err.response.text, diff --git a/tests/unit/llms/vertex_ai/context_caching/conftest.py b/tests/unit/llms/vertex_ai/context_caching/conftest.py index 75bcb5cde8a..9c80c964608 100644 --- a/tests/unit/llms/vertex_ai/context_caching/conftest.py +++ b/tests/unit/llms/vertex_ai/context_caching/conftest.py @@ -1,8 +1,6 @@ import pytest -from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( - known_cached_contents, -) +from litellm.llms.vertex_ai.context_caching.known_cached_contents import known_cached_contents @pytest.fixture(autouse=True) 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 index 33052341756..8f9c4ff05d9 100644 --- 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 @@ -1,21 +1,24 @@ import asyncio +import json from datetime import datetime, timedelta, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx +import openai import pytest +import respx + +import litellm 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, -) +from litellm.llms.vertex_ai.context_caching.known_cached_contents import known_cached_contents, parse_expire_time +from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ContextCachingEndpoints _MODULE: Final = "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching" +_REGISTRY_MODULE: Final = "litellm.llms.vertex_ai.context_caching.known_cached_contents" _SCOPE: Final = ("https://vertex.test/v1/projects/p/locations/us/cachedContents", "") @@ -163,7 +166,7 @@ def test_remembered_cache_is_listed_again_once_its_reuse_window_closes(): 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): + with patch(f"{_REGISTRY_MODULE}.time.time", return_value=datetime.now(timezone.utc).timestamp() + 55): assert _sync_call(client)[2] == "cachedContents/listed" @@ -290,18 +293,18 @@ async def test_waiters_resolve_on_their_own_when_the_first_request_fails(): def test_parse_expire_time_accepts_nanosecond_precision(): - assert _parse_expire_time("2026-10-14T18:06:21.123456789Z") == pytest.approx( + 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() + 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 + assert parse_expire_time("not a time") is None + assert parse_expire_time(None) is None def test_parse_expire_time_reads_a_missing_offset_as_utc(): - assert _parse_expire_time("2026-10-14T18:06:21") == datetime(2026, 10, 14, 18, 6, 21, tzinfo=timezone.utc).timestamp() + assert parse_expire_time("2026-10-14T18:06:21") == datetime(2026, 10, 14, 18, 6, 21, tzinfo=timezone.utc).timestamp() def test_cache_expiring_within_the_margin_is_not_remembered(): @@ -311,9 +314,59 @@ def test_cache_expiring_within_the_margin_is_not_remembered(): def test_registry_is_bounded(): - with patch(f"{_MODULE}.MAX_KNOWN_CACHED_CONTENTS", 3): + with patch(f"{_REGISTRY_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" + + +async def _complete(sync_mode: bool, stream: bool) -> None: + if sync_mode: + response = litellm.completion(model="gemini/gemini-3.1-pro", messages=_messages(), stream=stream) + for _chunk in response if stream else (): + pass + return + response = await litellm.acompletion(model="gemini/gemini-3.1-pro", messages=_messages(), stream=stream) + if stream: + async for _chunk in response: + pass + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize( + ("status_code", "relisted"), + [(400, True), (403, True), (404, True), (429, False), (500, False)], +) +@pytest.mark.asyncio +async def test_cache_rejected_by_generate_content_is_looked_up_again( + respx_mock: respx.MockRouter, + fake_provider_credentials: None, + monkeypatch: pytest.MonkeyPatch, + sync_mode: bool, + stream: bool, + status_code: int, + relisted: bool, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + list_route: Final = respx_mock.get(url__regex=r".*/cachedContents.*").mock( + return_value=httpx.Response(200, json={}) + ) + respx_mock.post(url__regex=r".*/cachedContents$").mock( + return_value=httpx.Response(200, json=_created("cachedContents/deleted", _expire_time(300))) + ) + generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock( + return_value=httpx.Response(status_code, json={"error": {"code": status_code, "message": "rejected"}}) + ) + + for _ in range(2): + with pytest.raises(openai.APIStatusError): + await _complete(sync_mode, stream) + + assert [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls] == [ + "cachedContents/deleted", + "cachedContents/deleted", + ] + assert list_route.call_count == (2 if relisted else 1)