mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(datadog): split LLM Obs span batches proactively under intake limits
Closes LIT-4358. DataDogLLMObsLogger.async_send_batch posted the entire log_queue to /api/intake/llm-obs/v1/trace/spans as one payload with no byte or event-count cap and no failure handling. Worse, the base CustomBatchLogger.flush_queue cleared the queue even when the send failed, so any failed flush window silently dropped every span. Port the _send_with_413_split strategy from the log-intake logger: - Proactively split batches exceeding DD_MAX_BATCH_SIZE (1000) events or DD_MAX_PAYLOAD_SIZE_BYTES (4 MB serialised) before POSTing - On a 413 response, halve and retry; drop single un-splittable spans - On transient errors, re-queue undelivered spans for the next flush - Override flush_queue so the base class cannot wipe re-queued spans - Run the proactive size check inside the send loop's try so a serialization failure re-queues only undelivered chunks Adds pytest cases covering proactive count and byte splits, 413 halving, single-span drop, transient and HTTP 500 re-queue, flush_queue preservation, and the no-duplicate guarantee on size check failure.
This commit is contained in:
Yucheng Zhu
parent
2c1d62ce2b
commit
c80ef283a6
2 changed files with 381 additions and 46 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue