From 9ba216d521cb992fab16209315f2e1eb69cf6077 Mon Sep 17 00:00:00 2001 From: quietjules <349048734+quietjules@users.noreply.github.com> Date: Sat, 14 Nov 2026 00:12:09 -0800 Subject: [PATCH] fix(vertex_ai): reuse one Google token across RAG Engine searches and cold bursts (#51239) * fix(vertex_ai): reuse one Google token across vector store searches * fix(vertex_ai): load Vertex credentials once when a cold worker gets a burst * test(vertex_ai): count Google token fetches for RAG Engine searches through the proxy * test(vertex_ai): type the fixtures and token callback in the concurrent search test * fix(vertex_ai): lock Google token loads per credential and reload a key Google rejects A shared RAG Engine config made one store's slow token fetch hold up every other store, and kept a rotated key file's deleted key until a restart. First loads and expired refreshes now take a lock per credentials and project, refreshes take a lock per credentials object, and any RefreshError reloads the credentials once. * fix(vertex_ai): share one Google token load or failure per credential and reload rejected keys on async calls * test(vertex_ai): check async callers sharing an expired credential refresh it once * fix(vertex_ai): reload Google credentials only when a refresh is rejected for good A retryable refresh error such as a 503 now raises without reloading the key, so an outage does not double google-auth's retries. Credentials that cannot be weak-referenced share the config's refresh lock again, as on main. The burst test for a rejected key moves to VertexBase and holds the token response until the other callers have started * fix(vertex_ai): reload Google credentials only when Google answers invalid_grant google-auth leaves its retryable flag False on metadata server, impersonation, and workload identity outages, and on a 502, so the previous check reloaded the key file on those outages and doubled the token tries. Now only an invalid_grant answer or "Reauthentication is needed" reloads The sync reload also refills the cache under the key the caller asked with, so a call without vertex_project keeps one cached copy, as the async reload already did * fix(vertex_ai): pick up a swapped Google key file on the next call The credential cache key now carries the key file's inode, mtime and size, for a vertex_credentials path and for GOOGLE_APPLICATION_CREDENTIALS. A key file that is replaced or rewritten, such as a rotated Kubernetes secret, is loaded on the next call instead of after Google rejects the old key or the proxy restarts * fix(vertex_ai): keep the earlier key while a replaced key file does not load A key file caught halfway through a non-atomic write, or replaced by one that does not load, now keeps the credentials loaded from its earlier version in use and logs a warning, as the cache did before key files were re-read. The next complete file is loaded on the next call. The pub/sub test no longer writes a key file it only reaches through the seeded cache * fix(vertex_ai): try a replaced key file again when its first token fetch fails The earlier key now stands in only when the replaced file does not load. A token fetch that fails on the new key raises like any first load, so the next call tries the new key again instead of keeping the old one * fix(vertex): raise when an async reload keeps an earlier key Google rejected --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/vertex_ai/credentials_source.py | 8 + .../vector_stores/rag_api/transformation.py | 3 + litellm/llms/vertex_ai/vertex_llm_base.py | 239 +++++++++++---- litellm/utils.py | 4 +- .../test_vector_store_search_wire.py | 63 ++++ .../integrations/gcs_pubsub/test_pub_sub.py | 17 +- .../llms/vertex_ai/test_vertex_llm_base.py | 284 +++++++++++++++++- tests/unit/vector_stores/test_main.py | 164 ++++++++++ 8 files changed, 707 insertions(+), 75 deletions(-) diff --git a/litellm/llms/vertex_ai/credentials_source.py b/litellm/llms/vertex_ai/credentials_source.py index ad51fd1e69d..4e481f1ca70 100644 --- a/litellm/llms/vertex_ai/credentials_source.py +++ b/litellm/llms/vertex_ai/credentials_source.py @@ -159,6 +159,14 @@ def load_vertex_credentials_source(credentials: str) -> VertexCredentialsSource: return _read_json_file(credentials) +def vertex_credentials_file_version(path: str) -> str | None: + try: + stat: Final = os.stat(path) + except (OSError, ValueError): + return None + return f"{stat.st_ino}:{stat.st_mtime_ns}:{stat.st_size}" + + def build_vertex_credentials( source: VertexCredentialsJson, build: Callable[[Mapping[str, object]], _Credentials] ) -> _Credentials: diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index 440604688bb..c929d2b1a00 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -361,3 +361,6 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): status_code=response.status_code, headers=response.headers, ) + + +vertex_rag_api_vector_store_config: Final = VertexVectorStoreConfig() diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 7dd98212f76..aaded19417e 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -6,10 +6,13 @@ Handles Authentication and generating request urls for Vertex AI and Google AI S import asyncio import json +import os import re import threading +import weakref from collections import OrderedDict -from collections.abc import Mapping +from collections.abc import Callable, Mapping +from concurrent.futures import Future from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final, Literal, Protocol from urllib.parse import urlparse @@ -35,6 +38,7 @@ from .credentials_source import ( build_vertex_credentials, load_vertex_credentials_source, raise_vertex_credentials_failure, + vertex_credentials_file_version, ) @@ -100,6 +104,32 @@ class _VertexCredentialsObject(Protocol): _VertexCredentialCacheKey = tuple[str | None, str | None] +_KEY_FILE_VERSION_SEPARATOR: Final = "\0" + + +def _credential_cache_identity(credentials: VERTEX_CREDENTIALS_TYPES | None) -> str | None: + if isinstance(credentials, dict): + return json.dumps(credentials) + key_file: Final = credentials if credentials is not None else os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") + version: Final = None if key_file is None else vertex_credentials_file_version(key_file) + if version is None: + return credentials + return f"{credentials or ''}{_KEY_FILE_VERSION_SEPARATOR}{version}" + + +def _is_other_version_of_key_file( + candidate: _VertexCredentialCacheKey, credential_cache_key: _VertexCredentialCacheKey +) -> bool: + candidate_identity, candidate_project_id = candidate + identity, project_id = credential_cache_key + if candidate_identity is None or identity is None or candidate_identity == identity: + return False + key_file, separator, _ = candidate_identity.partition(_KEY_FILE_VERSION_SEPARATOR) + return ( + bool(separator) + and candidate_project_id == project_id + and key_file == identity.partition(_KEY_FILE_VERSION_SEPARATOR)[0] + ) class VertexBase: @@ -128,11 +158,13 @@ class VertexBase: self._async_refresh_lock_refcounts: dict[_VertexCredentialCacheKey, int] = {} # Tracks in-flight background refresh tasks to avoid duplicate refreshes. 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 - # re-acquire it without deadlocking the current thread. - self._sync_refresh_lock = threading.RLock() + self._sync_token_flights: weakref.WeakValueDictionary[_VertexCredentialCacheKey, Future[tuple[str, str]]] = ( + weakref.WeakValueDictionary() + ) + self._credential_refresh_locks: weakref.WeakKeyDictionary[_VertexCredentialsObject, threading.RLock] = ( + weakref.WeakKeyDictionary() + ) + self._unkeyable_credentials_refresh_lock: Final = threading.RLock() @staticmethod def get_vertex_region(vertex_region: str | None, model: str) -> str: @@ -399,18 +431,64 @@ class VertexBase: except ImportError: raise ImportError(GOOGLE_IMPORT_ERROR_MESSAGE) - # Serialize all refreshes on this VertexBase across threads. # ``credentials.refresh()`` is not safe to call concurrently on the # same credentials object, and this method is invoked from three # places that can run on different threads: - # - sync ``get_access_token`` (already holds ``_sync_refresh_lock``) + # - sync ``get_access_token`` # - the async slow path (via ``asyncify`` in a worker thread) # - the background proactive refresh task (via ``asyncify``) - # ``_sync_refresh_lock`` is an ``RLock`` so reentrant acquisition - # from the sync path is safe. - with self._sync_refresh_lock: + with self._credential_refresh_lock(credentials): credentials.refresh(Request()) + def _credential_refresh_lock(self, credentials: _VertexCredentialsObject) -> threading.RLock: + with self._credentials_cache_lock: + try: + return self._credential_refresh_locks.setdefault(credentials, threading.RLock()) + except TypeError: + return self._unkeyable_credentials_refresh_lock + + def _refresh_if_still_needed( + self, + credentials: _VertexCredentialsObject, + needs_refresh: Callable[[_VertexCredentialsObject], object], + ) -> None: + with self._credential_refresh_lock(credentials): + if needs_refresh(credentials): + verbose_logger.debug("Credentials expired, refreshing") + self.refresh_auth(credentials) + + def _needs_blocking_refresh(self, credentials: _VertexCredentialsObject) -> bool: + from google.auth.credentials import TokenState + + return not isinstance(credentials.token, str) or self._get_token_state(credentials) == TokenState.INVALID + + def _join_sync_token_flight( + self, credential_cache_key: _VertexCredentialCacheKey + ) -> tuple[Future[tuple[str, str]], bool]: + with self._credentials_cache_lock: + in_flight: Final = self._sync_token_flights.get(credential_cache_key) + if in_flight is not None: + return in_flight, False + flight: Final[Future[tuple[str, str]]] = Future() + self._sync_token_flights[credential_cache_key] = flight + return flight, True + + @staticmethod + def _refresh_rejected(error: Exception) -> bool: + from google.auth.exceptions import RefreshError + + match error: + case RefreshError(args=(str(), {"error": "invalid_grant"})): + return True + case _: + return "Reauthentication is needed" in str(error) + + @staticmethod + def _token_fetch_failed(error: Exception) -> bool: + from google.auth.exceptions import RefreshError, TransportError + + return isinstance(error, (RefreshError, TransportError)) + def _acquire_async_refresh_lock(self, credential_cache_key: _VertexCredentialCacheKey) -> asyncio.Lock: """Increment the refcount and return the lock for ``credential_cache_key``. @@ -506,6 +584,51 @@ class VertexBase: return cached_entry return cached_entry, cached_entry.quota_project_id or getattr(cached_entry, "project_id", None) + def _load_auth_or_earlier_key_file( + self, + credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + credential_cache_key: _VertexCredentialCacheKey, + ) -> tuple[_VertexCredentialsObject | None, str | None]: + try: + return self.load_auth(credentials=credentials, project_id=project_id) + except Exception as e: + earlier: Final = ( + None + if self._token_fetch_failed(e) + else self._credentials_from_an_earlier_key_file(credential_cache_key, e) + ) + if earlier is None: + verbose_logger.exception( + "Failed to load vertex credentials. Check to see if credentials containing partial/invalid information. Error: %s", + e, + ) + raise + return earlier + + def _credentials_from_an_earlier_key_file( + self, credential_cache_key: _VertexCredentialCacheKey, error: Exception + ) -> tuple[_VertexCredentialsObject, str | None] | None: + with self._credentials_cache_lock: + earlier_key: Final = next( + ( + key + for key in reversed(self._credentials_project_mapping) + if _is_other_version_of_key_file(key, credential_cache_key) + ), + None, + ) + if earlier_key is None: + return None + earlier_credentials, earlier_project_id = self._unpack_cached_credentials(earlier_key) + if earlier_credentials is None: + return None + verbose_logger.warning( + "Vertex: the replaced credentials file did not load, so the credentials loaded before it stay in use: %s", + error, + ) + return earlier_credentials, earlier_project_id + def _cache_credentials( self, credential_cache_key: _VertexCredentialCacheKey, @@ -548,14 +671,9 @@ class VertexBase: credential_cache_key: _VertexCredentialCacheKey, ) -> tuple[_VertexCredentialsObject, str | None]: """Load credentials via load_auth (in thread) and cache the result.""" - try: - _credentials, credential_project_id = await asyncify(self.load_auth)( - credentials=credentials, - project_id=project_id, - ) - except Exception as e: - verbose_logger.exception("Failed to load vertex credentials: %s", str(e)) - raise + _credentials, credential_project_id = await asyncify(self._load_auth_or_earlier_key_file)( + credentials, project_id, credential_cache_key + ) if _credentials is None: raise ValueError("Could not resolve credentials") self._cache_credentials(credential_cache_key, _credentials, credential_project_id) @@ -842,7 +960,8 @@ class VertexBase: Handle reauthentication when credentials refresh fails. This method clears the cached credentials and attempts to reload them once. - It should only be called when "Reauthentication is needed" error occurs. + It should only be called when Google rejects the cached credentials' refresh, + e.g. "Reauthentication is needed" or a key file rotated to a new key. Args: credentials: The original credentials @@ -863,12 +982,13 @@ class VertexBase: # Clear the cached credentials with self._credentials_cache_lock: self._credentials_project_mapping.pop(credential_cache_key, None) + self._credentials_project_mapping.pop((credential_cache_key[0], project_id), None) # Retry once with _retry_reauth=True to prevent infinite recursion try: return self.get_access_token( credentials=credentials, - project_id=project_id, + project_id=credential_cache_key[1], _retry_reauth=True, ) except Exception as retry_error: @@ -897,6 +1017,7 @@ class VertexBase: with self._credentials_cache_lock: self._credentials_project_mapping.pop(credential_cache_key, None) + self._credentials_project_mapping.pop((credential_cache_key[0], project_id), None) try: ( @@ -909,14 +1030,14 @@ class VertexBase: ) if project_id is None and isinstance(credential_project_id, str): project_id = credential_project_id - cache_credentials: Final = json.dumps(credentials) if isinstance(credentials, dict) else credentials - resolved_cache_key: Final = (cache_credentials, project_id) + resolved_cache_key: Final = (credential_cache_key[0], project_id) # Always overwrite — any pre-existing entry at the resolved key # 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._cache_credentials(resolved_cache_key, _credentials, credential_project_id) + await asyncify(self._refresh_if_still_needed)(_credentials, self._needs_blocking_refresh) if _credentials.token is None or not isinstance(_credentials.token, str): raise ValueError( f"Could not resolve credentials token. Got None or non-string token (type={type(_credentials.token).__name__})" @@ -958,11 +1079,32 @@ class VertexBase: Tuple of (access_token, project_id) """ - # Convert dict credentials to string for caching - cache_credentials: Final = json.dumps(credentials) if isinstance(credentials, dict) else credentials + cache_credentials: Final = _credential_cache_identity(credentials) credential_cache_key: Final = (cache_credentials, project_id) verbose_logger.debug("Checking cached credentials for project_id: %s", project_id) + if _retry_reauth: + return self._resolve_access_token(credentials, project_id, credential_cache_key, _retry_reauth=True) + flight, leads = self._join_sync_token_flight(credential_cache_key) + if not leads: + return flight.result() + try: + flight.set_result(self._resolve_access_token(credentials, project_id, credential_cache_key)) + except BaseException as e: + flight.set_exception(e) + raise + finally: + with self._credentials_cache_lock: + self._sync_token_flights.pop(credential_cache_key, None) + return flight.result() + def _resolve_access_token( + self, + credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + credential_cache_key: _VertexCredentialCacheKey, + _retry_reauth: bool = False, + ) -> tuple[str, str]: + cache_credentials: Final = credential_cache_key[0] _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) @@ -976,14 +1118,9 @@ class VertexBase: "Credential cache key not found for project_id: %s, loading new credentials", project_id ) - try: - _credentials, credential_project_id = self.load_auth(credentials=credentials, project_id=project_id) - except Exception as e: - verbose_logger.exception( - "Failed to load vertex credentials. Check to see if credentials containing partial/invalid information. Error: %s", - e, - ) - raise e + _credentials, credential_project_id = self._load_auth_or_earlier_key_file( + credentials, project_id, credential_cache_key + ) if _credentials is None: raise ValueError( @@ -1005,23 +1142,19 @@ class VertexBase: raise ValueError("Credentials are None after loading") if _credentials.expired: - with self._sync_refresh_lock: - # Double-check after acquiring lock - if _credentials.expired: - try: - verbose_logger.debug("Credentials expired, refreshing") - self.refresh_auth(_credentials) - 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 - if "Reauthentication is needed" in str(e) and not _retry_reauth: - return self._handle_reauthentication( - credentials=credentials, - project_id=project_id, - credential_cache_key=credential_cache_key, - error=e, - ) - raise e + try: + self._refresh_if_still_needed(_credentials, lambda loaded: loaded.expired) + except Exception as e: + # if refresh fails, it's possible the user has re-authenticated via `gcloud auth application-default login` + # or rotated the key file, so we should try to reload the credentials by clearing the cache and retrying + if self._refresh_rejected(e) and not _retry_reauth: + return self._handle_reauthentication( + credentials=credentials, + project_id=project_id, + credential_cache_key=credential_cache_key, + error=e, + ) + raise e ## VALIDATION STEP if _credentials.token is None or not isinstance(_credentials.token, str): @@ -1048,7 +1181,7 @@ class VertexBase: """ from google.auth.credentials import TokenState - cache_credentials: Final = json.dumps(credentials) if isinstance(credentials, dict) else credentials + cache_credentials: Final = _credential_cache_identity(credentials) credential_cache_key: Final = (cache_credentials, project_id) # === FAST PATH (no lock) === @@ -1128,9 +1261,9 @@ class VertexBase: # Token is expired or missing — must block until refresh completes. try: verbose_logger.debug("Credentials expired, refreshing") - await asyncify(self.refresh_auth)(_credentials) + await asyncify(self._refresh_if_still_needed)(_credentials, self._needs_blocking_refresh) except Exception as e: - if "Reauthentication is needed" in str(e): + if self._refresh_rejected(e): verbose_logger.debug("Reauthentication needed, clearing cache and retrying") return await self._handle_reauthentication_async( credentials=credentials, diff --git a/litellm/utils.py b/litellm/utils.py index 494a98de98b..205dec2f9f9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -11031,10 +11031,10 @@ class ProviderConfigManager: def _get_vertex_ai_vector_stores_config(api_type: str | None) -> BaseVectorStoreConfig | None: if api_type == "rag_api" or api_type is None: from litellm.llms.vertex_ai.vector_stores.rag_api.transformation import ( - VertexVectorStoreConfig, + vertex_rag_api_vector_store_config, ) - return VertexVectorStoreConfig() + return vertex_rag_api_vector_store_config if api_type == "search_api": from litellm.llms.vertex_ai.vector_stores.search_api.transformation import ( vertex_search_api_vector_store_config, diff --git a/tests/integration/providers/test_vector_store_search_wire.py b/tests/integration/providers/test_vector_store_search_wire.py index accefe1bcf8..e20f906616b 100644 --- a/tests/integration/providers/test_vector_store_search_wire.py +++ b/tests/integration/providers/test_vector_store_search_wire.py @@ -4,7 +4,9 @@ from typing import Final from urllib.parse import urlsplit import httpx +import pytest from integration._support.client import Gateway +from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, wire_server from openai import OpenAI from pydantic import BaseModel, JsonValue, TypeAdapter @@ -30,6 +32,9 @@ _ALIAS_BODY: Final[dict[str, JsonValue]] = { "rewrite_query": None, } _ALIAS_REQUEST_BODY: Final[dict[str, JsonValue]] = {"query": "c"} +_VERTEX_PROJECT: Final = "scripted-search-project" +_VERTEX_TOKEN: Final = "scripted-search-token" +_VERTEX_SEARCH_COUNT: Final = 3 _SEARCH_RESULT: Final[dict[str, JsonValue]] = { "file_id": "file-search-wire", "filename": "search.txt", @@ -139,3 +144,61 @@ def test_vector_store_search_forwards_all_search_fields_on_both_routes(gateway: ("POST", f"/v1{search_path}"), ("POST", f"/v1{search_path}"), ] + + +@pytest.mark.parametrize( + ("api_type", "search_target", "search_reply"), + ( + ( + "search_api", + "/store:search", + {"results": [{"id": "doc-1", "document": {"id": "doc-1", "derivedStructData": {"title": "hit"}}}]}, + ), + ("rag_api", "/store:retrieveContexts", {"contexts": {"contexts": []}}), + ), +) +def test_vertex_vector_store_reuses_one_google_token_across_searches( + gateway: Gateway, api_type: str, search_target: str, search_reply: dict[str, JsonValue] +) -> None: + store_id: Final = f"vs_vertex_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" + if request.target == "/_oauth/token": + return Reply( + body=json.dumps({"access_token": _VERTEX_TOKEN, "token_type": "Bearer", "expires_in": 3600}).encode() + ) + assert request.target == search_target, request.target + assert request.headers["authorization"] == f"Bearer {_VERTEX_TOKEN}" + return Reply(body=json.dumps(search_reply).encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + registered: Final = gateway.request( + "POST", + "/vector_store/new", + { + "vector_store_id": store_id, + "custom_llm_provider": f"vertex_ai/{api_type}", + "litellm_params": { + "api_base": f"{wire.url}/store", + "vertex_project": _VERTEX_PROJECT, + "vertex_location": "global", + "vertex_credentials": service_account_json(_VERTEX_PROJECT, wire.url), + }, + }, + ) + assert registered.status_code == 200, registered.text + scenario.cleanups.callback(gateway.post, "/vector_store/delete", {"vector_store_id": store_id}) + + responses: Final = tuple( + gateway.request("POST", f"/v1/vector_stores/{store_id}/search", {"query": "hit"}) + for _ in range(_VERTEX_SEARCH_COUNT) + ) + assert [response.status_code for response in responses] == [200] * _VERTEX_SEARCH_COUNT, [ + response.text for response in responses + ] + + assert [request.target for request in wire.drain()] == [ + "/_oauth/token", + *(search_target,) * _VERTEX_SEARCH_COUNT, + ] diff --git a/tests/unit/integrations/gcs_pubsub/test_pub_sub.py b/tests/unit/integrations/gcs_pubsub/test_pub_sub.py index 77fc51a9026..81b8ed4c860 100644 --- a/tests/unit/integrations/gcs_pubsub/test_pub_sub.py +++ b/tests/unit/integrations/gcs_pubsub/test_pub_sub.py @@ -93,22 +93,7 @@ async def test_construct_request_headers_project_id_from_env(monkeypatch): def _authorized_user_credentials(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, project_id: str) -> str: - credentials_path: Final = tmp_path / "google-auth.json" - credentials_path.write_text( - json.dumps( - { - "type": "authorized_user", - "client_id": "unit-client", - "client_secret": "unit-secret", - "refresh_token": "unit-refresh-token", - "token": "unit-access-token", - "token_uri": "https://oauth2.googleapis.com/token", - "quota_project_id": "unit-project", - "expiry": "2099-01-01T00:00:00Z", - } - ) - ) - credentials_path_string: Final = str(credentials_path) + credentials_path_string: Final = str(tmp_path / "google-auth.json") credentials: Final = _CachedVertexCredentials(token="unit-access-token", quota_project_id=project_id) vertex_chat_completion: Final = litellm.vertex_chat_completion monkeypatch.setitem( 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 fa332cef5d8..89e2596d13c 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/unit/llms/vertex_ai/test_vertex_llm_base.py @@ -5,10 +5,14 @@ import logging import os import pathlib import re +import time import traceback +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta -from threading import Event -from typing import Final, Literal +from functools import partial +from threading import Event, Semaphore +from typing import Final, Literal, NoReturn from unittest.mock import MagicMock, call, patch import google.auth.exceptions @@ -2595,14 +2599,286 @@ def test_get_access_token_reraises_malformed_service_account_credentials_without ] -def _authorized_user_cache_credentials(identity: str) -> str: +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.asyncio +async def test_expired_credentials_shared_by_an_inferred_and_an_explicit_project_refresh_once(is_async: bool) -> None: + vertex: Final = VertexBase() + now: Final = datetime(2025, 1, 1) + credentials: Final = _authorized_user_cache_credentials("shared") + 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]: + if len(oauth.calls) == 1: + time.sleep(0.3) + 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) + assert vertex.get_access_token(credentials, None) == ("token-1", "shared") + clock.return_value = now + timedelta(seconds=7200) + fetch: Final = vertex.get_access_token_async if is_async else partial(asyncio.to_thread, vertex.get_access_token) + tokens: Final = tuple(await asyncio.gather(*(fetch(credentials, project) for project in (None, "shared")))) + + assert tokens == (("token-2", "shared"), ("token-2", "shared")) + assert len(oauth.calls) == 2 + + +def test_callers_that_join_a_rejected_google_token_fetch_share_its_error_and_the_next_call_retries() -> None: + vertex: Final = VertexBase() + credentials: Final = _authorized_user_cache_credentials("revoked") + leader_fetching: Final = Event() + followers_calling: Final = Semaphore(0) + + def follow() -> tuple[str, str]: + followers_calling.release() + return vertex.get_access_token(credentials, "unit-project") + + def revoked_token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]: + if not leader_fetching.is_set(): + leader_fetching.set() + assert all(followers_calling.acquire(timeout=5) for _ in range(4)) + time.sleep(0.2) + return 400, {}, json.dumps({"error": "invalid_grant", "error_description": "Token has been expired or revoked."}) + + with responses.RequestsMock() as oauth, ThreadPoolExecutor(max_workers=5) as pool: + oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=revoked_token_response) + leader: Final = pool.submit(vertex.get_access_token, credentials, "unit-project") + assert leader_fetching.wait(timeout=5) + followers: Final = tuple(pool.submit(follow) for _ in range(4)) + errors: Final = tuple(call.exception(timeout=10) for call in (leader, *followers)) + + assert len(oauth.calls) == 1 + assert all(isinstance(error, google.auth.exceptions.RefreshError) for error in errors) + with pytest.raises(google.auth.exceptions.RefreshError, match="invalid_grant"): + vertex.get_access_token(credentials, "unit-project") + assert len(oauth.calls) == 2 + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("project", ["unit-project", None], ids=["explicit", "inferred"]) +@pytest.mark.asyncio +async def test_callers_reload_a_rotated_key_file_once_when_google_rejects_the_old_key( + tmp_path: pathlib.Path, is_async: bool, project: str | None +) -> None: + vertex: Final = VertexBase() + now: Final = datetime(2025, 1, 1) + key_file: Final = tmp_path / "vertex-key.json" + key_file.write_text(_authorized_user_cache_credentials("old-key", quota_project="unit-project")) + old_key_revoked: Final = Event() + + def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]: + form: Final = request.body.decode() if isinstance(request.body, bytes) else str(request.body) + is_old_key: Final = "old-key" in form + if is_old_key and old_key_revoked.is_set(): + return 400, {}, json.dumps({"error": "invalid_grant", "error_description": "Token has been expired or revoked."}) + return 200, {}, json.dumps({"access_token": "old-token" if is_old_key else "new-token", "expires_in": 3600}) + + fetch: Final = vertex.get_access_token_async if is_async else partial(asyncio.to_thread, vertex.get_access_token) + with patch("google.auth._helpers.utcnow", return_value=now) as clock, responses.RequestsMock() as oauth: + oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response) + assert await fetch(str(key_file), project) == ("old-token", "unit-project") + _rewrite_keeping_its_version( + key_file, _authorized_user_cache_credentials("new-key", quota_project="unit-project") + ) + old_key_revoked.set() + clock.return_value = now + timedelta(seconds=7200) + assert await fetch(str(key_file), project) == ("new-token", "unit-project") + assert await fetch(str(key_file), project) == ("new-token", "unit-project") + assert len(oauth.calls) == 3 + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + ("outage_status", "outage_body"), + [ + (503, json.dumps({"error": "server_error", "error_description": "Service unavailable."})), + (502, "Bad Gateway"), + ], + ids=["503", "502"], +) +@pytest.mark.asyncio +async def test_a_google_outage_during_refresh_keeps_the_cached_key_instead_of_reloading_it( + tmp_path: pathlib.Path, is_async: bool, outage_status: int, outage_body: str +) -> None: + vertex: Final = VertexBase() + now: Final = datetime(2025, 1, 1) + key_file: Final = tmp_path / "vertex-key.json" + key_file.write_text(_authorized_user_cache_credentials("old-key")) + google_down: Final = Event() + + def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]: + if google_down.is_set(): + return outage_status, {}, outage_body + return 200, {}, json.dumps({"access_token": "old-token", "expires_in": 3600}) + + fetch: Final = vertex.get_access_token_async if is_async else partial(asyncio.to_thread, vertex.get_access_token) + with ( + patch("google.auth._helpers.utcnow", return_value=now) as clock, + patch("google.auth._exponential_backoff.time"), + responses.RequestsMock() as oauth, + ): + oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response) + assert await fetch(str(key_file), "unit-project") == ("old-token", "unit-project") + _rewrite_keeping_its_version(key_file, _authorized_user_cache_credentials("new-key")) + google_down.set() + clock.return_value = now + timedelta(seconds=7200) + with pytest.raises(google.auth.exceptions.RefreshError): + await fetch(str(key_file), "unit-project") + assert not any("new-key" in str(call.request.body) for call in oauth.calls) + + +def _rewrite_keeping_its_version(key_file: pathlib.Path, contents: str) -> None: + before: Final = key_file.stat() + key_file.write_text(contents) + os.utime(key_file, ns=(before.st_atime_ns, before.st_mtime_ns)) + + +def _move_into_place(key_file: pathlib.Path, contents: str) -> None: + before: Final = key_file.stat() + staged: Final = key_file.with_name("staged-key.json") + staged.write_text(contents) + os.utime(staged, ns=(before.st_atime_ns, before.st_mtime_ns)) + os.replace(staged, key_file) + + +def _rewrite_later(key_file: pathlib.Path, contents: str) -> None: + before: Final = key_file.stat() + key_file.write_text(contents) + os.utime(key_file, ns=(before.st_atime_ns, before.st_mtime_ns + 1_000_000_000)) + + +def _rewrite_longer(key_file: pathlib.Path, contents: str) -> None: + _rewrite_keeping_its_version(key_file, contents.replace("new-key", "new-key-v2")) + + +def _old_or_new_token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]: + form: Final = request.body.decode() if isinstance(request.body, bytes) else str(request.body) + return 200, {}, json.dumps({"access_token": "old-token" if "old-key" in form else "new-token", "expires_in": 3600}) + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("from_adc_env", [False, True], ids=["vertex_credentials", "GOOGLE_APPLICATION_CREDENTIALS"]) +@pytest.mark.parametrize( + "swap_key_file", + [_move_into_place, _rewrite_later, _rewrite_longer], + ids=["moved_into_place", "rewritten_later", "rewritten_longer"], +) +@pytest.mark.asyncio +async def test_a_swapped_key_file_is_used_on_the_next_call_while_the_old_key_still_works( + tmp_path: pathlib.Path, + monkeypatch: pytest.MonkeyPatch, + is_async: bool, + from_adc_env: bool, + swap_key_file: Callable[[pathlib.Path, str], None], +) -> None: + vertex: Final = VertexBase() + key_file: Final = tmp_path / "vertex-key.json" + key_file.write_text(_authorized_user_cache_credentials("old-key", quota_project="unit-project")) + if from_adc_env: + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", str(key_file)) + credentials: Final = None if from_adc_env else str(key_file) + + fetch: Final = vertex.get_access_token_async if is_async else partial(asyncio.to_thread, vertex.get_access_token) + with responses.RequestsMock() as oauth: + oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=_old_or_new_token_response) + assert await fetch(credentials, "unit-project") == ("old-token", "unit-project") + swap_key_file(key_file, _authorized_user_cache_credentials("new-key", quota_project="unit-project")) + assert await fetch(credentials, "unit-project") == ("new-token", "unit-project") + assert await fetch(credentials, "unit-project") == ("new-token", "unit-project") + assert len(oauth.calls) == 2 + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.asyncio +async def test_a_key_file_replaced_by_one_that_does_not_load_keeps_the_earlier_key_until_it_is_fixed( + tmp_path: pathlib.Path, is_async: bool +) -> None: + vertex: Final = VertexBase() + key_file: Final = tmp_path / "vertex-key.json" + key_file.write_text(_authorized_user_cache_credentials("old-key", quota_project="unit-project")) + + fetch: Final = vertex.get_access_token_async if is_async else partial(asyncio.to_thread, vertex.get_access_token) + with responses.RequestsMock() as oauth: + oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=_old_or_new_token_response) + assert await fetch(str(key_file), "unit-project") == ("old-token", "unit-project") + _move_into_place(key_file, '{"type": "authorized_user", "client_id": "half-wri') + assert await fetch(str(key_file), "unit-project") == ("old-token", "unit-project") + _move_into_place(key_file, _authorized_user_cache_credentials("repaired-key", quota_project="unit-project")) + assert await fetch(str(key_file), "unit-project") == ("new-token", "unit-project") + assert len(oauth.calls) == 2 + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.asyncio +async def test_a_key_file_that_does_not_load_raises_once_google_rejects_the_earlier_key( + tmp_path: pathlib.Path, is_async: bool +) -> None: + vertex: Final = VertexBase() + now: Final = datetime(2025, 1, 1) + key_file: Final = tmp_path / "vertex-key.json" + key_file.write_text(_authorized_user_cache_credentials("old-key", quota_project="unit-project")) + old_key_revoked: Final = Event() + + def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]: + if old_key_revoked.is_set(): + return 400, {}, json.dumps({"error": "invalid_grant", "error_description": "Token has been expired or revoked."}) + return _old_or_new_token_response(request) + + fetch: Final = vertex.get_access_token_async if is_async else partial(asyncio.to_thread, vertex.get_access_token) + with patch("google.auth._helpers.utcnow", return_value=now) as clock, responses.RequestsMock() as oauth: + oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response) + assert await fetch(str(key_file), "unit-project") == ("old-token", "unit-project") + _move_into_place(key_file, '{"type": "authorized_user", "client_id": "half-wri') + old_key_revoked.set() + clock.return_value = now + timedelta(seconds=7200) + with pytest.raises(google.auth.exceptions.RefreshError, match="invalid_grant"): + await fetch(str(key_file), "unit-project") + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + "google_unavailable", + [ + lambda: (503, {}, json.dumps({"error": "server_error", "error_description": "Service unavailable."})), + lambda: _raise(requests.exceptions.ConnectionError("oauth2.googleapis.com is unreachable")), + ], + ids=["503", "unreachable"], +) +@pytest.mark.asyncio +async def test_a_swapped_key_file_whose_first_token_fetch_fails_is_tried_again_on_the_next_call( + tmp_path: pathlib.Path, is_async: bool, google_unavailable: Callable[[], tuple[int, dict[str, str], str]] +) -> None: + vertex: Final = VertexBase() + key_file: Final = tmp_path / "vertex-key.json" + key_file.write_text(_authorized_user_cache_credentials("old-key", quota_project="unit-project")) + google_down: Final = Event() + + def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]: + return google_unavailable() if google_down.is_set() else _old_or_new_token_response(request) + + fetch: Final = vertex.get_access_token_async if is_async else partial(asyncio.to_thread, vertex.get_access_token) + with patch("google.auth._exponential_backoff.time"), responses.RequestsMock() as oauth: + oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response) + assert await fetch(str(key_file), "unit-project") == ("old-token", "unit-project") + _move_into_place(key_file, _authorized_user_cache_credentials("new-key", quota_project="unit-project")) + google_down.set() + with pytest.raises((google.auth.exceptions.RefreshError, google.auth.exceptions.TransportError)): + await fetch(str(key_file), "unit-project") + google_down.clear() + assert await fetch(str(key_file), "unit-project") == ("new-token", "unit-project") + assert await fetch(str(key_file), "unit-project") == ("new-token", "unit-project") + + +def _raise(error: Exception) -> NoReturn: + raise error + + +def _authorized_user_cache_credentials(identity: str, quota_project: str | None = None) -> str: return json.dumps( { "type": "authorized_user", "client_id": identity, "client_secret": "synthetic-secret", "refresh_token": identity, - "quota_project_id": identity, + "quota_project_id": quota_project or identity, } ) diff --git a/tests/unit/vector_stores/test_main.py b/tests/unit/vector_stores/test_main.py index b89806e6377..05398c49dd6 100644 --- a/tests/unit/vector_stores/test_main.py +++ b/tests/unit/vector_stores/test_main.py @@ -7,13 +7,24 @@ executor, and it must never leak into litellm_params/kwargs where logging would model_dump() it (the #19550 serialization trap). """ +from __future__ import annotations + +import asyncio import json +import pathlib +import threading +import time +import uuid from collections.abc import Mapping from typing import Final from unittest.mock import MagicMock, patch import httpx import pytest +import requests +import requests_mock as requests_mock_lib +import responses +import respx import litellm import litellm.vector_stores.main as vector_stores_main @@ -384,3 +395,156 @@ async def test_retrieving_a_missing_openai_vector_store_raises_not_found(is_asyn assert caught.value.status_code == 404 assert "No vector store found with id 'vs_missing'." in caught.value.message + + +GOOGLE_TOKEN_URL: Final = "https://oauth2.googleapis.com/token" + + +def _authorized_user_credentials(refresh_token: str) -> str: + return json.dumps( + { + "type": "authorized_user", + "client_id": "unit-client-id", + "client_secret": "unit-client-secret", + "refresh_token": refresh_token, + } + ) + + +def _search_rag_store(api_base: str, vertex_credentials: str): + return litellm.vector_stores.asearch( + query="what is LiteLLM?", + vector_store_id="unit-store", + custom_llm_provider="vertex_ai/rag_api", + api_base=api_base, + vertex_project="unit-project", + vertex_location="us-central1", + vertex_credentials=vertex_credentials, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("api_type", "search_method", "search_response"), + ( + ("search_api", ":search", {"results": []}), + ("rag_api", ":retrieveContexts", {"contexts": {"contexts": []}}), + ), +) +async def test_concurrent_vertex_searches_with_the_same_credentials_fetch_one_google_token( + api_type: str, + search_method: str, + search_response: dict[str, object], + respx_mock: respx.MockRouter, + requests_mock: requests_mock_lib.Mocker, + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + credentials: Final = _authorized_user_credentials(f"unit-refresh-{uuid.uuid4().hex}") + + def slow_token_response( + request: requests_mock_lib.Request, context: requests_mock_lib.Context + ) -> dict[str, object]: + time.sleep(0.2) + return {"access_token": "ya29.unit-token", "expires_in": 3600, "token_type": "Bearer"} + + token_route: Final = requests_mock.post(GOOGLE_TOKEN_URL, json=slow_token_response) + search_route: Final = respx_mock.post(f"https://vertex.test/{api_type}{search_method}").mock( + return_value=httpx.Response(200, json=search_response) + ) + + await asyncio.gather( + *( + litellm.vector_stores.asearch( + query="what is LiteLLM?", + vector_store_id="unit-store", + custom_llm_provider=f"vertex_ai/{api_type}", + api_base=f"https://vertex.test/{api_type}", + vertex_project="unit-project", + vertex_location="us-central1", + vertex_credentials=credentials, + ) + for _ in range(5) + ) + ) + + assert search_route.call_count == 5 + assert token_route.call_count == 1 + assert search_route.calls.last.request.headers["authorization"] == "Bearer ya29.unit-token" + + +@pytest.mark.asyncio +async def test_a_slow_google_token_fetch_for_one_rag_store_does_not_hold_up_another( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + slow_fetch_started: Final = threading.Event() + other_store_searched: Final = threading.Event() + + def access_token_for(form: str) -> str: + if "slow-store" not in form: + return "ya29.fast" + slow_fetch_started.set() + return f"ya29.slow-waited-{other_store_searched.wait(timeout=5)}" + + def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]: + form: Final = request.body.decode() if isinstance(request.body, bytes) else str(request.body) + token: Final = {"access_token": access_token_for(form), "expires_in": 3600, "token_type": "Bearer"} + return 200, {}, json.dumps(token) + + slow_route: Final = respx_mock.post("https://vertex.test/slow:retrieveContexts").mock( + return_value=httpx.Response(200, json={"contexts": {"contexts": []}}) + ) + respx_mock.post("https://vertex.test/fast:retrieveContexts").mock( + return_value=httpx.Response(200, json={"contexts": {"contexts": []}}) + ) + + with responses.RequestsMock() as google: + google.add_callback(responses.POST, GOOGLE_TOKEN_URL, callback=token_response) + slow_search: Final = asyncio.ensure_future( + _search_rag_store("https://vertex.test/slow", _authorized_user_credentials(f"slow-store-{uuid.uuid4().hex}")) + ) + assert await asyncio.to_thread(slow_fetch_started.wait, 5) + await _search_rag_store("https://vertex.test/fast", _authorized_user_credentials(f"fast-store-{uuid.uuid4().hex}")) + other_store_searched.set() + await slow_search + + assert slow_route.calls.last.request.headers["authorization"] == "Bearer ya29.slow-waited-True" + + +@pytest.mark.asyncio +async def test_a_rag_store_reloads_a_rotated_key_file_when_google_rejects_the_old_key( + tmp_path: pathlib.Path, + respx_mock: respx.MockRouter, + requests_mock: requests_mock_lib.Mocker, + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + key_file: Final = tmp_path / "vertex-key.json" + key_file.write_text(_authorized_user_credentials(f"old-key-{uuid.uuid4().hex}")) + old_key_deleted: Final = threading.Event() + + def token_response( + request: requests_mock_lib.Request, context: requests_mock_lib.Context + ) -> dict[str, object]: + is_old_key: Final = "old-key" in request.text + if is_old_key and old_key_deleted.is_set(): + context.status_code = 400 + return {"error": "invalid_grant", "error_description": "Token has been expired or revoked."} + return {"access_token": "ya29.old" if is_old_key else "ya29.new", "expires_in": 0, "token_type": "Bearer"} + + requests_mock.post(GOOGLE_TOKEN_URL, json=token_response) + search_route: Final = respx_mock.post("https://vertex.test/rotated:retrieveContexts").mock( + return_value=httpx.Response(200, json={"contexts": {"contexts": []}}) + ) + + await _search_rag_store("https://vertex.test/rotated", str(key_file)) + key_file.write_text(_authorized_user_credentials(f"new-key-{uuid.uuid4().hex}")) + old_key_deleted.set() + await _search_rag_store("https://vertex.test/rotated", str(key_file)) + + assert [call.request.headers["authorization"] for call in search_route.calls] == [ + "Bearer ya29.old", + "Bearer ya29.new", + ]