scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

fix(vertex_ai): reuse known context caches instead of listing on every request

Browse source 
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 <noreply@anthropic.com>
This commit is contained in:
Ondra Zahradnik 2026-10-15 15:52:56 +02:00
parent cda19d6a29
commit 755efd48a2
3 changed files with 472 additions and 63 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

285
litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py
View file

@ -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

12
tests/unit/llms/vertex_ai/context_caching/conftest.py Normal file
View file

@ -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()

238
tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py Normal file
View file

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