diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 1078f05165a..8f13b5602b0 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -9,6 +9,7 @@ API Reference: https://docs.datadoghq.com/llm_observability/setup/api/?tab=examp import asyncio import json import os +import time from litellm._uuid import uuid from datetime import datetime from typing import Any, Dict, List, Literal, Optional, Union @@ -35,6 +36,10 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.types.integrations.datadog import ( + DD_MAX_BATCH_SIZE, + DD_MAX_PAYLOAD_SIZE_BYTES, +) from litellm.types.integrations.datadog_llm_obs import * from litellm.types.utils import ( CallTypes, @@ -160,62 +165,158 @@ class DataDogLLMObsLogger(CustomBatchLogger): verbose_logger.exception(f"DataDogLLMObs: Error logging failure event - {str(e)}") async def async_send_batch(self): - try: - if not self.log_queue: - return + """Send queued LLM Obs spans to Datadog, splitting oversized batches. - verbose_logger.debug(f"DataDogLLMObs: Flushing {len(self.log_queue)} events") + Mirrors the proactive-split + 413-retry strategy from the log-intake + logger (``DataDogLogger._send_with_413_split``) so every POST stays + within the same conservative bounds (1000 events / 4 MB serialised). + """ + if not self.log_queue: + return + + batch_to_send = self.log_queue[:] + self.log_queue = [] + + try: + verbose_logger.debug(f"DataDogLLMObs: Flushing {len(batch_to_send)} events") if self.is_mock_mode: verbose_logger.debug("[DATADOG MOCK] Mock mode enabled - API calls will be intercepted") - # Prepare the payload - payload = { - "data": DDIntakePayload( - type="span", - attributes=DDSpanAttributes( - ml_app=get_datadog_service(), - tags=get_datadog_tags(), - spans=self.log_queue, - ), - ), - } - - # serialize datetime objects - for budget reset time in spend metrics - from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - - try: - verbose_logger.debug("payload %s", safe_dumps(payload)) - except Exception as debug_error: - verbose_logger.debug("payload serialization failed: %s", str(debug_error)) - - json_payload = safe_dumps(payload) - - headers = {"Content-Type": "application/json"} - if self.DD_API_KEY: - headers["DD-API-KEY"] = self.DD_API_KEY - - response = await self.async_client.post( - url=self.intake_url, - content=json_payload, - headers=headers, - ) - - if response.status_code != 202: - raise Exception( - f"DataDogLLMObs: Unexpected response - status_code: {response.status_code}, text: {response.text}" - ) + undelivered = await self._send_with_413_split(batch_to_send) + if undelivered: + self.log_queue = undelivered + self.log_queue if self.is_mock_mode: - verbose_logger.debug(f"[DATADOG MOCK] Batch of {len(self.log_queue)} events successfully mocked") - else: - verbose_logger.debug(f"DataDogLLMObs: Successfully sent batch - status_code: {response.status_code}") - self.log_queue.clear() - except httpx.HTTPStatusError as e: - verbose_logger.exception(f"DataDogLLMObs: Error sending batch - {e.response.text}") + verbose_logger.debug(f"[DATADOG MOCK] Batch of {len(batch_to_send)} events successfully mocked") + except Exception as e: + self.log_queue = batch_to_send + self.log_queue verbose_logger.exception(f"DataDogLLMObs: Error sending batch - {str(e)}") + async def flush_queue(self): + """Flush under the lock without clearing re-queued spans. + + The base ``CustomBatchLogger.flush_queue`` clears ``log_queue`` after + ``async_send_batch`` returns, which would erase the undelivered spans + ``async_send_batch`` just re-queued for retry. Mirrors the same + override on the log-intake ``DataDogLogger``. + """ + if self.flush_lock is None: + return + + async with self.flush_lock: + if self.log_queue: + verbose_logger.debug("DataDogLLMObs: Flushing batch of %s events", len(self.log_queue)) + await self.async_send_batch() + if not self.log_queue: + self.last_flush_time = time.time() + + async def _send_with_413_split(self, batch: List[LLMObsPayload]) -> List[LLMObsPayload]: + """Send *batch*, halving any sub-batch that exceeds intake limits. + + Proactively splits before serializing when the chunk is too large, + and retries with smaller halves on a 413 response. A lone span + that still 413s is dropped to prevent wedging the queue. + + Returns spans that could not be delivered due to a transient + (non-413) error so the caller can re-queue them. + """ + pending: List[List[LLMObsPayload]] = [batch] + while pending: + chunk = pending.pop() + if not chunk: + continue + try: + if len(chunk) > 1 and self._exceeds_intake_limits(chunk): + mid = len(chunk) // 2 + pending.append(chunk[mid:]) + pending.append(chunk[:mid]) + continue + response = await self._post_spans(chunk) + except httpx.HTTPStatusError as e: + if e.response.status_code == 413: + response = e.response + else: + verbose_logger.exception("DataDogLLMObs: Error sending batch - %s", str(e)) + return self._undelivered(chunk, pending) + except Exception as e: + verbose_logger.exception("DataDogLLMObs: Error sending batch - %s", str(e)) + return self._undelivered(chunk, pending) + + if response.status_code == 413: + if len(chunk) == 1: + verbose_logger.error("DataDogLLMObs: single span exceeds intake limit, dropping") + continue + mid = len(chunk) // 2 + pending.append(chunk[mid:]) + pending.append(chunk[:mid]) + continue + + if response.status_code != 202: + verbose_logger.error( + "DataDogLLMObs: unexpected response status_code=%s, text=%s", + response.status_code, + response.text, + ) + return self._undelivered(chunk, pending) + + verbose_logger.debug( + "DataDogLLMObs: delivered %s spans, status_code=%s", + len(chunk), + response.status_code, + ) + return [] + + async def _post_spans(self, spans: List[LLMObsPayload]) -> httpx.Response: + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + + payload = { + "data": DDIntakePayload( + type="span", + attributes=DDSpanAttributes( + ml_app=get_datadog_service(), + tags=get_datadog_tags(), + spans=spans, + ), + ), + } + + json_payload = safe_dumps(payload) + + headers: Dict[str, str] = {"Content-Type": "application/json"} + if self.DD_API_KEY: + headers["DD-API-KEY"] = self.DD_API_KEY + + response = await self.async_client.post( + url=self.intake_url, + content=json_payload, + headers=headers, + ) + response.raise_for_status() + return response + + @staticmethod + def _exceeds_intake_limits(chunk: List[LLMObsPayload]) -> bool: + """True when *chunk* would breach Datadog's span intake limits. + + Uses the same conservative bounds as the log intake logger: 1000 + events max and 4 MB serialised payload. The spans intake accepts + oversized payloads with a 202 and enforces limits asynchronously, + so bounded chunks are the reliable protection; the 413 path covers + DD Agent mode and intermediaries that enforce body limits. + """ + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + + if len(chunk) > DD_MAX_BATCH_SIZE: + return True + payload_size_bytes = len(safe_dumps(chunk).encode("utf-8")) + return payload_size_bytes > DD_MAX_PAYLOAD_SIZE_BYTES + + @staticmethod + def _undelivered(chunk: List[LLMObsPayload], pending: List[List[LLMObsPayload]]) -> List[LLMObsPayload]: + return chunk + [span for remaining in reversed(pending) for span in remaining] + def create_llm_obs_payload(self, kwargs: Dict, start_time: datetime, end_time: datetime) -> LLMObsPayload: standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") if standard_logging_payload is None: diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py b/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py index 1cc3591392b..2fd25efa447 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py @@ -5,6 +5,7 @@ from datetime import datetime, timedelta, timezone from typing import Optional from unittest.mock import MagicMock, Mock, patch +import httpx import pytest # Adds the grandparent directory to sys.path to allow importing project modules @@ -1193,3 +1194,236 @@ async def test_spend_metrics_in_datadog_payload(mock_env_vars): now = datetime.now(timezone.utc) time_diff = (budget_reset_dt - now).total_seconds() / 86400 # days assert 9.5 <= time_diff <= 10.5 # Should be close to 10 days + + +# ------------------------------------------------------------------- +# LIT-4358: async_send_batch splits oversized payloads +# ------------------------------------------------------------------- + +class TestDataDogLLMObsBatchSplit: + """Verify that ``async_send_batch`` proactively splits batches that + exceed the intake byte / event-count limits, and handles 413 retries.""" + + @staticmethod + def _make_logger() -> DataDogLLMObsLogger: + """Create a logger instance in mock mode so no real HTTP is used.""" + with patch.dict( + os.environ, + {"DD_API_KEY": "fake", "DD_SITE": "us5.datadoghq.com"}, + ): + return DataDogLLMObsLogger() + + @staticmethod + def _make_span(**overrides) -> dict: + """Return a minimal LLMObsPayload-like dict for queue testing.""" + base = { + "trace_id": "t1", + "span_id": "s1", + "parent_id": "undefined", + "name": "test", + "start_ns": 0, + "duration": 1_000_000, + "status": "ok", + "meta": {"kind": "llm", "input": {}, "output": {}}, + "metrics": {}, + "tags": [], + } + base.update(overrides) + return base + + @pytest.mark.asyncio + async def test_proactive_split_on_oversized_batch(self): + """When the batch exceeds DD_MAX_BATCH_SIZE, it should be split + and posted in multiple smaller requests.""" + logger = self._make_logger() + + from litellm.types.integrations.datadog import DD_MAX_BATCH_SIZE + + # Queue more than DD_MAX_BATCH_SIZE spans + num_spans = DD_MAX_BATCH_SIZE + 10 + logger.log_queue = [self._make_span(span_id=str(i)) for i in range(num_spans)] + + post_call_count = 0 + posted_span_counts: list = [] + + async def _mock_post_spans(spans): + nonlocal post_call_count + post_call_count += 1 + posted_span_counts.append(len(spans)) + resp = MagicMock() + resp.status_code = 202 + resp.text = "OK" + return resp + + with patch.object(logger, "_post_spans", side_effect=_mock_post_spans): + await logger.async_send_batch() + + # Should have split into >1 POST + assert post_call_count >= 2, f"Expected multiple POSTs, got {post_call_count}" + assert all(count <= DD_MAX_BATCH_SIZE for count in posted_span_counts) + # Total delivered should equal the original batch + assert sum(posted_span_counts) == num_spans + # Queue should be empty after successful delivery + assert len(logger.log_queue) == 0 + + @pytest.mark.asyncio + async def test_413_retries_with_smaller_chunks(self): + """On a 413 response, the chunk should be halved and retried.""" + logger = self._make_logger() + + logger.log_queue = [self._make_span(span_id=str(i)) for i in range(4)] + + call_sizes: list = [] + + async def _mock_post_spans(spans): + call_sizes.append(len(spans)) + resp = MagicMock() + if len(spans) > 2: + # Simulate 413 for large batches + resp.status_code = 413 + resp.text = "Payload Too Large" + err = httpx.HTTPStatusError( + "413", request=MagicMock(), response=resp + ) + raise err + resp.status_code = 202 + resp.text = "OK" + return resp + + with patch.object(logger, "_post_spans", side_effect=_mock_post_spans): + await logger.async_send_batch() + + # First attempt with 4 -> 413 -> split to 2+2 -> both succeed + assert 4 in call_sizes, f"Expected initial attempt of 4, got {call_sizes}" + assert call_sizes.count(2) >= 2, f"Expected two chunks of 2, got {call_sizes}" + assert len(logger.log_queue) == 0 + + @pytest.mark.asyncio + async def test_single_oversized_span_dropped(self): + """A single span that 413s should be dropped, not wedge the queue.""" + logger = self._make_logger() + logger.log_queue = [self._make_span(span_id="big")] + + async def _mock_post_spans(spans): + resp = MagicMock() + resp.status_code = 413 + resp.text = "Payload Too Large" + err = httpx.HTTPStatusError( + "413", request=MagicMock(), response=resp + ) + raise err + + with patch.object(logger, "_post_spans", side_effect=_mock_post_spans): + await logger.async_send_batch() + + assert len(logger.log_queue) == 0 + + @pytest.mark.asyncio + async def test_transient_error_requeues(self): + """A non-413 error should re-queue undelivered spans.""" + logger = self._make_logger() + logger.log_queue = [self._make_span(span_id=str(i)) for i in range(3)] + + async def _mock_post_spans(spans): + raise Exception("Connection refused") + + with patch.object(logger, "_post_spans", side_effect=_mock_post_spans): + await logger.async_send_batch() + + # All 3 spans should be back in the queue for retry + assert len(logger.log_queue) == 3 + + @pytest.mark.asyncio + async def test_periodic_flush_preserves_requeued_spans(self): + """The flush_queue path must not clear spans that async_send_batch + re-queued after a transient error.""" + logger = self._make_logger() + logger.log_queue = [self._make_span(span_id=str(i)) for i in range(3)] + + async def _mock_post_spans(spans): + raise Exception("Connection refused") + + with patch.object(logger, "_post_spans", side_effect=_mock_post_spans): + await logger.flush_queue() + + assert len(logger.log_queue) == 3 + + @pytest.mark.asyncio + async def test_proactive_split_on_oversized_bytes(self): + """A batch over DD_MAX_PAYLOAD_SIZE_BYTES is split before any POST, + so every posted chunk serializes under the byte limit.""" + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + from litellm.types.integrations.datadog import DD_MAX_PAYLOAD_SIZE_BYTES + + logger = self._make_logger() + logger.log_queue = [ + self._make_span( + span_id=str(i), + meta={"kind": "llm", "input": "x" * 3_000_000, "output": {}}, + ) + for i in range(3) + ] + + posted_chunks: list = [] + + async def _mock_post_spans(spans): + posted_chunks.append(list(spans)) + resp = MagicMock() + resp.status_code = 202 + resp.text = "OK" + return resp + + with patch.object(logger, "_post_spans", side_effect=_mock_post_spans): + await logger.async_send_batch() + + assert len(posted_chunks) == 3 + assert all( + len(safe_dumps(chunk).encode("utf-8")) <= DD_MAX_PAYLOAD_SIZE_BYTES + for chunk in posted_chunks + ) + assert [s["span_id"] for chunk in posted_chunks for s in chunk] == ["0", "1", "2"] + assert len(logger.log_queue) == 0 + + @pytest.mark.asyncio + async def test_http_500_requeues_all_spans(self): + """A non-413 HTTP error is transient: every undelivered span is + re-queued in order for the next flush.""" + logger = self._make_logger() + logger.log_queue = [self._make_span(span_id=str(i)) for i in range(4)] + + async def _mock_post_spans(spans): + resp = MagicMock() + resp.status_code = 500 + resp.text = "Internal Server Error" + raise httpx.HTTPStatusError("500", request=MagicMock(), response=resp) + + with patch.object(logger, "_post_spans", side_effect=_mock_post_spans): + await logger.async_send_batch() + + assert [s["span_id"] for s in logger.log_queue] == ["0", "1", "2", "3"] + + @pytest.mark.asyncio + async def test_size_check_failure_requeues_only_undelivered(self): + """If the size check itself raises mid-loop, chunks already delivered + must not be re-queued as duplicates.""" + logger = self._make_logger() + logger.log_queue = [self._make_span(span_id=str(i)) for i in range(4)] + + delivered: list = [] + + async def _mock_post_spans(spans): + delivered.extend(s["span_id"] for s in spans) + resp = MagicMock() + resp.status_code = 202 + resp.text = "OK" + return resp + + with patch.object( + logger, + "_exceeds_intake_limits", + side_effect=[True, False, Exception("serialization failed")], + ), patch.object(logger, "_post_spans", side_effect=_mock_post_spans): + await logger.async_send_batch() + + assert delivered == ["0", "1"] + assert [s["span_id"] for s in logger.log_queue] == ["2", "3"]