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

@@ -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):
         """

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
 
 

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