From ff98dcdc556468862d37d5edc431965ce876da48 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 13 Sep 2026 01:50:04 +0000 Subject: [PATCH] fix(proxy): judge /key/health GCS status by the health check's own event The flush result now carries the ids it uploaded and the ids it failed, so /key/health reports unhealthy only when its own event failed to upload. Another team's broken bucket or a leftover retry in the shared queue no longer marks a working key unhealthy, and the flush walks every pending batch so an event queued behind more than batch_size items is still judged Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/gcs_bucket/gcs_bucket.py | 18 ++- .../key_management_endpoints.py | 34 ++++-- litellm/types/integrations/gcs_bucket.py | 12 +- .../gcs_bucket/test_gcs_bucket.py | 10 +- .../test_key_management_endpoints.py | 103 ++++++++++++++---- 5 files changed, 134 insertions(+), 43 deletions(-) diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index bf748037f89..3dcbea7dbf9 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -242,10 +242,14 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): Send each log individually as separate GCS objects (legacy behavior). This is used when GCS_USE_BATCHED_LOGGING is disabled. """ - failed_items: Final = tuple([item for item in items if not await self._send_single_log_item(item)]) + outcomes: Final = tuple([(item, await self._send_single_log_item(item)) for item in items]) + failed_items: Final = tuple(item for item, sent in outcomes if not sent) if failed_items: self._requeue(failed_items) - return GCSFlushResult(sent=len(items) - len(failed_items), failed=len(failed_items)) + return GCSFlushResult( + sent_ids=tuple(item["payload"]["id"] for item, sent in outcomes if sent), + failed_ids=tuple(item["payload"]["id"] for item in failed_items), + ) async def _send_single_log_item(self, item: GCSLogQueueItem) -> bool: """ @@ -288,8 +292,12 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): if group_failed: self._requeue(group_items) return GCSFlushResult( - sent=sum(group_sent for _, (group_sent, _) in results), - failed=sum(group_failed for _, (_, group_failed) in results), + sent_ids=tuple( + item["payload"]["id"] for group_items, (_, failed) in results if not failed for item in group_items + ), + failed_ids=tuple( + item["payload"]["id"] for group_items, (_, failed) in results if failed for item in group_items + ), ) async def async_send_batch(self) -> None: @@ -307,7 +315,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): items_to_process: Final = self._drain_queue_batch() if not items_to_process: - return GCSFlushResult(sent=0, failed=0) + return GCSFlushResult(sent_ids=(), failed_ids=()) if self.use_batched_logging: return await self._send_grouped_batches(items_to_process) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 684230f06fd..5828b9466e9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -157,6 +157,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import ( from litellm.types.router import Deployment from litellm.types.utils import ( BudgetConfig, + ModelResponse, PersonalUIKeyGenerationConfig, TeamUIKeyGenerationConfig, ) @@ -7140,15 +7141,20 @@ def _callback_entry_error(entry: Mapping[str, object]) -> str | None: return None -async def flush_gcs_and_describe_failures(gcs_logger: CustomLogger | None) -> str | None: +async def flush_gcs_and_describe_failures(gcs_logger: CustomLogger | None, health_check_event_id: str) -> str | None: from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger if not isinstance(gcs_logger, GCSBucketLogger): return "gcs_bucket callback was selected but no GCS logger was initialized" - flush_result: Final = await gcs_logger.flush_queue_and_report() - if flush_result.failed == 0: - return None - return f"GCS upload failed for {flush_result.failed} event(s), {flush_result.sent} uploaded" + flush_rounds: Final = max(1, math.ceil(gcs_logger.log_queue.qsize() / gcs_logger.batch_size)) + for _ in range(flush_rounds): + flush_result = await gcs_logger.flush_queue_and_report() + if health_check_event_id in flush_result.failed_ids: + return ( + f"GCS upload failed for the /key/health event and {flush_result.failed - 1} other event(s), " + f"{flush_result.sent} uploaded" + ) + return None async def test_key_logging( @@ -7194,7 +7200,7 @@ async def test_key_logging( request=request, ) data["mock_response"] = "test response" - await litellm.acompletion(**data) # make mock completion call to trigger key based callbacks + health_check_response: Final = await litellm.acompletion(**data) except Exception as e: return LoggingCallbackStatus( callbacks=logging_callbacks, @@ -7203,20 +7209,26 @@ async def test_key_logging( ) await asyncio.sleep(2) # wait for callbacks to run, callbacks use batching so wait for the flush event + callback_log_contents: Final = log_capture_string.getvalue() + health_check_event_id: Final = health_check_response.id if isinstance(health_check_response, ModelResponse) else "" gcs_failure: Final = ( - await flush_gcs_and_describe_failures(get_custom_logger_compatible_class("gcs_bucket")) + await flush_gcs_and_describe_failures(get_custom_logger_compatible_class("gcs_bucket"), health_check_event_id) if "gcs_bucket" in logging_callbacks else None ) - - log_contents: Final = log_capture_string.getvalue() logger.removeHandler(ch) - if gcs_failure is not None or log_contents: + flush_log_contents: Final = ( + log_capture_string.getvalue()[len(callback_log_contents) :] if gcs_failure is not None else "" + ) + if gcs_failure is not None or callback_log_contents: return LoggingCallbackStatus( callbacks=logging_callbacks, status="unhealthy", - details=f"Logger exceptions triggered, system is unhealthy: {gcs_failure or ''} {log_contents}".strip(), + details=( + "Logger exceptions triggered, system is unhealthy: " + f"{gcs_failure or ''} {callback_log_contents}{flush_log_contents}" + ).strip(), ) return LoggingCallbackStatus( callbacks=logging_callbacks, diff --git a/litellm/types/integrations/gcs_bucket.py b/litellm/types/integrations/gcs_bucket.py index 7f8ea7caee4..1ae3191a1d7 100644 --- a/litellm/types/integrations/gcs_bucket.py +++ b/litellm/types/integrations/gcs_bucket.py @@ -38,5 +38,13 @@ class GCSLogQueueItem(TypedDict): @dataclass(frozen=True, slots=True) class GCSFlushResult: - sent: int - failed: int + sent_ids: tuple[str, ...] + failed_ids: tuple[str, ...] + + @property + def sent(self) -> int: + return len(self.sent_ids) + + @property + def failed(self) -> int: + return len(self.failed_ids) diff --git a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket.py b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket.py index cee457fab50..b6e43fff6b7 100644 --- a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket.py +++ b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket.py @@ -61,14 +61,14 @@ async def test_failed_batch_stays_queued_and_is_retried_on_the_next_flush(): failed_flush = await logger.flush_queue_and_report() - assert failed_flush == GCSFlushResult(sent=0, failed=2) + assert failed_flush == GCSFlushResult(sent_ids=(), failed_ids=("req-1", "req-2")) assert logger.log_queue.qsize() == 2 assert logger.uploaded == [] logger.failing_ids = frozenset() retried_flush = await logger.flush_queue_and_report() - assert retried_flush == GCSFlushResult(sent=2, failed=0) + assert retried_flush == GCSFlushResult(sent_ids=("req-1", "req-2"), failed_ids=()) assert logger.log_queue.qsize() == 0 assert logger.uploaded == [["req-1", "req-2"]] @@ -83,7 +83,7 @@ async def test_individual_mode_requeues_only_the_failed_items(): result = await logger.flush_queue_and_report() - assert result == GCSFlushResult(sent=1, failed=1) + assert result == GCSFlushResult(sent_ids=("req-ok",), failed_ids=("req-fail",)) assert logger.uploaded == [["req-ok"]] assert logger.queued_ids() == ["req-fail"] @@ -110,7 +110,7 @@ async def test_failed_batch_is_dropped_when_new_events_filled_the_queue_during_t result = await logger.flush_queue_and_report() - assert result == GCSFlushResult(sent=0, failed=2) + assert result == GCSFlushResult(sent_ids=(), failed_ids=("req-1", "req-2")) assert logger.queued_ids() == ["req-3", "req-4"] @@ -118,4 +118,4 @@ async def test_failed_batch_is_dropped_when_new_events_filled_the_queue_during_t async def test_empty_queue_flush_reports_nothing_sent_or_failed(): logger = _FakeUploadGCSLogger() - assert await logger.flush_queue_and_report() == GCSFlushResult(sent=0, failed=0) + assert await logger.flush_queue_and_report() == GCSFlushResult(sent_ids=(), failed_ids=()) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 41419f5adba..448203b57d7 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -18415,33 +18415,72 @@ async def test_key_health_rejects_key_logging_entries_without_a_callback_name(): assert "callback_name is required" in exc.value.message -def _gcs_logger_whose_flush_reports(sent: int, failed: int): +def _fake_upload_gcs_logger(broken_bucket: str | None = None, batch_size: int = 2048, enqueue_error: str | None = None): from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger - from litellm.types.integrations.gcs_bucket import GCSFlushResult + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + from litellm.types.integrations.gcs_bucket import GCSLoggingConfig, GCSLogQueueItem + from litellm.types.utils import StandardLoggingPayload + + class _FakeUploadGCSLogger(GCSBucketLogger): + """Skips GCP auth; an upload to `broken_bucket` raises, every other upload records its bucket""" - class _FixedFlushGCSLogger(GCSBucketLogger): def __init__(self) -> None: with patch("litellm.proxy.proxy_server.premium_user", True): # test-quality-ok: GCS logging is premium-gated - super().__init__(bucket_name="test-bucket") + super().__init__(bucket_name="team-bucket") + self.batch_size = batch_size + self.uploaded_buckets: list[str] = [] - async def flush_queue_and_report(self) -> GCSFlushResult: - return GCSFlushResult(sent=sent, failed=failed) + async def _enqueue(self, item: GCSLogQueueItem) -> None: + if enqueue_error is not None: + raise RuntimeError(enqueue_error) + await super()._enqueue(item) - return _FixedFlushGCSLogger() + async def enqueue(self, request_id: str, bucket_name: str) -> None: + payload: Final = StandardLoggingPayload(id=request_id) # pyright: ignore[reportCallIssue] # partial payload is enough for queueing + kwargs: Final = {"standard_callback_dynamic_params": {"gcs_bucket_name": bucket_name}} + await self._enqueue(GCSLogQueueItem(payload=payload, kwargs=kwargs, response_obj=None)) + + async def get_gcs_logging_config(self, kwargs: dict | None = None) -> GCSLoggingConfig: + dynamic_params: Final = (kwargs or {}).get("standard_callback_dynamic_params") or {} + return GCSLoggingConfig( + bucket_name=dynamic_params.get("gcs_bucket_name") or "team-bucket", + vertex_instance=None, + path_service_account=None, + ) + + async def construct_request_headers( + self, service_account_json: str | None, vertex_instance: VertexBase | None = None + ) -> dict[str, str]: + return {} + + async def _log_json_data_on_gcs( + self, headers: dict[str, str], bucket_name: str, object_name: str, logging_payload: StandardLoggingPayload | str + ) -> None: + if bucket_name == broken_bucket: + raise RuntimeError("storage.googleapis.com returned 403") + self.uploaded_buckets.append(bucket_name) + + return _FakeUploadGCSLogger() @pytest.mark.asyncio -async def test_flush_gcs_reports_the_failed_upload_count_from_the_registered_logger(): +async def test_flush_gcs_reports_only_when_the_health_event_itself_failed_to_upload(): from litellm.proxy.management_endpoints.key_management_endpoints import flush_gcs_and_describe_failures + logger = _fake_upload_gcs_logger(broken_bucket="team-bucket") + await logger.enqueue("req-1", "team-bucket") + await logger.enqueue("req-2", "team-bucket") + await logger.enqueue("req-3", "ok-bucket") + await logger.enqueue("health-event", "team-bucket") + assert ( - await flush_gcs_and_describe_failures(_gcs_logger_whose_flush_reports(sent=0, failed=3)) - == "GCS upload failed for 3 event(s), 0 uploaded" + await flush_gcs_and_describe_failures(logger, "health-event") + == "GCS upload failed for the /key/health event and 2 other event(s), 1 uploaded" ) - assert await flush_gcs_and_describe_failures(_gcs_logger_whose_flush_reports(sent=2, failed=0)) is None + assert await flush_gcs_and_describe_failures(logger, "health-event-already-uploaded") is None -async def _key_logging_status_after_gcs_flush(sent: int, failed: int) -> LoggingCallbackStatus: +async def _key_logging_status_with_gcs_logger(gcs_logger) -> LoggingCallbackStatus: from starlette.requests import Request as StarletteRequest from litellm.proxy.management_endpoints.key_management_endpoints import test_key_logging @@ -18454,28 +18493,52 @@ async def _key_logging_status_after_gcs_flush(sent: int, failed: int) -> Logging patch("litellm.proxy.proxy_server.proxy_config", _default_team_gcs_proxy_config("team-gcs")), # test-quality-ok: test_key_logging reads the module-level proxy config patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: the mock completion's GCS success event is premium-gated patch( # test-quality-ok: the mock completion and the flush both look the logger up in this process-wide registry - "litellm.litellm_core_utils.litellm_logging._in_memory_loggers", - [_gcs_logger_whose_flush_reports(sent=sent, failed=failed)], + "litellm.litellm_core_utils.litellm_logging._in_memory_loggers", [gcs_logger] ), ): return await test_key_logging(user_api_key_dict=caller, request=request, logging_callbacks=("gcs_bucket",)) @pytest.mark.asyncio -async def test_key_logging_marks_the_key_unhealthy_when_the_gcs_flush_leaves_events_undelivered(): - status = await _key_logging_status_after_gcs_flush(sent=1, failed=3) +async def test_key_logging_marks_the_key_unhealthy_when_its_own_gcs_upload_fails(): + status = await _key_logging_status_with_gcs_logger(_fake_upload_gcs_logger(broken_bucket="team-bucket")) assert status["status"] == "unhealthy" - assert "GCS upload failed for 3 event(s), 1 uploaded" in (status["details"] or "") + assert "GCS upload failed for the /key/health event and 0 other event(s), 0 uploaded" in (status["details"] or "") + assert "storage.googleapis.com returned 403" in (status["details"] or "") @pytest.mark.asyncio -async def test_key_logging_stays_healthy_when_the_gcs_flush_delivers_every_event(): - status = await _key_logging_status_after_gcs_flush(sent=1, failed=0) +async def test_key_logging_reports_the_callback_error_when_the_event_never_reaches_the_gcs_queue(): + status = await _key_logging_status_with_gcs_logger(_fake_upload_gcs_logger(enqueue_error="queue closed")) + + assert status["status"] == "unhealthy" + assert "GCS Bucket logging error: queue closed" in (status["details"] or "") + assert "GCS upload failed" not in (status["details"] or "") + + +@pytest.mark.asyncio +async def test_key_logging_stays_healthy_when_only_another_teams_queued_upload_fails(): + logger = _fake_upload_gcs_logger(broken_bucket="other-team-bucket") + await logger.enqueue("req-other-team", "other-team-bucket") + + status = await _key_logging_status_with_gcs_logger(logger) assert status["status"] == "healthy" assert status["callbacks"] == ("gcs_bucket",) assert "Manually check if logs were sent to gcs_bucket" in (status["details"] or "") + assert logger.uploaded_buckets == ["team-bucket"] + + +@pytest.mark.asyncio +async def test_key_logging_flushes_past_the_first_batch_to_reach_its_own_event(): + logger = _fake_upload_gcs_logger(broken_bucket="team-bucket", batch_size=1) + await logger.enqueue("req-other-team", "other-team-bucket") + + status = await _key_logging_status_with_gcs_logger(logger) + + assert status["status"] == "unhealthy" + assert logger.uploaded_buckets == ["other-team-bucket"] @pytest.mark.asyncio @@ -18483,6 +18546,6 @@ async def test_flush_gcs_names_a_missing_logger_when_the_callback_never_initiali from litellm.proxy.management_endpoints.key_management_endpoints import flush_gcs_and_describe_failures assert ( - await flush_gcs_and_describe_failures(None) + await flush_gcs_and_describe_failures(None, "health-event") == "gcs_bucket callback was selected but no GCS logger was initialized" )