From caa0a66760b7db98b454406ba59f41e3c40d312a Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 13 Sep 2026 02:08:47 +0000 Subject: [PATCH] fix(gcs_bucket): serialize the /key/health flush with the periodic flush so an in-flight failure is not missed Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/gcs_bucket/gcs_bucket.py | 19 +++++++++--- .../key_management_endpoints.py | 14 ++++----- .../test_key_management_endpoints.py | 29 ++++++++++++++++++- 3 files changed, 49 insertions(+), 13 deletions(-) diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 3dcbea7dbf9..373726da29a 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -1,6 +1,7 @@ import asyncio import hashlib import json +import math import os import time from collections.abc import Sequence @@ -403,12 +404,22 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): """ Override flush_queue to work with asyncio.Queue. """ - await self.flush_queue_and_report() + async with self.flush_lock: + await self._send_queued_events() + self.last_flush_time = time.time() async def flush_queue_and_report(self) -> GCSFlushResult: - result: Final = await self._send_queued_events() - self.last_flush_time = time.time() - return result + """ + Flush everything queued at call time, waiting for any in-flight periodic flush first, and report every event id. + """ + async with self.flush_lock: + batch_count: Final = math.ceil(self.log_queue.qsize() / self.batch_size) + results: Final = tuple([await self._send_queued_events() for _ in range(batch_count)]) + self.last_flush_time = time.time() + return GCSFlushResult( + sent_ids=tuple(event_id for result in results for event_id in result.sent_ids), + failed_ids=tuple(event_id for result in results for event_id in result.failed_ids), + ) async def periodic_flush(self): """ diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 5828b9466e9..fce4e7a83fd 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -7146,14 +7146,12 @@ async def flush_gcs_and_describe_failures(gcs_logger: CustomLogger | None, healt if not isinstance(gcs_logger, GCSBucketLogger): return "gcs_bucket callback was selected but no GCS logger was initialized" - 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" - ) + flush_result: Final = 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 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 448203b57d7..4c5a5072e9f 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 @@ -1,4 +1,5 @@ from typing import Final +import asyncio import json from datetime import datetime, timedelta, timezone @@ -18415,7 +18416,12 @@ async def test_key_health_rejects_key_logging_entries_without_a_callback_name(): assert "callback_name is required" in exc.value.message -def _fake_upload_gcs_logger(broken_bucket: str | None = None, batch_size: int = 2048, enqueue_error: str | None = None): +def _fake_upload_gcs_logger( + broken_bucket: str | None = None, + batch_size: int = 2048, + enqueue_error: str | None = None, + upload_gate: asyncio.Event | None = None, +): from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.integrations.gcs_bucket import GCSLoggingConfig, GCSLogQueueItem @@ -18456,6 +18462,8 @@ def _fake_upload_gcs_logger(broken_bucket: str | None = None, batch_size: int = async def _log_json_data_on_gcs( self, headers: dict[str, str], bucket_name: str, object_name: str, logging_payload: StandardLoggingPayload | str ) -> None: + if upload_gate is not None: + await upload_gate.wait() if bucket_name == broken_bucket: raise RuntimeError("storage.googleapis.com returned 403") self.uploaded_buckets.append(bucket_name) @@ -18463,6 +18471,25 @@ def _fake_upload_gcs_logger(broken_bucket: str | None = None, batch_size: int = return _FakeUploadGCSLogger() +@pytest.mark.asyncio +async def test_flush_gcs_waits_for_an_in_flight_periodic_flush_that_took_the_health_event(): + from litellm.proxy.management_endpoints.key_management_endpoints import flush_gcs_and_describe_failures + + upload_gate = asyncio.Event() + logger = _fake_upload_gcs_logger(broken_bucket="team-bucket", upload_gate=upload_gate) + await logger.enqueue("health-event", "team-bucket") + periodic_flush = asyncio.create_task(logger.flush_queue()) + await asyncio.sleep(0) + assert logger.log_queue.empty() + + health_flush = asyncio.create_task(flush_gcs_and_describe_failures(logger, "health-event")) + await asyncio.sleep(0) + upload_gate.set() + await periodic_flush + + assert await health_flush == "GCS upload failed for the /key/health event and 0 other event(s), 0 uploaded" + + @pytest.mark.asyncio 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