diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 6cff9729189..0cdd92bb071 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -7,12 +7,14 @@ Handles Authentication and generating request urls for Vertex AI and Google AI S import asyncio import json import threading +from collections import OrderedDict from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Literal, Protocol from urllib.parse import urlparse import litellm from litellm._logging import verbose_logger +from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -70,16 +72,20 @@ class _VertexCredentialsObject(Protocol): def refresh(self, request: object) -> None: ... +_VertexCredentialCacheKey = tuple[str | None, str | None] + + class VertexBase: def __init__(self) -> None: super().__init__() self.access_token: str | None = None self.refresh_token: str | None = None self._credentials: GoogleCredentialsObject | None = None - self._credentials_project_mapping: dict[ - tuple[VERTEX_CREDENTIALS_TYPES | None, str | None], + self._credentials_project_mapping: OrderedDict[ + _VertexCredentialCacheKey, tuple[_VertexCredentialsObject, str | None], - ] = {} + ] = OrderedDict() + self._credentials_cache_lock = threading.RLock() self.project_id: str | None = None self.async_handler: AsyncHTTPHandler | None = None # Per-credential-key asyncio.Lock for single-flight async refresh. @@ -91,10 +97,10 @@ class VertexBase: # each lock; the entry is pruned when the count reaches zero, so the # dict stays bounded even in long-running high-cardinality deployments # without depending on any private asyncio internals. - self._async_refresh_locks: dict[tuple, asyncio.Lock] = {} - self._async_refresh_lock_refcounts: dict[tuple, int] = {} + self._async_refresh_locks: dict[_VertexCredentialCacheKey, asyncio.Lock] = {} + self._async_refresh_lock_refcounts: dict[_VertexCredentialCacheKey, int] = {} # Tracks in-flight background refresh tasks to avoid duplicate refreshes. - self._background_refresh_tasks: dict[tuple, asyncio.Task] = {} + self._background_refresh_tasks: dict[_VertexCredentialCacheKey, asyncio.Task] = {} # Protects the sync get_access_token refresh path. # Use RLock so that the reauthentication retry path (which calls # back into get_access_token while still holding the lock) can @@ -378,7 +384,7 @@ class VertexBase: with self._sync_refresh_lock: credentials.refresh(Request()) - def _acquire_async_refresh_lock(self, credential_cache_key: tuple) -> asyncio.Lock: + def _acquire_async_refresh_lock(self, credential_cache_key: _VertexCredentialCacheKey) -> asyncio.Lock: """Increment the refcount and return the lock for ``credential_cache_key``. Every call must be paired with ``_release_async_refresh_lock`` once the @@ -391,7 +397,7 @@ class VertexBase: ) return lock - def _release_async_refresh_lock(self, credential_cache_key: tuple, lock: asyncio.Lock) -> None: + def _release_async_refresh_lock(self, credential_cache_key: _VertexCredentialCacheKey, lock: asyncio.Lock) -> None: """Decrement the refcount and drop the lock entry when it reaches zero. Must be called only after the caller has released ``lock`` (i.e. once @@ -409,7 +415,7 @@ class VertexBase: def _try_get_cached_token( self, - credential_cache_key: tuple, + credential_cache_key: _VertexCredentialCacheKey, project_id: str | None, ) -> tuple[str, str] | None: """ @@ -432,7 +438,7 @@ class VertexBase: def _try_get_usable_cached_token( self, - credential_cache_key: tuple, + credential_cache_key: _VertexCredentialCacheKey, project_id: str | None, ) -> tuple[str, str, "TokenState", _VertexCredentialsObject, str | None] | None: """ @@ -458,19 +464,37 @@ class VertexBase: return creds.token, resolved_project, token_state, creds, cached_project_id def _unpack_cached_credentials( - self, credential_cache_key: tuple + self, credential_cache_key: _VertexCredentialCacheKey ) -> tuple[_VertexCredentialsObject | None, str | None]: """ Return (credentials, project_id) from the cache, or (None, None) if not cached. Handles both tuple and legacy cache formats. """ - if credential_cache_key not in self._credentials_project_mapping: - return None, None - cached_entry: Final = self._credentials_project_mapping[credential_cache_key] + with self._credentials_cache_lock: + cached_entry: Final = self._credentials_project_mapping.get(credential_cache_key) + if cached_entry is None: + return None, None + self._credentials_project_mapping.move_to_end(credential_cache_key) if isinstance(cached_entry, tuple): return cached_entry return cached_entry, cached_entry.quota_project_id or getattr(cached_entry, "project_id", None) + def _cache_credentials( + self, + credential_cache_key: _VertexCredentialCacheKey, + credentials: _VertexCredentialsObject, + project_id: str | None, + *, + replace: bool = True, + ) -> None: + with self._credentials_cache_lock: + if not replace and credential_cache_key in self._credentials_project_mapping: + return + self._credentials_project_mapping[credential_cache_key] = (credentials, project_id) + self._credentials_project_mapping.move_to_end(credential_cache_key) + while len(self._credentials_project_mapping) > max(0, DEFAULT_MAX_LRU_CACHE_SIZE): + self._credentials_project_mapping.popitem(last=False) + def _get_token_state(self, credentials: _VertexCredentialsObject) -> "TokenState": """ Return the token state using google-auth's TokenState enum. @@ -494,7 +518,7 @@ class VertexBase: self, credentials: VERTEX_CREDENTIALS_TYPES | None, project_id: str | None, - credential_cache_key: tuple, + credential_cache_key: _VertexCredentialCacheKey, ) -> tuple[_VertexCredentialsObject, str | None]: """Load credentials via load_auth (in thread) and cache the result.""" try: @@ -507,17 +531,12 @@ class VertexBase: raise if _credentials is None: raise ValueError("Could not resolve credentials") - self._credentials_project_mapping[credential_cache_key] = ( - _credentials, - credential_project_id, - ) + self._cache_credentials(credential_cache_key, _credentials, credential_project_id) return _credentials, credential_project_id async def _background_refresh_credentials( self, credentials: _VertexCredentialsObject, - credential_cache_key: tuple, - credential_project_id: str | None, ) -> None: """ Refresh credentials in the background without blocking the calling request. @@ -528,25 +547,13 @@ class VertexBase: try: verbose_logger.debug("Background proactive credential refresh") await asyncify(self.refresh_auth)(credentials) - # Only update the cache if it still points at the credentials - # object we just refreshed. The per-key async lock is not held - # here, so a concurrent INVALID path may have already replaced - # this entry (e.g. via _handle_reauthentication_async, which - # creates a fresh credentials object). In that case our write - # would clobber the newer entry with a stale reference. - cached_creds, _ = self._unpack_cached_credentials(credential_cache_key) - if cached_creds is credentials: - self._credentials_project_mapping[credential_cache_key] = ( - credentials, - credential_project_id, - ) except Exception: verbose_logger.debug( "Background credential refresh failed, will retry on next request", exc_info=True, ) - async def _await_in_flight_background_refresh(self, credential_cache_key: tuple) -> None: + async def _await_in_flight_background_refresh(self, credential_cache_key: _VertexCredentialCacheKey) -> None: """Wait for an in-flight background refresh to finish, if any. google-auth's ``Credentials.refresh()`` is not safe to invoke @@ -568,8 +575,7 @@ class VertexBase: def _schedule_background_refresh( self, credentials: _VertexCredentialsObject, - credential_cache_key: tuple, - credential_project_id: str | None, + credential_cache_key: _VertexCredentialCacheKey, ) -> None: """Kick off a single background refresh for ``credential_cache_key``. @@ -581,9 +587,7 @@ class VertexBase: if existing is not None and not existing.done(): return self._background_refresh_tasks.pop(credential_cache_key, None) - task: Final = asyncio.create_task( - self._background_refresh_credentials(credentials, credential_cache_key, credential_project_id) - ) + task: Final = asyncio.create_task(self._background_refresh_credentials(credentials)) def _drop_background_refresh_task(_fut: asyncio.Future[None]) -> None: if self._background_refresh_tasks.get(credential_cache_key) is _fut: @@ -796,7 +800,7 @@ class VertexBase: self, credentials: VERTEX_CREDENTIALS_TYPES | None, project_id: str | None, - credential_cache_key: tuple, + credential_cache_key: _VertexCredentialCacheKey, error: Exception, ) -> tuple[str, str]: """ @@ -822,8 +826,8 @@ class VertexBase: ) # Clear the cached credentials - if credential_cache_key in self._credentials_project_mapping: - del self._credentials_project_mapping[credential_cache_key] + with self._credentials_cache_lock: + self._credentials_project_mapping.pop(credential_cache_key, None) # Retry once with _retry_reauth=True to prevent infinite recursion try: @@ -846,7 +850,7 @@ class VertexBase: self, credentials: VERTEX_CREDENTIALS_TYPES | None, project_id: str | None, - credential_cache_key: tuple, + credential_cache_key: _VertexCredentialCacheKey, error: Exception, ) -> tuple[str, str]: """ @@ -856,7 +860,8 @@ class VertexBase: "Handling async reauthentication for project_id: %s. Clearing cache and retrying once.", project_id ) - self._credentials_project_mapping.pop(credential_cache_key, None) + with self._credentials_cache_lock: + self._credentials_project_mapping.pop(credential_cache_key, None) try: ( @@ -875,10 +880,7 @@ class VertexBase: # references the OLD credentials object we just replaced, and # leaving it would force the next request to do a redundant # refresh/reauth before realizing the cached creds are stale. - self._credentials_project_mapping[resolved_cache_key] = ( - _credentials, - credential_project_id, - ) + self._cache_credentials(resolved_cache_key, _credentials, credential_project_id) if _credentials.token is None or not isinstance(_credentials.token, str): raise ValueError( @@ -924,21 +926,11 @@ class VertexBase: # Convert dict credentials to string for caching cache_credentials: Final = json.dumps(credentials) if isinstance(credentials, dict) else credentials credential_cache_key: Final = (cache_credentials, project_id) - _credentials: _VertexCredentialsObject | None = None - verbose_logger.debug("Checking cached credentials for project_id: %s", project_id) - if credential_cache_key in self._credentials_project_mapping: + _credentials, credential_project_id = self._unpack_cached_credentials(credential_cache_key) + if _credentials is not None: verbose_logger.debug("Cached credentials found for project_id: %s.", project_id) - # Retrieve both credentials and cached project_id - cached_entry: Final = self._credentials_project_mapping[credential_cache_key] - verbose_logger.debug("cached_entry: %s", cached_entry) - if isinstance(cached_entry, tuple): - _credentials, credential_project_id = cached_entry - else: - # Backward compatibility with old cache format - _credentials = cached_entry - credential_project_id = _credentials.quota_project_id or getattr(_credentials, "project_id", None) verbose_logger.debug( "Using cached credentials for project_id: %s", credential_project_id, @@ -963,10 +955,7 @@ class VertexBase: f"Could not resolve credentials - either dynamically or from environment, for project_id: {project_id}" ) # Cache the project_id and credentials from load_auth result (resolved project_id) - self._credentials_project_mapping[credential_cache_key] = ( - _credentials, - credential_project_id, - ) + self._cache_credentials(credential_cache_key, _credentials, credential_project_id) ## VALIDATE CREDENTIALS verbose_logger.debug("Validating credentials") @@ -974,11 +963,7 @@ class VertexBase: project_id = credential_project_id # Update cache with resolved project_id for future lookups resolved_cache_key: Final = (cache_credentials, project_id) - if resolved_cache_key not in self._credentials_project_mapping: - self._credentials_project_mapping[resolved_cache_key] = ( - _credentials, - credential_project_id, - ) + self._cache_credentials(resolved_cache_key, _credentials, credential_project_id, replace=False) # Check if credentials are None before accessing attributes if _credentials is None: @@ -991,10 +976,6 @@ class VertexBase: try: verbose_logger.debug("Credentials expired, refreshing") self.refresh_auth(_credentials) - self._credentials_project_mapping[credential_cache_key] = ( - _credentials, - credential_project_id, - ) except Exception as e: # if refresh fails, it's possible the user has re-authenticated via `gcloud auth application-default login` # in this case, we should try to reload the credentials by clearing the cache and retrying @@ -1043,9 +1024,9 @@ class VertexBase: # callers on the lock just to schedule that refresh. usable: Final = self._try_get_usable_cached_token(credential_cache_key, project_id) if usable is not None: - cached_token, resolved_project, token_state, creds, cached_project_id = usable + cached_token, resolved_project, token_state, creds, _cached_project_id = usable if token_state == TokenState.STALE: - self._schedule_background_refresh(creds, credential_cache_key, cached_project_id) + self._schedule_background_refresh(creds, credential_cache_key) return cached_token, resolved_project # === SLOW PATH (per-key lock) === @@ -1075,10 +1056,7 @@ class VertexBase: # reauth that only repopulated the unresolved key), which # would force the next request through an unnecessary # refresh/reauth cycle. - self._credentials_project_mapping[resolved_cache_key] = ( - _credentials, - credential_project_id, - ) + self._cache_credentials(resolved_cache_key, _credentials, credential_project_id) # Use google-auth's token_state to decide refresh strategy: # - STALE: token is usable but within REFRESH_THRESHOLD (3:45) of @@ -1098,7 +1076,6 @@ class VertexBase: self._schedule_background_refresh( _credentials, credential_cache_key, - credential_project_id, ) return current_token, project_id @@ -1117,10 +1094,6 @@ class VertexBase: try: verbose_logger.debug("Credentials expired, refreshing") await asyncify(self.refresh_auth)(_credentials) - self._credentials_project_mapping[credential_cache_key] = ( - _credentials, - credential_project_id, - ) except Exception as e: if "Reauthentication is needed" in str(e): verbose_logger.debug("Reauthentication needed, clearing cache and retrying") diff --git a/tests/unit/llms/vertex_ai/test_vertex_llm_base.py b/tests/unit/llms/vertex_ai/test_vertex_llm_base.py index dd2b86b6f81..fa332cef5d8 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/unit/llms/vertex_ai/test_vertex_llm_base.py @@ -6,13 +6,17 @@ import os import pathlib import re import traceback -from typing import Final +from datetime import datetime, timedelta +from threading import Event +from typing import Final, Literal from unittest.mock import MagicMock, call, patch import google.auth.exceptions import pytest +import requests +import responses -from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE, DEFAULT_MAX_RECURSE_DEPTH import litellm @@ -2589,3 +2593,120 @@ def test_get_access_token_reraises_malformed_service_account_credentials_without assert [record.getMessage().split(". Error: ")[0] for record in caplog.records] == [ "Failed to load vertex credentials. Check to see if credentials containing partial/invalid information" ] + + +def _authorized_user_cache_credentials(identity: str) -> str: + return json.dumps( + { + "type": "authorized_user", + "client_id": identity, + "client_secret": "synthetic-secret", + "refresh_token": identity, + "quota_project_id": identity, + } + ) + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("inferred_project", [False, True], ids=["explicit", "inferred"]) +@pytest.mark.asyncio +async def test_credential_cache_bounds_retention_and_reloads_least_recently_used( + is_async: bool, inferred_project: bool +) -> None: + vertex: Final = VertexBase() + limit: Final = max(0, DEFAULT_MAX_LRU_CACHE_SIZE) + with ( + patch("google.auth._helpers.utcnow", return_value=datetime(2025, 1, 1)), + responses.RequestsMock() as oauth, + ): + def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]: + return 200, {}, json.dumps({"access_token": f"token-{len(oauth.calls) + 1}", "expires_in": 3600}) + + oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response) + + async def access(index: int, infer: bool = False) -> tuple[str, str]: + identity: Final = f"project-{index}" + credentials: Final = _authorized_user_cache_credentials(identity) + project: Final = None if infer else identity + result: Final = ( + await vertex.get_access_token_async(credentials, project) + if is_async + else vertex.get_access_token(credentials, project) + ) + assert result[1] == identity + assert len(vertex._credentials_project_mapping) <= limit + return result + + if limit == 0: + assert await access(0, inferred_project) == ("token-1", "project-0") + assert await access(0, inferred_project) == ("token-2", "project-0") + return + + for index in range(limit): + assert await access(index) == (f"token-{index + 1}", f"project-{index}") + assert await access(0) == ("token-1", "project-0") + assert len(oauth.calls) == limit + assert await access(limit, inferred_project) == (f"token-{limit + 1}", f"project-{limit}") + assert len(oauth.calls) == limit + 1 + if limit > (2 if inferred_project else 1): + assert await access(0) == ("token-1", "project-0") + assert len(oauth.calls) == limit + 1 + evicted: Final = 1 if limit > 1 else 0 + assert await access(evicted) == (f"token-{limit + 2}", f"project-{evicted}") + assert len(oauth.calls) == limit + 2 + + +@pytest.mark.parametrize("mode", ["sync", "async", "background"]) +@pytest.mark.skipif(DEFAULT_MAX_LRU_CACHE_SIZE < 2, reason="Eviction during refresh requires two retained entries") +@pytest.mark.asyncio +async def test_evicted_credentials_finish_refresh_without_being_readmitted( + mode: Literal["sync", "async", "background"], +) -> None: + vertex: Final = VertexBase() + limit: Final = DEFAULT_MAX_LRU_CACHE_SIZE + now: Final = datetime(2025, 1, 1) + refresh_started: Final = Event() + release_refresh: Final = Event() + donor: Final = _authorized_user_cache_credentials("donor") + target: Final = _authorized_user_cache_credentials("target") + with patch("google.auth._helpers.utcnow", return_value=now) as clock, responses.RequestsMock() as oauth: + def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]: + body: Final = request.body.decode() if isinstance(request.body, bytes) else request.body + assert isinstance(body, str) + is_target: Final = "refresh_token=target" in body + if is_target and len(oauth.calls) == limit: + refresh_started.set() + assert release_refresh.wait(5), "OAuth refresh was not released" + return 200, {}, json.dumps( + {"access_token": f"token-{len(oauth.calls) + 1}", "expires_in": 3600 if is_target else 86400} + ) + + oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response) + assert vertex.get_access_token(donor, None) == ("token-1", "donor") + assert vertex.get_access_token(donor, None) == ("token-1", "donor") + for index in range(limit - 2): + assert vertex.get_access_token(donor, f"filler-{index}") == (f"token-{index + 2}", f"filler-{index}") + assert vertex.get_access_token(target, "target") == (f"token-{limit}", "target") + clock.return_value = now + timedelta(seconds=3550 if mode == "background" else 7200) + refresh: Final = asyncio.create_task( + asyncio.to_thread(vertex.get_access_token, target, "target") + if mode == "sync" + else vertex.get_access_token_async(target, "target") + ) + try: + assert await asyncio.to_thread(refresh_started.wait, 5), "OAuth refresh did not start" + for index in range(limit - 2): + assert vertex.get_access_token(donor, f"filler-{index}") == (f"token-{index + 2}", f"filler-{index}") + assert vertex.get_access_token(donor, None) == ("token-1", "donor") + assert len(vertex._credentials_project_mapping) == limit + assert (target, "target") not in vertex._credentials_project_mapping + finally: + release_refresh.set() + completed: Final = await refresh + await asyncio.gather(*tuple(vertex._background_refresh_tasks.values())) + expected_token: Final = f"token-{limit if mode == 'background' else limit + 1}" + assert completed == (expected_token, "target") + assert len(oauth.calls) == limit + 1 + assert (target, "target") not in vertex._credentials_project_mapping + assert vertex.get_access_token(target, "target") == (f"token-{limit + 2}", "target") + assert len(vertex._credentials_project_mapping) == limit diff --git a/tests/unit/secret_managers/test_google_secret_manager.py b/tests/unit/secret_managers/test_google_secret_manager.py index 2051be08f7c..7effde4fb30 100644 --- a/tests/unit/secret_managers/test_google_secret_manager.py +++ b/tests/unit/secret_managers/test_google_secret_manager.py @@ -1,6 +1,7 @@ import base64 import binascii import json +from collections import OrderedDict from dataclasses import dataclass from typing import Final @@ -120,7 +121,7 @@ def manager(monkeypatch: pytest.MonkeyPatch) -> GoogleSecretManager: monkeypatch.setattr(litellm, "secret_manager_client", None) monkeypatch.setattr(litellm, "_key_management_system", None) monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings()) - monkeypatch.setattr(litellm.vertex_chat_completion, "_credentials_project_mapping", {}) + monkeypatch.setattr(litellm.vertex_chat_completion, "_credentials_project_mapping", OrderedDict()) monkeypatch.setattr(google.auth, "default", _default_credentials) monkeypatch.setenv("GOOGLE_SECRET_MANAGER_PROJECT_ID", PROJECT_ID) monkeypatch.delenv("GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER", raising=False)