diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 26259bc3e21..c4b504070a1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14741,6 +14741,13 @@ async def _with_resolved_api_key(litellm_params: object) -> Mapping[str, object] param="model", code=401, ) from e + except litellm.ServiceUnavailableError as e: + raise ProxyException( + message=e.message, + type="service_unavailable", + param="model", + code=503, + ) from e async def _with_resolved_credentials(deployment: Mapping[str, object] | None) -> dict[str, object] | None: diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index f4d9439d735..c7fbcc80229 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -14,6 +14,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.env_utils import get_env_float from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, @@ -170,7 +171,7 @@ LOCAL_OIDC_SECRET_PREFIXES: Final = ("oidc/file/", "oidc/env/") GOOGLE_OIDC_SECRET_PREFIX: Final = "oidc/google/" API_KEY_OIDC_SECRET_PREFIXES: Final = (*LOCAL_OIDC_SECRET_PREFIXES, GOOGLE_OIDC_SECRET_PREFIX) _GOOGLE_METADATA_DEFAULT_HOST: Final = "metadata.google.internal" -_GOOGLE_METADATA_DEFAULT_TIMEOUT_SECONDS: Final = "3" +_GOOGLE_METADATA_DEFAULT_TIMEOUT_SECONDS: Final = 3.0 _GOOGLE_TOKEN_UNAVAILABLE: Final = "the metadata server did not return an ID token" @@ -180,7 +181,8 @@ def _google_metadata_identity_url() -> str: def _google_metadata_timeout() -> float: - return float(os.getenv("GCE_METADATA_TIMEOUT") or _GOOGLE_METADATA_DEFAULT_TIMEOUT_SECONDS) + seconds: Final = get_env_float("GCE_METADATA_TIMEOUT", _GOOGLE_METADATA_DEFAULT_TIMEOUT_SECONDS) + return seconds if seconds > 0 else _GOOGLE_METADATA_DEFAULT_TIMEOUT_SECONDS class GoogleMetadataStatusError(ValueError): diff --git a/tests/integration/security/test_api_key_secret_reference.py b/tests/integration/security/test_api_key_secret_reference.py index 4fda6365665..86a7d58e0d3 100644 --- a/tests/integration/security/test_api_key_secret_reference.py +++ b/tests/integration/security/test_api_key_secret_reference.py @@ -2,6 +2,7 @@ from __future__ import annotations import base64 import json +import threading import time import uuid from collections.abc import Callable, Iterator, Mapping, Sequence @@ -58,6 +59,7 @@ _GOOGLE_AUDIENCES: Final = MappingProxyType( } ) _UNAVAILABLE_AUDIENCE: Final = "https://vllm-unavailable.example.test" +_HUNG_AUDIENCE: Final = "https://vllm-hung.example.test" _IDENTITY_PATH: Final = "/computeMetadata/v1/instance/service-accounts/default/identity" _RUN_SECRETS_TOKEN: Final = "/run/secrets/wif-absent-token" _SOURCES: Final = ("config_credential", "config_inline", "db_credential", "db_inline") @@ -152,6 +154,8 @@ def _metadata_server(google: _Google) -> Callable[[Request], Reply]: return Reply(status=404, body=b"not found", content_type="text/plain") if audience == _UNAVAILABLE_AUDIENCE: return Reply(status=503, body=b"metadata server unavailable", content_type="text/plain") + if audience == _HUNG_AUDIENCE: + return Reply(chunks=(b"never sent",), gate_before_first=threading.Event(), content_type="text/plain") return Reply(body=google.token(audience).encode(), content_type="text/plain") return respond @@ -286,6 +290,14 @@ def _config(directory: Path, secrets: _Secrets, provider: str) -> Path: "wif-google-unavailable", {"model": _GOOGLE_MODEL, "api_key": _google_reference(_UNAVAILABLE_AUDIENCE), "api_base": provider}, ), + _deployment( + "wif-google-unavailable-count", + {"model": _CHAT_MODEL, "api_key": _google_reference(_UNAVAILABLE_AUDIENCE), "api_base": provider}, + ), + _deployment( + "wif-google-hung", + {"model": _GOOGLE_MODEL, "api_key": _google_reference(_HUNG_AUDIENCE), "api_base": provider}, + ), _deployment( "wif-mistral-passthrough", { @@ -420,6 +432,7 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: **secrets.overrides(), "MISTRAL_API_BASE": provider_wire.url, "GCE_METADATA_HOST": httpx.URL(metadata_wire.url).netloc.decode(), + "GCE_METADATA_TIMEOUT": "1", }, config=_config(directory, secrets, provider_base), workers=2, @@ -847,6 +860,28 @@ def test_google_reference_fails_closed_when_the_metadata_server_fails(rig: _Rig) assert _UNAVAILABLE_AUDIENCE in _asked_audiences(rig) +def test_google_reference_answers_503_when_the_metadata_server_never_finishes_its_reply(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _chat(rig.owned.gateway, "wif-google-hung", marker) + reached: Final = _carrying(rig.provider.requests(), marker) + assert (response.status_code, reached) == (503, ()), (response.text, _bearers(reached)) + assert "the metadata server did not return an ID token" in response.text, response.text + assert _HUNG_AUDIENCE in _asked_audiences(rig) + + +def test_input_token_count_answers_503_when_the_metadata_server_fails(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.owned.gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": "wif-google-unavailable-count", "input": f"say {marker}"}, + headers=_CLOSE, + ) + reached: Final = _carrying(rig.provider.requests(), marker) + assert (response.status_code, reached) == (503, ()), (response.text, _bearers(reached)) + assert "oidc/google secret reference could not be resolved" in response.text, response.text + + def _side_call(rig: _Rig, seam: str, marker: str) -> tuple[httpx.Response, Callable[[Request], bool]]: gateway: Final = rig.owned.gateway match seam: diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 0e8fc6823bd..378f86fff26 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -279,6 +279,39 @@ async def test_token_counter_reads_an_oidc_api_key_only_when_it_calls_the_provid assert str(tmp_path) not in excinfo.value.message +@pytest.mark.asyncio +async def test_token_counter_answers_503_when_the_metadata_server_behind_an_oidc_google_api_key_fails( + monkeypatch: pytest.MonkeyPatch, +): + audience: Final = f"https://count.test/{uuid.uuid4().hex}" + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "claude", + "litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": f"oidc/google/{audience}"}, + } + ] + ), + ) + + with respx.mock(assert_all_called=True) as upstream: + upstream.get(url__regex=r"/computeMetadata/v1/instance/service-accounts/default/identity").respond( + 500, text="metadata server unavailable" + ) + with pytest.raises(ProxyException) as excinfo: + await token_counter( + request=TokenCountRequest(model="claude", messages=[{"role": "user", "content": "hello"}]), + call_endpoint=True, + ) + + assert excinfo.value.code == "503" + assert "oidc/google secret reference could not be resolved" in excinfo.value.message + assert audience.removeprefix("https://") not in excinfo.value.message + + @pytest.mark.asyncio async def test_native_openai_responses_count_sends_the_token_an_oidc_api_key_points_to( tmp_path: Path, monkeypatch: pytest.MonkeyPatch diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index e1d97fa40f7..5b783e0a7f1 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -10,7 +10,6 @@ import json import logging import os import queue -import socket import threading import time import uuid @@ -7576,32 +7575,68 @@ async def test_an_oidc_google_api_key_token_cached_when_the_call_starts_is_sent_ @pytest.mark.asyncio -async def test_a_hung_metadata_server_fails_an_oidc_google_api_key_call_with_a_503_and_leaves_the_default_executor_free( +async def test_a_pending_oidc_google_metadata_fetch_leaves_the_default_executor_free( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) _, reference = _google_reference() monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) - monkeypatch.setenv("GCE_METADATA_TIMEOUT", "1") loop: Final = asyncio.get_running_loop() loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) + fetch_started: Final[queue.SimpleQueue[int]] = queue.SimpleQueue() + release: Final = threading.Event() - with socket.create_server(("127.0.0.1", 0)) as never_answering, respx.mock(assert_all_called=False) as upstream: - monkeypatch.setenv("GCE_METADATA_HOST", f"127.0.0.1:{never_answering.getsockname()[1]}") - upstream.route(host="127.0.0.1").pass_through() + def pending_metadata(_request: httpx.Request) -> httpx.Response: + fetch_started.put(threading.get_ident()) + release.wait(timeout=30) + return httpx.Response(500, text="metadata server unavailable") + + with ThreadPoolExecutor(max_workers=1) as observer, respx.mock(assert_all_called=False) as upstream: + upstream.get(_GOOGLE_METADATA_IDENTITY).mock(side_effect=pending_metadata) chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) call: Final = asyncio.ensure_future(_ping_vllm_through("async", "vllm-sa")) - await asyncio.sleep(0.2) - other_work: Final = await asyncio.wait_for(loop.run_in_executor(None, lambda: "done"), timeout=0.5) + try: + await asyncio.wait_for(loop.run_in_executor(observer, fetch_started.get), timeout=30) + other_work: Final = await asyncio.wait_for(loop.run_in_executor(None, lambda: "done"), timeout=30) + finally: + release.set() with pytest.raises(litellm.ServiceUnavailableError) as excinfo: - await asyncio.wait_for(call, timeout=10) + await call assert other_work == "done" assert excinfo.value.status_code == 503 - assert "the metadata server did not return an ID token" in excinfo.value.message assert chat.call_count == 0 +@pytest.mark.parametrize( + ("configured", "bound"), + [(None, 3.0), ("1.5", 1.5), ("three seconds", 3.0), ("0", 3.0), ("-2", 3.0)], + ids=["unset", "configured", "malformed", "zero", "negative"], +) +def test_oidc_google_api_key_metadata_fetch_waits_at_most_gce_metadata_timeout_seconds( + monkeypatch: pytest.MonkeyPatch, configured: str | None, bound: float +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + if configured is None: + monkeypatch.delenv("GCE_METADATA_TIMEOUT", raising=False) + else: + monkeypatch.setenv("GCE_METADATA_TIMEOUT", configured) + monkeypatch.setattr(time, "time", lambda: _CLOCK_START) + audience, reference = _google_reference() + monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) + token: Final = _google_id_token(audience, _CLOCK_START) + + with respx.mock(assert_all_called=True) as upstream: + metadata: Final = upstream.get(_GOOGLE_METADATA_IDENTITY).respond(200, text=token) + chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) + litellm.completion(model=_VLLM_MODEL, messages=_PING, litellm_credential_name="vllm-sa") + sent: Final = chat.calls.last.request.headers["Authorization"] + waits: Final = metadata.calls.last.request.extensions["timeout"] + + assert sent == f"Bearer {token}" + assert waits == {"connect": bound, "read": bound, "write": bound, "pool": bound} + + def test_an_oidc_google_api_key_sent_with_the_call_goes_upstream_as_written_without_a_metadata_fetch() -> None: _, reference = _google_reference()