mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
feat(control_plane): otel server and dataplane client spans with w3c propagation
Some checks are pending
Cost map pins recheck / rerun-cost-map-guard (push) Waiting to run
LiteLLM Rust DB / db-queries (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run
Terraform Modules / fmt, validate, test (aws) (push) Waiting to run
Terraform Modules / fmt, validate, test (gcp) (push) Waiting to run
Some checks are pending
Cost map pins recheck / rerun-cost-map-guard (push) Waiting to run
LiteLLM Rust DB / db-queries (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run
Terraform Modules / fmt, validate, test (aws) (push) Waiting to run
Terraform Modules / fmt, validate, test (gcp) (push) Waiting to run
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
537919ab1d
commit
dbff9ad0aa
13 changed files with 859 additions and 20 deletions
|
|
@ -364,7 +364,7 @@ class OpenTelemetryV2Config(BaseSettings):
|
|||
traces_endpoint=self.traces_endpoint,
|
||||
headers=self.headers,
|
||||
)
|
||||
if not self.model_fields_set.isdisjoint(("exporter", "endpoint", "headers"))
|
||||
if not self.model_fields_set.isdisjoint(("exporter", "endpoint", "traces_endpoint", "headers"))
|
||||
else ExporterSpec()
|
||||
]
|
||||
# Ensure ``genai`` is always present and first.
|
||||
|
|
|
|||
|
|
@ -9,7 +9,6 @@ from typing import TYPE_CHECKING, Final
|
|||
|
||||
from opentelemetry import baggage
|
||||
from opentelemetry.context import Context, get_current
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
from opentelemetry.trace import (
|
||||
INVALID_SPAN,
|
||||
Link,
|
||||
|
|
@ -97,6 +96,8 @@ def request_root_http_route() -> str | None:
|
|||
where the FastAPI instrumentation did not mount.
|
||||
"""
|
||||
span: Final = request_root_span()
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
|
||||
route: Final = span.attributes.get(HTTP.ROUTE) if isinstance(span, ReadableSpan) and span.attributes else None
|
||||
return route if isinstance(route, str) and route else None
|
||||
|
||||
|
|
@ -279,6 +280,8 @@ def post_response_root(span: Span) -> Generator[None]:
|
|||
|
||||
|
||||
def _is_post_response(parent: Span, end_time_ns: int | None) -> bool:
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
|
||||
if not isinstance(parent, ReadableSpan):
|
||||
return False
|
||||
if in_post_response_phase():
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
"""App factory for the LiteLLM global control plane. Enterprise only: refuses to build without a valid license."""
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
from contextvars import ContextVar
|
||||
from typing import Final, Protocol
|
||||
|
||||
|
|
@ -8,7 +10,7 @@ from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
|||
|
||||
from litellm.proxy.control_plane.audit.dependencies import get_audit_log
|
||||
from litellm.proxy.control_plane.audit.router import router as audit_router
|
||||
from litellm.proxy.control_plane.auth.dependencies import get_token_verifier
|
||||
from litellm.proxy.control_plane.auth.dependencies import get_telemetry, get_token_verifier
|
||||
from litellm.proxy.control_plane.auth.router import get_login_service
|
||||
from litellm.proxy.control_plane.auth.router import router as auth_router
|
||||
from litellm.proxy.control_plane.common.exceptions import (
|
||||
|
|
@ -82,14 +84,25 @@ class ControlPlaneRequestSizeLimitMiddleware:
|
|||
def create_app(container: Container, license_check: License) -> FastAPI:
|
||||
if not license_check.is_premium():
|
||||
raise LicenseRequired()
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
container.telemetry.shutdown()
|
||||
|
||||
app: Final = FastAPI(
|
||||
title=TITLE,
|
||||
version="1",
|
||||
docs_url=f"{container.settings.api_prefix}/docs",
|
||||
openapi_url=f"{container.settings.api_prefix}/openapi.json",
|
||||
lifespan=lifespan,
|
||||
)
|
||||
app.add_middleware(ControlPlaneRequestSizeLimitMiddleware)
|
||||
container.telemetry.instrument(app)
|
||||
app.dependency_overrides[get_token_verifier] = lambda: container.token_verifier
|
||||
app.dependency_overrides[get_telemetry] = lambda: container.telemetry
|
||||
app.dependency_overrides[get_login_service] = lambda: container.login_service
|
||||
app.dependency_overrides[get_registry] = lambda: container.registry
|
||||
app.dependency_overrides[get_control_plane_settings] = lambda: container.settings
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from starlette.concurrency import run_in_threadpool
|
|||
from litellm.proxy.control_plane.auth.jwt import Rejected, TokenVerifier
|
||||
from litellm.proxy.control_plane.auth.rbac import Principal, Scope
|
||||
from litellm.proxy.control_plane.common.exceptions import ControlPlaneException, forbidden, unauthorized
|
||||
from litellm.proxy.control_plane.common.tracing import UNTRACED, Telemetry, annotate_principal
|
||||
|
||||
bearer_scheme: Final = HTTPBearer(auto_error=False, scheme_name="ControlPlaneBearer")
|
||||
|
||||
|
|
@ -18,15 +19,21 @@ def get_token_verifier() -> TokenVerifier:
|
|||
raise RuntimeError("get_token_verifier is not bound, build the app with create_app")
|
||||
|
||||
|
||||
def get_telemetry() -> Telemetry:
|
||||
return UNTRACED
|
||||
|
||||
|
||||
async def get_principal(
|
||||
credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(bearer_scheme)],
|
||||
verifier: Annotated[TokenVerifier, Depends(get_token_verifier)],
|
||||
telemetry: Annotated[Telemetry, Depends(get_telemetry)],
|
||||
) -> Principal:
|
||||
if credentials is None:
|
||||
raise ControlPlaneException(unauthorized("Missing bearer credentials."))
|
||||
verified: Final = await run_in_threadpool(verifier.verify, credentials.credentials)
|
||||
if isinstance(verified, Rejected):
|
||||
raise ControlPlaneException(unauthorized("Invalid bearer credentials."))
|
||||
annotate_principal(verified, telemetry.identity_attributes)
|
||||
return verified
|
||||
|
||||
|
||||
|
|
|
|||
309
litellm/proxy/control_plane/common/tracing.py
Normal file
309
litellm/proxy/control_plane/common/tracing.py
Normal file
|
|
@ -0,0 +1,309 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from contextlib import AbstractContextManager
|
||||
from dataclasses import dataclass
|
||||
from types import TracebackType
|
||||
from typing import TYPE_CHECKING, Final, Literal, Protocol
|
||||
from urllib.parse import urlsplit, urlunsplit
|
||||
|
||||
from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
|
||||
from opentelemetry.semconv.attributes.http_attributes import HTTP_REQUEST_METHOD, HTTP_RESPONSE_STATUS_CODE
|
||||
from opentelemetry.semconv.attributes.server_attributes import SERVER_ADDRESS, SERVER_PORT
|
||||
from opentelemetry.semconv.attributes.url_attributes import URL_PATH, URL_SCHEME
|
||||
from opentelemetry.trace import Span, SpanKind, StatusCode, Tracer, get_current_span, use_span
|
||||
from pydantic import AliasChoices
|
||||
|
||||
from litellm.integrations.otel.model.config import OpenTelemetryV2Config
|
||||
from litellm.integrations.otel.plumbing.context import inject_trace_context
|
||||
from litellm.proxy.control_plane.audit.log import redact_path
|
||||
from litellm.proxy.control_plane.auth.rbac import Principal
|
||||
from litellm.proxy.control_plane.settings import ControlPlaneSettings
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi import FastAPI
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
from opentelemetry.sdk.trace.export import SpanExporter
|
||||
|
||||
|
||||
def _setting_aliases(field_name: Literal["endpoint", "traces_endpoint"]) -> tuple[str, ...]:
|
||||
alias: Final = OpenTelemetryV2Config.model_fields[field_name].validation_alias
|
||||
return (
|
||||
tuple(choice for choice in alias.choices if isinstance(choice, str)) if isinstance(alias, AliasChoices) else ()
|
||||
)
|
||||
|
||||
|
||||
TRACER_NAME: Final = "litellm.control_plane"
|
||||
ENDUSER_ID: Final = "enduser.id"
|
||||
ENDUSER_ROLE: Final = "enduser.role"
|
||||
HTTP_REQUEST_RESEND_COUNT: Final = "http.request.resend_count"
|
||||
PEER_SERVICE: Final = "peer.service"
|
||||
TOKEN_KIND_ATTRIBUTE: Final = "litellm.control_plane.token_kind"
|
||||
SCOPES_ATTRIBUTE: Final = "litellm.control_plane.scopes"
|
||||
CREDENTIAL_ATTRIBUTE: Final = "litellm.control_plane.credential"
|
||||
DATAPLANE_ATTRIBUTE: Final = "litellm.control_plane.dataplane"
|
||||
UNTRACED_PATHS: Final = (
|
||||
"/healthz",
|
||||
"/readyz",
|
||||
)
|
||||
OTEL_ENDPOINT_ENV_ALIASES: Final = _setting_aliases("endpoint") + _setting_aliases("traces_endpoint")
|
||||
DEFAULT_PORTS: Final = {"https": 443, "http": 80}
|
||||
PATH_ATTRIBUTES: Final = ("http.target", "url.path")
|
||||
FULL_URL_ATTRIBUTES: Final = ("http.url", "url.full")
|
||||
QUERY_ATTRIBUTES: Final = ("url.query",)
|
||||
HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS_ENV: Final = "OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS"
|
||||
HTTP_CAPTURE_HEADERS_SERVER_REQUEST_ENV: Final = "OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST"
|
||||
SENSITIVE_HTTP_HEADERS: Final = ("authorization", "cookie", "set-cookie", "x-api-key", "x-litellm-api-key")
|
||||
|
||||
|
||||
def _http_capture_headers_sanitize_fields(environ: Mapping[str, str]) -> tuple[str, ...]:
|
||||
configured: Final = tuple(
|
||||
field.strip().casefold() for field in environ.get(HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS_ENV, "").split(",")
|
||||
)
|
||||
return tuple(dict.fromkeys((*SENSITIVE_HTTP_HEADERS, *(field for field in configured if field))))
|
||||
|
||||
|
||||
def _http_capture_headers_server_request(environ: Mapping[str, str]) -> tuple[str, ...] | None:
|
||||
configured: Final = environ.get(HTTP_CAPTURE_HEADERS_SERVER_REQUEST_ENV)
|
||||
if configured is None:
|
||||
return None
|
||||
return tuple(header.strip() for header in configured.split(",") if header.strip())
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DataplaneTarget:
|
||||
dataplane_id: str
|
||||
url: str
|
||||
method: str
|
||||
path: str
|
||||
credential: str
|
||||
|
||||
@property
|
||||
def server(self) -> tuple[str, int | None]:
|
||||
parts: Final = urlsplit(self.url)
|
||||
return parts.hostname or self.url, parts.port or DEFAULT_PORTS.get(parts.scheme)
|
||||
|
||||
|
||||
class SpanRecorder(Protocol):
|
||||
def client_span(self, target: DataplaneTarget, attempt: int) -> AbstractSpan: ...
|
||||
|
||||
def outgoing_headers(self, headers: Mapping[str, str]) -> dict[str, str]: ... # mutable-ok: httpx takes dict
|
||||
|
||||
|
||||
class AbstractSpan(Protocol):
|
||||
def __enter__(self) -> AbstractSpan: ...
|
||||
|
||||
def __exit__(self, exc_type: object, exc: object, tb: object) -> None: ...
|
||||
|
||||
def replied(self, status: int) -> None: ...
|
||||
|
||||
def failed(self, error_type: str) -> None: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class NoSpan:
|
||||
def __enter__(self) -> NoSpan:
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type: object, exc: object, tb: object) -> None:
|
||||
return None
|
||||
|
||||
def replied(self, status: int) -> None:
|
||||
return None
|
||||
|
||||
def failed(self, error_type: str) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Untraced:
|
||||
def client_span(self, target: DataplaneTarget, attempt: int) -> NoSpan:
|
||||
return NoSpan()
|
||||
|
||||
def outgoing_headers(self, headers: Mapping[str, str]) -> dict[str, str]: # mutable-ok: httpx takes dict
|
||||
return dict(headers)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ClientSpan:
|
||||
span: Span
|
||||
scope: AbstractContextManager[Span]
|
||||
|
||||
def __enter__(self) -> ClientSpan:
|
||||
self.scope.__enter__()
|
||||
return self
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc: BaseException | None,
|
||||
tb: TracebackType | None,
|
||||
) -> bool | None:
|
||||
if exc is not None:
|
||||
error_type: Final = f"{type(exc).__module__}.{type(exc).__qualname__}"
|
||||
self.span.set_status(StatusCode.ERROR)
|
||||
self.span.set_attribute(ERROR_TYPE, error_type)
|
||||
return self.scope.__exit__(exc_type, exc, tb)
|
||||
|
||||
def replied(self, status: int) -> None:
|
||||
self.span.set_attribute(HTTP_RESPONSE_STATUS_CODE, status)
|
||||
if status >= 400:
|
||||
self.span.set_status(StatusCode.ERROR)
|
||||
self.span.set_attribute(ERROR_TYPE, str(status))
|
||||
|
||||
def failed(self, error_type: str) -> None:
|
||||
self.span.set_status(StatusCode.ERROR, error_type)
|
||||
self.span.set_attribute(ERROR_TYPE, error_type)
|
||||
|
||||
|
||||
def _client_attributes(target: DataplaneTarget, attempt: int) -> Mapping[str, str | int]:
|
||||
host, port = target.server
|
||||
base: Final[Mapping[str, str | int]] = {
|
||||
HTTP_REQUEST_METHOD: target.method,
|
||||
URL_PATH: target.path,
|
||||
URL_SCHEME: urlsplit(target.url).scheme,
|
||||
SERVER_ADDRESS: host,
|
||||
PEER_SERVICE: target.dataplane_id,
|
||||
DATAPLANE_ATTRIBUTE: target.dataplane_id,
|
||||
CREDENTIAL_ATTRIBUTE: target.credential,
|
||||
}
|
||||
with_port: Final = base if port is None else {**base, SERVER_PORT: port}
|
||||
return with_port if attempt == 0 else {**with_port, HTTP_REQUEST_RESEND_COUNT: attempt}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Traced:
|
||||
tracer: Tracer
|
||||
|
||||
def client_span(self, target: DataplaneTarget, attempt: int) -> ClientSpan:
|
||||
span: Final = self.tracer.start_span(
|
||||
f"{target.method} {target.path}", kind=SpanKind.CLIENT, attributes=_client_attributes(target, attempt)
|
||||
)
|
||||
return ClientSpan(span, use_span(span, end_on_exit=True))
|
||||
|
||||
def outgoing_headers(self, headers: Mapping[str, str]) -> dict[str, str]: # mutable-ok: httpx takes dict
|
||||
return inject_trace_context(headers)
|
||||
|
||||
|
||||
def annotate_principal(principal: Principal, identity: bool) -> None:
|
||||
span: Final = get_current_span()
|
||||
if not span.is_recording():
|
||||
return
|
||||
span.set_attribute(TOKEN_KIND_ATTRIBUTE, principal.token_kind)
|
||||
span.set_attribute(SCOPES_ATTRIBUTE, tuple(sorted(scope.value for scope in principal.scopes)))
|
||||
if not identity:
|
||||
return
|
||||
span.set_attribute(ENDUSER_ID, principal.subject)
|
||||
span.set_attribute(ENDUSER_ROLE, tuple(sorted(principal.roles)))
|
||||
|
||||
|
||||
def _url_with_path(attributes: Mapping[str, object], path: str) -> str | None:
|
||||
url: Final = next(
|
||||
(value for key in FULL_URL_ATTRIBUTES if isinstance(value := attributes.get(key), str)),
|
||||
None,
|
||||
)
|
||||
if url is None:
|
||||
return None
|
||||
parts: Final = urlsplit(url)
|
||||
if not parts.scheme or not parts.netloc:
|
||||
return None
|
||||
return urlunsplit((parts.scheme, parts.netloc, path, "", ""))
|
||||
|
||||
|
||||
def redact_server_span(span: Span | None, scope: Mapping[str, object]) -> None:
|
||||
if span is None or not span.is_recording():
|
||||
return
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
|
||||
raw_path: Final = scope.get("path")
|
||||
path: Final = redact_path(raw_path) if isinstance(raw_path, str) else ""
|
||||
for key in PATH_ATTRIBUTES:
|
||||
span.set_attribute(key, path)
|
||||
attributes: Final = span.attributes if isinstance(span, ReadableSpan) else None
|
||||
if full_url := (_url_with_path(attributes, path) if attributes is not None else None):
|
||||
for key in FULL_URL_ATTRIBUTES:
|
||||
span.set_attribute(key, full_url)
|
||||
for key in QUERY_ATTRIBUTES:
|
||||
span.set_attribute(key, "")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Telemetry:
|
||||
provider: TracerProvider | None
|
||||
recorder: SpanRecorder
|
||||
identity_attributes: bool
|
||||
http_capture_headers_server_request: tuple[str, ...] | None = None
|
||||
http_capture_headers_sanitize_fields: tuple[str, ...] = SENSITIVE_HTTP_HEADERS
|
||||
|
||||
def instrument(self, app: FastAPI) -> None:
|
||||
if self.provider is None:
|
||||
return
|
||||
from opentelemetry.instrumentation import fastapi # pyright: ignore[reportMissingTypeStubs] # untyped contrib
|
||||
|
||||
paths: Final = (*UNTRACED_PATHS, app.docs_url, app.openapi_url)
|
||||
excluded_urls: Final = ",".join(f"^https?://[^/]+{re.escape(path)}$" for path in paths if path is not None)
|
||||
fastapi.FastAPIInstrumentor.instrument_app( # pyright: ignore[reportUnknownMemberType] # untyped contrib
|
||||
app,
|
||||
tracer_provider=self.provider,
|
||||
server_request_hook=redact_server_span,
|
||||
http_capture_headers_server_request=(
|
||||
None
|
||||
if self.http_capture_headers_server_request is None
|
||||
else list(self.http_capture_headers_server_request)
|
||||
),
|
||||
http_capture_headers_sanitize_fields=list(self.http_capture_headers_sanitize_fields),
|
||||
excluded_urls=excluded_urls,
|
||||
exclude_spans=["receive", "send"],
|
||||
)
|
||||
|
||||
def shutdown(self) -> None:
|
||||
if self.provider is None:
|
||||
return
|
||||
self.provider.shutdown()
|
||||
|
||||
|
||||
UNTRACED: Final = Telemetry(provider=None, recorder=Untraced(), identity_attributes=False)
|
||||
|
||||
|
||||
def tracing_enabled(settings: ControlPlaneSettings, environ: Mapping[str, str] = os.environ) -> bool:
|
||||
return (
|
||||
settings.otel_enabled
|
||||
if settings.otel_enabled is not None
|
||||
else any(environ.get(alias) for alias in OTEL_ENDPOINT_ENV_ALIASES)
|
||||
)
|
||||
|
||||
|
||||
def build_provider(settings: ControlPlaneSettings, exporter: SpanExporter | None = None) -> TracerProvider:
|
||||
from opentelemetry.sdk.trace import SpanProcessor
|
||||
|
||||
from litellm.integrations.otel.plumbing.providers import build_tracer_provider
|
||||
|
||||
config: Final = OpenTelemetryV2Config.from_env().model_copy(update={"service_name": settings.otel_service_name})
|
||||
return build_tracer_provider(
|
||||
config,
|
||||
exporter=exporter,
|
||||
baggage_processor=SpanProcessor(),
|
||||
use_simple_processor=None if exporter is None else True,
|
||||
console_default=exporter is None,
|
||||
)
|
||||
|
||||
|
||||
def build_telemetry(
|
||||
settings: ControlPlaneSettings,
|
||||
provider: TracerProvider | None = None,
|
||||
environ: Mapping[str, str] = os.environ,
|
||||
) -> Telemetry:
|
||||
chosen: Final = (
|
||||
provider if provider is not None else build_provider(settings) if tracing_enabled(settings, environ) else None
|
||||
)
|
||||
if chosen is None:
|
||||
return UNTRACED
|
||||
return Telemetry(
|
||||
provider=chosen,
|
||||
recorder=Traced(chosen.get_tracer(TRACER_NAME)),
|
||||
identity_attributes=settings.otel_identity_attributes,
|
||||
http_capture_headers_server_request=_http_capture_headers_server_request(environ),
|
||||
http_capture_headers_sanitize_fields=_http_capture_headers_sanitize_fields(environ),
|
||||
)
|
||||
|
|
@ -1,12 +1,14 @@
|
|||
"""Long-lived objects built once per process and handed to the app. Tests build their own."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from collections.abc import Mapping, Sequence
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -27,12 +29,16 @@ from litellm.proxy.control_plane.auth.rbac import (
|
|||
parse_role_definition,
|
||||
)
|
||||
from litellm.proxy.control_plane.common.fanout import Fanout
|
||||
from litellm.proxy.control_plane.common.tracing import Telemetry, build_telemetry
|
||||
from litellm.proxy.control_plane.dataplanes.client import DataplaneClient, HttpTransport, RetryPolicy
|
||||
from litellm.proxy.control_plane.dataplanes.credentials import TokenExchanger, http_token_exchanger
|
||||
from litellm.proxy.control_plane.dataplanes.registry import DataplaneRegistry, build_registry
|
||||
from litellm.proxy.control_plane.settings import ControlPlaneSettings
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Container:
|
||||
|
|
@ -44,6 +50,7 @@ class Container:
|
|||
dataplane_client: DataplaneClient
|
||||
fanout: Fanout
|
||||
audit_log: AuditLog
|
||||
telemetry: Telemetry
|
||||
|
||||
|
||||
class InvalidRolePolicy(ValueError):
|
||||
|
|
@ -113,6 +120,7 @@ def build_container(
|
|||
jwks: JwksKeyLookup | None = None,
|
||||
transport: HttpTransport | None = None,
|
||||
exchanger: TokenExchanger | None = None,
|
||||
tracer_provider: TracerProvider | None = None,
|
||||
) -> Container:
|
||||
master_key: Final = settings.master_key.get_secret_value()
|
||||
policy: Final = build_role_policy(settings)
|
||||
|
|
@ -130,6 +138,7 @@ def build_container(
|
|||
policy=policy,
|
||||
)
|
||||
password: Final = settings.ui_password
|
||||
telemetry: Final = build_telemetry(settings, tracer_provider)
|
||||
return Container(
|
||||
settings=settings,
|
||||
policy=policy,
|
||||
|
|
@ -147,7 +156,9 @@ def build_container(
|
|||
transport=http,
|
||||
timeout_seconds=settings.dataplane_timeout_seconds,
|
||||
retry=RetryPolicy(attempts=settings.dataplane_retry_attempts),
|
||||
tracing=telemetry.recorder,
|
||||
),
|
||||
fanout=Fanout(concurrency=settings.fanout_concurrency),
|
||||
audit_log=build_audit_log(settings.audit_capacity),
|
||||
telemetry=telemetry,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -8,13 +8,15 @@ import asyncio
|
|||
import random
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence
|
||||
from contextlib import AbstractAsyncContextManager, AsyncExitStack
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Final, Literal, Protocol, TypeAlias
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy.control_plane.audit.log import redact_path
|
||||
from litellm.proxy.control_plane.common.tracing import DataplaneTarget, SpanRecorder, Untraced
|
||||
from litellm.proxy.control_plane.dataplanes.credentials import Caller, CredentialKind, CredentialUnavailable
|
||||
from litellm.proxy.control_plane.dataplanes.registry import Dataplane
|
||||
|
||||
|
|
@ -135,6 +137,7 @@ class DataplaneClient:
|
|||
retry: RetryPolicy = RetryPolicy()
|
||||
sleep: Sleeper = asyncio.sleep
|
||||
jitter: Callable[[], float] = random.random
|
||||
tracing: SpanRecorder = field(default_factory=Untraced)
|
||||
|
||||
async def get(
|
||||
self,
|
||||
|
|
@ -172,11 +175,15 @@ class DataplaneClient:
|
|||
credential: Final = await dataplane.credential.headers(caller)
|
||||
if isinstance(credential, CredentialUnavailable):
|
||||
return credential
|
||||
url: Final = f"{dataplane.url}/{path.lstrip('/')}"
|
||||
attempts: Final = 1 if method in WRITE_METHODS else self.retry.attempts
|
||||
outcome: Final = await self._attempt(
|
||||
method, url, credential.headers, params, json, attempt=0, attempts=attempts
|
||||
target: Final = DataplaneTarget(
|
||||
dataplane_id=dataplane.id,
|
||||
url=f"{dataplane.url}/{path.lstrip('/')}",
|
||||
method=method,
|
||||
path=redact_path(f"/{path.lstrip('/')}"),
|
||||
credential=credential.kind,
|
||||
)
|
||||
attempts: Final = 1 if method in WRITE_METHODS else self.retry.attempts
|
||||
outcome: Final = await self._attempt(target, credential.headers, params, json, attempt=0, attempts=attempts)
|
||||
if isinstance(outcome, DataplaneResponseTooLarge):
|
||||
return outcome
|
||||
return (
|
||||
|
|
@ -195,27 +202,32 @@ class DataplaneClient:
|
|||
|
||||
async def _attempt(
|
||||
self,
|
||||
method: Method,
|
||||
url: str,
|
||||
target: DataplaneTarget,
|
||||
headers: Mapping[str, str],
|
||||
params: Sequence[tuple[str, str]] | None,
|
||||
json: Mapping[str, object] | None,
|
||||
attempt: int,
|
||||
attempts: int,
|
||||
) -> _RawReply | DataplaneUnreachable | DataplaneResponseTooLarge:
|
||||
outcome: Final = await self._send(method, url, headers, params, json)
|
||||
with self.tracing.client_span(target, attempt) as span:
|
||||
outcome: Final = await self._send(target, headers, params, json)
|
||||
if isinstance(outcome, _RawReply):
|
||||
span.replied(outcome.status)
|
||||
elif isinstance(outcome, DataplaneResponseTooLarge):
|
||||
span.failed("dataplane response too large")
|
||||
else:
|
||||
span.failed(outcome.detail.split(":", 1)[0])
|
||||
if isinstance(outcome, DataplaneResponseTooLarge):
|
||||
return outcome
|
||||
last: Final = attempt + 1 >= attempts
|
||||
if last or not _retryable(outcome):
|
||||
return outcome if isinstance(outcome, _RawReply) else DataplaneUnreachable(outcome.detail, attempt + 1)
|
||||
await self.sleep(self.retry.delay(attempt, self.jitter()))
|
||||
return await self._attempt(method, url, headers, params, json, attempt + 1, attempts)
|
||||
return await self._attempt(target, headers, params, json, attempt + 1, attempts)
|
||||
|
||||
async def _send(
|
||||
self,
|
||||
method: Method,
|
||||
url: str,
|
||||
target: DataplaneTarget,
|
||||
headers: Mapping[str, str],
|
||||
params: Sequence[tuple[str, str]] | None,
|
||||
json: Mapping[str, object] | None,
|
||||
|
|
@ -224,11 +236,11 @@ class DataplaneClient:
|
|||
try:
|
||||
response: Final = await stack.enter_async_context(
|
||||
self.transport.stream(
|
||||
method,
|
||||
url,
|
||||
target.method,
|
||||
target.url,
|
||||
params=params,
|
||||
json=None if json is None else dict(json),
|
||||
headers=dict(headers),
|
||||
headers=self.tracing.outgoing_headers(headers),
|
||||
timeout=self.timeout_seconds,
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -114,6 +114,9 @@ class ControlPlaneSettings(BaseSettings):
|
|||
usage_max_window_days: int = Field(default=366, ge=1)
|
||||
audit_capacity: int = Field(default=10_000, ge=1)
|
||||
api_prefix: str = "/api/v1"
|
||||
otel_enabled: bool | None = None
|
||||
otel_service_name: str = "litellm-control-plane"
|
||||
otel_identity_attributes: bool = True
|
||||
|
||||
@field_validator("master_key")
|
||||
@classmethod
|
||||
|
|
|
|||
|
|
@ -63,6 +63,10 @@ proxy = [
|
|||
"granian>=2.8.0,<3.0; python_version >= '3.15'",
|
||||
"uvloop>=0.22.1,<1.0; sys_platform != 'win32'",
|
||||
"fastapi>=0.136.3,<1.0",
|
||||
"opentelemetry-api==1.45.2",
|
||||
"opentelemetry-sdk==1.45.2",
|
||||
"opentelemetry-exporter-otlp==1.45.2",
|
||||
"opentelemetry-instrumentation-fastapi==0.66b2",
|
||||
"starlette>=1.0.1,<2.0",
|
||||
"backoff>=2.2.1,<3.0",
|
||||
"pyyaml>=6.0.3,<7.0",
|
||||
|
|
|
|||
|
|
@ -932,6 +932,13 @@ def _export_one_span(cfg: OpenTelemetryV2Config) -> None:
|
|||
provider.shutdown()
|
||||
|
||||
|
||||
def test_traces_endpoint_alone_creates_an_otlp_http_exporter_spec():
|
||||
endpoint: Final = "https://collector.example.test/custom/traces"
|
||||
config: Final = OpenTelemetryV2Config(traces_endpoint=endpoint)
|
||||
(spec,) = config.exporters
|
||||
assert spec.kind == "otlp_http" and spec.traces_endpoint == endpoint
|
||||
|
||||
|
||||
def test_traces_endpoint_env_posts_to_the_configured_url_verbatim(monkeypatch, otlp_collector):
|
||||
base_url, received_paths = otlp_collector
|
||||
for var in ("OTEL_EXPORTER", "OTEL_EXPORTER_OTLP_PROTOCOL", "OTEL_EXPORTER_OTLP_ENDPOINT"):
|
||||
|
|
|
|||
457
tests/unit/proxy/control_plane/common/test_tracing.py
Normal file
457
tests/unit/proxy/control_plane/common/test_tracing.py
Normal file
|
|
@ -0,0 +1,457 @@
|
|||
import asyncio
|
||||
import importlib.abc
|
||||
import sys
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from importlib.machinery import ModuleSpec
|
||||
from types import ModuleType
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
|
||||
from opentelemetry.sdk.trace.export import BatchSpanProcessor
|
||||
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
||||
from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
|
||||
from opentelemetry.trace import Span, SpanKind, StatusCode, format_trace_id, use_span
|
||||
from opentelemetry.trace.span import SpanContext
|
||||
|
||||
from litellm.proxy.control_plane.common.tracing import (
|
||||
ClientSpan,
|
||||
build_provider,
|
||||
build_telemetry,
|
||||
redact_server_span,
|
||||
tracing_enabled,
|
||||
)
|
||||
from litellm.proxy.control_plane.settings import ControlPlaneSettings
|
||||
from tests.unit.proxy.control_plane.support import DATAPLANES, MASTER_KEY, Licensed, response, rig
|
||||
|
||||
PROXY_ROUTE: Final = "GET /api/v1/dataplanes/{dataplane_id}/proxy/{path:path}"
|
||||
USER_INFO: Final = "https://eu.test/user/info"
|
||||
|
||||
|
||||
def _settings(**overrides: object) -> ControlPlaneSettings:
|
||||
return ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": MASTER_KEY, **overrides})
|
||||
|
||||
|
||||
def _by_kind(spans: Sequence[ReadableSpan], kind: SpanKind) -> list[ReadableSpan]:
|
||||
return [span for span in spans if span.kind == kind]
|
||||
|
||||
|
||||
def _ctx(span: ReadableSpan) -> SpanContext:
|
||||
assert span.context is not None
|
||||
return span.context
|
||||
|
||||
|
||||
def _attr(span: ReadableSpan, key: str) -> object:
|
||||
return None if span.attributes is None else span.attributes.get(key)
|
||||
|
||||
|
||||
def _recording_server_span(attributes: Mapping[str, str] | None = None) -> tuple[Span, InMemorySpanExporter]:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
span: Final = build_provider(_settings(), exporter).get_tracer("test").start_span("server")
|
||||
if attributes is not None:
|
||||
for key, value in attributes.items():
|
||||
span.set_attribute(key, value)
|
||||
return span, exporter
|
||||
|
||||
|
||||
def test_one_proxy_call_is_one_trace_from_server_span_to_dataplane_header() -> None:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
client, transport, minter = rig({USER_INFO: [response(200, {"ok": True})]}, build_provider(_settings(), exporter))
|
||||
reply: Final = client.get("/api/v1/dataplanes/eu/proxy/user/info?secret=1", headers=minter.headers("eu"))
|
||||
assert reply.status_code == 200, reply.text
|
||||
spans: Final = exporter.get_finished_spans()
|
||||
(server,), (dataplane_call,) = _by_kind(spans, SpanKind.SERVER), _by_kind(spans, SpanKind.CLIENT)
|
||||
assert server.name == PROXY_ROUTE and dataplane_call.name == "GET /user/info"
|
||||
assert dataplane_call.parent is not None and dataplane_call.parent.span_id == _ctx(server).span_id
|
||||
_, _, headers = transport.calls[0]
|
||||
assert headers is not None and headers["traceparent"].split("-")[1:3] == [
|
||||
format_trace_id(_ctx(dataplane_call).trace_id),
|
||||
format(_ctx(dataplane_call).span_id, "016x"),
|
||||
]
|
||||
assert headers["Authorization"] == "Bearer sk-eu", "the credential survives header injection"
|
||||
|
||||
|
||||
def test_client_span_carries_the_dataplane_and_never_the_query_string() -> None:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
client, _, minter = rig({USER_INFO: [response(200, {})]}, build_provider(_settings(), exporter))
|
||||
client.get("/api/v1/dataplanes/eu/proxy/user/info?token=sk-leak", headers=minter.headers("eu"))
|
||||
(dataplane_call,) = _by_kind(exporter.get_finished_spans(), SpanKind.CLIENT)
|
||||
attributes: Final = dict(dataplane_call.attributes or {})
|
||||
assert attributes["peer.service"] == "eu" and attributes["server.address"] == "eu.test"
|
||||
assert attributes["url.path"] == "/user/info" and attributes["http.response.status_code"] == 200
|
||||
assert attributes["litellm.control_plane.credential"] == "service_token"
|
||||
assert "sk-leak" not in repr(attributes) and "url.full" not in attributes and "url.query" not in attributes
|
||||
|
||||
|
||||
def test_a_virtual_key_in_the_forwarded_path_never_reaches_the_span() -> None:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
url: Final = "https://eu.test/key/sk-live-123/reset_spend"
|
||||
client, _, minter = rig({url: [response(200, {})]}, build_provider(_settings(), exporter))
|
||||
client.get("/api/v1/dataplanes/eu/proxy/key/sk-live-123/reset_spend", headers=minter.headers("eu"))
|
||||
(dataplane_call,) = _by_kind(exporter.get_finished_spans(), SpanKind.CLIENT)
|
||||
assert dataplane_call.name == "GET /key/sk-[redacted]/reset_spend"
|
||||
assert _attr(dataplane_call, "url.path") == "/key/sk-[redacted]/reset_spend"
|
||||
assert "sk-live-123" not in repr(dataplane_call.attributes)
|
||||
|
||||
|
||||
def test_the_server_span_keeps_neither_the_virtual_key_nor_the_query_string() -> None:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
url: Final = "https://eu.test/key/sk-live-123/reset_spend"
|
||||
client, _, minter = rig({url: [response(200, {})]}, build_provider(_settings(), exporter))
|
||||
client.get("/api/v1/dataplanes/eu/proxy/key/sk-live-123/reset_spend?token=sk-q", headers=minter.headers("eu"))
|
||||
spans: Final = exporter.get_finished_spans()
|
||||
(server,) = _by_kind(spans, SpanKind.SERVER)
|
||||
attributes: Final = repr(tuple(span.attributes for span in spans))
|
||||
assert "sk-live-123" not in attributes and "sk-q" not in attributes and "token=" not in attributes
|
||||
redacted_path: Final = "/api/v1/dataplanes/eu/proxy/key/sk-[redacted]/reset_spend"
|
||||
assert _attr(server, "http.target") == redacted_path
|
||||
assert _attr(server, "url.path") == redacted_path
|
||||
assert _attr(server, "url.query") == ""
|
||||
expected_url: Final = "http://testserver/api/v1/dataplanes/eu/proxy/key/sk-[redacted]/reset_spend"
|
||||
assert _attr(server, "http.url") == expected_url
|
||||
assert _attr(server, "url.full") == expected_url
|
||||
|
||||
|
||||
def test_server_span_redacts_double_encoded_virtual_keys() -> None:
|
||||
span, exporter = _recording_server_span(
|
||||
{
|
||||
"http.url": "https://instrumented.test/original?token=sk-live-123",
|
||||
"url.full": "https://instrumented.test/original?token=sk-live-123",
|
||||
}
|
||||
)
|
||||
redact_server_span(
|
||||
span,
|
||||
{
|
||||
"path": "/api/v1/dataplanes/eu/proxy/key/%73k-live-123/secret/reset_spend",
|
||||
"scheme": "https",
|
||||
"server": ("proxy.test", 443),
|
||||
},
|
||||
)
|
||||
span.end()
|
||||
(recorded,) = exporter.get_finished_spans()
|
||||
attributes: Final = recorded.attributes or {}
|
||||
redacted_path: Final = "/api/v1/dataplanes/eu/proxy/key/sk-[redacted]/reset_spend"
|
||||
expected_url: Final = f"https://instrumented.test{redacted_path}"
|
||||
assert attributes["http.target"] == redacted_path
|
||||
assert attributes["url.path"] == redacted_path
|
||||
assert attributes["http.url"] == expected_url
|
||||
assert attributes["url.full"] == expected_url
|
||||
assert "sk-live-123" not in repr(attributes)
|
||||
|
||||
|
||||
def test_server_span_without_an_instrumented_url_leaves_full_url_unset() -> None:
|
||||
span, exporter = _recording_server_span()
|
||||
redact_server_span(span, {"path": "/healthz", "scheme": "https", "server": ("proxy.test", 443)})
|
||||
span.end()
|
||||
(recorded,) = exporter.get_finished_spans()
|
||||
attributes: Final = recorded.attributes or {}
|
||||
assert "http.url" not in attributes and "url.full" not in attributes
|
||||
|
||||
|
||||
def test_server_span_drops_fragments_from_redacted_urls() -> None:
|
||||
instrumented_url: Final = "https://instrumented.test/original#%23&key=sk-live-123"
|
||||
span, exporter = _recording_server_span({"http.url": instrumented_url, "url.full": instrumented_url})
|
||||
redact_server_span(span, {"path": "/healthz"})
|
||||
span.end()
|
||||
(recorded,) = exporter.get_finished_spans()
|
||||
attributes: Final = recorded.attributes or {}
|
||||
expected_url: Final = "https://instrumented.test/healthz"
|
||||
assert attributes["http.url"] == expected_url
|
||||
assert attributes["url.full"] == expected_url
|
||||
assert "key=" not in attributes["http.url"] and "key=" not in attributes["url.full"]
|
||||
|
||||
|
||||
def test_double_encoded_key_is_redacted_on_the_real_server_span() -> None:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
target: Final = "https://eu.test/key/sk-live-123/reset_spend"
|
||||
client, _, minter = rig({target: [response(200, {})]}, build_provider(_settings(), exporter))
|
||||
reply: Final = client.get(
|
||||
"/api/v1/dataplanes/eu/proxy/key/%2573k-live-123/reset_spend",
|
||||
headers=minter.headers("eu"),
|
||||
)
|
||||
assert reply.status_code == 200, reply.text
|
||||
(server,) = _by_kind(exporter.get_finished_spans(), SpanKind.SERVER)
|
||||
assert "k-live-123" not in repr(server.attributes)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("error", [ValueError("boom"), asyncio.CancelledError()])
|
||||
def test_client_span_records_exceptions(error: BaseException) -> None:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
span: Final = build_provider(_settings(), exporter).get_tracer("test").start_span("client")
|
||||
with pytest.raises(type(error)):
|
||||
with ClientSpan(span, use_span(span, end_on_exit=True)):
|
||||
raise error
|
||||
(recorded,) = exporter.get_finished_spans()
|
||||
assert recorded.status.status_code == StatusCode.ERROR
|
||||
assert _attr(recorded, ERROR_TYPE) == f"{type(error).__module__}.{type(error).__qualname__}"
|
||||
if isinstance(error, asyncio.CancelledError):
|
||||
assert recorded.events == ()
|
||||
else:
|
||||
assert tuple(event.name for event in recorded.events) == ("exception",)
|
||||
|
||||
|
||||
def test_server_instrumentation_redacts_captured_request_headers() -> None:
|
||||
environ: Final = {
|
||||
"OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST": ".*",
|
||||
"OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS": "X-Request-Id",
|
||||
}
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
provider: Final = build_provider(_settings(), exporter)
|
||||
telemetry: Final = build_telemetry(_settings(otel_enabled=True), provider, environ=environ)
|
||||
app: Final = FastAPI()
|
||||
|
||||
@app.get("/headers")
|
||||
def headers() -> dict[str, bool]:
|
||||
return {"ok": True}
|
||||
|
||||
telemetry.instrument(app)
|
||||
response: Final = TestClient(app).get(
|
||||
"/headers",
|
||||
headers={
|
||||
"Authorization": "Bearer auth-secret",
|
||||
"Cookie": "session=cookie-secret",
|
||||
"X-Api-Key": "api-key-secret",
|
||||
"X-Request-Id": "request-id-secret",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
(server,) = _by_kind(exporter.get_finished_spans(), SpanKind.SERVER)
|
||||
attributes: Final = server.attributes or {}
|
||||
header_attributes: Final = (
|
||||
"http.request.header.authorization",
|
||||
"http.request.header.cookie",
|
||||
"http.request.header.x_api_key",
|
||||
"http.request.header.x_request_id",
|
||||
)
|
||||
assert all("[REDACTED]" in repr(attributes.get(name)) for name in header_attributes)
|
||||
assert all(
|
||||
secret not in repr(attributes)
|
||||
for secret in ("auth-secret", "cookie-secret", "api-key-secret", "request-id-secret")
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def without_otel_sdk() -> Iterator[None]:
|
||||
modules: Final = sys.modules.copy()
|
||||
meta_path: Final = tuple(sys.meta_path)
|
||||
|
||||
class BlockSdk(importlib.abc.MetaPathFinder):
|
||||
def find_spec(
|
||||
self, fullname: str, path: Sequence[str] | None, target: ModuleType | None = None
|
||||
) -> ModuleSpec | None:
|
||||
if fullname == "opentelemetry.sdk" or fullname.startswith("opentelemetry.sdk."):
|
||||
raise ImportError(fullname)
|
||||
return None
|
||||
|
||||
sys.meta_path = [BlockSdk(), *meta_path]
|
||||
for name in tuple(sys.modules):
|
||||
if (
|
||||
name == "opentelemetry.sdk"
|
||||
or name.startswith("opentelemetry.sdk.")
|
||||
or name == "litellm.proxy.control_plane"
|
||||
or name.startswith("litellm.proxy.control_plane.")
|
||||
or name == "litellm.integrations.otel.plumbing.context"
|
||||
):
|
||||
sys.modules.pop(name)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
sys.modules.clear()
|
||||
sys.modules.update(modules)
|
||||
sys.meta_path = list(meta_path)
|
||||
|
||||
|
||||
def test_control_plane_app_builds_without_the_otel_sdk_when_tracing_is_disabled(
|
||||
without_otel_sdk: None,
|
||||
) -> None:
|
||||
from litellm.proxy.control_plane.app import create_app
|
||||
from litellm.proxy.control_plane.container import build_container
|
||||
|
||||
settings: Final = _settings(otel_enabled=False)
|
||||
app: Final = create_app(build_container(settings), Licensed())
|
||||
reply: Final = TestClient(app).get("/healthz")
|
||||
assert reply.status_code == 200, reply.text
|
||||
assert reply.json() == {"status": "ok"}
|
||||
|
||||
|
||||
def test_inbound_baggage_cannot_stamp_control_plane_spans_and_traceparent_is_honored() -> None:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
client, transport, minter = rig({USER_INFO: [response(200, {})]}, build_provider(_settings(), exporter))
|
||||
trace_id: Final = "1234567890abcdef1234567890abcdef"
|
||||
span_id: Final = "0123456789abcdef"
|
||||
reply: Final = client.get(
|
||||
"/api/v1/dataplanes/eu/proxy/user/info",
|
||||
headers={
|
||||
**minter.headers("eu"),
|
||||
"baggage": "litellm.team.id=spoofed",
|
||||
"traceparent": f"00-{trace_id}-{span_id}-01",
|
||||
},
|
||||
)
|
||||
assert reply.status_code == 200
|
||||
spans: Final = exporter.get_finished_spans()
|
||||
(server,) = _by_kind(spans, SpanKind.SERVER)
|
||||
(dataplane_call,) = _by_kind(spans, SpanKind.CLIENT)
|
||||
assert format_trace_id(_ctx(server).trace_id) == trace_id
|
||||
assert _attr(server, "litellm.team.id") is None
|
||||
assert _attr(dataplane_call, "litellm.team.id") is None
|
||||
_, _, outbound_headers = transport.calls[0]
|
||||
assert outbound_headers is not None and outbound_headers["traceparent"].split("-")[1] == trace_id
|
||||
|
||||
|
||||
def test_server_span_names_the_caller_but_not_the_email() -> None:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
client, _, minter = rig({USER_INFO: [response(200, {})]}, build_provider(_settings(), exporter))
|
||||
client.get("/api/v1/dataplanes/eu/proxy/user/info", headers=minter.headers("eu"))
|
||||
(server,) = _by_kind(exporter.get_finished_spans(), SpanKind.SERVER)
|
||||
assert _attr(server, "enduser.id") == "alice" and _attr(server, "enduser.role") == ("viewer",)
|
||||
assert _attr(server, "litellm.control_plane.token_kind") == "session"
|
||||
assert "alice@x.test" not in repr(server.attributes)
|
||||
|
||||
|
||||
def test_identity_attributes_can_be_turned_off() -> None:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
client, _, minter = rig(
|
||||
{USER_INFO: [response(200, {})]}, build_provider(_settings(), exporter), otel_identity_attributes=False
|
||||
)
|
||||
client.get("/api/v1/dataplanes/eu/proxy/user/info", headers=minter.headers("eu"))
|
||||
(server,) = _by_kind(exporter.get_finished_spans(), SpanKind.SERVER)
|
||||
assert _attr(server, "enduser.id") is None and _attr(server, "litellm.control_plane.token_kind") == "session"
|
||||
|
||||
|
||||
def test_each_retry_is_its_own_client_span_with_a_resend_count() -> None:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
outcomes: Final = [httpx.ConnectError("boom"), response(503), response(200, {})]
|
||||
client, _, minter = rig({USER_INFO: outcomes}, build_provider(_settings(), exporter))
|
||||
reply: Final = client.get("/api/v1/dataplanes/eu/proxy/user/info", headers=minter.headers("eu"))
|
||||
assert reply.status_code == 200, reply.text
|
||||
attempts: Final = _by_kind(exporter.get_finished_spans(), SpanKind.CLIENT)
|
||||
assert [_attr(span, "http.request.resend_count") for span in attempts] == [None, 1, 2]
|
||||
assert [span.status.status_code for span in attempts] == [StatusCode.ERROR, StatusCode.ERROR, StatusCode.UNSET]
|
||||
assert [_attr(span, "error.type") for span in attempts] == ["ConnectError", "503", None]
|
||||
assert len({_ctx(span).trace_id for span in attempts}) == 1
|
||||
|
||||
|
||||
def test_client_span_marks_a_4xx_reply_as_an_error() -> None:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
client, _, minter = rig({USER_INFO: [response(400)]}, build_provider(_settings(), exporter))
|
||||
|
||||
reply: Final = client.get("/api/v1/dataplanes/eu/proxy/user/info", headers=minter.headers("eu"))
|
||||
|
||||
assert reply.status_code == 400
|
||||
(dataplane_call,) = _by_kind(exporter.get_finished_spans(), SpanKind.CLIENT)
|
||||
assert dataplane_call.status.status_code == StatusCode.ERROR
|
||||
assert _attr(dataplane_call, "error.type") == "400"
|
||||
|
||||
|
||||
def test_health_is_not_traced() -> None:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
client, _, _ = rig({}, build_provider(_settings(), exporter))
|
||||
assert client.get("/healthz").status_code == 200
|
||||
assert exporter.get_finished_spans() == ()
|
||||
|
||||
|
||||
def test_server_exclusions_match_only_the_exact_url_path() -> None:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
docs_url: Final = "https://docs-prod.test/user/info"
|
||||
dataplanes: Final = {
|
||||
**DATAPLANES,
|
||||
"docs-prod": {"url": "https://docs-prod.test", "service_token": {"token": "sk-docs"}},
|
||||
}
|
||||
client, _, minter = rig(
|
||||
{docs_url: [response(200, {"ok": True})]},
|
||||
build_provider(_settings(), exporter),
|
||||
dataplanes=dataplanes,
|
||||
)
|
||||
|
||||
assert client.get("/api/v1/docs").status_code == 200
|
||||
assert _by_kind(exporter.get_finished_spans(), SpanKind.SERVER) == []
|
||||
assert client.get("/api/v1/openapi.json").status_code == 200
|
||||
assert _by_kind(exporter.get_finished_spans(), SpanKind.SERVER) == []
|
||||
reply: Final = client.get(
|
||||
"/api/v1/dataplanes/docs-prod/proxy/user/info?probe=1",
|
||||
headers=minter.headers("docs-prod"),
|
||||
)
|
||||
|
||||
assert reply.status_code == 200, reply.text
|
||||
(server,) = _by_kind(exporter.get_finished_spans(), SpanKind.SERVER)
|
||||
assert server.name == PROXY_ROUTE
|
||||
assert _attr(server, "http.url") == "http://testserver/api/v1/dataplanes/docs-prod/proxy/user/info"
|
||||
|
||||
|
||||
def test_without_a_provider_nothing_is_recorded_and_no_trace_header_is_sent() -> None:
|
||||
client, transport, minter = rig({USER_INFO: [response(200, {})]}, otel_enabled=False)
|
||||
assert client.get("/api/v1/dataplanes/eu/proxy/user/info", headers=minter.headers("eu")).status_code == 200
|
||||
_, _, headers = transport.calls[0]
|
||||
assert headers is not None and "traceparent" not in headers
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"env_name",
|
||||
(
|
||||
"OTEL_ENDPOINT",
|
||||
"OTEL_EXPORTER_OTLP_ENDPOINT",
|
||||
"OTEL_TRACES_ENDPOINT",
|
||||
"OTEL_EXPORTER_OTLP_TRACES_ENDPOINT",
|
||||
),
|
||||
)
|
||||
def test_tracing_enables_for_each_endpoint_alias(env_name: str) -> None:
|
||||
assert tracing_enabled(_settings(), {env_name: "http://collector:4318"}) is True
|
||||
|
||||
|
||||
def test_tracing_is_disabled_without_an_endpoint() -> None:
|
||||
assert tracing_enabled(_settings(), {}) is False
|
||||
|
||||
|
||||
def test_explicitly_disabled_tracing_wins_over_an_endpoint() -> None:
|
||||
assert tracing_enabled(_settings(otel_enabled=False), {"OTEL_ENDPOINT": "http://collector:4318"}) is False
|
||||
|
||||
|
||||
def test_explicitly_enabled_tracing_needs_no_endpoint() -> None:
|
||||
assert tracing_enabled(_settings(otel_enabled=True), {}) is True
|
||||
|
||||
|
||||
def test_explicit_enable_without_an_endpoint_exports_spans_to_console(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
capfd: pytest.CaptureFixture[str],
|
||||
) -> None:
|
||||
endpoint_aliases: Final = (
|
||||
"OTEL_EXPORTER",
|
||||
"OTEL_EXPORTER_OTLP_PROTOCOL",
|
||||
"OTEL_ENDPOINT",
|
||||
"OTEL_EXPORTER_OTLP_ENDPOINT",
|
||||
"OTEL_TRACES_ENDPOINT",
|
||||
"OTEL_EXPORTER_OTLP_TRACES_ENDPOINT",
|
||||
"OTEL_LOGS_ENDPOINT",
|
||||
"OTEL_EXPORTER_OTLP_LOGS_ENDPOINT",
|
||||
)
|
||||
for alias in endpoint_aliases:
|
||||
monkeypatch.delenv(alias, raising=False)
|
||||
|
||||
telemetry: Final = build_telemetry(_settings(otel_enabled=True), environ={})
|
||||
assert telemetry.provider is not None
|
||||
telemetry.provider.get_tracer("test").start_span("console-default").end()
|
||||
telemetry.shutdown()
|
||||
|
||||
assert "console-default" in capfd.readouterr().out
|
||||
|
||||
|
||||
def test_disabled_telemetry_has_no_provider_to_instrument_with() -> None:
|
||||
assert build_telemetry(_settings(otel_enabled=False)).provider is None
|
||||
|
||||
|
||||
def test_app_shutdown_flushes_the_server_span() -> None:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
tracer_provider: Final = TracerProvider()
|
||||
tracer_provider.add_span_processor(BatchSpanProcessor(exporter, schedule_delay_millis=60_000))
|
||||
client, _, minter = rig({USER_INFO: [response(200, {"ok": True})]}, tracer_provider)
|
||||
with client:
|
||||
reply: Final = client.get(
|
||||
"/api/v1/dataplanes/eu/proxy/user/info",
|
||||
headers=minter.headers("eu"),
|
||||
)
|
||||
assert reply.status_code == 200, reply.text
|
||||
server_spans: Final = _by_kind(exporter.get_finished_spans(), SpanKind.SERVER)
|
||||
assert tuple(span.name for span in server_spans) == (PROXY_ROUTE,)
|
||||
|
|
@ -9,6 +9,7 @@ from typing import Final
|
|||
|
||||
import httpx
|
||||
from fastapi.testclient import TestClient
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
|
||||
from litellm.litellm_core_utils.oauth_token_exchange import OAuthTokenExchangeConfig
|
||||
from litellm.proxy.control_plane.app import create_app
|
||||
|
|
@ -83,7 +84,9 @@ class Minter:
|
|||
|
||||
|
||||
def rig(
|
||||
routes: Mapping[str, list[httpx.Response | Exception]], **overrides: object
|
||||
routes: Mapping[str, list[httpx.Response | Exception]],
|
||||
tracer_provider: TracerProvider | None = None,
|
||||
**overrides: object,
|
||||
) -> tuple[TestClient, RoutedTransport, Minter]:
|
||||
settings: Final = ControlPlaneSettings.model_validate(
|
||||
{"LITELLM_MASTER_KEY": MASTER_KEY, "dataplanes": DATAPLANES, **overrides}
|
||||
|
|
@ -91,6 +94,8 @@ def rig(
|
|||
queues: defaultdict[str, Iterator[httpx.Response | Exception]] = defaultdict(lambda: iter(()))
|
||||
queues.update({url: iter(outcomes) for url, outcomes in routes.items()})
|
||||
transport: Final = RoutedTransport(queues)
|
||||
container: Final = build_container(settings, transport=transport, exchanger=no_exchange)
|
||||
container: Final = build_container(
|
||||
settings, transport=transport, exchanger=no_exchange, tracer_provider=tracer_provider
|
||||
)
|
||||
codec: Final = SessionTokenCodec(MASTER_KEY, settings.session_issuer, timedelta(hours=1), container.policy)
|
||||
return TestClient(create_app(container, Licensed())), transport, Minter(codec, container.policy)
|
||||
|
|
|
|||
8
uv.lock
generated
8
uv.lock
generated
|
|
@ -5134,6 +5134,10 @@ proxy = [
|
|||
{ name = "litellm-enterprise" },
|
||||
{ name = "litellm-proxy-extras" },
|
||||
{ name = "mcp" },
|
||||
{ name = "opentelemetry-api" },
|
||||
{ name = "opentelemetry-exporter-otlp" },
|
||||
{ name = "opentelemetry-instrumentation-fastapi" },
|
||||
{ name = "opentelemetry-sdk" },
|
||||
{ name = "orjson" },
|
||||
{ name = "polars" },
|
||||
{ name = "pydantic", version = "2.13.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.15'" },
|
||||
|
|
@ -5387,9 +5391,13 @@ requires-dist = [
|
|||
{ name = "numpy", marker = "extra == 'stt-nvidia-riva'", specifier = ">=1.26.0" },
|
||||
{ name = "numpydoc", marker = "extra == 'utils'", specifier = ">=1.8.0,<2.0" },
|
||||
{ name = "openai", specifier = ">=2.20.0,<4.0.0" },
|
||||
{ name = "opentelemetry-api", marker = "extra == 'proxy'", specifier = "==1.45.2" },
|
||||
{ name = "opentelemetry-api", marker = "extra == 'proxy-runtime'", specifier = "==1.45.2" },
|
||||
{ name = "opentelemetry-exporter-otlp", marker = "extra == 'proxy'", specifier = "==1.45.2" },
|
||||
{ name = "opentelemetry-exporter-otlp", marker = "extra == 'proxy-runtime'", specifier = "==1.45.2" },
|
||||
{ name = "opentelemetry-instrumentation-fastapi", marker = "extra == 'proxy'", specifier = "==0.66b2" },
|
||||
{ name = "opentelemetry-instrumentation-fastapi", marker = "extra == 'proxy-runtime'", specifier = "==0.66b2" },
|
||||
{ name = "opentelemetry-sdk", marker = "extra == 'proxy'", specifier = "==1.45.2" },
|
||||
{ name = "opentelemetry-sdk", marker = "extra == 'proxy-runtime'", specifier = "==1.45.2" },
|
||||
{ name = "orjson", marker = "python_full_version >= '3.15' and extra == 'proxy'", specifier = ">=3.12.0,<4.0" },
|
||||
{ name = "orjson", marker = "python_full_version < '3.15' and extra == 'proxy'", specifier = ">=3.11.6,<4.0" },
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue