fix(azure_sentinel): split batches under the 1MB ingestion cap

litellm/integrations/azure_sentinel/azure_sentinel.py

@@ -16,23 +16,27 @@
 import os
 import time
 import traceback
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from types import MappingProxyType
-from typing import Final
+from typing import Final, TypeVar
 from urllib.parse import urlparse
 
 from litellm._logging import verbose_logger
+from litellm.integrations.batch_utils import send_batch_with_413_split
 from litellm.integrations.custom_batch_logger import CustomBatchLogger
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,
     httpxSpecialProvider,
 )
+from litellm.types.integrations.azure_sentinel import AZURE_SENTINEL_MAX_PAYLOAD_SIZE_BYTES
 from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
 
 DEFAULT_AZURE_AUTHORITY_HOST: Final = "https://login.microsoftonline.com"
 DEFAULT_AZURE_MONITOR_SCOPE: Final = "https://monitor.azure.com/.default"
 
+_QueuedPayload = TypeVar("_QueuedPayload", StandardLoggingPayload, StandardAuditLogPayload)
+
 MONITOR_SCOPE_BY_AUTHORITY_HOST: Final[Mapping[str, str]] = MappingProxyType(
     {
         "login.microsoftonline.com": DEFAULT_AZURE_MONITOR_SCOPE,
@@ -153,6 +157,8 @@ def __init__(
         asyncio.create_task(self.periodic_flush())
         self.log_queue: list[StandardLoggingPayload] = []
         self.audit_log_queue: list[StandardAuditLogPayload] = []
+        self.logs_awaiting_retry = False
+        self.audit_logs_awaiting_retry = False
 
     @staticmethod
     def _normalize_authority_host(authority_host: str) -> str:
@@ -245,8 +251,8 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti
 
             self.log_queue.append(standard_logging_payload)
 
-            if len(self.log_queue) >= self.batch_size:
-                await self.async_send_batch()
+            if len(self.log_queue) >= self.batch_size and not self.logs_awaiting_retry:
+                await self.flush_queue()
 
         except Exception as e:
             verbose_logger.exception("Azure Sentinel Layer Error - %s\n%s", e, traceback.format_exc())
@@ -275,8 +281,8 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti
 
             self.log_queue.append(standard_logging_payload)
 
-            if len(self.log_queue) >= self.batch_size:
-                await self.async_send_batch()
+            if len(self.log_queue) >= self.batch_size and not self.logs_awaiting_retry:
+                await self.flush_queue()
 
         except Exception as e:
             verbose_logger.exception("Azure Sentinel Layer Error - %s\n%s", e, traceback.format_exc())
@@ -298,8 +304,8 @@ async def async_log_audit_log_event(self, audit_log: StandardAuditLogPayload) ->
 
             self.audit_log_queue.append(audit_log)
 
-            if len(self.audit_log_queue) >= self.batch_size:
-                await self.async_send_audit_batch()
+            if len(self.audit_log_queue) >= self.batch_size and not self.audit_logs_awaiting_retry:
+                await self.flush_queue()
 
         except Exception as e:
             verbose_logger.exception("Azure Sentinel Audit Log Layer Error - %s\n%s", e, traceback.format_exc())
@@ -311,67 +317,86 @@ async def async_send_batch(self):
         Raises:
             Raises a NON Blocking verbose_logger.exception if an error occurs
         """
-        await self._async_send_batch_to_api(
-            log_queue=self.log_queue,
+        batch_to_send: Final = tuple(self.log_queue)
+        self.log_queue = []
+        undelivered: Final = await self._async_send_batch_to_api(
+            log_queue=batch_to_send,
             api_endpoint=self.api_endpoint,
             log_type="logs",
         )
+        self.logs_awaiting_retry = bool(undelivered)
+        self.log_queue = self._requeue(undelivered, self.log_queue, "logs")
 
     async def async_send_audit_batch(self):
         """
         Sends the batch of audit logs to Azure Monitor Logs Ingestion API
         """
-        await self._async_send_batch_to_api(
-            log_queue=self.audit_log_queue,
+        batch_to_send: Final = tuple(self.audit_log_queue)
+        self.audit_log_queue = []
+        undelivered: Final = await self._async_send_batch_to_api(
+            log_queue=batch_to_send,
             api_endpoint=self.audit_api_endpoint,
             log_type="audit logs",
         )
+        self.audit_logs_awaiting_retry = bool(undelivered)
+        self.audit_log_queue = self._requeue(undelivered, self.audit_log_queue, "audit logs")
+
+    def _requeue(
+        self,
+        undelivered: tuple[_QueuedPayload, ...],
+        queue: list[_QueuedPayload],
+        log_type: str,
+    ) -> list[_QueuedPayload]:
+        merged: Final = [*undelivered, *queue]
+        overflow: Final = len(merged) - self.max_queue_size
+        if overflow <= 0:
+            return merged
+
+        verbose_logger.warning(
+            "Azure Sentinel: %s queue exceeded max_queue_size=%s, dropped %s oldest records",
+            log_type,
+            self.max_queue_size,
+            overflow,
+        )
+        return merged[overflow:]
 
     async def _async_send_batch_to_api(
         self,
-        log_queue: list[StandardLoggingPayload | StandardAuditLogPayload],
+        log_queue: tuple[_QueuedPayload, ...],
         api_endpoint: str,
         log_type: str,
-    ) -> None:
-        try:
-            if not log_queue:
-                return
+    ) -> tuple[_QueuedPayload, ...]:
+        if not log_queue:
+            return ()
 
-            verbose_logger.debug("Azure Sentinel - about to flush %s %s", len(log_queue), log_type)
-
-            # Get OAuth2 token
+        verbose_logger.debug("Azure Sentinel - about to flush %s %s", len(log_queue), log_type)
+        try:
             bearer_token: Final = await self._get_oauth_token()
+        except Exception as e:
+            verbose_logger.exception("Azure Sentinel Error getting OAuth token - %s", e)
+            return tuple(log_queue)
 
-            # Convert log queue to JSON array format expected by Logs Ingestion API
-            # Each log entry should be a JSON object in the array
-            body: Final = safe_dumps(log_queue)
-
-            # Set headers for Logs Ingestion API
-            headers: Final = {
-                "Authorization": f"Bearer {bearer_token}",
-                "Content-Type": "application/json",
-            }
-
-            # Send the request
-            response = await self.async_httpx_client.post(url=api_endpoint, data=body.encode("utf-8"), headers=headers)
-
-            if response.status_code not in [200, 204]:
-                verbose_logger.error(
-                    "Azure Sentinel API error: status_code=%s, response=%s",
-                    response.status_code,
-                    response.text,
-                )
-                raise Exception(f"Failed to send logs to Azure Sentinel: {response.status_code} - {response.text}")
+        headers: Final = {
+            "Authorization": f"Bearer {bearer_token}",
+            "Content-Type": "application/json",
+        }
 
-            verbose_logger.debug(
-                "Azure Sentinel: Response from API status_code: %s",
-                response.status_code,
+        async def _send_batch(batch: Sequence[_QueuedPayload]):
+            body: Final = safe_dumps(batch)
+            return await self.async_httpx_client.post(
+                url=api_endpoint,
+                data=body.encode("utf-8"),
+                headers=headers,
             )
 
-        except Exception as e:
-            verbose_logger.exception("Azure Sentinel Error sending batch API - %s\n%s", e, traceback.format_exc())
-        finally:
-            log_queue.clear()
+        return await send_batch_with_413_split(
+            batch=log_queue,
+            send_batch=_send_batch,
+            exceeds_limits=lambda batch: len(safe_dumps(batch).encode("utf-8")) > AZURE_SENTINEL_MAX_PAYLOAD_SIZE_BYTES,
+            success_status_codes=frozenset({200, 204}),
+            integration_name="Azure Sentinel",
+            drop_error_message="Azure Sentinel API Error - Payload too large for a single record",
+        )
 
     async def flush_queue(self):
         if self.flush_lock is None:

litellm/integrations/batch_utils.py

@@ -0,0 +1,83 @@
+from collections.abc import Awaitable, Callable, Sequence
+from typing import Final, TypeVar
+
+import httpx
+
+from litellm._logging import verbose_logger
+from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError
+
+_BatchItem = TypeVar("_BatchItem")
+
+
+async def send_batch_with_413_split(
+    batch: Sequence[_BatchItem],
+    send_batch: Callable[[Sequence[_BatchItem]], Awaitable[httpx.Response]],
+    exceeds_limits: Callable[[Sequence[_BatchItem]], bool],
+    success_status_codes: frozenset[int],
+    integration_name: str,
+    drop_error_message: str,
+) -> tuple[_BatchItem, ...]:
+    async def _halve() -> tuple[_BatchItem, ...]:
+        midpoint: Final = len(batch) // 2
+        left_undelivered: Final = await send_batch_with_413_split(
+            batch=batch[:midpoint],
+            send_batch=send_batch,
+            exceeds_limits=exceeds_limits,
+            success_status_codes=success_status_codes,
+            integration_name=integration_name,
+            drop_error_message=drop_error_message,
+        )
+        if left_undelivered:
+            return left_undelivered + tuple(batch[midpoint:])
+        return await send_batch_with_413_split(
+            batch=batch[midpoint:],
+            send_batch=send_batch,
+            exceeds_limits=exceeds_limits,
+            success_status_codes=success_status_codes,
+            integration_name=integration_name,
+            drop_error_message=drop_error_message,
+        )
+
+    async def _handle_413() -> tuple[_BatchItem, ...]:
+        if len(batch) == 1:
+            verbose_logger.error(drop_error_message)
+            return ()
+        return await _halve()
+
+    if not batch:
+        return ()
+
+    try:
+        oversized: Final = exceeds_limits(batch)
+    except (TypeError, ValueError) as e:
+        if len(batch) > 1:
+            return await _halve()
+        verbose_logger.exception("%s dropped a record that cannot be serialized - %s", integration_name, e)
+        return ()
+    if oversized and len(batch) > 1:
+        return await _halve()
+
+    try:
+        response: Final = await send_batch(batch)
+    except MaskedHTTPStatusError as e:
+        if e.status_code == 413:
+            return await _handle_413()
+        verbose_logger.exception("%s Error sending batch API - %s", integration_name, e)
+        return tuple(batch)
+    except Exception as e:
+        verbose_logger.exception("%s Error sending batch API - %s", integration_name, e)
+        return tuple(batch)
+
+    if response.status_code == 413:
+        return await _handle_413()
+    if response.status_code not in success_status_codes:
+        verbose_logger.error(
+            "%s API error: status_code=%s, response=%s",
+            integration_name,
+            response.status_code,
+            response.text,
+        )
+        return tuple(batch)
+
+    verbose_logger.debug("%s delivered %s records, status_code=%s", integration_name, len(batch), response.status_code)
+    return ()

litellm/integrations/datadog/datadog.py

@@ -29,6 +29,7 @@
 import litellm
 from litellm._logging import verbose_logger
 from litellm._uuid import uuid
+from litellm.integrations.batch_utils import send_batch_with_413_split
 from litellm.integrations.custom_batch_logger import CustomBatchLogger
 from litellm.integrations.datadog.datadog_handler import (
     get_datadog_base_url_from_env,
@@ -43,7 +44,6 @@
 )
 from litellm.litellm_core_utils.dd_tracing import tracer
 from litellm.llms.custom_httpx.http_handler import (
-    MaskedHTTPStatusError,
     _get_httpx_client,
     get_async_httpx_client,
     httpxSpecialProvider,
@@ -413,53 +413,15 @@ async def _send_with_413_split(self, batch: list) -> list:
         that could not be delivered because of a non-413 (transient) error, so the caller
         re-queues only those and never the events already accepted by Datadog.
         """
-        pending: Final[list[list]] = [batch]
-        while pending:
-            chunk = pending.pop()
-            if not chunk:
-                continue
-            if len(chunk) > 1 and self._exceeds_intake_limits(chunk):
-                mid = len(chunk) // 2
-                pending.append(chunk[mid:])
-                pending.append(chunk[:mid])
-                continue
-            try:
-                response = await self.async_send_compressed_data(chunk)
-            except Exception as e:
-                if isinstance(e, MaskedHTTPStatusError) and e.status_code == 413:
-                    response = e.response
-                else:
-                    verbose_logger.exception("Datadog Error sending batch API - %s", e)
-                    return self._undelivered(chunk, pending)
-
-            if response.status_code == 413:
-                if len(chunk) == 1:
-                    verbose_logger.error(DD_ERRORS.DATADOG_413_ERROR.value)
-                    continue
-                mid = len(chunk) // 2
-                pending.append(chunk[mid:])
-                pending.append(chunk[:mid])
-                continue
-
-            if response.status_code != 202:
-                verbose_logger.error(
-                    "Datadog: unexpected response status_code=%s, text=%s",
-                    response.status_code,
-                    response.text,
-                )
-                return self._undelivered(chunk, pending)
-
-            verbose_logger.debug(
-                "Datadog: delivered %s events, status_code=%s, text=%s",
-                len(chunk),
-                response.status_code,
-                response.text,
-            )
-        return []
-
-    @staticmethod
-    def _undelivered(chunk: list, pending: list[list]) -> list:
-        return chunk + [event for remaining in reversed(pending) for event in remaining]
+        undelivered: Final = await send_batch_with_413_split(
+            batch=batch,
+            send_batch=self.async_send_compressed_data,
+            exceeds_limits=self._exceeds_intake_limits,
+            success_status_codes=frozenset({202}),
+            integration_name="Datadog",
+            drop_error_message=DD_ERRORS.DATADOG_413_ERROR.value,
+        )
+        return list(undelivered)
 
     @staticmethod
     def _exceeds_intake_limits(chunk: Sequence[DatadogPayload]) -> bool:
@@ -606,7 +568,7 @@ def create_datadog_logging_payload(
         )
         return dd_payload
 
-    async def async_send_compressed_data(self, data: list) -> Response:
+    async def async_send_compressed_data(self, data: Sequence[DatadogPayload]) -> Response:
         """
         Async helper to send compressed data to datadog self.intake_url
 

litellm/types/integrations/azure_sentinel.py

@@ -1,5 +1,9 @@
+from typing import Final
+
 from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams
 
+AZURE_SENTINEL_MAX_PAYLOAD_SIZE_BYTES: Final = 1_000_000
+
 
 class AzureSentinelInitParams(StandardCustomLoggerInitParams):
     """