diff --git a/litellm/integrations/datadog/datadog_metrics.py b/litellm/integrations/datadog/datadog_metrics.py index 5dda336dc94..6c461eaa075 100644 --- a/litellm/integrations/datadog/datadog_metrics.py +++ b/litellm/integrations/datadog/datadog_metrics.py @@ -2,8 +2,10 @@ import asyncio import gzip import os import time +import zlib +from collections.abc import Sequence from datetime import datetime -from typing import Final +from typing import Final, Literal from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger @@ -21,6 +23,8 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus from litellm.types.integrations.datadog_metrics import ( + DatadogDistributionPayload, + DatadogDistributionSeries, DatadogMetricPoint, DatadogMetricSeries, DatadogMetricsPayload, @@ -29,6 +33,9 @@ from litellm.types.utils import StandardLoggingPayload class DatadogMetricsLogger(CustomBatchLogger): + preserve_events_added_during_flush = True + log_queue: list[DatadogMetricSeries | DatadogDistributionSeries] + def __init__(self, start_periodic_flush: bool = True, **kwargs): self.dd_api_key = os.getenv("DD_API_KEY") self.dd_app_key = os.getenv("DD_APP_KEY") @@ -38,6 +45,7 @@ class DatadogMetricsLogger(CustomBatchLogger): verbose_logger.warning("Datadog Metrics: DD_API_KEY is required. Integration will not work.") self.upload_url = f"https://api.{self.dd_site}/api/v2/series" + self.distribution_upload_url = f"https://api.{self.dd_site}/api/v1/distribution_points" self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) @@ -54,6 +62,10 @@ class DatadogMetricsLogger(CustomBatchLogger): super().__init__(**kwargs) + # Distribution points at the head of the queue that Datadog rejected; they ride the next flush + # but do not count toward `batch_size` + self.unsent_distribution_count = 0 + # Start periodic flush task only if instructed if start_periodic_flush: asyncio.create_task(self.periodic_flush()) @@ -71,10 +83,12 @@ class DatadogMetricsLogger(CustomBatchLogger): f"env:{get_datadog_env()}", f"service:{get_datadog_service()}", f"version:{os.getenv('DD_VERSION', 'unknown')}", - f"HOSTNAME:{get_datadog_hostname()}", f"POD_NAME:{get_datadog_pod_name()}", ] + if hostname := get_datadog_hostname(): + tags.append(f"HOSTNAME:{hostname}") + # Add metric-specific tags if provider := log.get("custom_llm_provider"): tags.append(f"provider:{provider}") @@ -102,6 +116,27 @@ class DatadogMetricsLogger(CustomBatchLogger): return tags + def _add_latency_metric( + self, gauge_metric: str, distribution_metric: str, seconds: float, timestamp: int, tags: Sequence[str] + ) -> None: + """ + Queues a latency sample as a gauge under its `litellm.*` name and as a distribution + under its `litellm_proxy.*` name, which Datadog can compute percentiles over + """ + gauge: Final[DatadogMetricSeries] = { + "metric": gauge_metric, + "type": 3, # gauge + "points": [{"timestamp": timestamp, "value": seconds}], + "tags": tags, + } + distribution: Final[DatadogDistributionSeries] = { + "metric": distribution_metric, + "type": "distribution", + "points": ((timestamp, (seconds,)),), + "tags": tuple(tags), + } + self.log_queue.extend((gauge, distribution)) + def _add_metrics_from_log( self, log: StandardLoggingPayload, @@ -121,43 +156,30 @@ class DatadogMetricsLogger(CustomBatchLogger): start_time_dt: Final = kwargs.get("start_time") if start_time_dt and end_time_dt: total_duration: Final = (end_time_dt - start_time_dt).total_seconds() - series_total_latency: Final[DatadogMetricSeries] = { - "metric": "litellm.request.total_latency", - "type": 3, # gauge - "points": [{"timestamp": timestamp, "value": total_duration}], - "tags": tags, - } - self.log_queue.append(series_total_latency) + self._add_latency_metric( + "litellm.request.total_latency", "litellm_proxy.request.total_latency", total_duration, timestamp, tags + ) # 2. LLM API Latency Metric (Provider alone) api_call_start_time: Final = kwargs.get("api_call_start_time") if api_call_start_time and end_time_dt: llm_api_duration: Final = (end_time_dt - api_call_start_time).total_seconds() - series_llm_latency: Final[DatadogMetricSeries] = { - "metric": "litellm.llm_api.latency", - "type": 3, # gauge - "points": [{"timestamp": timestamp, "value": llm_api_duration}], - "tags": tags, - } - self.log_queue.append(series_llm_latency) + self._add_latency_metric( + "litellm.llm_api.latency", "litellm_proxy.llm_api.latency", llm_api_duration, timestamp, tags + ) # 3. LiteLLM Overhead Latency Metric (total - llm_api time) hidden_params: Final = log.get("hidden_params", {}) or {} litellm_overhead_time_ms: Final = hidden_params.get("litellm_overhead_time_ms") if litellm_overhead_time_ms is not None: overhead_tags: Final = self._extract_tags(log) # no status_code on latency metric - series_overhead: Final[DatadogMetricSeries] = { - "metric": "litellm.overhead.latency", - "type": 3, # gauge - "points": [ - { - "timestamp": timestamp, - "value": litellm_overhead_time_ms / 1000, # convert ms → seconds - } - ], - "tags": overhead_tags, - } - self.log_queue.append(series_overhead) + self._add_latency_metric( + "litellm.overhead.latency", + "litellm_proxy.overhead.latency", + litellm_overhead_time_ms / 1000, + timestamp, + overhead_tags, + ) # 4. Request Count / Status Code series_count: Final[DatadogMetricSeries] = { @@ -178,7 +200,7 @@ class DatadogMetricsLogger(CustomBatchLogger): self._add_metrics_from_log(log=standard_logging_object, kwargs=kwargs, status_code="200") - if len(self.log_queue) >= self.batch_size: + if len(self.log_queue) - self.unsent_distribution_count >= self.batch_size: await self.flush_queue() except Exception as e: @@ -200,7 +222,7 @@ class DatadogMetricsLogger(CustomBatchLogger): self._add_metrics_from_log(log=standard_logging_object, kwargs=kwargs, status_code=status_code) - if len(self.log_queue) >= self.batch_size: + if len(self.log_queue) - self.unsent_distribution_count >= self.batch_size: await self.flush_queue() except Exception as e: @@ -210,41 +232,69 @@ class DatadogMetricsLogger(CustomBatchLogger): if not self.log_queue: return - batch: Final = self.log_queue.copy() - payload_data: Final[DatadogMetricsPayload] = {"series": batch} + batch: Final[tuple[DatadogMetricSeries | DatadogDistributionSeries, ...]] = tuple(self.log_queue) + series: Final[tuple[DatadogMetricSeries, ...]] = tuple(s for s in batch if s["type"] != "distribution") + distributions: Final[tuple[DatadogDistributionSeries, ...]] = tuple( + s for s in batch if s["type"] == "distribution" + ) try: - await self._upload_to_datadog(payload_data) + if series: + await self._upload_to_datadog({"series": series}) except Exception as e: verbose_logger.exception("Datadog Metrics: Error in async_send_batch: %s", e) raise + try: + if distributions: + await self._upload_distributions_to_datadog({"series": distributions}) + except (Exception, asyncio.CancelledError) as e: + unsent: Final = distributions[-self.batch_size :] + self.log_queue[: len(batch)] = unsent + self.unsent_distribution_count = len(unsent) + if dropped := len(distributions) - len(unsent): + verbose_logger.warning( + "Datadog Metrics: keeping the newest %s unsent distribution points; dropped %s older ones.", + len(unsent), + dropped, + ) + verbose_logger.exception("Datadog Metrics: Error in async_send_batch: %s", e) + raise + + self.unsent_distribution_count = 0 + async def _upload_to_datadog(self, payload: DatadogMetricsPayload): + await self._post_compressed(self.upload_url, safe_dumps(payload), len(payload["series"]), encoding="gzip") + + async def _upload_distributions_to_datadog(self, payload: DatadogDistributionPayload) -> None: + await self._post_compressed( + self.distribution_upload_url, safe_dumps(payload), len(payload["series"]), encoding="deflate" + ) + + async def _post_compressed( + self, url: str, json_data: str, point_count: int, encoding: Literal["gzip", "deflate"] + ) -> None: if not self.dd_api_key: return headers: Final = { "Content-Type": "application/json", + "Content-Encoding": encoding, "DD-API-KEY": self.dd_api_key, } if self.dd_app_key: headers["DD-APPLICATION-KEY"] = self.dd_app_key - json_data: Final = safe_dumps(payload) - compressed_data: Final = gzip.compress(json_data.encode("utf-8")) - headers["Content-Encoding"] = "gzip" + raw: Final = json_data.encode("utf-8") + compressed_data: Final = gzip.compress(raw) if encoding == "gzip" else zlib.compress(raw) - response: Final = await self.async_client.post( - self.upload_url, - content=compressed_data, - headers=headers, - ) + response: Final = await self.async_client.post(url, content=compressed_data, headers=headers) response.raise_for_status() verbose_logger.debug( - "Datadog Metrics: Uploaded %s metric points. Status: %s", len(payload["series"]), response.status_code + "Datadog Metrics: Uploaded %s metric points. Status: %s", point_count, response.status_code ) async def async_health_check(self) -> IntegrationHealthCheckStatus: diff --git a/litellm/types/integrations/datadog_metrics.py b/litellm/types/integrations/datadog_metrics.py index c294fdd2522..3e5ddf97cf8 100644 --- a/litellm/types/integrations/datadog_metrics.py +++ b/litellm/types/integrations/datadog_metrics.py @@ -1,4 +1,7 @@ -from typing_extensions import TypedDict +from collections.abc import Sequence +from typing import Literal, TypeAlias + +from typing_extensions import NotRequired, ReadOnly, TypedDict class DatadogMetricPoint(TypedDict): @@ -6,13 +9,27 @@ class DatadogMetricPoint(TypedDict): value: float # The metric value -class DatadogMetricSeries(TypedDict, total=False): - metric: str - type: int # 0=unspecified, 1=count, 2=rate, 3=gauge - points: list[DatadogMetricPoint] - tags: list[str] - interval: int | None # Required for count (type=1) and rate (type=2) metrics +class DatadogMetricSeries(TypedDict): + metric: ReadOnly[str] + type: ReadOnly[Literal[0, 1, 2, 3]] # 0=unspecified, 1=count, 2=rate, 3=gauge + points: ReadOnly[Sequence[DatadogMetricPoint]] + tags: ReadOnly[Sequence[str]] + interval: ReadOnly[NotRequired[int]] # Required for count (type=1) and rate (type=2) metrics class DatadogMetricsPayload(TypedDict): - series: list[DatadogMetricSeries] + series: ReadOnly[Sequence[DatadogMetricSeries]] + + +DatadogDistributionPoint: TypeAlias = tuple[int, tuple[float, ...]] + + +class DatadogDistributionSeries(TypedDict): + metric: ReadOnly[str] + type: ReadOnly[Literal["distribution"]] + points: ReadOnly[tuple[DatadogDistributionPoint, ...]] + tags: ReadOnly[tuple[str, ...]] + + +class DatadogDistributionPayload(TypedDict): + series: ReadOnly[tuple[DatadogDistributionSeries, ...]] diff --git a/tests/integration/observability/test_datadog_metrics_distributions.py b/tests/integration/observability/test_datadog_metrics_distributions.py new file mode 100644 index 00000000000..864ff715367 --- /dev/null +++ b/tests/integration/observability/test_datadog_metrics_distributions.py @@ -0,0 +1,616 @@ +"""Datadog custom metrics: every request's latency reaches Datadog as a distribution sample. + +Datadog keeps one gauge value per second per tag set, so concurrent requests with the same tags collapse +into one point and `p95:litellm.request.total_latency` is not a per-request percentile. The +``datadog_metrics`` callback therefore also sends each latency sample to ``/api/v1/distribution_points`` +under a ``litellm_proxy.*`` name, since Datadog refuses to enable percentiles on its integration's ``litellm.*``. +The owned proxy points ``DD_SITE`` at a TLS double for ``api.localhost`` that inflates every upload, answers +like the real intake (202, or 503 while an outage is scripted) and records what each flush carried, so the +cases count samples per request instead of reading Datadog's aggregates. +""" + +from __future__ import annotations + +import asyncio +import gzip +import json +import threading +import time +import uuid +import zlib +from collections.abc import Callable, Iterator, Mapping +from contextlib import AbstractContextManager, contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit + +import httpx +import pytest +import yaml +from _s3_v2_support import SURFACES, call_surface, surface_reply +from integration._support.client import Gateway, eventually +from integration._support.process import owned_proxy +from integration._support.tls import server_context, write_self_signed_cert +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +SERIES: Final = "/api/v2/series" +DISTRIBUTIONS: Final = "/api/v1/distribution_points" +LATENCIES: Final = ("litellm.request.total_latency", "litellm.llm_api.latency", "litellm.overhead.latency") +DISTRIBUTION_OF: Final = { + "litellm.request.total_latency": "litellm_proxy.request.total_latency", + "litellm.llm_api.latency": "litellm_proxy.llm_api.latency", + "litellm.overhead.latency": "litellm_proxy.overhead.latency", +} +GAUGE_OF: Final = {distribution: gauge for gauge, distribution in DISTRIBUTION_OF.items()} +REQUEST_COUNT_METRIC: Final = "litellm.llm_api.request_count" +SITE_HOSTS: Final = ("api.localhost",) +DD_API_KEY: Final = "synthetic-datadog-api-key" +WORKERS: Final = 2 +FLUSH_WAIT_SECONDS: Final = 60 +LOGGER_FLUSH_INTERVAL_SECONDS: Final = 5 +SERIES_LIST: Final = TypeAdapter(list[dict[str, JsonValue]]) + + +@dataclass(frozen=True, slots=True) +class Upload: + target: str + encoding: str | None + api_key: str | None + status: int + series: tuple[dict[str, JsonValue], ...] + received: float + + +class DatadogIntake: + """Answers `/api/v2/series` and `/api/v1/distribution_points` like Datadog and keeps every inflated body. + + `failing` scripts an outage as the number of uploads to answer 503 per path; `delay` holds every answer + so a flush stays in flight while the proxy keeps serving requests.""" + + def __init__(self, failing: Mapping[str, int] | None = None, delay: float = 0) -> None: + self._failing: Final = dict(failing or {}) + self.delay: Final = delay + self._lock: Final = threading.Lock() + self._uploads: Final[list[Upload]] = [] + + def respond(self, request: Request) -> Reply: + assert request.method == "POST" and request.target in (SERIES, DISTRIBUTIONS), request.target + encoding: Final = request.headers.get("content-encoding") + body: Final = ( + gzip.decompress(request.body) + if encoding == "gzip" + else zlib.decompress(request.body) + if encoding == "deflate" + else request.body + ) + payload: Final = json.loads(body) + with self._lock: + remaining: Final = self._failing.get(request.target, 0) + if remaining: + self._failing[request.target] = remaining - 1 + status: Final = 503 if remaining else 202 + self._uploads.append( + Upload( + request.target, + encoding, + request.headers.get("dd-api-key"), + status, + tuple(SERIES_LIST.validate_python(payload["series"])), + time.monotonic(), + ) + ) + if self.delay: + time.sleep(self.delay) + if status == 503: + return Reply(status=503, body=b'{"errors": ["Service Unavailable"]}') + return Reply(status=202, body=b'{"status": "ok"}' if request.target == DISTRIBUTIONS else b'{"errors": []}') + + def uploads(self, target: str | None = None) -> tuple[Upload, ...]: + with self._lock: + return tuple(upload for upload in self._uploads if target is None or upload.target == target) + + def accepted(self, target: str) -> tuple[Upload, ...]: + return tuple(upload for upload in self.uploads(target) if upload.status == 202) + + def rejected(self, target: str) -> tuple[Upload, ...]: + return tuple(upload for upload in self.uploads(target) if upload.status != 202) + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + intake: DatadogIntake + provider_url: str + + +def _tags(series: Mapping[str, JsonValue]) -> tuple[str, ...]: + tags: Final = series["tags"] + assert isinstance(tags, list), series + return tuple(str(tag) for tag in tags) + + +def _for_model(series: Mapping[str, JsonValue], metric: str, model: str) -> bool: + return series["metric"] == metric and f"model_group:{model}" in _tags(series) + + +def distribution_samples( + intake: DatadogIntake, metric: str, model: str, status_code: str | None = None +) -> tuple[float, ...]: + """Every sample value Datadog accepted for the distribution of gauge `metric` from requests to `model`.""" + samples: Final[list[float]] = [] + for upload in intake.accepted(DISTRIBUTIONS): + for series in upload.series: + if not _for_model(series, DISTRIBUTION_OF[metric], model): + continue + if status_code is not None and f"status_code:{status_code}" not in _tags(series): + continue + assert series["type"] == "distribution", series + points: Final = series["points"] + assert isinstance(points, list), series + for point in points: + assert isinstance(point, list) and len(point) == 2 and isinstance(point[1], list), series + samples.extend(float(str(value)) for value in point[1]) + return tuple(samples) + + +def series_points(intake: DatadogIntake, metric: str, model: str, status_code: str | None = None) -> tuple[float, ...]: + """Every point value Datadog accepted on `/api/v2/series` for `metric` from requests to `model`.""" + values: Final[list[float]] = [] + for upload in intake.accepted(SERIES): + for series in upload.series: + if not _for_model(series, metric, model): + continue + if status_code is not None and f"status_code:{status_code}" not in _tags(series): + continue + points: Final = series["points"] + assert isinstance(points, list), series + for point in points: + assert isinstance(point, dict), series + values.append(float(str(point["value"]))) + return tuple(values) + + +def every_series(intake: DatadogIntake) -> tuple[dict[str, JsonValue], ...]: + return tuple(series for upload in intake.uploads() for series in upload.series) + + +def _config(path: Path, callback_setting: str | None) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + if callback_setting is not None: + config["litellm_settings"].update({callback_setting: ["datadog_metrics"]}) + target: Final = path / "datadog_metrics.yaml" + target.write_text(yaml.safe_dump(config)) + return target + + +def upstream_failing_marked_requests(request: Request) -> Reply: + body: Final = json.loads(request.body) if request.body else {} + messages: Final = body.get("messages") if isinstance(body, dict) else None + content: Final = messages[0].get("content") if isinstance(messages, list) and messages else None + if isinstance(content, str) and content.startswith("fail-"): + return Reply(status=500, body=json.dumps({"error": {"message": f"upstream rejected {content}"}}).encode()) + return surface_reply(request) + + +RigFactory = Callable[..., AbstractContextManager[Rig]] + + +@pytest.fixture +def rig_factory(gateway: Gateway, tmp_path: Path) -> RigFactory: + """Boots one owned proxy wired to a fresh Datadog intake double per case.""" + + @contextmanager + def build( + intake: DatadogIntake, + *, + hostname: str | None = None, + callback_setting: str | None = "callbacks", + upstream: Callable[[Request], Reply] = surface_reply, + batch_size: int | None = None, + ) -> Iterator[Rig]: + cert, key = write_self_signed_cert(tmp_path, SITE_HOSTS) + with ( + wire_server(upstream) as provider, + wire_server(intake.respond, tls=server_context(cert, key), keep_alive=True) as site, + ): + environment: Final = { + "DD_API_KEY": DD_API_KEY, + "DD_SITE": f"localhost:{urlsplit(site.url).port}", + "DD_ENV": "integration", + "DD_SERVICE": "litellm-integration", + "SSL_CERT_FILE": str(cert), + **({"HOSTNAME": hostname} if hostname is not None else {}), + **({"DEFAULT_BATCH_SIZE": str(batch_size)} if batch_size is not None else {}), + } + removed: Final = ("DD_APP_KEY", "POD_NAME", *(("HOSTNAME",) if hostname is None else ())) + with owned_proxy( + gateway, + tmp_path, + environment, + config=_config(tmp_path, callback_setting), + remove_environment=removed, + workers=WORKERS, + ) as candidate: + yield Rig(candidate, intake, provider.url) + provider.drain() + + return build + + +def _chat(base_url: str, key: str, model: str, markers: tuple[str, ...]) -> tuple[int, ...]: + async def send_all() -> tuple[int, ...]: + async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client: + responses: Final = await asyncio.gather( + *( + client.post( + "/v1/chat/completions", + json={"model": model, "messages": [{"role": "user", "content": marker}]}, + headers={"Authorization": f"Bearer {key}"}, + ) + for marker in markers + ) + ) + return tuple(response.status_code for response in responses) + + return asyncio.run(send_all()) + + +def _wait_for_samples(intake: DatadogIntake, model: str, count: int, metrics: tuple[str, ...] = LATENCIES) -> None: + eventually( + lambda: tuple(len(distribution_samples(intake, metric, model)) for metric in metrics), + lambda counts: all(total >= count for total in counts), + seconds=FLUSH_WAIT_SECONDS, + return_last_on_timeout=True, + ) + + +def test_concurrent_chat_requests_each_reach_datadog_as_a_distribution_sample(rig_factory: RigFactory) -> None: + """The reporter's flow: a burst of concurrent `/v1/chat/completions` to one model, all 200. Every request + must add one sample to each latency distribution, while the legacy gauges and the request count keep + arriving on `/api/v2/series` unchanged, and an unset HOSTNAME adds no empty `HOSTNAME:` tag.""" + burst: Final = 8 + intake: Final = DatadogIntake() + with rig_factory(intake) as rig: + with rig.proxy.scenario() as scenario: + model: Final = scenario.model(api_base=rig.provider_url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + marker: Final = f"dd-{uuid.uuid4().hex[:8]}" + statuses: Final = _chat( + str(rig.proxy.client.base_url), key, model, tuple(f"{marker}-{index}" for index in range(burst)) + ) + assert statuses == (200,) * burst, statuses + _wait_for_samples(intake, model, burst) + + for metric in LATENCIES: + samples = distribution_samples(intake, metric, model) + assert len(samples) == burst, ( + f"{DISTRIBUTION_OF[metric]} got {len(samples)} samples for {burst} requests: " + f"{[(upload.target, upload.status) for upload in intake.uploads()]}" + ) + assert all(value >= 0 for value in samples), samples + for metric in LATENCIES: + assert len(series_points(intake, metric, model)) == burst, metric + assert sum(series_points(intake, REQUEST_COUNT_METRIC, model, "200")) == burst + + assert {(upload.encoding, upload.api_key) for upload in intake.uploads(DISTRIBUTIONS)} == { + ("deflate", DD_API_KEY) + } + assert {(upload.encoding, upload.api_key) for upload in intake.uploads(SERIES)} == {("gzip", DD_API_KEY)} + hostname_tags: Final = { + tag for series in every_series(intake) for tag in _tags(series) if "HOSTNAME" in tag + } + assert hostname_tags == set(), hostname_tags + + +def test_every_surface_reports_its_latency_as_distribution_samples(rig_factory: RigFactory) -> None: + """Chat, Anthropic messages and Responses, streaming and not, through the sync and async SDKs, each on + its own model: every call adds one total and one provider latency sample, every overhead gauge point + the surface reports gets its distribution twin, each with the same provider, model and status tags as + its gauge, and the overhead distribution carries no status code, like the overhead gauge.""" + rounds: Final = 2 + intake: Final = DatadogIntake() + with rig_factory(intake) as rig: + with rig.proxy.scenario() as scenario: + models: Final = { + surface: ( + scenario.model( + model="anthropic/claude-sonnet-4-5-20250929", + api_base=rig.provider_url, + api_key="synthetic-provider-key", + ) + if surface.startswith("messages") + else scenario.model(api_base=rig.provider_url + "/v1", api_key="synthetic-provider-key") + ) + for surface in SURFACES + } + key: Final = scenario.key(models=list(models.values())) + marker: Final = f"dd-{uuid.uuid4().hex[:8]}" + for index in range(rounds): + for surface, model in models.items(): + call_surface(rig.proxy, surface, model, model, key, f"{marker}-{surface}-{index}") + for model in models.values(): + _wait_for_samples(intake, model, rounds, LATENCIES[:2]) + + for surface, model in models.items(): + for metric in LATENCIES[:2]: + assert len(distribution_samples(intake, metric, model, "200")) == rounds, (surface, metric) + assert len(series_points(intake, metric, model, "200")) == rounds, (surface, metric) + overhead_gauges = series_points(intake, LATENCIES[2], model) + assert len(distribution_samples(intake, LATENCIES[2], model)) == len(overhead_gauges), surface + + for upload in intake.accepted(DISTRIBUTIONS): + for series in upload.series: + tags = _tags(series) + provider = ( + "anthropic" + if any( + f"model_group:{models[surface]}" in tags for surface in SURFACES if "messages" in surface + ) + else "openai" + ) + assert f"provider:{provider}" in tags, series + assert {"env:integration", "service:litellm-integration"} <= set(tags), series + if series["metric"] == "litellm_proxy.overhead.latency": + assert not any(tag.startswith("status_code:") for tag in tags), series + gauge_tags: Final = { + (str(series["metric"]), _tags(series)) + for upload in intake.accepted(SERIES) + for series in upload.series + if series["metric"] in LATENCIES + } + distribution_tags: Final = { + (GAUGE_OF[str(series["metric"])], _tags(series)) + for upload in intake.accepted(DISTRIBUTIONS) + for series in upload.series + } + assert distribution_tags == gauge_tags + + +def test_a_set_hostname_tags_every_series_and_distribution(rig_factory: RigFactory) -> None: + intake: Final = DatadogIntake() + with rig_factory(intake, hostname="pod-lit6804") as rig: + with rig.proxy.scenario() as scenario: + model: Final = scenario.model(api_base=rig.provider_url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + assert _chat(str(rig.proxy.client.base_url), key, model, (f"dd-{uuid.uuid4().hex[:8]}",)) == (200,) + _wait_for_samples(intake, model, 1) + + for metric in LATENCIES: + assert len(distribution_samples(intake, metric, model)) == 1, metric + uploaded: Final = tuple( + series for series in every_series(intake) if f"model_group:{model}" in _tags(series) + ) + assert {str(series["metric"]) for series in uploaded} == { + *LATENCIES, + *DISTRIBUTION_OF.values(), + REQUEST_COUNT_METRIC, + } + assert {tuple(tag for tag in _tags(series) if tag.startswith("HOSTNAME")) for series in uploaded} == { + ("HOSTNAME:pod-lit6804",) + } + + +def test_an_empty_hostname_adds_no_hostname_tag(rig_factory: RigFactory) -> None: + """`HOSTNAME=""` is treated like an unset HOSTNAME: no series or distribution carries a `HOSTNAME` tag.""" + intake: Final = DatadogIntake() + with rig_factory(intake, hostname="") as rig: + with rig.proxy.scenario() as scenario: + model: Final = scenario.model(api_base=rig.provider_url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + assert _chat(str(rig.proxy.client.base_url), key, model, (f"dd-{uuid.uuid4().hex[:8]}",)) == (200,) + _wait_for_samples(intake, model, 1) + + uploaded: Final = tuple( + series for series in every_series(intake) if f"model_group:{model}" in _tags(series) + ) + assert len(uploaded) == 2 * len(LATENCIES) + 1, uploaded + assert not any(tag.startswith("HOSTNAME") for series in uploaded for tag in _tags(series)), uploaded + + +@pytest.mark.parametrize(("callback_setting", "reports_failures"), (("callbacks", True), ("success_callback", False))) +def test_failed_requests_report_latency_under_their_status_code( + rig_factory: RigFactory, callback_setting: str, reports_failures: bool +) -> None: + """A provider 500 reaches the caller as an error; with the callback on both events its latency lands as + distribution samples tagged `status_code:500` next to the successes, and in success-only mode nothing + of the failures is sent.""" + per_outcome: Final = 3 + intake: Final = DatadogIntake() + with rig_factory(intake, callback_setting=callback_setting, upstream=upstream_failing_marked_requests) as rig: + with rig.proxy.scenario() as scenario: + model: Final = scenario.model(api_base=rig.provider_url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + marker: Final = uuid.uuid4().hex[:8] + statuses: Final = _chat( + str(rig.proxy.client.base_url), + key, + model, + (*(f"ok-{marker}-{index}" for index in range(per_outcome)),) + + tuple(f"fail-{marker}-{index}" for index in range(per_outcome)), + ) + assert sorted(statuses) == [200] * per_outcome + [500] * per_outcome, statuses + _wait_for_samples(intake, model, per_outcome) + + def failed_samples() -> int: + return len(distribution_samples(intake, LATENCIES[0], model, "500")) + + if reports_failures: + eventually( + failed_samples, + lambda total: total >= per_outcome, + seconds=FLUSH_WAIT_SECONDS, + return_last_on_timeout=True, + ) + else: + # The failures finished with the successes, so two more periodic flushes would have carried them + eventually( + failed_samples, + lambda total: total > 0, + seconds=2 * LOGGER_FLUSH_INTERVAL_SECONDS, + return_last_on_timeout=True, + ) + + assert len(distribution_samples(intake, LATENCIES[0], model, "200")) == per_outcome + failed: Final = distribution_samples(intake, LATENCIES[0], model, "500") + assert len(failed) == (per_outcome if reports_failures else 0), failed + assert sum(series_points(intake, REQUEST_COUNT_METRIC, model, "500")) == ( + per_outcome if reports_failures else 0 + ) + assert len(series_points(intake, LATENCIES[0], model, "500")) == len(failed) + + +@pytest.mark.parametrize( + ("outage", "rejected_path"), + (({DISTRIBUTIONS: 2}, DISTRIBUTIONS), ({SERIES: 2}, SERIES)), + ids=("distribution-endpoint-down", "series-endpoint-down"), +) +def test_an_endpoint_outage_delivers_every_sample_exactly_once_after_recovery( + rig_factory: RigFactory, outage: Mapping[str, int], rejected_path: str +) -> None: + """Datadog answers 503 to the first uploads of one endpoint. After it recovers, every request's gauges, + count and distribution samples have been accepted exactly once: the retry neither drops the rejected + half nor resends the half Datadog already took.""" + burst: Final = 6 + intake: Final = DatadogIntake(failing=outage) + with rig_factory(intake) as rig: + with rig.proxy.scenario() as scenario: + model: Final = scenario.model(api_base=rig.provider_url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + marker: Final = f"dd-{uuid.uuid4().hex[:8]}" + statuses: Final = _chat( + str(rig.proxy.client.base_url), key, model, tuple(f"{marker}-{index}" for index in range(burst)) + ) + assert statuses == (200,) * burst, statuses + _wait_for_samples(intake, model, burst) + eventually( + lambda: sum(series_points(intake, REQUEST_COUNT_METRIC, model)), + lambda total: total >= burst, + seconds=FLUSH_WAIT_SECONDS, + return_last_on_timeout=True, + ) + + assert intake.rejected(rejected_path), "the scripted outage never answered an upload" + for metric in LATENCIES: + assert len(distribution_samples(intake, metric, model)) == burst, metric + assert len(series_points(intake, metric, model)) == burst, metric + assert sum(series_points(intake, REQUEST_COUNT_METRIC, model)) == burst + + +def test_a_lasting_distribution_outage_retries_one_bounded_batch_per_flush(rig_factory: RigFactory) -> None: + """Datadog rejects every distribution upload while requests keep finishing. The rejected samples ride the + next flush, but they neither make every request flush nor pile up into one ever-growing retry: each retry + carries at most two batches of samples, and every gauge and count still lands on `/api/v2/series` once.""" + batch_size: Final = 21 # the metrics of three requests + requests: Final = 120 + intake: Final = DatadogIntake(failing={DISTRIBUTIONS: 10**6}) + with rig_factory(intake, batch_size=batch_size) as rig: + with rig.proxy.scenario() as scenario: + model: Final = scenario.model(api_base=rig.provider_url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + marker: Final = f"dd-{uuid.uuid4().hex[:8]}" + statuses: Final = tuple( + status + for index in range(requests) + for status in _chat(str(rig.proxy.client.base_url), key, model, (f"{marker}-{index}",)) + ) + assert statuses == (200,) * requests, statuses + eventually( + lambda: sum(series_points(intake, REQUEST_COUNT_METRIC, model)), + lambda total: total >= requests, + seconds=FLUSH_WAIT_SECONDS, + return_last_on_timeout=True, + ) + + retries: Final = tuple(len(upload.series) for upload in intake.rejected(DISTRIBUTIONS)) + assert retries, "the scripted outage never answered an upload" + assert len(retries) <= requests // 2, retries + assert max(retries) <= 2 * batch_size, retries + for metric in LATENCIES: + assert len(series_points(intake, metric, model)) == requests, metric + assert sum(series_points(intake, REQUEST_COUNT_METRIC, model)) == requests + + +def test_requests_finishing_while_a_slow_flush_is_in_flight_are_not_lost(rig_factory: RigFactory) -> None: + """Datadog takes 1.5 s to answer each upload while requests keep finishing. Samples queued during an + upload must ride the next flush rather than be cleared with the batch that was being sent.""" + requests: Final = 40 + intake: Final = DatadogIntake(delay=1.5) + with rig_factory(intake) as rig: + with rig.proxy.scenario() as scenario: + model: Final = scenario.model(api_base=rig.provider_url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + marker: Final = f"dd-{uuid.uuid4().hex[:8]}" + finished: Final[list[float]] = [] + for index in range(requests): + assert _chat(str(rig.proxy.client.base_url), key, model, (f"{marker}-{index}",)) == (200,) + finished.append(time.monotonic()) + time.sleep(0.3) + _wait_for_samples(intake, model, requests) + eventually( + lambda: len(series_points(intake, LATENCIES[0], model)), + lambda total: total >= requests, + seconds=FLUSH_WAIT_SECONDS, + return_last_on_timeout=True, + ) + + in_flight: Final = tuple( + done + for done in finished + if any(upload.received < done < upload.received + intake.delay for upload in intake.uploads()) + ) + assert in_flight, "no request finished while an upload was in flight" + for metric in LATENCIES: + assert len(distribution_samples(intake, metric, model)) == requests, metric + assert len(series_points(intake, metric, model)) == requests, metric + + +def _mock_health_request_series(intake: DatadogIntake, target: str, metric: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + series + for upload in intake.accepted(target) + for series in upload.series + if series["metric"] == metric and any("litellm-mock-response-model" in tag for tag in _tags(series)) + ) + + +def test_the_health_check_test_request_reaches_datadog_as_a_distribution_sample(rig_factory: RigFactory) -> None: + """With the callback configured, `GET /health/services?service=datadog_metrics` answers with a mock + completion, and that completion's latency reaches Datadog like any request: one gauge point and one + distribution sample with the same tags.""" + intake: Final = DatadogIntake() + with rig_factory(intake) as rig: + response: Final = rig.proxy.request("GET", "/health/services", params={"service": "datadog_metrics"}) + assert response.status_code == 200, response.text + assert response.json()["status"] == "success", response.text + metric: Final = DISTRIBUTION_OF[LATENCIES[0]] + eventually( + lambda: len(_mock_health_request_series(intake, DISTRIBUTIONS, metric)), + lambda total: total >= 1, + seconds=FLUSH_WAIT_SECONDS, + return_last_on_timeout=True, + ) + + distributions: Final = _mock_health_request_series(intake, DISTRIBUTIONS, metric) + gauges: Final = _mock_health_request_series(intake, SERIES, LATENCIES[0]) + assert len(distributions) == 1 and len(gauges) == 1, intake.uploads() + assert _tags(distributions[0]) == _tags(gauges[0]) + + +def test_the_health_check_probe_without_the_callback_still_posts_only_to_the_series_endpoint( + rig_factory: RigFactory, +) -> None: + """Without the callback configured, the same health check sends its own `litellm.health_check` probe, + which still goes only to `/api/v2/series`, so a distribution intake outage does not show up there.""" + intake: Final = DatadogIntake() + with rig_factory(intake, callback_setting=None) as rig: + response: Final = rig.proxy.request("GET", "/health/services", params={"service": "datadog_metrics"}) + assert response.status_code == 200, response.text + assert response.json() == {"status": "healthy", "message": "Datadog Metrics is healthy"}, response.text + probes: Final = tuple( + (upload.target, upload.encoding, str(series["metric"])) + for upload in intake.uploads() + for series in upload.series + ) + assert probes == ((SERIES, "gzip", "litellm.health_check"),), intake.uploads() diff --git a/tests/unit/integrations/datadog/test_datadog_metrics.py b/tests/unit/integrations/datadog/test_datadog_metrics.py index eade92d6672..474bc2a2007 100644 --- a/tests/unit/integrations/datadog/test_datadog_metrics.py +++ b/tests/unit/integrations/datadog/test_datadog_metrics.py @@ -1,4 +1,8 @@ +import asyncio +import gzip +import json import time +import zlib from datetime import datetime, timedelta from unittest.mock import AsyncMock @@ -125,9 +129,9 @@ async def test_add_metrics_from_log(clean_env): logger._add_metrics_from_log(log=payload, kwargs=kwargs, status_code="200") - # Should have 3 series: total_latency, llm_api_latency, request_count + # total_latency and llm_api_latency (each as gauge + distribution) plus request_count # (no overhead metric because payload has no hidden_params litellm_overhead_time_ms) - assert len(logger.log_queue) == 3 + assert len(logger.log_queue) == 5 metrics = {s["metric"]: s for s in logger.log_queue} @@ -148,6 +152,63 @@ async def test_add_metrics_from_log(clean_env): assert "status_code:200" in count["tags"] +@pytest.mark.asyncio +async def test_latency_metrics_also_emitted_as_distributions(clean_env): + """Each latency sample is queued as a `litellm_proxy.*` distribution point so Datadog can compute percentiles.""" + logger = DatadogMetricsLogger(batch_size=100, start_periodic_flush=False) + + now = datetime.now() + payload = StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4o", + hidden_params={"litellm_overhead_time_ms": 250}, + ) + kwargs = { + "start_time": now - timedelta(seconds=2), + "api_call_start_time": now - timedelta(seconds=1), + "end_time": now, + } + + logger._add_metrics_from_log(log=payload, kwargs=kwargs, status_code="200") + + distributions = {s["metric"]: s for s in logger.log_queue if s["type"] == "distribution"} + assert set(distributions) == { + "litellm_proxy.request.total_latency", + "litellm_proxy.llm_api.latency", + "litellm_proxy.overhead.latency", + } + + expected_seconds = { + "litellm_proxy.request.total_latency": 2.0, + "litellm_proxy.llm_api.latency": 1.0, + "litellm_proxy.overhead.latency": 0.25, + } + for metric, seconds in expected_seconds.items(): + ((timestamp, values),) = distributions[metric]["points"] + assert timestamp == int(now.timestamp()) + assert len(values) == 1 + assert abs(values[0] - seconds) < 0.1 + assert "provider:openai" in distributions[metric]["tags"] + assert "model_name:gpt-4o" in distributions[metric]["tags"] + + assert "status_code:200" in distributions["litellm_proxy.request.total_latency"]["tags"] + assert not any(tag.startswith("status_code:") for tag in distributions["litellm_proxy.overhead.latency"]["tags"]) + + +@pytest.mark.asyncio +async def test_extract_tags_omits_hostname_when_unset(clean_env, monkeypatch: pytest.MonkeyPatch): + """An unset HOSTNAME must not produce an empty `HOSTNAME:` tag.""" + monkeypatch.delenv("HOSTNAME", raising=False) + logger = DatadogMetricsLogger(start_periodic_flush=False) + + tags = logger._extract_tags(log=StandardLoggingPayload(model="gpt-4o")) + + assert not any(tag.startswith("HOSTNAME") for tag in tags) + + monkeypatch.setenv("HOSTNAME", "pod-abc") + assert "HOSTNAME:pod-abc" in logger._extract_tags(log=StandardLoggingPayload(model="gpt-4o")) + + @pytest.mark.asyncio async def test_overhead_latency_metric_emitted(clean_env): """Test that litellm.overhead.latency is emitted when hidden_params contains litellm_overhead_time_ms.""" @@ -176,9 +237,9 @@ async def test_overhead_latency_metric_emitted(clean_env): metrics = {s["metric"]: s for s in logger.log_queue} # Overhead metric must be present - assert ( - "litellm.overhead.latency" in metrics - ), f"Expected 'litellm.overhead.latency' in emitted metrics, got: {list(metrics.keys())}" + assert "litellm.overhead.latency" in metrics, ( + f"Expected 'litellm.overhead.latency' in emitted metrics, got: {list(metrics.keys())}" + ) overhead = metrics["litellm.overhead.latency"] assert overhead["type"] == 3 # gauge # 250 ms → 0.25 s @@ -321,9 +382,7 @@ async def test_async_send_batch(clean_env): logger = DatadogMetricsLogger(start_periodic_flush=False) logger.async_client = AsyncMock() mock_request = Request("POST", "https://api.test.datadoghq.com/api/v2/series") - logger.async_client.post.return_value = Response( - 202, json={"status": "ok"}, request=mock_request - ) + logger.async_client.post.return_value = Response(202, json={"status": "ok"}, request=mock_request) # Manually add a metric series to the queue logger.log_queue = [ @@ -351,6 +410,317 @@ async def test_async_send_batch(clean_env): assert payload["series"][0]["metric"] == "litellm.request.total_latency" +@pytest.mark.asyncio +async def test_async_send_batch_routes_distributions_to_v1_endpoint(clean_env): + """Distribution series go to /api/v1/distribution_points (deflate), gauges/counts stay on /api/v2/series.""" + logger = DatadogMetricsLogger(start_periodic_flush=False) + logger.async_client = AsyncMock() + logger.async_client.post.return_value = Response( + 202, json={"status": "ok"}, request=Request("POST", "https://api.test.datadoghq.com") + ) + + timestamp = int(time.time()) + logger.log_queue = [ + { + "metric": "litellm.request.total_latency", + "type": 3, + "points": [{"timestamp": timestamp, "value": 1.5}], + "tags": ["env:test"], + }, + { + "metric": "litellm_proxy.request.total_latency", + "type": "distribution", + "points": ((timestamp, (1.5,)),), + "tags": ("env:test",), + }, + ] + + await logger.async_send_batch() + + calls = {call.args[0]: call.kwargs for call in logger.async_client.post.call_args_list} + assert set(calls) == { + "https://api.test.datadoghq.com/api/v2/series", + "https://api.test.datadoghq.com/api/v1/distribution_points", + } + + series_call = calls["https://api.test.datadoghq.com/api/v2/series"] + assert series_call["headers"]["Content-Encoding"] == "gzip" + series_payload = json.loads(gzip.decompress(series_call["content"])) + assert [s["metric"] for s in series_payload["series"]] == ["litellm.request.total_latency"] + + distribution_call = calls["https://api.test.datadoghq.com/api/v1/distribution_points"] + assert distribution_call["headers"]["Content-Encoding"] == "deflate" + assert distribution_call["headers"]["DD-API-KEY"] == "test_api_key" + distribution_payload = json.loads(zlib.decompress(distribution_call["content"])) + assert distribution_payload == { + "series": [ + { + "metric": "litellm_proxy.request.total_latency", + "type": "distribution", + "points": [[timestamp, [1.5]]], + "tags": ["env:test"], + } + ] + } + + +@pytest.mark.asyncio +async def test_async_send_batch_skips_v2_when_only_distributions_queued(clean_env): + logger = DatadogMetricsLogger(start_periodic_flush=False) + logger.async_client = AsyncMock() + logger.async_client.post.return_value = Response( + 202, json={"status": "ok"}, request=Request("POST", "https://api.test.datadoghq.com") + ) + logger.log_queue = [ + { + "metric": "litellm_proxy.llm_api.latency", + "type": "distribution", + "points": ((int(time.time()), (0.4,)),), + "tags": ("env:test",), + } + ] + + await logger.async_send_batch() + + assert [call.args[0] for call in logger.async_client.post.call_args_list] == [ + "https://api.test.datadoghq.com/api/v1/distribution_points" + ] + + +@pytest.mark.asyncio +async def test_flush_retries_only_distributions_after_v1_failure(clean_env): + logger = DatadogMetricsLogger(start_periodic_flush=False) + logger.async_client = AsyncMock() + v2_url = "https://api.test.datadoghq.com/api/v2/series" + v1_url = "https://api.test.datadoghq.com/api/v1/distribution_points" + responses = { + v2_url: Response(202, json={"status": "ok"}, request=Request("POST", v2_url)), + v1_url: Response(503, json={"errors": ["down"]}, request=Request("POST", v1_url)), + } + timestamp = int(time.time()) + late_count = { + "metric": "litellm.llm_api.request_count", + "type": 1, + "points": [{"timestamp": timestamp, "value": 1}], + "tags": ["env:test"], + } + + def post(url, **_): + if url == v1_url and late_count not in logger.log_queue: + logger.log_queue.append(late_count) + return responses[url] + + logger.async_client.post.side_effect = post + + gauge = { + "metric": "litellm.request.total_latency", + "type": 3, + "points": [{"timestamp": timestamp, "value": 1.5}], + "tags": ["env:test"], + } + distribution = { + "metric": "litellm_proxy.request.total_latency", + "type": "distribution", + "points": ((timestamp, (1.5,)),), + "tags": ("env:test",), + } + logger.log_queue = [gauge, distribution] + + await logger.flush_queue() + + assert logger.log_queue == [distribution, late_count] + + responses[v1_url] = Response(202, json={"status": "ok"}, request=Request("POST", v1_url)) + await logger.flush_queue() + + assert logger.log_queue == [] + assert [call.args[0] for call in logger.async_client.post.call_args_list] == [v2_url, v1_url, v2_url, v1_url] + second_v2 = json.loads(gzip.decompress(logger.async_client.post.call_args_list[2].kwargs["content"])) + assert [s["metric"] for s in second_v2["series"]] == ["litellm.llm_api.request_count"] + + +@pytest.mark.asyncio +async def test_flush_keeps_only_the_newest_batch_of_rejected_distributions(clean_env): + logger = DatadogMetricsLogger(batch_size=2, start_periodic_flush=False) + logger.async_client = AsyncMock() + v1_url = "https://api.test.datadoghq.com/api/v1/distribution_points" + logger.async_client.post.return_value = Response(503, json={"errors": ["down"]}, request=Request("POST", v1_url)) + timestamp = int(time.time()) + distributions = [ + { + "metric": "litellm_proxy.request.total_latency", + "type": "distribution", + "points": ((timestamp, (float(seconds),)),), + "tags": ("env:test",), + } + for seconds in range(3) + ] + logger.log_queue = list(distributions) + + await logger.flush_queue() + + assert logger.log_queue == distributions[1:] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("log_event", ("async_log_success_event", "async_log_failure_event")) +async def test_rejected_distributions_do_not_count_toward_the_batch_size(clean_env, log_event): + """While Datadog rejects distributions, a request triggers a flush only once a batch of new metrics has + queued behind them, and once Datadog takes them every queued metric counts again.""" + logger = DatadogMetricsLogger(batch_size=8, start_periodic_flush=False) + logger.async_client = AsyncMock() + v2_url = "https://api.test.datadoghq.com/api/v2/series" + v1_url = "https://api.test.datadoghq.com/api/v1/distribution_points" + responses = { + v2_url: Response(202, json={"status": "ok"}, request=Request("POST", v2_url)), + v1_url: Response(503, json={"errors": ["down"]}, request=Request("POST", v1_url)), + } + logger.async_client.post.side_effect = lambda url, **_: responses[url] + now = datetime.now() + kwargs = { + "standard_logging_object": StandardLoggingPayload(custom_llm_provider="openai", model="gpt-4o"), + "start_time": now - timedelta(seconds=1), + "api_call_start_time": now - timedelta(seconds=0.5), + "end_time": now, + } + + async def posts_after_each_of(requests: int) -> list[int]: + posts = [] + for _ in range(requests): + await getattr(logger, log_event)(kwargs, None, kwargs["start_time"], now) + posts.append(logger.async_client.post.call_count) + return posts + + assert await posts_after_each_of(4) == [0, 2, 2, 4] + + responses[v1_url] = Response(202, json={"status": "ok"}, request=Request("POST", v1_url)) + + assert await posts_after_each_of(4) == [4, 6, 6, 8] + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_flush_keeps_the_whole_batch_when_the_series_upload_fails(clean_env): + logger = DatadogMetricsLogger(start_periodic_flush=False) + logger.async_client = AsyncMock() + v2_url = "https://api.test.datadoghq.com/api/v2/series" + v1_url = "https://api.test.datadoghq.com/api/v1/distribution_points" + responses = { + v2_url: Response(503, json={"errors": ["down"]}, request=Request("POST", v2_url)), + v1_url: Response(202, json={"status": "ok"}, request=Request("POST", v1_url)), + } + logger.async_client.post.side_effect = lambda url, **_: responses[url] + timestamp = int(time.time()) + gauge = { + "metric": "litellm.request.total_latency", + "type": 3, + "points": [{"timestamp": timestamp, "value": 1.5}], + "tags": ["env:test"], + } + distribution = { + "metric": "litellm_proxy.request.total_latency", + "type": "distribution", + "points": ((timestamp, (1.5,)),), + "tags": ("env:test",), + } + logger.log_queue = [gauge, distribution] + + await logger.flush_queue() + + assert logger.log_queue == [gauge, distribution] + + responses[v2_url] = Response(202, json={"status": "ok"}, request=Request("POST", v2_url)) + await logger.flush_queue() + + assert logger.log_queue == [] + assert [call.args[0] for call in logger.async_client.post.call_args_list] == [v2_url, v2_url, v1_url] + + +@pytest.mark.asyncio +async def test_flush_cancelled_during_the_distribution_upload_resends_only_distributions(clean_env): + """A size-triggered flush runs under the logging worker's deadline; when it is cancelled after Datadog + took the series, the next flush must not count those series a second time.""" + logger = DatadogMetricsLogger(start_periodic_flush=False) + logger.async_client = AsyncMock() + v2_url = "https://api.test.datadoghq.com/api/v2/series" + v1_url = "https://api.test.datadoghq.com/api/v1/distribution_points" + v1_hangs = asyncio.Event() + v1_hangs.set() + + async def post(url, **_): + if url == v1_url and v1_hangs.is_set(): + await asyncio.sleep(3600) + return Response(202, json={"status": "ok"}, request=Request("POST", url)) + + logger.async_client.post.side_effect = post + timestamp = int(time.time()) + count = { + "metric": "litellm.llm_api.request_count", + "type": 1, + "points": [{"timestamp": timestamp, "value": 1}], + "tags": ["env:test"], + } + distribution = { + "metric": "litellm_proxy.request.total_latency", + "type": "distribution", + "points": ((timestamp, (1.5,)),), + "tags": ("env:test",), + } + logger.log_queue = [count, distribution] + + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(logger.flush_queue(), timeout=0.1) + + assert logger.log_queue == [distribution] + + v1_hangs.clear() + await logger.flush_queue() + + assert logger.log_queue == [] + assert [call.args[0] for call in logger.async_client.post.call_args_list] == [v2_url, v1_url, v1_url] + + +@pytest.mark.asyncio +async def test_flush_keeps_metrics_queued_while_the_upload_was_in_flight(clean_env): + """A request finishing while a flush awaits Datadog rides the next flush instead of being cleared with the batch.""" + logger = DatadogMetricsLogger(start_periodic_flush=False) + logger.async_client = AsyncMock() + timestamp = int(time.time()) + late = { + "metric": "litellm_proxy.request.total_latency", + "type": "distribution", + "points": ((timestamp, (2.0,)),), + "tags": ("env:late",), + } + + def post(url, **_): + if late not in logger.log_queue: + logger.log_queue.append(late) + return Response(202, json={"status": "ok"}, request=Request("POST", url)) + + logger.async_client.post.side_effect = post + logger.log_queue = [ + { + "metric": "litellm.request.total_latency", + "type": 3, + "points": [{"timestamp": timestamp, "value": 1.5}], + "tags": ["env:test"], + } + ] + + await logger.flush_queue() + + assert logger.log_queue == [late] + + await logger.flush_queue() + + assert logger.log_queue == [] + assert [call.args[0] for call in logger.async_client.post.call_args_list] == [ + "https://api.test.datadoghq.com/api/v2/series", + "https://api.test.datadoghq.com/api/v1/distribution_points", + ] + + @pytest.mark.asyncio async def test_async_send_batch_empty_queue(clean_env): """Test that async_send_batch does nothing when queue is empty."""