mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
feat(otel v2): excluded_services opt-out for datastore spans on tenant destinations (#43278)
* feat(otel): excluded_services opt-out for datastore spans on tenant destinations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): keep upstream support unchanged Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): excluded_services resolves from the otel callback config only Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): name the otel callback logger so excluded_services owner lookup matches Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): assert no aux datastore traces reach the tenant sink Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): read bogus-start proxy log from the results dir Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): read only this invocation's bogus-start proxy log Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): assert operator kept db spans over the whole recorded window Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): split operator db-span asserts by trace scope Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): build the otel logger after preset callbacks and validate the exclusion env at boot Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): tolerate a bogus exclusion env when callback config wins Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): keep bogus exclusion env fatal when a preset parses it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): hoist the preset check out of the callback loop Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(otel): use a rule-scoped pyright suppression Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): log and drop unknown excluded_services instead of failing boot Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): wait for the operator spend-writer span before checking the tenant for postgres Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel v2): leave callback init and boot untouched when excluded_services is unset Read callback_settings.otel.excluded_services directly instead of making the otel callback build its own logger, and drop the new boot-time parse of callback_settings.otel, so a proxy without the setting behaves exactly as on main Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel v2): normalize callback_settings excluded_services without rereading OTel env vars Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel v2): log and ignore malformed excluded_services instead of failing startup Lowercase and trim names, drop non-string items, and add an integration matrix over endpoints, clients, cache hits, destination outages and setting shapes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(otel v2): cover failed upstream calls in the excluded_services matrix Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(otel v2): pin operator Langfuse credentials in preset-only excluded_services tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng <yucheng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mrinal <mrinal@berri.ai>
This commit is contained in:
parent
e78845afcf
commit
50f5cc9bbb
10 changed files with 2018 additions and 21 deletions
|
|
@ -213,6 +213,15 @@ nothing here imports outside it:
|
|||
`config.yaml` — the latter reach the config through the logger's constructor
|
||||
kwargs. `baggage_team_metadata_keys` is empty by default, so none of a team's
|
||||
free-form metadata is promoted until each sub-key is explicitly allowlisted.
|
||||
`excluded_services` withholds datastore spans from key/team `callback_vars`
|
||||
destinations while the operator's own exporters keep them: set
|
||||
`LITELLM_OTEL_EXCLUDED_SERVICES` (comma-separated) or `excluded_services`
|
||||
(a YAML list) under `callback_settings.otel`, naming the datastore services
|
||||
to withhold (`redis`, `postgres`, `batch_write_to_db`, `redis_*`, or their
|
||||
`db.system.name` spellings `redis` / `postgresql`). Unknown names are logged
|
||||
as an error and ignored. A span is withheld when its `db.system.name` /
|
||||
`db.system` attribute is in the set, so request root, auth, guardrail and
|
||||
model spans can never be excluded.
|
||||
- [`baggage.py`](./model/baggage.py) — the single definition of which request-identity
|
||||
values are promoted into Baggage (so child spans inherit them) and under which
|
||||
attribute keys.
|
||||
|
|
|
|||
|
|
@ -30,7 +30,7 @@ from litellm.integrations.custom_logger import CustomLogger
|
|||
from litellm.integrations.otel.emitter import SpanEmitter, stamp_error
|
||||
from litellm.integrations.otel.mappers import resolve_mappers
|
||||
from litellm.integrations.otel.model.baggage import promoted_baggage
|
||||
from litellm.integrations.otel.model.config import OpenTelemetryV2Config
|
||||
from litellm.integrations.otel.model.config import OpenTelemetryV2Config, excluded_db_systems_from
|
||||
from litellm.integrations.otel.model.metadata import (
|
||||
LLMCallEvent,
|
||||
RequestIdentity,
|
||||
|
|
@ -898,12 +898,29 @@ def publish_global_otel_v2_provider(
|
|||
"""
|
||||
global _published_v2_provider
|
||||
logger: Final = select_global_otel_v2_logger(in_memory_loggers, registered=registered)
|
||||
attach_tenant_fan_out(logger.tracer_provider, *_v2_configs(in_memory_loggers, logger))
|
||||
attach_tenant_fan_out(
|
||||
logger.tracer_provider,
|
||||
*_v2_configs(in_memory_loggers, logger),
|
||||
excluded_db_systems=_excluded_db_systems(logger),
|
||||
)
|
||||
set_global_provider(logger.tracer_provider)
|
||||
_published_v2_provider = logger.tracer_provider # rebind-ok: startup records the one provider carrying the fan-out
|
||||
return logger
|
||||
|
||||
|
||||
def _excluded_db_systems(logger: "OpenTelemetryV2") -> frozenset[str]:
|
||||
"""The datastore services withheld from tenant destinations.
|
||||
|
||||
``callback_settings.otel.excluded_services`` wins over the env var whichever
|
||||
logger got published: with ``callbacks: [langfuse_otel, otel]`` the ``otel``
|
||||
callback folds into the preset, whose config is env-only.
|
||||
"""
|
||||
configured: Final = litellm.callback_settings.get("otel", {}).get("excluded_services")
|
||||
if configured is None:
|
||||
return logger.config.excluded_services
|
||||
return excluded_db_systems_from(configured)
|
||||
|
||||
|
||||
def _v2_configs(in_memory_loggers: Sequence[object], logger: "OpenTelemetryV2") -> tuple[OpenTelemetryV2Config, ...]:
|
||||
"""Every v2 logger's config, the published logger's first.
|
||||
|
||||
|
|
@ -963,7 +980,11 @@ def fan_out_provider() -> ApiTracerProvider:
|
|||
return published
|
||||
logger: Final = _registered_v2_logger()
|
||||
if logger is not None:
|
||||
attach_tenant_fan_out(logger.tracer_provider, logger.config)
|
||||
attach_tenant_fan_out(
|
||||
logger.tracer_provider,
|
||||
logger.config,
|
||||
excluded_db_systems=_excluded_db_systems(logger),
|
||||
)
|
||||
return logger.tracer_provider
|
||||
return get_tracer_provider()
|
||||
|
||||
|
|
|
|||
|
|
@ -4,14 +4,16 @@ from enum import Enum
|
|||
from functools import lru_cache
|
||||
from typing import Annotated, Any, Final
|
||||
|
||||
from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator
|
||||
from pydantic import AliasChoices, BaseModel, Field, TypeAdapter, ValidationError, field_validator, model_validator
|
||||
from pydantic_settings import BaseSettings, NoDecode, SettingsConfigDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.otel.model.baggage import (
|
||||
BAGGAGE_PROMOTED_KEYS,
|
||||
DEFAULT_BAGGAGE_METADATA_KEYS,
|
||||
DEFAULT_BAGGAGE_TEAM_METADATA_KEYS,
|
||||
)
|
||||
from litellm.integrations.otel.model.spans import POSTGRESQL, db_system
|
||||
from litellm.types.utils import OtelSpanScope
|
||||
|
||||
#: Master feature-flag env var. The logger is inert until this is truthy.
|
||||
|
|
@ -174,6 +176,19 @@ class OpenTelemetryV2Config(BaseSettings):
|
|||
"key/team destinations are not affected."
|
||||
),
|
||||
)
|
||||
excluded_services: Annotated[frozenset[str], NoDecode] = Field(
|
||||
default_factory=frozenset,
|
||||
validation_alias=AliasChoices("excluded_services", "LITELLM_OTEL_EXCLUDED_SERVICES"),
|
||||
description=(
|
||||
"Datastore services whose spans are withheld from key/team ``callback_vars`` "
|
||||
"OTel destinations (the operator's own exporters still receive them). Accepted "
|
||||
"values are the datastore ``ServiceTypes`` names (``redis``, ``postgres``, "
|
||||
"``batch_write_to_db``, ``redis_*``) or their ``db.system.name`` spellings "
|
||||
"(``redis``, ``postgresql``); stored normalized to ``db.system.name`` values. "
|
||||
"Configure via the ``LITELLM_OTEL_EXCLUDED_SERVICES`` env var (comma-separated) "
|
||||
"or ``callback_settings.otel.excluded_services`` in config.yaml (a YAML list)."
|
||||
),
|
||||
)
|
||||
|
||||
# ----- explicit multi-destination / vocabulary configuration ------------ #
|
||||
|
||||
|
|
@ -284,6 +299,11 @@ class OpenTelemetryV2Config(BaseSettings):
|
|||
return [item.strip() for item in value.split(",") if item.strip()]
|
||||
return value
|
||||
|
||||
@field_validator("excluded_services", mode="before")
|
||||
@classmethod
|
||||
def _read_excluded_services(cls, value: object) -> frozenset[str]:
|
||||
return excluded_service_names(value)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _normalize(self) -> "OpenTelemetryV2Config":
|
||||
# An endpoint with the default exporter kind implies OTLP/HTTP.
|
||||
|
|
@ -316,6 +336,7 @@ class OpenTelemetryV2Config(BaseSettings):
|
|||
if self.legacy_compat and "legacy" not in names:
|
||||
names.append("legacy")
|
||||
self.mapper_names = names
|
||||
self.excluded_services = _normalize_excluded_services(self.excluded_services)
|
||||
return self
|
||||
|
||||
@property
|
||||
|
|
@ -334,3 +355,55 @@ class OpenTelemetryV2Config(BaseSettings):
|
|||
@classmethod
|
||||
def from_env(cls) -> "OpenTelemetryV2Config":
|
||||
return cls()
|
||||
|
||||
|
||||
_EXCLUDED_SERVICES_INPUT: Final[TypeAdapter[str | tuple[object, ...]]] = TypeAdapter(str | tuple[object, ...])
|
||||
|
||||
|
||||
def excluded_db_systems_from(value: object) -> frozenset[str]:
|
||||
"""Normalize a raw ``excluded_services`` value without building a settings model that rereads the env"""
|
||||
return _normalize_excluded_services(excluded_service_names(value))
|
||||
|
||||
|
||||
def excluded_service_names(value: object) -> frozenset[str]:
|
||||
"""Read a YAML list or comma-separated string of service names, logging and dropping unusable input
|
||||
so a malformed value cannot stop the OTel logger from being built"""
|
||||
if value is None:
|
||||
return frozenset()
|
||||
try:
|
||||
parsed: Final = _EXCLUDED_SERVICES_INPUT.validate_python(value)
|
||||
except ValidationError:
|
||||
verbose_logger.error("excluded_services must be a list or comma-separated string; %r ignored", value)
|
||||
return frozenset()
|
||||
items: Final = tuple(parsed.split(",")) if isinstance(parsed, str) else parsed
|
||||
return frozenset(name for item in items if (name := _service_name(item)))
|
||||
|
||||
|
||||
def _service_name(item: object) -> str:
|
||||
if not isinstance(item, str):
|
||||
verbose_logger.error("excluded_services must be a list of service names; %r ignored", item)
|
||||
return ""
|
||||
return item.strip().lower()
|
||||
|
||||
|
||||
def _normalize_excluded_services(services: frozenset[str]) -> frozenset[str]:
|
||||
"""Fold each accepted spelling to its ``db.system.name`` value.
|
||||
|
||||
``postgres`` and ``postgresql`` name the same system, as do every
|
||||
``ServiceTypes`` member that ``db_system`` maps. Anything else means the
|
||||
operator pointed the setting at a span family it cannot cover; those names
|
||||
are logged and dropped so a typo cannot take the proxy down.
|
||||
"""
|
||||
resolved: Final = frozenset(
|
||||
system for service in services if (system := _db_system_for_excluded_service(service)) is not None
|
||||
)
|
||||
return resolved
|
||||
|
||||
|
||||
def _db_system_for_excluded_service(service: str) -> str | None:
|
||||
resolved: Final = db_system(service) if service != POSTGRESQL else POSTGRESQL
|
||||
if resolved is None:
|
||||
verbose_logger.error(
|
||||
"excluded_services: %r is not a datastore service; ignored. Allowed: postgres, redis", service
|
||||
)
|
||||
return resolved
|
||||
|
|
|
|||
|
|
@ -418,6 +418,13 @@ def _is_database_span(attributes: Mapping[str, AttributeValue]) -> bool:
|
|||
return any(key in attributes for key in _DB_SYSTEM_KEYS)
|
||||
|
||||
|
||||
def _is_excluded_database_span(attributes: Mapping[str, AttributeValue], excluded: frozenset[str]) -> bool:
|
||||
if not excluded:
|
||||
return False
|
||||
system: Final = attributes.get(DB.SYSTEM_NAME) or attributes.get(DB.SYSTEM_LEGACY)
|
||||
return isinstance(system, str) and system in excluded
|
||||
|
||||
|
||||
def _is_tenant_owned_span(attributes: Mapping[str, AttributeValue]) -> bool:
|
||||
return any(key in attributes for key in _TENANT_OWNED_KEYS)
|
||||
|
||||
|
|
@ -549,10 +556,12 @@ class TenantFanOutSpanProcessor(SpanProcessor):
|
|||
processor_factory: 'Callable[["OtelDestination"], SpanProcessor | None] | None' = None,
|
||||
shutdown_drain_seconds: float = _SHUTDOWN_DRAIN_SECONDS,
|
||||
operator_sinks: 'Mapping[_SinkKey, "OtelSpanScope"]' = MappingProxyType({}),
|
||||
excluded_db_systems: frozenset[str] = frozenset(),
|
||||
pending_drains: int = _MAX_PENDING_DRAINS,
|
||||
drain_pool: _DrainPool | None = None,
|
||||
) -> None:
|
||||
self._operator_sinks: Final = operator_sinks
|
||||
self._excluded_db_systems: Final = excluded_db_systems
|
||||
self._drain_seconds: Final = shutdown_drain_seconds
|
||||
self._lock: Final = threading.Condition()
|
||||
self._closed = False # guarded by ``_lock``: an unlocked read races the teardown it gates
|
||||
|
|
@ -567,9 +576,12 @@ class TenantFanOutSpanProcessor(SpanProcessor):
|
|||
|
||||
def on_end(self, span: ReadableSpan) -> None:
|
||||
suppressed: Final = suppressed_backends()
|
||||
attributes: Final = span.attributes or _NO_ATTRIBUTES
|
||||
for destination in request_destinations():
|
||||
if self._operator_already_writes(span, destination, suppressed) or not _in_scope(
|
||||
span, destination.span_scope
|
||||
if (
|
||||
self._operator_already_writes(span, destination, suppressed)
|
||||
or not _in_scope(span, destination.span_scope)
|
||||
or _is_excluded_database_span(attributes, self._excluded_db_systems)
|
||||
):
|
||||
continue
|
||||
processor = self._acquire(destination)
|
||||
|
|
@ -1155,7 +1167,9 @@ def build_tracer_provider(
|
|||
_FAN_OUT_ATTACH_LOCK: Final = threading.Lock()
|
||||
|
||||
|
||||
def attach_tenant_fan_out(provider: TracerProvider, *configs: OpenTelemetryV2Config) -> None:
|
||||
def attach_tenant_fan_out(
|
||||
provider: TracerProvider, *configs: OpenTelemetryV2Config, excluded_db_systems: frozenset[str] = frozenset()
|
||||
) -> None:
|
||||
"""Give ``provider`` the fan-out that delivers spans to key/team destinations.
|
||||
|
||||
Called on the one provider published as the OTel global, and idempotent so a
|
||||
|
|
@ -1164,12 +1178,18 @@ def attach_tenant_fan_out(provider: TracerProvider, *configs: OpenTelemetryV2Con
|
|||
so exactly one fan-out lands. ``configs`` name the operator's own exporters, one
|
||||
config per v2 logger since each keeps its own provider and still writes its
|
||||
account, so an additive destination pointing at any of them is delivered once
|
||||
rather than twice.
|
||||
rather than twice. ``excluded_db_systems`` only filters what the fan-out
|
||||
delivers, never the operator's own exporters.
|
||||
"""
|
||||
with _FAN_OUT_ATTACH_LOCK:
|
||||
if any(isinstance(processor, TenantFanOutSpanProcessor) for processor in _attached_processors(provider)):
|
||||
return
|
||||
provider.add_span_processor(TenantFanOutSpanProcessor(operator_sinks=operator_sink_scopes(*configs)))
|
||||
provider.add_span_processor(
|
||||
TenantFanOutSpanProcessor(
|
||||
operator_sinks=operator_sink_scopes(*configs),
|
||||
excluded_db_systems=excluded_db_systems,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def deliverable_destinations(
|
||||
|
|
|
|||
528
tests/integration/_support/otlp_sink.py
Normal file
528
tests/integration/_support/otlp_sink.py
Normal file
|
|
@ -0,0 +1,528 @@
|
|||
"""OTLP/HTTP trace sink: records exported spans and exposes them over a control API.
|
||||
|
||||
Accepts ``application/x-protobuf`` ``ExportTraceServiceRequest`` bodies and OTLP
|
||||
``http/json`` bodies on any path. Tests read spans through ``recorded_spans`` and
|
||||
steer the sink through ``configure``; the process can also be frozen with
|
||||
``SIGSTOP``/``SIGCONT`` after reading its pid from ``/__pid``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import datetime
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import socket
|
||||
import ssl
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
import psutil
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
INTERNAL_MARKERS: Final = ("gen_ai.operation.name", "mcp.method.name", "litellm.guardrail_name")
|
||||
|
||||
|
||||
class Span(TypedDict):
|
||||
trace_id: ReadOnly[str]
|
||||
span_id: ReadOnly[str]
|
||||
parent_span_id: ReadOnly[str]
|
||||
kind: ReadOnly[int]
|
||||
name: ReadOnly[str]
|
||||
attributes: ReadOnly[Mapping[str, JsonValue]]
|
||||
resource: ReadOnly[Mapping[str, JsonValue]]
|
||||
|
||||
|
||||
class _SpanListing(TypedDict):
|
||||
next: ReadOnly[int]
|
||||
spans: ReadOnly[list[Span]]
|
||||
|
||||
|
||||
_SPAN_LISTING: Final = TypeAdapter(_SpanListing)
|
||||
|
||||
|
||||
def _proto_spans(body: bytes) -> list[Span]:
|
||||
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest
|
||||
from opentelemetry.proto.common.v1.common_pb2 import AnyValue
|
||||
|
||||
def scalar(value: AnyValue) -> JsonValue:
|
||||
match value.WhichOneof("value"):
|
||||
case "string_value":
|
||||
return value.string_value
|
||||
case "bool_value":
|
||||
return value.bool_value
|
||||
case "int_value":
|
||||
return int(value.int_value)
|
||||
case "double_value":
|
||||
return value.double_value
|
||||
case "bytes_value":
|
||||
return value.bytes_value.decode("utf-8", errors="replace")
|
||||
case "array_value":
|
||||
return [scalar(item) for item in value.array_value.values]
|
||||
case "kvlist_value":
|
||||
return {pair.key: scalar(pair.value) for pair in value.kvlist_value.values}
|
||||
case _:
|
||||
return None
|
||||
|
||||
request: Final = ExportTraceServiceRequest()
|
||||
request.ParseFromString(body)
|
||||
return [
|
||||
Span(
|
||||
trace_id=span.trace_id.hex(),
|
||||
span_id=span.span_id.hex(),
|
||||
parent_span_id=span.parent_span_id.hex(),
|
||||
kind=span.kind,
|
||||
name=span.name,
|
||||
attributes={attribute.key: scalar(attribute.value) for attribute in span.attributes},
|
||||
resource={attribute.key: scalar(attribute.value) for attribute in resource.resource.attributes},
|
||||
)
|
||||
for resource in request.resource_spans
|
||||
for scope in resource.scope_spans
|
||||
for span in scope.spans
|
||||
]
|
||||
|
||||
|
||||
def _json_spans(body: bytes) -> list[Span]:
|
||||
payload: Final = json.loads(body)
|
||||
|
||||
def scalar(value: object) -> JsonValue:
|
||||
if not isinstance(value, dict):
|
||||
return value if isinstance(value, (str, int, float, bool)) or value is None else str(value)
|
||||
for key in ("stringValue", "intValue", "doubleValue", "boolValue", "bytesValue"):
|
||||
if key in value:
|
||||
return value[key]
|
||||
if "arrayValue" in value:
|
||||
return [scalar(item) for item in value["arrayValue"].get("values", [])]
|
||||
if "kvlistValue" in value:
|
||||
return {pair["key"]: scalar(pair["value"]) for pair in value["kvlistValue"].get("values", [])}
|
||||
return None
|
||||
|
||||
return [
|
||||
Span(
|
||||
trace_id=str(span.get("traceId", "")),
|
||||
span_id=str(span.get("spanId", "")),
|
||||
parent_span_id=str(span.get("parentSpanId", "")),
|
||||
kind=int(span.get("kind", 0)),
|
||||
name=str(span.get("name", "")),
|
||||
attributes={attribute["key"]: scalar(attribute.get("value")) for attribute in span.get("attributes", [])},
|
||||
resource={
|
||||
attribute["key"]: scalar(attribute.get("value"))
|
||||
for attribute in resource.get("resource", {}).get("attributes", [])
|
||||
},
|
||||
)
|
||||
for resource in payload.get("resourceSpans", [])
|
||||
for scope in resource.get("scopeSpans", [])
|
||||
for span in scope.get("spans", [])
|
||||
]
|
||||
|
||||
|
||||
def decode_spans(body: bytes, content_type: str) -> list[Span]:
|
||||
if "protobuf" in content_type:
|
||||
return _proto_spans(body)
|
||||
return _json_spans(body)
|
||||
|
||||
|
||||
def span_class(span: Span) -> str:
|
||||
if span["kind"] == 2:
|
||||
return "root"
|
||||
if any(marker in span["attributes"] for marker in INTERNAL_MARKERS):
|
||||
return "tenant"
|
||||
return "internal"
|
||||
|
||||
|
||||
def spans_for_trace(spans: tuple[Span, ...], trace_id: str) -> tuple[Span, ...]:
|
||||
return tuple(span for span in spans if span["trace_id"] == trace_id)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _State:
|
||||
spans: list[Span] = field(default_factory=list)
|
||||
requests: list[dict[str, JsonValue]] = field(default_factory=list)
|
||||
status: int = 200
|
||||
delay_seconds: float = 0.0
|
||||
pause: threading.Event = field(default_factory=threading.Event)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.pause.set()
|
||||
|
||||
|
||||
class _Handler(BaseHTTPRequestHandler):
|
||||
state: _State
|
||||
protocol_version = "HTTP/1.1"
|
||||
|
||||
def _read_body(self) -> bytes:
|
||||
return self.rfile.read(int(self.headers.get("content-length", "0")))
|
||||
|
||||
def _send_json(self, payload: object, status: int = 200) -> None:
|
||||
body: Final = json.dumps(payload).encode()
|
||||
self.send_response(status)
|
||||
self.send_header("content-type", "application/json")
|
||||
self.send_header("content-length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def _record(self) -> None:
|
||||
body: Final = self._read_body()
|
||||
self.state.pause.wait(timeout=120)
|
||||
if self.state.delay_seconds > 0:
|
||||
time.sleep(self.state.delay_seconds)
|
||||
recorded: Final = decode_spans(body, self.headers.get("content-type", ""))
|
||||
self.state.spans.extend(recorded)
|
||||
self.state.requests.append(
|
||||
{
|
||||
"path": self.path,
|
||||
"count": len(recorded),
|
||||
"host": self.headers.get("host", ""),
|
||||
"headers": dict(self.headers),
|
||||
}
|
||||
)
|
||||
self._send_json({"recorded": len(recorded)}, status=self.state.status)
|
||||
|
||||
do_POST = _record
|
||||
do_PUT = _record
|
||||
|
||||
def do_GET(self) -> None:
|
||||
parsed: Final = urlparse(self.path)
|
||||
if parsed.path == "/__spans":
|
||||
since: Final = int(dict(part.split("=", 1) for part in parsed.query.split("&") if part).get("since", "0"))
|
||||
self._send_json({"next": len(self.state.spans), "spans": self.state.spans[since:]})
|
||||
return
|
||||
if parsed.path == "/__pid":
|
||||
self._send_json({"pid": os.getpid()})
|
||||
return
|
||||
if parsed.path == "/__requests":
|
||||
self._send_json({"requests": self.state.requests})
|
||||
return
|
||||
self._send_json({"error": "unknown"}, status=404)
|
||||
|
||||
def do_DELETE(self) -> None:
|
||||
if urlparse(self.path).path == "/__spans":
|
||||
self.state.spans.clear()
|
||||
self.state.requests.clear()
|
||||
self._send_json({"cleared": True})
|
||||
return
|
||||
self._send_json({"error": "unknown"}, status=404)
|
||||
|
||||
def do_PATCH(self) -> None:
|
||||
if urlparse(self.path).path != "/__control":
|
||||
self._send_json({"error": "unknown"}, status=404)
|
||||
return
|
||||
fields: Final = json.loads(self._read_body() or b"{}")
|
||||
if "status" in fields:
|
||||
self.state.status = int(fields["status"])
|
||||
if "delay_seconds" in fields:
|
||||
self.state.delay_seconds = float(fields["delay_seconds"])
|
||||
if fields.get("paused") is True:
|
||||
self.state.pause.clear()
|
||||
if fields.get("paused") is False:
|
||||
self.state.pause.set()
|
||||
self._send_json({"status": self.state.status, "delay_seconds": self.state.delay_seconds})
|
||||
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class _ConnectHandler(_Handler):
|
||||
tunnel_context: ssl.SSLContext
|
||||
|
||||
def do_CONNECT(self) -> None:
|
||||
self.state.requests.append({"connect": self.path})
|
||||
self.connection.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n")
|
||||
wrapped: Final = self.tunnel_context.wrap_socket(self.connection, server_side=True)
|
||||
self.close_connection = True
|
||||
type(self)(wrapped, self.client_address, self.server)
|
||||
|
||||
|
||||
_MITM_HOSTS: Final = ("otlp.nr-data.net", "otlp.eu01.nr-data.net")
|
||||
|
||||
|
||||
def _mitm_context(directory: Path) -> ssl.SSLContext:
|
||||
from cryptography import x509
|
||||
from cryptography.hazmat.primitives import hashes, serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
from cryptography.x509.oid import NameOID
|
||||
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
now: Final = datetime.datetime.now(datetime.timezone.utc)
|
||||
window: Final = datetime.timedelta(days=2)
|
||||
ca_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
ca_name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "otlp-sink test CA")])
|
||||
ca_cert: Final = (
|
||||
x509.CertificateBuilder()
|
||||
.subject_name(ca_name)
|
||||
.issuer_name(ca_name)
|
||||
.public_key(ca_key.public_key())
|
||||
.serial_number(x509.random_serial_number())
|
||||
.not_valid_before(now - window)
|
||||
.not_valid_after(now + window)
|
||||
.add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True)
|
||||
.sign(ca_key, hashes.SHA256())
|
||||
)
|
||||
leaf_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
leaf_cert: Final = (
|
||||
x509.CertificateBuilder()
|
||||
.subject_name(x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, _MITM_HOSTS[0])]))
|
||||
.issuer_name(ca_cert.subject)
|
||||
.public_key(leaf_key.public_key())
|
||||
.serial_number(x509.random_serial_number())
|
||||
.not_valid_before(now - window)
|
||||
.not_valid_after(now + window)
|
||||
.add_extension(x509.SubjectAlternativeName([x509.DNSName(host) for host in _MITM_HOSTS]), critical=False)
|
||||
.sign(ca_key, hashes.SHA256())
|
||||
)
|
||||
ca_pem: Final = directory / "ca.pem"
|
||||
ca_pem.write_bytes(ca_cert.public_bytes(serialization.Encoding.PEM))
|
||||
leaf_pem: Final = directory / "leaf.pem"
|
||||
leaf_pem.write_bytes(leaf_cert.public_bytes(serialization.Encoding.PEM))
|
||||
leaf_key_pem: Final = directory / "leaf-key.pem"
|
||||
leaf_key_pem.write_bytes(
|
||||
leaf_key.private_bytes(
|
||||
serialization.Encoding.PEM,
|
||||
serialization.PrivateFormat.TraditionalOpenSSL,
|
||||
serialization.NoEncryption(),
|
||||
)
|
||||
)
|
||||
context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
|
||||
context.load_cert_chain(str(leaf_pem), str(leaf_key_pem))
|
||||
return context
|
||||
|
||||
|
||||
def _grpc_trace_server(state: _State, port: int) -> object:
|
||||
from concurrent import futures
|
||||
|
||||
import grpc
|
||||
from opentelemetry.proto.collector.trace.v1 import trace_service_pb2, trace_service_pb2_grpc
|
||||
|
||||
class _TraceService(trace_service_pb2_grpc.TraceServiceServicer):
|
||||
def Export(self, request: object, context: grpc.ServicerContext) -> object:
|
||||
state.pause.wait(timeout=120)
|
||||
if state.delay_seconds > 0:
|
||||
time.sleep(state.delay_seconds)
|
||||
recorded: Final = _proto_spans(request.SerializeToString())
|
||||
state.spans.extend(recorded)
|
||||
state.requests.append(
|
||||
{
|
||||
"grpc": "Export",
|
||||
"metadata": {key: value for key, value in context.invocation_metadata()},
|
||||
"count": len(recorded),
|
||||
}
|
||||
)
|
||||
return trace_service_pb2.ExportTraceServiceResponse()
|
||||
|
||||
server: Final = grpc.server(futures.ThreadPoolExecutor(max_workers=4))
|
||||
trace_service_pb2_grpc.add_TraceServiceServicer_to_server(_TraceService(), server)
|
||||
server.add_insecure_port(f"127.0.0.1:{port}")
|
||||
server.start()
|
||||
return server
|
||||
|
||||
|
||||
def recorded_spans(url: str, since: int = 0) -> tuple[int, tuple[Span, ...]]:
|
||||
response: Final = httpx.get(f"{url}/__spans", params={"since": since}, trust_env=False, timeout=15)
|
||||
response.raise_for_status()
|
||||
listing: Final = _SPAN_LISTING.validate_python(response.json())
|
||||
return listing["next"], tuple(listing["spans"])
|
||||
|
||||
|
||||
def configure_sink(url: str, **fields: JsonValue) -> None:
|
||||
httpx.request("PATCH", f"{url}/__control", json=dict(fields), trust_env=False, timeout=15).raise_for_status()
|
||||
|
||||
|
||||
def reset_sink(url: str) -> None:
|
||||
httpx.delete(f"{url}/__spans", trust_env=False, timeout=15).raise_for_status()
|
||||
|
||||
|
||||
def sink_pid(url: str) -> int:
|
||||
return int(httpx.get(f"{url}/__pid", trust_env=False, timeout=15).json()["pid"])
|
||||
|
||||
|
||||
_REQUEST_LISTING: Final = TypeAdapter(list[dict[str, JsonValue]])
|
||||
|
||||
|
||||
def recorded_requests(url: str) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
response: Final = httpx.get(f"{url}/__requests", trust_env=False, timeout=15)
|
||||
response.raise_for_status()
|
||||
return tuple(_REQUEST_LISTING.validate_python(response.json()["requests"]))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SpanSinks:
|
||||
operator: str
|
||||
tenant: str
|
||||
arize: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class GrpcSink:
|
||||
url: str
|
||||
control_url: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ConnectSink:
|
||||
proxy_url: str
|
||||
control_url: str
|
||||
ca_pem: str
|
||||
|
||||
|
||||
def _free_port() -> int:
|
||||
with socket.socket() as reserve:
|
||||
reserve.bind(("127.0.0.1", 0))
|
||||
return int(reserve.getsockname()[1])
|
||||
|
||||
|
||||
def _pid_reachable(url: str) -> bool:
|
||||
try:
|
||||
return httpx.get(f"{url}/__pid", trust_env=False, timeout=2).status_code == 200
|
||||
except httpx.TransportError:
|
||||
return False
|
||||
|
||||
|
||||
@contextmanager
|
||||
def owned_sinks(directory: Path) -> Iterator[SpanSinks]:
|
||||
from integration._support.process import group_members, signal_group, stop_root_process
|
||||
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
ports: Final = tuple(_free_port() for _ in range(3))
|
||||
root: Final = Path(__file__).resolve().parents[3]
|
||||
with ExitStack() as stack:
|
||||
processes: Final = tuple(
|
||||
subprocess.Popen(
|
||||
[sys.executable, "-P", "-m", "integration._support.otlp_sink", "--port", str(port)],
|
||||
cwd=root,
|
||||
stdout=stack.enter_context((directory / f"otlp-sink-{port}.log").open("w")),
|
||||
stderr=subprocess.STDOUT,
|
||||
start_new_session=True,
|
||||
)
|
||||
for port in ports
|
||||
)
|
||||
try:
|
||||
urls: Final = tuple(f"http://127.0.0.1:{port}" for port in ports)
|
||||
deadline: Final = time.monotonic() + 30
|
||||
while True:
|
||||
alive: Final = all(process.poll() is None for process in processes)
|
||||
assert alive, "OTLP sink exited before readiness"
|
||||
if all(_pid_reachable(url) for url in urls):
|
||||
break
|
||||
assert time.monotonic() < deadline, "OTLP sink readiness deadline exceeded"
|
||||
time.sleep(0.05)
|
||||
yield SpanSinks(operator=urls[0], tenant=urls[1], arize=urls[2])
|
||||
finally:
|
||||
for process in processes:
|
||||
stopped: Final = stop_root_process(process)
|
||||
residual: Final = group_members(process.pid)
|
||||
if residual:
|
||||
signal_group(process.pid, signal.SIGKILL)
|
||||
psutil.wait_procs(residual, timeout=5)
|
||||
survivors: Final = group_members(process.pid)
|
||||
assert not survivors and stopped, "OTLP sink required forced cleanup"
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _spawn_sink(directory: Path, log_name: str, argv: Sequence[str]) -> Iterator[None]:
|
||||
from integration._support.process import group_members, signal_group, stop_root_process
|
||||
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
root: Final = Path(__file__).resolve().parents[3]
|
||||
with (directory / log_name).open("w") as log:
|
||||
process: Final = subprocess.Popen(
|
||||
[sys.executable, "-P", "-m", "integration._support.otlp_sink", *argv],
|
||||
cwd=root,
|
||||
stdout=log,
|
||||
stderr=subprocess.STDOUT,
|
||||
start_new_session=True,
|
||||
)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
stopped: Final = stop_root_process(process)
|
||||
residual: Final = group_members(process.pid)
|
||||
if residual:
|
||||
signal_group(process.pid, signal.SIGKILL)
|
||||
psutil.wait_procs(residual, timeout=5)
|
||||
survivors: Final = group_members(process.pid)
|
||||
assert not survivors and stopped, "OTLP sink required forced cleanup"
|
||||
|
||||
|
||||
def _await_sink(url: str) -> None:
|
||||
deadline: Final = time.monotonic() + 30
|
||||
while not _pid_reachable(url):
|
||||
assert time.monotonic() < deadline, "OTLP sink readiness deadline exceeded"
|
||||
time.sleep(0.05)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def owned_grpc_sink(directory: Path) -> Iterator[GrpcSink]:
|
||||
http_port: Final = _free_port()
|
||||
grpc_port: Final = _free_port()
|
||||
with _spawn_sink(
|
||||
directory, "otlp-grpc-sink.log", ["--port", str(http_port), "--grpc-port", str(grpc_port)]
|
||||
):
|
||||
control_url: Final = f"http://127.0.0.1:{http_port}"
|
||||
_await_sink(control_url)
|
||||
yield GrpcSink(url=f"http://127.0.0.1:{grpc_port}", control_url=control_url)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def owned_connect_sink(directory: Path) -> Iterator[ConnectSink]:
|
||||
http_port: Final = _free_port()
|
||||
tunnel_port: Final = _free_port()
|
||||
ca_dir: Final = directory / "mitm"
|
||||
with _spawn_sink(
|
||||
directory,
|
||||
"otlp-connect-sink.log",
|
||||
["--port", str(http_port), "--connect-port", str(tunnel_port), "--ca-dir", str(ca_dir)],
|
||||
):
|
||||
control_url: Final = f"http://127.0.0.1:{http_port}"
|
||||
_await_sink(control_url)
|
||||
yield ConnectSink(
|
||||
proxy_url=f"http://127.0.0.1:{tunnel_port}",
|
||||
control_url=control_url,
|
||||
ca_pem=str(ca_dir / "ca.pem"),
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser: Final = argparse.ArgumentParser()
|
||||
parser.add_argument("--port", type=int, required=True)
|
||||
parser.add_argument("--grpc-port", type=int, default=0)
|
||||
parser.add_argument("--connect-port", type=int, default=0)
|
||||
parser.add_argument("--ca-dir", type=Path, default=None)
|
||||
arguments: Final = parser.parse_args()
|
||||
bound_state: Final = _State()
|
||||
|
||||
class BoundHandler(_Handler):
|
||||
state = bound_state
|
||||
|
||||
if arguments.grpc_port:
|
||||
grpc_server: Final = _grpc_trace_server(bound_state, arguments.grpc_port)
|
||||
assert grpc_server is not None
|
||||
if arguments.connect_port:
|
||||
assert arguments.ca_dir is not None, "--connect-port needs --ca-dir"
|
||||
bound_context: Final = _mitm_context(arguments.ca_dir)
|
||||
|
||||
class BoundConnectHandler(_ConnectHandler):
|
||||
state = bound_state
|
||||
tunnel_context = bound_context # pyright: ignore[reportIncompatibleVariableOverride] # bound context, not a new field
|
||||
|
||||
tunnel: Final = ThreadingHTTPServer(("127.0.0.1", arguments.connect_port), BoundConnectHandler)
|
||||
tunnel.daemon_threads = True
|
||||
threading.Thread(target=tunnel.serve_forever, daemon=True).start()
|
||||
server: Final = ThreadingHTTPServer(("127.0.0.1", arguments.port), BoundHandler)
|
||||
server.daemon_threads = True
|
||||
server.serve_forever()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
53
tests/integration/observability/conftest.py
Normal file
53
tests/integration/observability/conftest.py
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.otlp_sink import SpanSinks, owned_sinks
|
||||
from pydantic import JsonValue
|
||||
|
||||
AuditConfigWriter = Callable[[Path, Mapping[str, JsonValue]], Path]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def audit_sinks(tmp_path_factory: pytest.TempPathFactory) -> Iterator[SpanSinks]:
|
||||
directory: Final = tmp_path_factory.mktemp("otel-audit-sinks")
|
||||
with owned_sinks(directory) as sinks:
|
||||
yield sinks
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def otel_audit_config(audit_sinks: SpanSinks) -> AuditConfigWriter:
|
||||
tenant_host: Final = urlparse(audit_sinks.tenant).netloc
|
||||
|
||||
def write(directory: Path, litellm_settings: Mapping[str, JsonValue] = {}) -> Path:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["litellm_settings"] = {
|
||||
**config.get("litellm_settings", {}),
|
||||
"callbacks": ["otel"],
|
||||
"provider_url_destination_allowed_hosts": [tenant_host],
|
||||
**dict(litellm_settings),
|
||||
}
|
||||
config["callback_settings"] = {
|
||||
"otel": {"exporter": "http/json", "endpoint": audit_sinks.operator, "use_simple_processor": True}
|
||||
}
|
||||
config["general_settings"] = {**config.get("general_settings", {}), "user_api_key_cache_ttl": 2}
|
||||
path: Final = directory / f"otel-audit-{uuid.uuid4().hex}.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
return write
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def langfuse_vars(audit_sinks: SpanSinks) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"langfuse_public_key": "pk-lf-audit",
|
||||
"langfuse_secret_key": "sk-lf-audit",
|
||||
"langfuse_host": audit_sinks.tenant,
|
||||
}
|
||||
392
tests/integration/observability/test_otel_excluded_services.py
Normal file
392
tests/integration/observability/test_otel_excluded_services.py
Normal file
|
|
@ -0,0 +1,392 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import (
|
||||
Gateway,
|
||||
eventually,
|
||||
gateway_from_environment,
|
||||
)
|
||||
from integration._support.otlp_sink import (
|
||||
Span,
|
||||
SpanSinks,
|
||||
recorded_spans,
|
||||
spans_for_trace,
|
||||
)
|
||||
from integration._support.process import owned_proxy, owned_proxy_process
|
||||
from pydantic import JsonValue
|
||||
|
||||
AuditConfigWriter = Callable[[Path, Mapping[str, JsonValue]], Path]
|
||||
|
||||
DB_SYSTEM_KEYS: Final = frozenset({"db.system.name", "db.system"})
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def gateway(audit_sinks: SpanSinks) -> Iterator[Gateway]:
|
||||
with gateway_from_environment() as base:
|
||||
yield base
|
||||
|
||||
|
||||
def _config_with(
|
||||
directory: Path,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
*,
|
||||
otel: Mapping[str, JsonValue] = MappingProxyType({}),
|
||||
extra: Callable[[dict[str, JsonValue]], None] | None = None,
|
||||
) -> Path:
|
||||
config: Final = yaml.safe_load(otel_audit_config(directory, {}).read_text())
|
||||
config["callback_settings"]["otel"].update(dict(otel))
|
||||
if extra is not None:
|
||||
extra(config)
|
||||
path: Final = directory / f"otel-excl-{uuid.uuid4().hex}.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
def _operator_langfuse(audit_sinks: SpanSinks) -> dict[str, str]:
|
||||
return {
|
||||
"LANGFUSE_HOST": audit_sinks.operator,
|
||||
"LANGFUSE_PUBLIC_KEY": "pk-lf-operator",
|
||||
"LANGFUSE_SECRET_KEY": "sk-lf-operator",
|
||||
"OTEL_EXPORTER": "http/json",
|
||||
"OTEL_ENDPOINT": audit_sinks.operator,
|
||||
}
|
||||
|
||||
|
||||
def _add_callback(gateway: Gateway, team_id: str, callback_vars: Mapping[str, JsonValue]) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
f"/team/{team_id}/callback",
|
||||
{"callback_name": "langfuse_otel", "callback_vars": dict(callback_vars)},
|
||||
)
|
||||
|
||||
|
||||
def _drive(candidate: Gateway, langfuse_vars: Mapping[str, JsonValue]) -> httpx.Response:
|
||||
with candidate.scenario() as scenario:
|
||||
model: Final = scenario.model(model="openai/audit-chat", api_base=f"{candidate.upstream_url}/v1")
|
||||
team_id: Final = scenario.team()
|
||||
callback: Final = _add_callback(candidate, team_id, langfuse_vars)
|
||||
assert callback.status_code == 200, callback.text
|
||||
key: Final = scenario.key(team_id=team_id)
|
||||
response: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": f"otel-excl-{uuid.uuid4().hex}"}]},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
return response
|
||||
|
||||
|
||||
def _trace_id(sink_url: str, response: httpx.Response, seconds: float = 40) -> str:
|
||||
call_id: Final = response.headers.get("x-litellm-call-id")
|
||||
response_id: Final = response.json().get("id")
|
||||
|
||||
def look() -> str | None:
|
||||
_, spans = recorded_spans(sink_url)
|
||||
return next(
|
||||
(
|
||||
str(span["trace_id"])
|
||||
for span in spans
|
||||
if (call_id is not None and span["attributes"].get("litellm.call_id") == call_id)
|
||||
or (response_id is not None and span["attributes"].get("gen_ai.response.id") == response_id)
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
found: Final = eventually(look, lambda value: value is not None, seconds=seconds)
|
||||
assert found is not None
|
||||
return found
|
||||
|
||||
|
||||
def _trace_spans(sink_url: str, trace_id: str, seconds: float = 30) -> tuple[Span, ...]:
|
||||
"""The trace's spans once the post-call tail has landed.
|
||||
|
||||
The spend-writer and other post-response spans flush after the request
|
||||
answers, so absence assertions poll for the whole window instead of
|
||||
settling at the first glimpse of the root span.
|
||||
"""
|
||||
deadline: Final = time.monotonic() + seconds
|
||||
group: tuple[Span, ...] = () # rebind-ok: drains samples until the post-call tail lands
|
||||
while time.monotonic() < deadline:
|
||||
_, spans = recorded_spans(sink_url)
|
||||
group = spans_for_trace(spans, trace_id)
|
||||
time.sleep(0.5)
|
||||
assert group, f"trace {trace_id} never reached {sink_url}"
|
||||
return group
|
||||
|
||||
|
||||
def _await_db_span(sink_url: str, trace_id: str | None, needle: str, seconds: float = 40, since: int = 0) -> None:
|
||||
def seen() -> bool:
|
||||
_, spans = recorded_spans(sink_url, since)
|
||||
group: Final = spans if trace_id is None else spans_for_trace(spans, trace_id)
|
||||
return any(
|
||||
needle in str(span["name"]) or needle in {str(span["attributes"].get(k)) for k in DB_SYSTEM_KEYS}
|
||||
for span in group
|
||||
)
|
||||
|
||||
landed: Final = eventually(seen, bool, seconds=seconds)
|
||||
assert landed, f"{needle} span never landed at {sink_url}"
|
||||
|
||||
|
||||
def _db_systems(spans: tuple[Span, ...]) -> set[str]:
|
||||
return {str(span["attributes"][key]) for span in spans for key in DB_SYSTEM_KEYS if key in span["attributes"]}
|
||||
|
||||
|
||||
def _assert_core_spans_present(spans: tuple[Span, ...]) -> None:
|
||||
attributes_by_span: Final = tuple(span["attributes"] for span in spans)
|
||||
assert any(span["kind"] == 2 for span in spans), "request root span missing"
|
||||
assert any("gen_ai.operation.name" in attrs for attrs in attributes_by_span), "model span missing"
|
||||
assert any("litellm.guardrail.name" in attrs for attrs in attributes_by_span), "guardrail span missing"
|
||||
names: Final = sorted(str(span["name"]) for span in spans)
|
||||
assert any(name.startswith("auth") for name in names), f"auth span missing in {names}"
|
||||
|
||||
|
||||
def _assert_tenant_keeps_redis_without_postgres(
|
||||
candidate: Gateway, audit_sinks: SpanSinks, langfuse_vars: Mapping[str, JsonValue]
|
||||
) -> None:
|
||||
tenant_start, _ = recorded_spans(audit_sinks.tenant)
|
||||
operator_start, _ = recorded_spans(audit_sinks.operator)
|
||||
traffic: Final = _drive(candidate, langfuse_vars)
|
||||
_await_db_span(audit_sinks.operator, None, "batch_write_to_db", seconds=60, since=operator_start)
|
||||
tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic)
|
||||
_await_db_span(audit_sinks.tenant, tenant_trace, "redis", seconds=60)
|
||||
systems: Final = _db_systems(_trace_spans(audit_sinks.tenant, tenant_trace, seconds=15))
|
||||
assert "redis" in systems, f"redis spans missing at tenant: {systems}"
|
||||
_, all_tenant = recorded_spans(audit_sinks.tenant, tenant_start)
|
||||
assert "postgresql" not in _db_systems(all_tenant), f"postgresql spans reached tenant: {_db_systems(all_tenant)}"
|
||||
|
||||
|
||||
def _guardrail_block(config: dict) -> None:
|
||||
config["guardrails"] = [
|
||||
{
|
||||
"guardrail_name": f"excl-filter-{uuid.uuid4().hex[:8]}",
|
||||
"litellm_params": {
|
||||
"guardrail": "litellm_content_filter",
|
||||
"mode": "pre_call",
|
||||
"default_on": True,
|
||||
"patterns": [
|
||||
{
|
||||
"pattern_type": "regex",
|
||||
"pattern_name": "excl_secret",
|
||||
"pattern": "TOPSECRET\\d{9}",
|
||||
"action": "BLOCK",
|
||||
}
|
||||
],
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
def test_excluded_services_drops_db_spans_at_tenant_only(
|
||||
gateway: Gateway,
|
||||
audit_sinks: SpanSinks,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
langfuse_vars: dict[str, JsonValue],
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
config: Final = _config_with(
|
||||
tmp_path, otel_audit_config, otel={"excluded_services": ["redis", "postgres"]}, extra=_guardrail_block
|
||||
)
|
||||
with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as candidate:
|
||||
ten_start, _ = recorded_spans(audit_sinks.tenant)
|
||||
op_start, _ = recorded_spans(audit_sinks.operator)
|
||||
traffic: Final = _drive(candidate, langfuse_vars)
|
||||
_await_db_span(audit_sinks.operator, None, "postgresql", seconds=60, since=op_start)
|
||||
tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic)
|
||||
tenant_spans: Final = _trace_spans(audit_sinks.tenant, tenant_trace)
|
||||
_assert_core_spans_present(tenant_spans)
|
||||
assert _db_systems(tenant_spans) == set(), (
|
||||
f"db spans reached tenant: {sorted(str(s['name']) for s in tenant_spans)}"
|
||||
)
|
||||
operator_trace: Final = _trace_id(audit_sinks.operator, traffic)
|
||||
assert operator_trace == tenant_trace
|
||||
trace_systems: Final = _db_systems(_trace_spans(audit_sinks.operator, operator_trace))
|
||||
assert "redis" in trace_systems, f"operator trace lost redis spans: {trace_systems}"
|
||||
_, all_operator = recorded_spans(audit_sinks.operator, op_start)
|
||||
operator_systems: Final = _db_systems(all_operator)
|
||||
assert "postgresql" in operator_systems, f"operator lost aux db spans: {operator_systems}"
|
||||
_, all_tenant = recorded_spans(audit_sinks.tenant, ten_start)
|
||||
names: Final = sorted(str(span["name"]) for span in all_tenant)
|
||||
assert _db_systems(all_tenant) == set(), f"aux db spans reached tenant: {names}"
|
||||
assert not any("batch_write_to_db" in name for name in names), f"spend writer reached tenant: {names}"
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
def test_without_excluded_services_the_tenant_still_gets_redis_and_postgres_spans(
|
||||
gateway: Gateway,
|
||||
audit_sinks: SpanSinks,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
langfuse_vars: dict[str, JsonValue],
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
config: Final = _config_with(tmp_path, otel_audit_config, extra=_guardrail_block)
|
||||
with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as candidate:
|
||||
tenant_start, _ = recorded_spans(audit_sinks.tenant)
|
||||
traffic: Final = _drive(candidate, langfuse_vars)
|
||||
_await_db_span(audit_sinks.tenant, None, "batch_write_to_db", seconds=60, since=tenant_start)
|
||||
tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic)
|
||||
_await_db_span(audit_sinks.tenant, tenant_trace, "redis", seconds=60)
|
||||
_assert_core_spans_present(_trace_spans(audit_sinks.tenant, tenant_trace, seconds=15))
|
||||
_, all_tenant = recorded_spans(audit_sinks.tenant, tenant_start)
|
||||
systems: Final = _db_systems(all_tenant)
|
||||
assert {"redis", "postgresql"} <= systems, f"datastore spans missing at tenant: {systems}"
|
||||
|
||||
|
||||
def test_env_excluded_services_drops_only_redis(
|
||||
gateway: Gateway,
|
||||
audit_sinks: SpanSinks,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
langfuse_vars: dict[str, JsonValue],
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
config: Final = _config_with(tmp_path, otel_audit_config)
|
||||
with owned_proxy(
|
||||
gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "redis"}, config=config, workers=2
|
||||
) as candidate:
|
||||
start, _ = recorded_spans(audit_sinks.tenant)
|
||||
_drive(candidate, langfuse_vars)
|
||||
_await_db_span(audit_sinks.tenant, None, "postgresql", seconds=60, since=start)
|
||||
_, tenant_spans = recorded_spans(audit_sinks.tenant, start)
|
||||
systems: Final = _db_systems(tenant_spans)
|
||||
assert "postgresql" in systems, f"postgresql spans missing at tenant: {systems}"
|
||||
assert "redis" not in systems, f"redis spans reached tenant: {sorted(str(s['name']) for s in tenant_spans)}"
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
def test_config_excluded_services_wins_over_env(
|
||||
gateway: Gateway,
|
||||
audit_sinks: SpanSinks,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
langfuse_vars: dict[str, JsonValue],
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
def with_langfuse_otel(config: dict) -> None:
|
||||
config["litellm_settings"]["callbacks"] = ["otel", "langfuse_otel"]
|
||||
|
||||
config: Final = _config_with(
|
||||
tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]}, extra=with_langfuse_otel
|
||||
)
|
||||
with owned_proxy(
|
||||
gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "redis"}, config=config, workers=2
|
||||
) as candidate:
|
||||
_assert_tenant_keeps_redis_without_postgres(candidate, audit_sinks, langfuse_vars)
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
def test_excluded_services_applies_with_preset_ordered_first(
|
||||
gateway: Gateway,
|
||||
audit_sinks: SpanSinks,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
langfuse_vars: dict[str, JsonValue],
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
def preset_first(config: dict) -> None:
|
||||
config["litellm_settings"]["callbacks"] = ["langfuse_otel", "otel"]
|
||||
|
||||
config: Final = _config_with(
|
||||
tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]}, extra=preset_first
|
||||
)
|
||||
overrides: Final = {"LITELLM_OTEL_V2": "1", **_operator_langfuse(audit_sinks)}
|
||||
with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate:
|
||||
_assert_tenant_keeps_redis_without_postgres(candidate, audit_sinks, langfuse_vars)
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
def test_bogus_excluded_service_logs_error_and_drops_at_proxy_start(
|
||||
gateway: Gateway,
|
||||
audit_sinks: SpanSinks,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
langfuse_vars: dict[str, JsonValue],
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
config: Final = _config_with(tmp_path, otel_audit_config, otel={"excluded_services": ["auth", "postgres"]})
|
||||
with owned_proxy_process(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as owned:
|
||||
assert "'auth' is not a datastore service; ignored" in owned.log.read_text(), owned.log.read_text()[-3000:]
|
||||
_assert_tenant_keeps_redis_without_postgres(owned.gateway, audit_sinks, langfuse_vars)
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
def test_valid_config_excluded_services_tolerates_bogus_env(
|
||||
gateway: Gateway,
|
||||
audit_sinks: SpanSinks,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
langfuse_vars: dict[str, JsonValue],
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
config: Final = _config_with(tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]})
|
||||
with owned_proxy(
|
||||
gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "auth"}, config=config, workers=2
|
||||
) as candidate:
|
||||
_assert_tenant_keeps_redis_without_postgres(candidate, audit_sinks, langfuse_vars)
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
def test_bogus_excluded_services_env_logs_and_drops_with_preset_alongside_otel(
|
||||
gateway: Gateway,
|
||||
audit_sinks: SpanSinks,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
langfuse_vars: dict[str, JsonValue],
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
def with_langfuse_otel(config: dict) -> None:
|
||||
config["litellm_settings"]["callbacks"] = ["otel", "langfuse_otel"]
|
||||
|
||||
config: Final = _config_with(tmp_path, otel_audit_config, extra=with_langfuse_otel)
|
||||
overrides: Final = {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "auth,postgres"}
|
||||
with owned_proxy_process(gateway, tmp_path, overrides, config=config, workers=2) as owned:
|
||||
assert "'auth' is not a datastore service; ignored" in owned.log.read_text(), owned.log.read_text()[-3000:]
|
||||
_assert_tenant_keeps_redis_without_postgres(owned.gateway, audit_sinks, langfuse_vars)
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
def test_bogus_excluded_services_env_logs_and_drops_without_otel_callback(
|
||||
gateway: Gateway,
|
||||
audit_sinks: SpanSinks,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
langfuse_vars: dict[str, JsonValue],
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
def presets_only(config: dict) -> None:
|
||||
config["litellm_settings"]["callbacks"] = ["langfuse_otel"]
|
||||
|
||||
config: Final = _config_with(tmp_path, otel_audit_config, extra=presets_only)
|
||||
overrides: Final = {
|
||||
"LITELLM_OTEL_V2": "1",
|
||||
"LITELLM_OTEL_EXCLUDED_SERVICES": "auth,postgres",
|
||||
**_operator_langfuse(audit_sinks),
|
||||
}
|
||||
with owned_proxy_process(gateway, tmp_path, overrides, config=config, workers=2) as owned:
|
||||
assert "'auth' is not a datastore service; ignored" in owned.log.read_text(), owned.log.read_text()[-3000:]
|
||||
_assert_tenant_keeps_redis_without_postgres(owned.gateway, audit_sinks, langfuse_vars)
|
||||
|
||||
|
||||
def test_postgres_exclusion_covers_batch_write_to_db(
|
||||
gateway: Gateway,
|
||||
audit_sinks: SpanSinks,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
langfuse_vars: dict[str, JsonValue],
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
config: Final = _config_with(tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]})
|
||||
with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as candidate:
|
||||
op_start, _ = recorded_spans(audit_sinks.operator)
|
||||
ten_start, _ = recorded_spans(audit_sinks.tenant)
|
||||
traffic: Final = _drive(candidate, langfuse_vars)
|
||||
_await_db_span(audit_sinks.operator, None, "batch_write_to_db", seconds=60, since=op_start)
|
||||
tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic)
|
||||
_await_db_span(audit_sinks.tenant, tenant_trace, "redis", seconds=60)
|
||||
tenant_spans: Final = _trace_spans(audit_sinks.tenant, tenant_trace, seconds=15)
|
||||
_, all_tenant = recorded_spans(audit_sinks.tenant, ten_start)
|
||||
names: Final = sorted(str(span["name"]) for span in all_tenant)
|
||||
assert "redis" in _db_systems(tenant_spans), f"redis spans missing at tenant: {names}"
|
||||
assert not any("batch_write_to_db" in name for name in names), f"spend writer reached tenant: {names}"
|
||||
|
|
@ -0,0 +1,713 @@
|
|||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import signal
|
||||
import uuid
|
||||
from collections.abc import Callable, Generator, Iterator, Mapping
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final, Literal
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
import psutil
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value
|
||||
from integration._support.otlp_sink import Span, SpanSinks, configure_sink, recorded_spans, spans_for_trace
|
||||
from integration._support.process import OwnedProxy, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
MARKER: Final = re.compile(rb"excl-[0-9a-f]{32}")
|
||||
FAILING: Final = re.compile(rb"excl-fail-[0-9a-f]{32}")
|
||||
JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
REPLY_TEXT: Final = "excluded ok"
|
||||
SERVER: Final = 2
|
||||
INVALID_NAME_LOG: Final = "is not a datastore service"
|
||||
INVALID_VALUE_LOG: Final = "excluded_services must be"
|
||||
Endpoint = Literal["chat", "responses", "messages"]
|
||||
Client = Literal["raw", "sdk", "async_sdk"]
|
||||
ENDPOINTS: Final[tuple[Endpoint, ...]] = ("chat", "responses", "messages")
|
||||
CLIENTS: Final[tuple[Client, ...]] = ("raw", "sdk", "async_sdk")
|
||||
AuditConfigWriter = Callable[[Path, Mapping[str, JsonValue]], Path]
|
||||
|
||||
|
||||
def _marker() -> str:
|
||||
return "excl-" + uuid.uuid4().hex
|
||||
|
||||
|
||||
def _chat_reply(identity: str, stream: bool) -> Reply:
|
||||
if not stream:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [
|
||||
{"index": 0, "message": {"role": "assistant", "content": REPLY_TEXT}, "finish_reason": "stop"}
|
||||
],
|
||||
"usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
chunk: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini"}
|
||||
first, _, rest = REPLY_TEXT.partition(" ")
|
||||
deltas: Final[tuple[dict[str, JsonValue], ...]] = (
|
||||
{**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": first}}]},
|
||||
{**chunk, "choices": [{"index": 0, "delta": {"content": " " + rest}, "finish_reason": "stop"}]},
|
||||
{**chunk, "choices": [], "usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9}},
|
||||
)
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=(*(b"data: " + json.dumps(delta).encode() + b"\n\n" for delta in deltas), b"data: [DONE]\n\n"),
|
||||
)
|
||||
|
||||
|
||||
def _responses_reply(identity: str, stream: bool) -> Reply:
|
||||
response: Final[dict[str, JsonValue]] = {
|
||||
"id": identity,
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-4o-mini",
|
||||
"output": [
|
||||
{
|
||||
"id": "msg_" + identity,
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{"type": "output_text", "text": REPLY_TEXT, "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9},
|
||||
}
|
||||
if not stream:
|
||||
return Reply(body=json.dumps(response).encode())
|
||||
events: Final[tuple[dict[str, JsonValue], ...]] = (
|
||||
{
|
||||
"type": "response.created",
|
||||
"sequence_number": 0,
|
||||
"response": {**response, "status": "in_progress", "output": []},
|
||||
},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 1,
|
||||
"item_id": "msg_" + identity,
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": REPLY_TEXT,
|
||||
},
|
||||
{"type": "response.completed", "sequence_number": 2, "response": response},
|
||||
)
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events),
|
||||
)
|
||||
|
||||
|
||||
def _upstream(request: Request) -> Reply:
|
||||
if FAILING.search(request.body) is not None:
|
||||
return Reply(status=500, body=b'{"error":{"message":"scripted upstream failure","type":"server_error"}}')
|
||||
found: Final = MARKER.search(request.body)
|
||||
if found is None:
|
||||
return Reply(status=404, body=b'{"error":"no marker"}')
|
||||
marker: Final = found.group(0).decode()
|
||||
stream: Final = object_value(JSON.validate_json(request.body)).get("stream") is True
|
||||
if request.target.endswith("/responses"):
|
||||
return _responses_reply(f"resp_{marker}", stream)
|
||||
return _chat_reply(f"chatcmpl-{marker}", stream)
|
||||
|
||||
|
||||
def _at(payload: JsonValue, *path: str | int) -> JsonValue:
|
||||
if not path:
|
||||
return payload
|
||||
step: Final = path[0]
|
||||
if isinstance(step, int):
|
||||
assert isinstance(payload, list), payload
|
||||
return _at(payload[step], *path[1:])
|
||||
return _at(object_value(payload)[step], *path[1:])
|
||||
|
||||
|
||||
def _sse(body: str) -> tuple[JsonValue, ...]:
|
||||
return tuple(
|
||||
JSON.validate_json(line[6:])
|
||||
for line in body.splitlines()
|
||||
if line.startswith("data: ") and line != "data: [DONE]"
|
||||
)
|
||||
|
||||
|
||||
def _raw_text(endpoint: Endpoint, stream: bool, body: str) -> str:
|
||||
if not stream:
|
||||
path: Final[tuple[str | int, ...]] = {
|
||||
"chat": ("choices", 0, "message", "content"),
|
||||
"responses": ("output", 0, "content", 0, "text"),
|
||||
"messages": ("content", 0, "text"),
|
||||
}[endpoint]
|
||||
return str(_at(JSON.validate_json(body), *path))
|
||||
events: Final = _sse(body)
|
||||
if endpoint == "chat":
|
||||
return "".join(
|
||||
str(object_value(_at(event, "choices", 0, "delta")).get("content") or "")
|
||||
for event in events
|
||||
if _at(event, "choices")
|
||||
)
|
||||
if endpoint == "responses":
|
||||
return "".join(
|
||||
str(_at(event, "delta")) for event in events if _at(event, "type") == "response.output_text.delta"
|
||||
)
|
||||
return "".join(
|
||||
str(_at(event, "delta", "text"))
|
||||
for event in events
|
||||
if _at(event, "type") == "content_block_delta" and _at(event, "delta", "type") == "text_delta"
|
||||
)
|
||||
|
||||
|
||||
def _body(model: str, endpoint: Endpoint, marker: str, stream: bool) -> tuple[str, dict[str, JsonValue]]:
|
||||
if endpoint == "chat":
|
||||
return "/v1/chat/completions", {
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": marker}],
|
||||
"stream": stream,
|
||||
}
|
||||
if endpoint == "responses":
|
||||
return "/v1/responses", {"model": model, "input": marker, "stream": stream}
|
||||
return "/v1/messages", {
|
||||
"model": model,
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": marker}],
|
||||
"stream": stream,
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Sent:
|
||||
call_id: str
|
||||
text: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Cursors:
|
||||
operator: int
|
||||
tenant: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
proxy: Gateway
|
||||
owned: OwnedProxy
|
||||
scenario: Scenario
|
||||
model: str
|
||||
key: str
|
||||
upstream: Wire
|
||||
sinks: SpanSinks
|
||||
|
||||
def cursors(self) -> Cursors:
|
||||
self.upstream.drain()
|
||||
return Cursors(recorded_spans(self.sinks.operator)[0], recorded_spans(self.sinks.tenant)[0])
|
||||
|
||||
def upstream_hits(self, marker: str) -> int:
|
||||
return sum(1 for request in self.upstream.drain() if marker.encode() in request.body)
|
||||
|
||||
def base_url(self) -> str:
|
||||
return str(self.proxy.client.base_url)
|
||||
|
||||
def raw(
|
||||
self, endpoint: Endpoint, marker: str, stream: bool, key: str | None = None, trace_id: str | None = None
|
||||
) -> Sent:
|
||||
path, body = _body(self.model, endpoint, marker, stream)
|
||||
auth: Final = {"Authorization": f"Bearer {key or self.key}"}
|
||||
parent: Final = {} if trace_id is None else {"traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01"}
|
||||
with self.proxy.client.stream("POST", path, json=body, headers={**auth, **parent}) as response:
|
||||
text: Final = response.read().decode()
|
||||
assert response.status_code == 200, text
|
||||
return Sent(response.headers["x-litellm-call-id"], _raw_text(endpoint, stream, text))
|
||||
|
||||
def sdk(self, endpoint: Endpoint, marker: str, stream: bool) -> Sent:
|
||||
if endpoint == "messages":
|
||||
messages: Final = anthropic.Anthropic(base_url=self.base_url(), api_key=self.key, max_retries=0).messages
|
||||
if not stream:
|
||||
reply: Final = messages.with_raw_response.create(
|
||||
model=self.model, max_tokens=16, messages=[{"role": "user", "content": marker}]
|
||||
)
|
||||
block: Final = reply.parse().content[0]
|
||||
assert isinstance(block, anthropic.types.TextBlock), block
|
||||
return Sent(reply.headers["x-litellm-call-id"], block.text)
|
||||
with messages.with_streaming_response.create(
|
||||
model=self.model, max_tokens=16, messages=[{"role": "user", "content": marker}], stream=True
|
||||
) as streamed:
|
||||
return Sent(
|
||||
streamed.headers["x-litellm-call-id"],
|
||||
"".join(
|
||||
event.delta.text
|
||||
for event in streamed.parse()
|
||||
if event.type == "content_block_delta" and event.delta.type == "text_delta"
|
||||
),
|
||||
)
|
||||
client: Final = openai.OpenAI(base_url=self.base_url() + "/v1", api_key=self.key, max_retries=0)
|
||||
if endpoint == "chat":
|
||||
if not stream:
|
||||
completion: Final = client.chat.completions.with_raw_response.create(
|
||||
model=self.model, messages=[{"role": "user", "content": marker}]
|
||||
)
|
||||
return Sent(
|
||||
completion.headers["x-litellm-call-id"], completion.parse().choices[0].message.content or ""
|
||||
)
|
||||
with client.chat.completions.with_streaming_response.create(
|
||||
model=self.model, messages=[{"role": "user", "content": marker}], stream=True
|
||||
) as chunks:
|
||||
return Sent(
|
||||
chunks.headers["x-litellm-call-id"],
|
||||
"".join(chunk.choices[0].delta.content or "" for chunk in chunks.parse() if chunk.choices),
|
||||
)
|
||||
if not stream:
|
||||
created: Final = client.responses.with_raw_response.create(model=self.model, input=marker)
|
||||
return Sent(created.headers["x-litellm-call-id"], created.parse().output_text)
|
||||
with client.responses.with_streaming_response.create(model=self.model, input=marker, stream=True) as events:
|
||||
return Sent(
|
||||
events.headers["x-litellm-call-id"],
|
||||
"".join(event.delta for event in events.parse() if event.type == "response.output_text.delta"),
|
||||
)
|
||||
|
||||
async def async_sdk(self, endpoint: Endpoint, marker: str, stream: bool) -> Sent:
|
||||
if endpoint == "messages":
|
||||
messages: Final = anthropic.AsyncAnthropic(
|
||||
base_url=self.base_url(), api_key=self.key, max_retries=0
|
||||
).messages
|
||||
if not stream:
|
||||
reply: Final = await messages.with_raw_response.create(
|
||||
model=self.model, max_tokens=16, messages=[{"role": "user", "content": marker}]
|
||||
)
|
||||
block: Final = reply.parse().content[0]
|
||||
assert isinstance(block, anthropic.types.TextBlock), block
|
||||
return Sent(reply.headers["x-litellm-call-id"], block.text)
|
||||
async with messages.with_streaming_response.create(
|
||||
model=self.model, max_tokens=16, messages=[{"role": "user", "content": marker}], stream=True
|
||||
) as streamed:
|
||||
pieces: Final = [
|
||||
event.delta.text
|
||||
async for event in await streamed.parse()
|
||||
if event.type == "content_block_delta" and event.delta.type == "text_delta"
|
||||
]
|
||||
return Sent(streamed.headers["x-litellm-call-id"], "".join(pieces))
|
||||
client: Final = openai.AsyncOpenAI(base_url=self.base_url() + "/v1", api_key=self.key, max_retries=0)
|
||||
if endpoint == "chat":
|
||||
if not stream:
|
||||
completion: Final = await client.chat.completions.with_raw_response.create(
|
||||
model=self.model, messages=[{"role": "user", "content": marker}]
|
||||
)
|
||||
return Sent(
|
||||
completion.headers["x-litellm-call-id"], completion.parse().choices[0].message.content or ""
|
||||
)
|
||||
async with client.chat.completions.with_streaming_response.create(
|
||||
model=self.model, messages=[{"role": "user", "content": marker}], stream=True
|
||||
) as chunks:
|
||||
deltas: Final = [
|
||||
chunk.choices[0].delta.content or "" async for chunk in await chunks.parse() if chunk.choices
|
||||
]
|
||||
return Sent(chunks.headers["x-litellm-call-id"], "".join(deltas))
|
||||
if not stream:
|
||||
created: Final = await client.responses.with_raw_response.create(model=self.model, input=marker)
|
||||
return Sent(created.headers["x-litellm-call-id"], created.parse().output_text)
|
||||
async with client.responses.with_streaming_response.create(
|
||||
model=self.model, input=marker, stream=True
|
||||
) as events:
|
||||
texts: Final = [
|
||||
event.delta async for event in await events.parse() if event.type == "response.output_text.delta"
|
||||
]
|
||||
return Sent(events.headers["x-litellm-call-id"], "".join(texts))
|
||||
|
||||
def send(self, endpoint: Endpoint, client: Client, marker: str, stream: bool) -> Sent:
|
||||
if client == "raw":
|
||||
return self.raw(endpoint, marker, stream)
|
||||
if client == "sdk":
|
||||
return self.sdk(endpoint, marker, stream)
|
||||
return asyncio.run(self.async_sdk(endpoint, marker, stream))
|
||||
|
||||
|
||||
def _db_systems(spans: tuple[Span, ...]) -> set[str]:
|
||||
return {
|
||||
str(system)
|
||||
for span in spans
|
||||
if (system := span["attributes"].get("db.system.name") or span["attributes"].get("db.system")) is not None
|
||||
}
|
||||
|
||||
|
||||
def _names(spans: tuple[Span, ...]) -> list[str]:
|
||||
return sorted(span["name"] for span in spans)
|
||||
|
||||
|
||||
def _has_root(spans: tuple[Span, ...]) -> bool:
|
||||
return any(span["kind"] == SERVER for span in spans)
|
||||
|
||||
|
||||
def _trace_of_call(sink: str, call_id: str, since: int) -> tuple[Span, ...]:
|
||||
_, spans = recorded_spans(sink, since)
|
||||
traces: Final = {span["trace_id"] for span in spans if span["attributes"].get("litellm.call_id") == call_id}
|
||||
return tuple(span for span in spans if span["trace_id"] in traces)
|
||||
|
||||
|
||||
def _operator_trace(rig: Rig, sent: Sent, cursors: Cursors) -> tuple[Span, ...]:
|
||||
trace: Final = eventually(
|
||||
lambda: _trace_of_call(rig.sinks.operator, sent.call_id, cursors.operator),
|
||||
lambda spans: _has_root(spans) and "redis" in _db_systems(spans),
|
||||
seconds=40,
|
||||
)
|
||||
assert len({span["trace_id"] for span in trace}) == 1, _names(trace)
|
||||
return trace
|
||||
|
||||
|
||||
def _traced_raw(rig: Rig, endpoint: Endpoint, marker: str) -> tuple[str, Sent]:
|
||||
trace_id: Final = uuid.uuid4().hex
|
||||
return trace_id, rig.raw(endpoint, marker, stream=False, trace_id=trace_id)
|
||||
|
||||
|
||||
def _operator_trace_by_id(rig: Rig, trace_id: str, cursors: Cursors) -> tuple[Span, ...]:
|
||||
return eventually(
|
||||
lambda: spans_for_trace(recorded_spans(rig.sinks.operator, cursors.operator)[1], trace_id),
|
||||
_has_root,
|
||||
seconds=40,
|
||||
)
|
||||
|
||||
|
||||
def _tenant_mirror(rig: Rig, operator: tuple[Span, ...], cursors: Cursors) -> tuple[Span, ...]:
|
||||
kept: Final = frozenset(span["name"] for span in operator if not _db_systems((span,)))
|
||||
return eventually(
|
||||
lambda: spans_for_trace(recorded_spans(rig.sinks.tenant, cursors.tenant)[1], operator[0]["trace_id"]),
|
||||
lambda spans: kept <= {span["name"] for span in spans},
|
||||
seconds=40,
|
||||
)
|
||||
|
||||
|
||||
def _assert_tenant_mirrors(rig: Rig, operator: tuple[Span, ...], cursors: Cursors) -> tuple[Span, ...]:
|
||||
tenant: Final = _tenant_mirror(rig, operator, cursors)
|
||||
assert _db_systems(tenant) == set(), f"datastore spans reached the tenant: {_names(tenant)}"
|
||||
assert sum(1 for span in tenant if span["kind"] == SERVER) == 1, _names(tenant)
|
||||
return tenant
|
||||
|
||||
|
||||
def _assert_withheld(rig: Rig, sent: Sent, cursors: Cursors) -> tuple[Span, ...]:
|
||||
tenant: Final = _assert_tenant_mirrors(rig, _operator_trace(rig, sent, cursors), cursors)
|
||||
assert any("gen_ai.operation.name" in span["attributes"] for span in tenant), _names(tenant)
|
||||
return tenant
|
||||
|
||||
|
||||
def _config(directory: Path, otel_audit_config: AuditConfigWriter, otel: Mapping[str, JsonValue], name: str) -> Path:
|
||||
written: Final = otel_audit_config(directory, {})
|
||||
loaded: Final = object_value(JSON.validate_python(yaml.safe_load(written.read_text())))
|
||||
settings: Final = object_value(loaded["callback_settings"])
|
||||
config: Final = {**loaded, "callback_settings": {**settings, "otel": {**object_value(settings["otel"]), **otel}}}
|
||||
path: Final = directory / f"{name}.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _started(
|
||||
provider: Wire,
|
||||
sinks: SpanSinks,
|
||||
config: Path,
|
||||
directory: Path,
|
||||
langfuse_vars: Mapping[str, JsonValue],
|
||||
workers: int,
|
||||
) -> Generator[Rig]:
|
||||
with (
|
||||
gateway_from_environment() as gateway,
|
||||
owned_proxy_process(
|
||||
gateway,
|
||||
directory,
|
||||
{"LITELLM_OTEL_V2": "1", "OTEL_BSP_SCHEDULE_DELAY": "300"},
|
||||
config=config,
|
||||
remove_environment=("LITELLM_OTEL_EXCLUDED_SERVICES",),
|
||||
workers=workers,
|
||||
) as owned,
|
||||
owned.gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
team: Final = scenario.team()
|
||||
attached: Final = owned.gateway.request(
|
||||
"POST", f"/team/{team}/callback", {"callback_name": "langfuse_otel", "callback_vars": dict(langfuse_vars)}
|
||||
)
|
||||
assert attached.status_code == 200, attached.text
|
||||
yield Rig(owned.gateway, owned, scenario, model, scenario.key(team_id=team), provider, sinks)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def provider() -> Iterator[Wire]:
|
||||
with wire_server(_upstream) as wire:
|
||||
yield wire
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(
|
||||
provider: Wire,
|
||||
audit_sinks: SpanSinks,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
langfuse_vars: dict[str, JsonValue],
|
||||
tmp_path_factory: pytest.TempPathFactory,
|
||||
) -> Iterator[Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("excluded-matrix")
|
||||
config: Final = _config(directory, otel_audit_config, {"excluded_services": ["redis", "postgres"]}, "matrix")
|
||||
with _started(provider, audit_sinks, config, directory, langfuse_vars, workers=2) as started:
|
||||
yield started
|
||||
|
||||
|
||||
@pytest.mark.timeout(120)
|
||||
@pytest.mark.parametrize("stream", [False, True], ids=["unary", "stream"])
|
||||
@pytest.mark.parametrize("client", CLIENTS)
|
||||
@pytest.mark.parametrize("endpoint", ENDPOINTS)
|
||||
def test_tenant_trace_keeps_request_spans_without_datastore_spans(
|
||||
rig: Rig, endpoint: Endpoint, client: Client, stream: bool
|
||||
) -> None:
|
||||
cursors: Final = rig.cursors()
|
||||
marker: Final = _marker()
|
||||
sent: Final = rig.send(endpoint, client, marker, stream)
|
||||
assert sent.text == REPLY_TEXT, sent
|
||||
assert rig.upstream_hits(marker) == 1
|
||||
_assert_withheld(rig, sent, cursors)
|
||||
|
||||
|
||||
@pytest.mark.timeout(120)
|
||||
@pytest.mark.parametrize("endpoint", ["chat", "messages"])
|
||||
def test_cache_hit_twin_keeps_datastore_spans_off_the_tenant(rig: Rig, endpoint: Endpoint) -> None:
|
||||
marker: Final = _marker()
|
||||
first: Final = rig.raw(endpoint, marker, stream=False)
|
||||
assert first.text == REPLY_TEXT, first
|
||||
assert rig.upstream_hits(marker) == 1
|
||||
cursors: Final = rig.cursors()
|
||||
trace_id, hit = eventually(
|
||||
lambda: _traced_raw(rig, endpoint, marker), lambda sent: rig.upstream_hits(marker) == 0, seconds=20
|
||||
)
|
||||
assert hit.text == REPLY_TEXT, hit
|
||||
_assert_tenant_mirrors(rig, _operator_trace_by_id(rig, trace_id, cursors), cursors)
|
||||
|
||||
|
||||
@pytest.mark.timeout(120)
|
||||
@pytest.mark.parametrize("endpoint", ENDPOINTS)
|
||||
def test_failed_upstream_call_keeps_datastore_spans_off_the_tenant(rig: Rig, endpoint: Endpoint) -> None:
|
||||
cursors: Final = rig.cursors()
|
||||
marker: Final = "excl-fail-" + uuid.uuid4().hex
|
||||
trace_id: Final = uuid.uuid4().hex
|
||||
path, body = _body(rig.model, endpoint, marker, stream=False)
|
||||
failed: Final = rig.proxy.client.post(
|
||||
path,
|
||||
json=body,
|
||||
headers={"Authorization": f"Bearer {rig.key}", "traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01"},
|
||||
)
|
||||
assert failed.status_code == 500, failed.text
|
||||
assert rig.upstream_hits(marker) >= 1
|
||||
operator: Final = eventually(
|
||||
lambda: spans_for_trace(recorded_spans(rig.sinks.operator, cursors.operator)[1], trace_id),
|
||||
lambda spans: _has_root(spans) and "redis" in _db_systems(spans),
|
||||
seconds=40,
|
||||
)
|
||||
_assert_tenant_mirrors(rig, operator, cursors)
|
||||
|
||||
|
||||
@pytest.mark.timeout(120)
|
||||
def test_key_level_callback_vars_destination_is_filtered_too(rig: Rig, langfuse_vars: dict[str, JsonValue]) -> None:
|
||||
key: Final = rig.scenario.key(
|
||||
metadata={
|
||||
"logging": [
|
||||
{"callback_name": "langfuse_otel", "callback_type": "success", "callback_vars": dict(langfuse_vars)}
|
||||
]
|
||||
}
|
||||
)
|
||||
cursors: Final = rig.cursors()
|
||||
marker: Final = _marker()
|
||||
sent: Final = rig.raw("chat", marker, stream=False, key=key)
|
||||
assert sent.text == REPLY_TEXT, sent
|
||||
assert rig.upstream_hits(marker) == 1
|
||||
_assert_withheld(rig, sent, cursors)
|
||||
|
||||
|
||||
@pytest.mark.timeout(120)
|
||||
@pytest.mark.parametrize("status", [403, 404])
|
||||
def test_rejecting_tenant_destination_leaves_serving_and_the_operator_trace_intact(rig: Rig, status: int) -> None:
|
||||
configure_sink(rig.sinks.tenant, status=status)
|
||||
try:
|
||||
cursors: Final = rig.cursors()
|
||||
marker: Final = _marker()
|
||||
sent: Final = rig.raw("chat", marker, stream=True)
|
||||
assert sent.text == REPLY_TEXT, sent
|
||||
assert rig.upstream_hits(marker) == 1
|
||||
_assert_withheld(rig, sent, cursors)
|
||||
finally:
|
||||
configure_sink(rig.sinks.tenant, status=200)
|
||||
after: Final = rig.cursors()
|
||||
_assert_withheld(rig, rig.raw("responses", _marker(), stream=False), after)
|
||||
|
||||
|
||||
def _burst(rig: Rig, count: int) -> tuple[Sent | str, ...]:
|
||||
def one(index: int) -> Sent | str:
|
||||
try:
|
||||
return rig.raw(ENDPOINTS[index % 3], _marker(), stream=index % 2 == 0)
|
||||
except (httpx.HTTPError, AssertionError) as error:
|
||||
return repr(error)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=10) as pool:
|
||||
return tuple(pool.map(one, range(count)))
|
||||
|
||||
|
||||
def _served(results: tuple[Sent | str, ...]) -> tuple[Sent, ...]:
|
||||
return tuple(result for result in results if isinstance(result, Sent))
|
||||
|
||||
|
||||
def _assert_operator_exactly_once(rig: Rig, served: tuple[Sent, ...], cursors: Cursors) -> set[str]:
|
||||
wanted: Final = {sent.call_id for sent in served}
|
||||
|
||||
def roots() -> dict[str, int]:
|
||||
_, spans = recorded_spans(rig.sinks.operator, cursors.operator)
|
||||
traced: Final = {
|
||||
span["trace_id"]: str(span["attributes"]["litellm.call_id"])
|
||||
for span in spans
|
||||
if span["attributes"].get("litellm.call_id") in wanted
|
||||
}
|
||||
counts: Final = {call: 0 for call in wanted}
|
||||
for span in spans:
|
||||
if span["kind"] == SERVER and span["trace_id"] in traced:
|
||||
counts[traced[span["trace_id"]]] += 1
|
||||
return counts
|
||||
|
||||
landed: Final = eventually(roots, lambda counts: all(count >= 1 for count in counts.values()), seconds=90)
|
||||
assert landed == {call: 1 for call in wanted}, landed
|
||||
_, spans = recorded_spans(rig.sinks.operator, cursors.operator)
|
||||
return {span["trace_id"] for span in spans if span["attributes"].get("litellm.call_id") in wanted}
|
||||
|
||||
|
||||
def _assert_tenant_never_saw_datastore_spans(rig: Rig, cursors: Cursors, traces: set[str]) -> None:
|
||||
tenant: Final = eventually(
|
||||
lambda: recorded_spans(rig.sinks.tenant, cursors.tenant)[1],
|
||||
lambda spans: traces <= {span["trace_id"] for span in spans if span["kind"] == SERVER},
|
||||
seconds=90,
|
||||
)
|
||||
assert _db_systems(tenant) == set(), _names(tenant)
|
||||
|
||||
|
||||
@pytest.mark.timeout(300)
|
||||
def test_tenant_outage_during_a_mixed_burst_keeps_serving_and_never_leaks_datastore_spans(rig: Rig) -> None:
|
||||
cursors: Final = rig.cursors()
|
||||
configure_sink(rig.sinks.tenant, status=503)
|
||||
try:
|
||||
results: Final = _burst(rig, 30)
|
||||
finally:
|
||||
configure_sink(rig.sinks.tenant, status=200)
|
||||
served: Final = _served(results)
|
||||
assert len(served) == 30, [result for result in results if isinstance(result, str)]
|
||||
assert all(sent.text == REPLY_TEXT for sent in served), served
|
||||
traces: Final = _assert_operator_exactly_once(rig, served, cursors)
|
||||
_assert_tenant_never_saw_datastore_spans(rig, cursors, traces)
|
||||
after: Final = rig.cursors()
|
||||
_assert_withheld(rig, rig.raw("messages", _marker(), stream=True), after)
|
||||
|
||||
|
||||
@pytest.mark.timeout(300)
|
||||
def test_stalled_tenant_destination_during_a_burst_does_not_block_responses(rig: Rig) -> None:
|
||||
cursors: Final = rig.cursors()
|
||||
configure_sink(rig.sinks.tenant, paused=True)
|
||||
try:
|
||||
results: Final = _burst(rig, 20)
|
||||
finally:
|
||||
configure_sink(rig.sinks.tenant, paused=False)
|
||||
served: Final = _served(results)
|
||||
assert len(served) == 20, [result for result in results if isinstance(result, str)]
|
||||
traces: Final = _assert_operator_exactly_once(rig, served, cursors)
|
||||
_assert_tenant_never_saw_datastore_spans(rig, cursors, traces)
|
||||
|
||||
|
||||
@pytest.mark.timeout(300)
|
||||
def test_killing_one_of_two_workers_mid_burst_keeps_the_filter_on_the_survivor(rig: Rig) -> None:
|
||||
root: Final = psutil.Process(rig.owned.process.pid)
|
||||
workers: Final = eventually(
|
||||
lambda: tuple(child for child in root.children() if "resource_tracker" not in " ".join(child.cmdline())),
|
||||
lambda found: len(found) == 2,
|
||||
seconds=30,
|
||||
)
|
||||
cursors: Final = rig.cursors()
|
||||
|
||||
def one(index: int) -> Sent | str:
|
||||
if index == 6:
|
||||
os.kill(workers[0].pid, signal.SIGKILL)
|
||||
try:
|
||||
return rig.raw("chat", _marker(), stream=index % 2 == 0)
|
||||
except (httpx.HTTPError, AssertionError) as error:
|
||||
return repr(error)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=6) as pool:
|
||||
results: Final = tuple(pool.map(one, range(18)))
|
||||
assert rig.owned.process.poll() is None, "Proxy root exited after a worker was killed"
|
||||
failures: Final = tuple(result for result in results if isinstance(result, str))
|
||||
assert all(failure.startswith(("ReadError(", "RemoteProtocolError(", "ConnectError(")) for failure in failures), (
|
||||
failures
|
||||
)
|
||||
assert len(failures) <= 6, failures
|
||||
settled: Final = tuple(result for index, result in enumerate(results) if index > 12 and isinstance(result, Sent))
|
||||
traces: Final = _assert_operator_exactly_once(rig, settled, cursors)
|
||||
_assert_tenant_never_saw_datastore_spans(rig, cursors, traces)
|
||||
after: Final = rig.cursors()
|
||||
_assert_withheld(rig, rig.raw("chat", _marker(), stream=False), after)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Setting:
|
||||
otel: Mapping[str, JsonValue]
|
||||
withholds_redis: bool
|
||||
logs: str | None
|
||||
|
||||
|
||||
SETTINGS: Final[dict[str, Setting]] = {
|
||||
"missing": Setting({}, False, None),
|
||||
"null": Setting({"excluded_services": None}, False, None),
|
||||
"empty_list": Setting({"excluded_services": []}, False, None),
|
||||
"empty_string": Setting({"excluded_services": ""}, False, None),
|
||||
"yaml_string": Setting({"excluded_services": "redis"}, True, None),
|
||||
"duplicates": Setting({"excluded_services": ["redis", "redis"]}, True, None),
|
||||
"case_and_space": Setting({"excluded_services": ["REDIS", " Postgres "]}, True, None),
|
||||
"integer": Setting({"excluded_services": 7}, False, INVALID_VALUE_LOG),
|
||||
"mapping": Setting({"excluded_services": {"redis": True}}, False, INVALID_VALUE_LOG),
|
||||
"non_string_item": Setting({"excluded_services": [7, "redis"]}, True, INVALID_VALUE_LOG),
|
||||
"oversized_name": Setting({"excluded_services": "x" * 5000}, False, INVALID_NAME_LOG),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
@pytest.mark.parametrize("name", SETTINGS)
|
||||
def test_excluded_services_setting_shapes_boot_and_resolve(
|
||||
name: str,
|
||||
provider: Wire,
|
||||
audit_sinks: SpanSinks,
|
||||
otel_audit_config: AuditConfigWriter,
|
||||
langfuse_vars: dict[str, JsonValue],
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
setting: Final = SETTINGS[name]
|
||||
config: Final = _config(tmp_path, otel_audit_config, setting.otel, name)
|
||||
with _started(provider, audit_sinks, config, tmp_path, langfuse_vars, workers=1) as started:
|
||||
cursors: Final = started.cursors()
|
||||
marker: Final = _marker()
|
||||
sent: Final = started.raw("chat", marker, stream=False)
|
||||
assert sent.text == REPLY_TEXT, sent
|
||||
assert started.upstream_hits(marker) == 1
|
||||
operator: Final = _operator_trace(started, sent, cursors)
|
||||
tenant: Final = _tenant_mirror(started, operator, cursors)
|
||||
if setting.withholds_redis:
|
||||
assert "redis" not in _db_systems(tenant), _names(tenant)
|
||||
else:
|
||||
eventually(
|
||||
lambda: _db_systems(
|
||||
spans_for_trace(recorded_spans(started.sinks.tenant, cursors.tenant)[1], tenant[0]["trace_id"])
|
||||
),
|
||||
lambda systems: "redis" in systems,
|
||||
seconds=30,
|
||||
)
|
||||
log: Final = started.owned.log.read_text()
|
||||
if setting.logs is None:
|
||||
assert INVALID_NAME_LOG not in log and INVALID_VALUE_LOG not in log, log[-2000:]
|
||||
else:
|
||||
assert setting.logs in log, log[-4000:]
|
||||
|
|
@ -11,6 +11,7 @@
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -22,17 +23,18 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( # noqa: E4
|
|||
)
|
||||
|
||||
from litellm.integrations.otel import LiteLLM, OpenTelemetryV2Config # noqa: E402
|
||||
from litellm.integrations.otel.plumbing import providers # noqa: E402
|
||||
from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402
|
||||
from litellm.integrations.otel.model.baggage import ( # noqa: E402
|
||||
BAGGAGE_PROMOTED_KEYS,
|
||||
DEFAULT_BAGGAGE_METADATA_KEYS,
|
||||
)
|
||||
from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402
|
||||
from litellm.integrations.otel.model.config import excluded_db_systems_from # noqa: E402
|
||||
from litellm.integrations.otel.model.payloads import GuardrailSpanData # noqa: E402
|
||||
from litellm.integrations.otel.model.spans import ( # noqa: E402
|
||||
LITELLM_PROXY_REQUEST_SPAN_NAME,
|
||||
SpanRole,
|
||||
)
|
||||
from litellm.integrations.otel.plumbing import providers # noqa: E402
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Area 1 — baggage allowlists configurable
|
||||
|
|
@ -74,13 +76,11 @@ def test_baggage_keys_from_config_yaml_kwargs():
|
|||
|
||||
|
||||
def test_baggage_processor_allowlist_uses_config_keys():
|
||||
cfg = OpenTelemetryV2Config(
|
||||
exporter="in_memory", baggage_promoted_keys=[LiteLLM.TEAM_ID]
|
||||
)
|
||||
cfg = OpenTelemetryV2Config(exporter="in_memory", baggage_promoted_keys=[LiteLLM.TEAM_ID])
|
||||
provider, exporter = providers.in_memory_provider(cfg)
|
||||
from litellm.integrations.otel.plumbing import context as ctx_mod
|
||||
from litellm.integrations.otel.emitter import SpanEmitter
|
||||
from litellm.integrations.otel.model.payloads import ServiceSpanData
|
||||
from litellm.integrations.otel.plumbing import context as ctx_mod
|
||||
|
||||
engine = SpanEmitter(providers.get_tracer(provider, "t"), cfg)
|
||||
ctx = ctx_mod.set_request_baggage({LiteLLM.TEAM_ID: "t1", LiteLLM.TEAM_ALIAS: "ta"})
|
||||
|
|
@ -90,6 +90,68 @@ def test_baggage_processor_allowlist_uses_config_keys():
|
|||
assert LiteLLM.TEAM_ALIAS not in span.attributes # not in this allowlist
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"given,expected",
|
||||
[
|
||||
(["redis"], frozenset({"redis"})),
|
||||
(["postgres"], frozenset({"postgresql"})),
|
||||
(["postgresql"], frozenset({"postgresql"})),
|
||||
(["batch_write_to_db"], frozenset({"postgresql"})),
|
||||
(["redis_spend_update_queue"], frozenset({"redis"})),
|
||||
(["redis", "postgres"], frozenset({"redis", "postgresql"})),
|
||||
],
|
||||
)
|
||||
def test_excluded_services_normalize_to_db_system_names(given, expected):
|
||||
assert OpenTelemetryV2Config(excluded_services=given).excluded_services == expected
|
||||
|
||||
|
||||
def test_excluded_services_from_env_csv(monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_OTEL_EXCLUDED_SERVICES", "redis, postgres")
|
||||
assert OpenTelemetryV2Config().excluded_services == frozenset({"redis", "postgresql"})
|
||||
|
||||
|
||||
def test_excluded_services_config_wins_over_env(monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_OTEL_EXCLUDED_SERVICES", "redis")
|
||||
assert OpenTelemetryV2Config(excluded_services=["postgres"]).excluded_services == frozenset({"postgresql"})
|
||||
|
||||
|
||||
def test_excluded_services_drops_a_non_datastore_service_and_logs(caplog):
|
||||
with caplog.at_level(logging.ERROR, logger="LiteLLM"):
|
||||
config = OpenTelemetryV2Config(excluded_services=["auth", "redis"])
|
||||
assert config.excluded_services == frozenset({"redis"})
|
||||
assert any("'auth' is not a datastore service; ignored" in record.message for record in caplog.records)
|
||||
|
||||
|
||||
def test_excluded_services_env_drops_a_bad_value_and_logs(monkeypatch, caplog):
|
||||
monkeypatch.setenv("LITELLM_OTEL_EXCLUDED_SERVICES", "auth,postgres")
|
||||
with caplog.at_level(logging.ERROR, logger="LiteLLM"):
|
||||
config = OpenTelemetryV2Config()
|
||||
assert config.excluded_services == frozenset({"postgresql"})
|
||||
assert any("'auth' is not a datastore service; ignored" in record.message for record in caplog.records)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"given,expected,logged",
|
||||
[
|
||||
(None, frozenset(), None),
|
||||
("", frozenset(), None),
|
||||
([], frozenset(), None),
|
||||
(["REDIS", " Postgres "], frozenset({"redis", "postgresql"}), None),
|
||||
(7, frozenset(), "excluded_services must be a list or comma-separated string; 7 ignored"),
|
||||
({"redis": True}, frozenset(), "excluded_services must be a list or comma-separated string"),
|
||||
([7, "redis"], frozenset({"redis"}), "excluded_services must be a list of service names; 7 ignored"),
|
||||
],
|
||||
)
|
||||
def test_malformed_excluded_services_logs_and_still_builds_the_config(given, expected, logged, caplog):
|
||||
with caplog.at_level(logging.ERROR, logger="LiteLLM"):
|
||||
config = OpenTelemetryV2Config(excluded_services=given)
|
||||
resolved = excluded_db_systems_from(given)
|
||||
assert config.excluded_services == expected
|
||||
assert resolved == expected
|
||||
messages = [record.message for record in caplog.records]
|
||||
assert (logged is None and messages == []) or any(logged in message for message in messages), messages
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Area 2 — pass-through LLM span parents to the ambient server span
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
|
@ -124,9 +186,7 @@ def test_passthrough_llm_span_parents_to_ambient_server_span():
|
|||
later (possibly detached) success callback only closes the already-parented
|
||||
span, so it never becomes a separate root trace."""
|
||||
logger, exporter = _logger()
|
||||
server = logger._emitter.start_span(
|
||||
SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
|
||||
)
|
||||
server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME)
|
||||
kwargs = {
|
||||
"standard_logging_object": _payload(),
|
||||
"litellm_params": {"metadata": {}},
|
||||
|
|
@ -150,9 +210,7 @@ def test_llm_span_unaffected_by_phase_span_active_at_close():
|
|||
successor to the old auth-failure-401 case where the LLM log nested under
|
||||
``auth``: the span is now born after auth, parented to the request root."""
|
||||
logger, exporter = _logger()
|
||||
server = logger._emitter.start_span(
|
||||
SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
|
||||
)
|
||||
server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME)
|
||||
kwargs = {
|
||||
"standard_logging_object": _payload(),
|
||||
"litellm_params": {"metadata": {}},
|
||||
|
|
|
|||
|
|
@ -514,6 +514,53 @@ class TestFanOut:
|
|||
for child in ("auth /v1/chat/completions", "chat gpt-4"):
|
||||
assert by_name[child].parent.span_id == root.context.span_id
|
||||
|
||||
def test_excluded_services_drop_only_the_datastore_spans_at_the_tenant(self):
|
||||
"""The exclusion is per ``db.system.*`` value: a span naming an excluded
|
||||
datastore never reaches the tenant, while every span of the request's
|
||||
own work (root, auth, guardrail, model) still does, and the operator's
|
||||
own exporter keeps the full tree."""
|
||||
dest_exporter, operator_exporter = InMemorySpanExporter(), InMemorySpanExporter()
|
||||
provider = TracerProvider()
|
||||
provider.add_span_processor(SimpleSpanProcessor(operator_exporter))
|
||||
provider.add_span_processor(
|
||||
TenantFanOutSpanProcessor(
|
||||
processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter),
|
||||
excluded_db_systems=frozenset({"redis", "postgresql"}),
|
||||
)
|
||||
)
|
||||
tracer = get_tracer(provider, "litellm")
|
||||
|
||||
def run():
|
||||
set_request_destinations((LANGFUSE_DEST,))
|
||||
with tracer.start_as_current_span("POST /v1/chat/completions"):
|
||||
with tracer.start_as_current_span("auth /v1/chat/completions"):
|
||||
pass
|
||||
with tracer.start_as_current_span("execute_guardrail pii"):
|
||||
pass
|
||||
with tracer.start_as_current_span("redis async_get_cache") as redis_span:
|
||||
redis_span.set_attribute("db.system.name", "redis")
|
||||
with tracer.start_as_current_span("batch_write_to_db _PROXY_track_cost_callback") as spend_span:
|
||||
spend_span.set_attribute("db.system", "postgresql")
|
||||
with tracer.start_as_current_span("chat gpt-4"):
|
||||
pass
|
||||
|
||||
in_fresh_context(run)
|
||||
|
||||
assert {s.name for s in dest_exporter.get_finished_spans()} == {
|
||||
"POST /v1/chat/completions",
|
||||
"auth /v1/chat/completions",
|
||||
"execute_guardrail pii",
|
||||
"chat gpt-4",
|
||||
}
|
||||
assert {s.name for s in operator_exporter.get_finished_spans()} == {
|
||||
"POST /v1/chat/completions",
|
||||
"auth /v1/chat/completions",
|
||||
"execute_guardrail pii",
|
||||
"redis async_get_cache",
|
||||
"batch_write_to_db _PROXY_track_cost_callback",
|
||||
"chat gpt-4",
|
||||
}
|
||||
|
||||
def test_a_team_naming_two_backends_gets_the_trace_at_both(self):
|
||||
"""The fan-out rides one provider, so it cannot skip a destination on the
|
||||
grounds that some other backend owns it: nothing else would deliver it."""
|
||||
|
|
@ -1022,6 +1069,89 @@ class TestProviderWiring:
|
|||
assert kinds(published).count("TenantFanOutSpanProcessor") == 1
|
||||
assert "TenantFanOutSpanProcessor" not in kinds(other)
|
||||
|
||||
@staticmethod
|
||||
def _fan_out_of(logger: OpenTelemetryV2) -> TenantFanOutSpanProcessor:
|
||||
return next(
|
||||
processor
|
||||
for processor in logger._tracer_provider._active_span_processor._span_processors
|
||||
if isinstance(processor, TenantFanOutSpanProcessor)
|
||||
)
|
||||
|
||||
def test_callback_settings_excluded_services_win_over_the_published_preset_env_config(self, monkeypatch):
|
||||
"""A preset builds its config env-only, so the fan-out must read
|
||||
``callback_settings.otel.excluded_services`` itself rather than the
|
||||
published logger's config, or the env value would win."""
|
||||
monkeypatch.setattr(litellm, "callback_settings", {"otel": {"excluded_services": ["postgres"]}}, raising=False)
|
||||
preset = OpenTelemetryV2(
|
||||
config=OpenTelemetryV2Config(exporters=[ExporterSpec(kind="in_memory")], excluded_services=["redis"]),
|
||||
callback_name="langfuse_otel",
|
||||
)
|
||||
|
||||
publish_global_otel_v2_provider([], lambda _p: None, registered=preset)
|
||||
|
||||
assert self._fan_out_of(preset)._excluded_db_systems == frozenset({"postgresql"})
|
||||
|
||||
def test_callback_settings_excluded_services_apply_even_when_other_otel_env_vars_are_malformed(self, monkeypatch):
|
||||
"""Reading the setting must not rebuild the whole settings model, or an unrelated bad env
|
||||
value the operator overrode in config would stop publication before the fan-out is attached"""
|
||||
preset = OpenTelemetryV2(
|
||||
config=OpenTelemetryV2Config(exporters=[ExporterSpec(kind="in_memory")]),
|
||||
callback_name="langfuse_otel",
|
||||
)
|
||||
monkeypatch.setenv("LITELLM_OTEL_LEGACY_COMPAT", "not-a-bool")
|
||||
monkeypatch.setattr(litellm, "callback_settings", {"otel": {"excluded_services": ["postgres"]}}, raising=False)
|
||||
|
||||
publish_global_otel_v2_provider([], lambda _p: None, registered=preset)
|
||||
|
||||
assert self._fan_out_of(preset)._excluded_db_systems == frozenset({"postgresql"})
|
||||
|
||||
def test_excluded_services_fall_back_to_the_published_logger_config_without_callback_settings(self, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "callback_settings", {"otel": {"exporter": "in_memory"}}, raising=False)
|
||||
preset = OpenTelemetryV2(
|
||||
config=OpenTelemetryV2Config(exporters=[ExporterSpec(kind="in_memory")], excluded_services=["redis"]),
|
||||
callback_name="langfuse_otel",
|
||||
)
|
||||
|
||||
publish_global_otel_v2_provider([], lambda _p: None, registered=preset)
|
||||
|
||||
assert self._fan_out_of(preset)._excluded_db_systems == frozenset({"redis"})
|
||||
|
||||
def test_otel_after_a_preset_reuses_it_and_still_takes_callback_settings_exclusions(self, monkeypatch):
|
||||
"""``callbacks: [langfuse_otel, otel]`` keeps one v2 logger, exactly as
|
||||
before ``excluded_services`` existed, and the exclusion still comes from
|
||||
``callback_settings.otel`` rather than the preset's env-only config."""
|
||||
from litellm.litellm_core_utils import litellm_logging as logging_module
|
||||
|
||||
logging_module._in_memory_loggers.clear()
|
||||
monkeypatch.setenv("LITELLM_OTEL_V2", "true")
|
||||
monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk")
|
||||
monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk")
|
||||
monkeypatch.setenv("LITELLM_OTEL_EXCLUDED_SERVICES", "redis")
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
monkeypatch.setattr(litellm, "callback_settings", {"otel": {"excluded_services": ["postgres"]}}, raising=False)
|
||||
try:
|
||||
|
||||
def init(name: str) -> CustomLogger | None:
|
||||
return logging_module._init_custom_logger_compatible_class(
|
||||
logging_integration=name, # pyright: ignore[reportArgumentType] # test passes a literal callback name
|
||||
internal_usage_cache=None,
|
||||
llm_router=None,
|
||||
custom_logger_init_args={},
|
||||
)
|
||||
|
||||
preset = init("langfuse_otel")
|
||||
otel_cb = init("otel")
|
||||
|
||||
assert isinstance(preset, OpenTelemetryV2)
|
||||
assert otel_cb is preset
|
||||
v2_loggers = [cb for cb in logging_module._in_memory_loggers if isinstance(cb, OpenTelemetryV2)]
|
||||
assert v2_loggers == [preset], v2_loggers
|
||||
publish_global_otel_v2_provider(logging_module._in_memory_loggers, lambda _p: None, registered=preset)
|
||||
assert self._fan_out_of(preset)._excluded_db_systems == frozenset({"postgresql"})
|
||||
finally:
|
||||
logging_module._in_memory_loggers.clear()
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
|
||||
@pytest.mark.parametrize("canonical", ["langfuse_otel", "arize"])
|
||||
def test_publishing_tells_the_fan_out_about_every_v2_loggers_account(self, monkeypatch, canonical):
|
||||
monkeypatch.setenv("LITELLM_OTEL_TENANT_DESTINATION_MODE", "additive")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue