mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix: reuse cached credentials in VertexAIPartnerModels (#26065)
* fix: reuse cached credentials in VertexAIPartnerModels instead of creating new VertexLLM per request VertexAIPartnerModels.completion() was creating a throwaway VertexLLM() instance on every call to get an access token, bypassing the credential cache inherited from VertexBase. This caused a fresh token fetch for every single request, adding significant latency overhead. Fix: call super().__init__() to initialize VertexBase's credential cache, and use self._ensure_access_token() instead of a new VertexLLM instance. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: apply same credential caching fix to VertexAIGemmaModels and VertexAIModelGardenModels Same bug as VertexAIPartnerModels: both classes had `pass` in __init__ instead of `super().__init__()`, and created throwaway VertexLLM() instances per request instead of using self._ensure_access_token(). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Matthew Lapointe
GitHub
parent
4583310313
commit
67e6a95cb0
4 changed files with 226 additions and 19 deletions
|
|
@ -45,7 +45,7 @@ class PartnerModelPrefixes(str, Enum):
|
|||
|
||||
class VertexAIPartnerModels(VertexBase):
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
super().__init__()
|
||||
|
||||
@staticmethod
|
||||
def is_vertex_partner_model(model: str):
|
||||
|
|
@ -116,9 +116,6 @@ class VertexAIPartnerModels(VertexBase):
|
|||
CodestralTextCompletion,
|
||||
)
|
||||
from litellm.llms.openai_like.chat.handler import OpenAILikeChatHandler
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexLLM,
|
||||
)
|
||||
except Exception as e:
|
||||
raise VertexAIError(
|
||||
status_code=400,
|
||||
|
|
@ -133,9 +130,7 @@ class VertexAIPartnerModels(VertexBase):
|
|||
message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""",
|
||||
)
|
||||
try:
|
||||
vertex_httpx_logic = VertexLLM()
|
||||
|
||||
access_token, project_id = vertex_httpx_logic._ensure_access_token(
|
||||
access_token, project_id = self._ensure_access_token(
|
||||
credentials=vertex_credentials,
|
||||
project_id=vertex_project,
|
||||
custom_llm_provider="vertex_ai",
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ from ..vertex_llm_base import VertexBase
|
|||
|
||||
class VertexAIGemmaModels(VertexBase):
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
super().__init__()
|
||||
|
||||
def completion(
|
||||
self,
|
||||
|
|
@ -62,9 +62,6 @@ class VertexAIGemmaModels(VertexBase):
|
|||
try:
|
||||
import vertexai
|
||||
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexLLM,
|
||||
)
|
||||
from litellm.llms.vertex_ai.vertex_gemma_models.transformation import (
|
||||
VertexGemmaConfig,
|
||||
)
|
||||
|
|
@ -83,9 +80,8 @@ class VertexAIGemmaModels(VertexBase):
|
|||
)
|
||||
try:
|
||||
model = get_vertex_base_model_name(model=model)
|
||||
vertex_httpx_logic = VertexLLM()
|
||||
|
||||
access_token, project_id = vertex_httpx_logic._ensure_access_token(
|
||||
access_token, project_id = self._ensure_access_token(
|
||||
credentials=vertex_credentials,
|
||||
project_id=vertex_project,
|
||||
custom_llm_provider="vertex_ai",
|
||||
|
|
|
|||
|
|
@ -41,7 +41,7 @@ def create_vertex_url(
|
|||
|
||||
class VertexAIModelGardenModels(VertexBase):
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
super().__init__()
|
||||
|
||||
def completion(
|
||||
self,
|
||||
|
|
@ -73,9 +73,6 @@ class VertexAIModelGardenModels(VertexBase):
|
|||
import vertexai
|
||||
|
||||
from litellm.llms.openai_like.chat.handler import OpenAILikeChatHandler
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexLLM,
|
||||
)
|
||||
except Exception as e:
|
||||
raise VertexAIError(
|
||||
status_code=400,
|
||||
|
|
@ -91,9 +88,8 @@ class VertexAIModelGardenModels(VertexBase):
|
|||
)
|
||||
try:
|
||||
model = get_vertex_base_model_name(model=model)
|
||||
vertex_httpx_logic = VertexLLM()
|
||||
|
||||
access_token, project_id = vertex_httpx_logic._ensure_access_token(
|
||||
access_token, project_id = self._ensure_access_token(
|
||||
credentials=vertex_credentials,
|
||||
project_id=vertex_project,
|
||||
custom_llm_provider="vertex_ai",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,220 @@
|
|||
"""
|
||||
Test that VertexBase subclasses (PartnerModels, Gemma, ModelGarden) reuse
|
||||
cached credentials instead of creating a new VertexLLM instance on every request.
|
||||
"""
|
||||
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
|
||||
VertexAIPartnerModels,
|
||||
)
|
||||
from litellm.llms.vertex_ai.vertex_gemma_models.main import VertexAIGemmaModels
|
||||
from litellm.llms.vertex_ai.vertex_model_garden.main import VertexAIModelGardenModels
|
||||
|
||||
|
||||
def _mock_vertexai():
|
||||
"""Return a MagicMock that satisfies the vertexai import guards."""
|
||||
m = MagicMock()
|
||||
m.preview = MagicMock()
|
||||
m.preview.language_models = MagicMock()
|
||||
return m
|
||||
|
||||
|
||||
class TestVertexBaseSubclassInit:
|
||||
"""All VertexBase subclasses must call super().__init__() so that
|
||||
the credential cache is initialized."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cls",
|
||||
[VertexAIPartnerModels, VertexAIGemmaModels, VertexAIModelGardenModels],
|
||||
ids=["PartnerModels", "Gemma", "ModelGarden"],
|
||||
)
|
||||
def test_init_calls_super(self, cls):
|
||||
instance = cls()
|
||||
assert hasattr(instance, "_credentials_project_mapping")
|
||||
assert isinstance(instance._credentials_project_mapping, dict)
|
||||
assert hasattr(instance, "access_token")
|
||||
assert hasattr(instance, "project_id")
|
||||
|
||||
|
||||
class TestPartnerModelsCredentialReuse:
|
||||
def test_completion_uses_self_ensure_access_token(self):
|
||||
"""completion() should call self._ensure_access_token, not create a
|
||||
throwaway VertexLLM instance."""
|
||||
partner = VertexAIPartnerModels()
|
||||
|
||||
with (
|
||||
patch.dict(sys.modules, {"vertexai": _mock_vertexai()}),
|
||||
patch.object(
|
||||
partner,
|
||||
"_ensure_access_token",
|
||||
return_value=("cached-token", "test-project"),
|
||||
) as mock_ensure,
|
||||
patch(
|
||||
"litellm.llms.vertex_ai.vertex_ai_partner_models.main.base_llm_http_handler"
|
||||
) as mock_handler,
|
||||
):
|
||||
mock_handler.completion.return_value = "response"
|
||||
|
||||
partner.completion(
|
||||
model="meta/llama-3.1-405b-instruct-maas",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
model_response=MagicMock(),
|
||||
print_verbose=lambda *a, **kw: None,
|
||||
encoding=MagicMock(),
|
||||
logging_obj=MagicMock(),
|
||||
api_base=None,
|
||||
optional_params={},
|
||||
custom_prompt_dict={},
|
||||
headers=None,
|
||||
timeout=30.0,
|
||||
litellm_params={},
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
vertex_credentials='{"type": "service_account"}',
|
||||
)
|
||||
|
||||
mock_ensure.assert_called_once_with(
|
||||
credentials='{"type": "service_account"}',
|
||||
project_id="test-project",
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
|
||||
def test_credential_cache_shared_across_calls(self):
|
||||
"""Two successive completion() calls should hit load_auth only once."""
|
||||
partner = VertexAIPartnerModels()
|
||||
|
||||
mock_creds = MagicMock()
|
||||
mock_creds.token = "my-token"
|
||||
mock_creds.expired = False
|
||||
mock_creds.project_id = "proj"
|
||||
mock_creds.quota_project_id = "proj"
|
||||
|
||||
with (
|
||||
patch.dict(sys.modules, {"vertexai": _mock_vertexai()}),
|
||||
patch.object(
|
||||
partner, "load_auth", return_value=(mock_creds, "proj")
|
||||
) as mock_load,
|
||||
patch(
|
||||
"litellm.llms.vertex_ai.vertex_ai_partner_models.main.base_llm_http_handler"
|
||||
) as mock_handler,
|
||||
):
|
||||
mock_handler.completion.return_value = "resp"
|
||||
|
||||
common_kwargs = dict(
|
||||
model="meta/llama-3.1-405b-instruct-maas",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
model_response=MagicMock(),
|
||||
print_verbose=lambda *a, **kw: None,
|
||||
encoding=MagicMock(),
|
||||
logging_obj=MagicMock(),
|
||||
api_base=None,
|
||||
optional_params={},
|
||||
custom_prompt_dict={},
|
||||
headers=None,
|
||||
timeout=30.0,
|
||||
litellm_params={},
|
||||
vertex_project="proj",
|
||||
vertex_location="us-central1",
|
||||
vertex_credentials='{"type": "service_account"}',
|
||||
)
|
||||
|
||||
partner.completion(**common_kwargs)
|
||||
partner.completion(**common_kwargs)
|
||||
|
||||
assert mock_load.call_count == 1
|
||||
|
||||
|
||||
class TestGemmaModelsCredentialReuse:
|
||||
def test_completion_uses_self_ensure_access_token(self):
|
||||
"""completion() should call self._ensure_access_token, not create a
|
||||
throwaway VertexLLM instance."""
|
||||
gemma = VertexAIGemmaModels()
|
||||
|
||||
mock_gemma_config = MagicMock()
|
||||
mock_gemma_config.return_value.completion.return_value = "response"
|
||||
|
||||
with (
|
||||
patch.dict(sys.modules, {"vertexai": _mock_vertexai()}),
|
||||
patch.object(
|
||||
gemma,
|
||||
"_ensure_access_token",
|
||||
return_value=("cached-token", "test-project"),
|
||||
) as mock_ensure,
|
||||
patch(
|
||||
"litellm.llms.vertex_ai.vertex_gemma_models.transformation.VertexGemmaConfig",
|
||||
mock_gemma_config,
|
||||
),
|
||||
):
|
||||
gemma.completion(
|
||||
model="gemma/gemma-3-12b-it-1234567890",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
model_response=MagicMock(),
|
||||
print_verbose=lambda *a, **kw: None,
|
||||
encoding=MagicMock(),
|
||||
logging_obj=MagicMock(),
|
||||
api_base="https://123.us-central1-1.prediction.vertexai.goog/v1/projects/proj/locations/us-central1/endpoints/456:predict",
|
||||
optional_params={},
|
||||
custom_prompt_dict={},
|
||||
headers=None,
|
||||
timeout=30.0,
|
||||
litellm_params={},
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
vertex_credentials='{"type": "service_account"}',
|
||||
)
|
||||
|
||||
mock_ensure.assert_called_once_with(
|
||||
credentials='{"type": "service_account"}',
|
||||
project_id="test-project",
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
|
||||
|
||||
class TestModelGardenCredentialReuse:
|
||||
def test_completion_uses_self_ensure_access_token(self):
|
||||
"""completion() should call self._ensure_access_token, not create a
|
||||
throwaway VertexLLM instance."""
|
||||
garden = VertexAIModelGardenModels()
|
||||
|
||||
mock_handler = MagicMock()
|
||||
mock_handler.return_value.completion.return_value = "response"
|
||||
|
||||
with (
|
||||
patch.dict(sys.modules, {"vertexai": _mock_vertexai()}),
|
||||
patch.object(
|
||||
garden,
|
||||
"_ensure_access_token",
|
||||
return_value=("cached-token", "test-project"),
|
||||
) as mock_ensure,
|
||||
patch(
|
||||
"litellm.llms.openai_like.chat.handler.OpenAILikeChatHandler",
|
||||
mock_handler,
|
||||
),
|
||||
):
|
||||
garden.completion(
|
||||
model="openai/5464397967697903616",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
model_response=MagicMock(),
|
||||
print_verbose=lambda *a, **kw: None,
|
||||
encoding=MagicMock(),
|
||||
logging_obj=MagicMock(),
|
||||
api_base=None,
|
||||
optional_params={},
|
||||
custom_prompt_dict={},
|
||||
headers=None,
|
||||
timeout=30.0,
|
||||
litellm_params={},
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
vertex_credentials='{"type": "service_account"}',
|
||||
)
|
||||
|
||||
mock_ensure.assert_called_once_with(
|
||||
credentials='{"type": "service_account"}',
|
||||
project_id="test-project",
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
Loading…
Add table
Reference in a new issue