diff --git a/litellm/integrations/otel/README.md b/litellm/integrations/otel/README.md index 1b97e159105..d9047b675ce 100644 --- a/litellm/integrations/otel/README.md +++ b/litellm/integrations/otel/README.md @@ -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. diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index e21711c2708..55eb8e8fb71 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -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() diff --git a/litellm/integrations/otel/model/config.py b/litellm/integrations/otel/model/config.py index 5a3965862e0..9eb29157d6f 100644 --- a/litellm/integrations/otel/model/config.py +++ b/litellm/integrations/otel/model/config.py @@ -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 diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py index 8bac36aad76..25878e8a302 100644 --- a/litellm/integrations/otel/plumbing/providers.py +++ b/litellm/integrations/otel/plumbing/providers.py @@ -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( diff --git a/tests/integration/_support/otlp_sink.py b/tests/integration/_support/otlp_sink.py new file mode 100644 index 00000000000..eeabe9d886f --- /dev/null +++ b/tests/integration/_support/otlp_sink.py @@ -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() diff --git a/tests/integration/observability/conftest.py b/tests/integration/observability/conftest.py new file mode 100644 index 00000000000..f09a703f047 --- /dev/null +++ b/tests/integration/observability/conftest.py @@ -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, + } diff --git a/tests/integration/observability/test_otel_excluded_services.py b/tests/integration/observability/test_otel_excluded_services.py new file mode 100644 index 00000000000..48b20651b97 --- /dev/null +++ b/tests/integration/observability/test_otel_excluded_services.py @@ -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}" diff --git a/tests/integration/observability/test_otel_excluded_services_matrix.py b/tests/integration/observability/test_otel_excluded_services_matrix.py new file mode 100644 index 00000000000..0d4b5d087c6 --- /dev/null +++ b/tests/integration/observability/test_otel_excluded_services_matrix.py @@ -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:] diff --git a/tests/unit/integrations/otel/test_otel_v2_config_baggage_parenting_guardrails.py b/tests/unit/integrations/otel/test_otel_v2_config_baggage_parenting_guardrails.py index dcaff3c911a..86837f7f46c 100644 --- a/tests/unit/integrations/otel/test_otel_v2_config_baggage_parenting_guardrails.py +++ b/tests/unit/integrations/otel/test_otel_v2_config_baggage_parenting_guardrails.py @@ -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": {}}, diff --git a/tests/unit/integrations/otel/test_otel_v2_destinations.py b/tests/unit/integrations/otel/test_otel_v2_destinations.py index 9cb3dbb9deb..5a7057203e4 100644 --- a/tests/unit/integrations/otel/test_otel_v2_destinations.py +++ b/tests/unit/integrations/otel/test_otel_v2_destinations.py @@ -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")