diff --git a/litellm/integrations/generic_api/generic_api_callback.py b/litellm/integrations/generic_api/generic_api_callback.py index 5f9851f595f..dc338722a36 100644 --- a/litellm/integrations/generic_api/generic_api_callback.py +++ b/litellm/integrations/generic_api/generic_api_callback.py @@ -197,6 +197,7 @@ class GenericAPILogger(CustomBatchLogger): # Init variables for batch flushing logs ######################################################### self.flush_lock = asyncio.Lock() + self._full_queue_flush: asyncio.Task[None] | None = None super().__init__(**kwargs, flush_lock=self.flush_lock) asyncio.create_task(self.periodic_flush()) self.log_queue: list[Mapping[str, object] | StandardLoggingPayload] = [] @@ -315,7 +316,7 @@ class GenericAPILogger(CustomBatchLogger): self.log_queue.append(standard_logging_payload) if len(self.log_queue) >= self.batch_size: - await self.async_send_batch() + self._flush_full_queue_in_background() except Exception as e: verbose_logger.exception("Generic API Logger Error - %s\n%s", e, traceback.format_exc()) @@ -346,7 +347,7 @@ class GenericAPILogger(CustomBatchLogger): self.log_queue.append(standard_logging_payload) if len(self.log_queue) >= self.batch_size: - await self.async_send_batch() + self._flush_full_queue_in_background() except Exception as e: verbose_logger.exception("Generic API Logger Error - %s\n%s", e, traceback.format_exc()) @@ -354,6 +355,11 @@ class GenericAPILogger(CustomBatchLogger): async def flush_queue(self) -> None: await self.async_send_batch() + def _flush_full_queue_in_background(self) -> None: + if self._full_queue_flush is not None and not self._full_queue_flush.done(): + return + self._full_queue_flush = asyncio.create_task(self.flush_queue()) + async def async_send_batch(self) -> None: async with self.flush_lock: await self._send_batch() diff --git a/tests/unit/integrations/generic_api/test_generic_api_callback.py b/tests/unit/integrations/generic_api/test_generic_api_callback.py index afaae02321d..99cafdb17ae 100644 --- a/tests/unit/integrations/generic_api/test_generic_api_callback.py +++ b/tests/unit/integrations/generic_api/test_generic_api_callback.py @@ -4,7 +4,7 @@ import json import logging import os from collections.abc import AsyncIterator, Iterator, Sequence -from typing import Awaitable, Final, TypedDict, TypeVar +from typing import Awaitable, Final, Literal, TypedDict, TypeVar from unittest.mock import AsyncMock import httpx @@ -19,7 +19,7 @@ from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE from litellm.integrations.gcs_pubsub.pub_sub import * from litellm.integrations.generic_api.generic_api_callback import LOG_FORMAT_TYPES, GenericAPILogger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker from litellm.types.utils import StandardLoggingPayload from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -795,3 +795,88 @@ async def test_generic_api_concurrent_flushes_deliver_each_event_once(log_format task.cancel() await asyncio.gather(*owned_tasks, return_exceptions=True) await handler.close() + + +class _RecoveringLogSink: + def __init__(self) -> None: + self.arrived: Final = asyncio.Event() + self.release: Final = asyncio.Event() + self.attempts: int = 0 + self.delivered: tuple[httpx.Request, ...] = () + + async def respond(self, request: httpx.Request) -> httpx.Response: + self.attempts += 1 + if self.attempts == 1: + self.arrived.set() + await self.release.wait() + return httpx.Response(503) + self.delivered = (*self.delivered, request) + return httpx.Response(200) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("event_type", ("success", "failure")) +async def test_generic_api_full_queue_send_survives_logging_worker_timeout( + event_type: Literal["success", "failure"], +) -> None: + sink: Final = _RecoveringLogSink() + handler: Final = AsyncHTTPHandler(transport=httpx.MockTransport(sink.respond)) + worker: Final = LoggingWorker(timeout=0.05) + before_tasks: Final = frozenset(asyncio.all_tasks()) + logger: Final = GenericAPILogger( + endpoint="https://owned-logger.invalid", + async_httpx_client=handler, + batch_size=1, + flush_interval=3600, + max_retries=1, + retry_delay=0, + ) + owned_tasks: Final = frozenset(asyncio.all_tasks()) - before_tasks + log_event: Final = logger.async_log_success_event if event_type == "success" else logger.async_log_failure_event + try: + worker.ensure_initialized_and_enqueue( + log_event({"standard_logging_object": {"id": "queued"}}, None, None, None) + ) + await worker.flush() + await sink.arrived.wait() + sink.release.set() + await logger.flush_queue() + assert tuple(_received_log_ids(sink.delivered, "json_array")) == ("queued",) + finally: + sink.release.set() + await worker.stop() + for task in owned_tasks: + task.cancel() + await asyncio.gather(*owned_tasks, return_exceptions=True) + await handler.close() + + +@pytest.mark.asyncio +async def test_generic_api_full_queue_starts_no_extra_send_while_one_is_stuck() -> None: + sink: Final = _InProcessLogSink(hold_first=True) + handler: Final = AsyncHTTPHandler(transport=httpx.MockTransport(sink.respond)) + before_tasks: Final = frozenset(asyncio.all_tasks()) + logger: Final = GenericAPILogger( + endpoint="https://owned-logger.invalid", + async_httpx_client=handler, + batch_size=1, + flush_interval=3600, + ) + owned_tasks: Final = frozenset(asyncio.all_tasks()) - before_tasks + later_ids: Final = tuple(f"later-{index}" for index in range(50)) + try: + await logger.async_log_success_event({"standard_logging_object": {"id": "stuck"}}, None, None, None) + await sink.arrived.wait() + tasks_while_stuck: Final = frozenset(asyncio.all_tasks()) + for log_id in later_ids: + await logger.async_log_success_event({"standard_logging_object": {"id": log_id}}, None, None, None) + assert frozenset(asyncio.all_tasks()) - tasks_while_stuck == frozenset() + sink.release.set() + await logger.flush_queue() + assert tuple(_received_log_ids(sink.requests, "json_array")) == ("stuck", *later_ids) + finally: + sink.release.set() + for task in owned_tasks: + task.cancel() + await asyncio.gather(*owned_tasks, return_exceptions=True) + await handler.close()