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
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
c6afb67f64
commit
140dc2be1a
8 changed files with 498 additions and 17 deletions
|
|
@ -1,12 +1,14 @@
|
|||
"""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 typing import Final, Protocol
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
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 install_exception_handlers
|
||||
|
|
@ -35,13 +37,24 @@ class LicenseRequired(RuntimeError):
|
|||
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,
|
||||
)
|
||||
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_dataplane_client] = lambda: container.dataplane_client
|
||||
|
|
|
|||
|
|
@ -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(verified.reason))
|
||||
annotate_principal(verified, telemetry.identity_attributes)
|
||||
return verified
|
||||
|
||||
|
||||
|
|
|
|||
255
litellm/proxy/control_plane/common/tracing.py
Normal file
255
litellm/proxy/control_plane/common/tracing.py
Normal file
|
|
@ -0,0 +1,255 @@
|
|||
import os
|
||||
from collections.abc import Mapping
|
||||
from contextlib import AbstractContextManager
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final, Literal, Protocol
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
from opentelemetry.sdk.trace.export import SpanExporter
|
||||
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.integrations.otel.plumbing.providers import build_tracer_provider
|
||||
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
|
||||
|
||||
|
||||
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", "/docs", "/openapi.json")
|
||||
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",)
|
||||
|
||||
|
||||
@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: object, exc: object, tb: object) -> None:
|
||||
self.scope.__exit__(None, None, None)
|
||||
|
||||
def replied(self, status: int) -> None:
|
||||
self.span.set_attribute(HTTP_RESPONSE_STATUS_CODE, status)
|
||||
if status >= 500:
|
||||
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)))
|
||||
|
||||
|
||||
class _ServerSpan(Protocol):
|
||||
def is_recording(self) -> bool: ...
|
||||
|
||||
def set_attribute(self, key: str, value: str) -> None: ...
|
||||
|
||||
|
||||
def _url_with_path(scope: Mapping[str, object], path: str) -> str:
|
||||
scheme: Final = scope.get("scheme")
|
||||
server: Final = scope.get("server")
|
||||
if not isinstance(scheme, str):
|
||||
return path
|
||||
match server:
|
||||
case (str() as host, int() as port):
|
||||
port_suffix: Final = "" if DEFAULT_PORTS.get(scheme) == port else f":{port}"
|
||||
return f"{scheme}://{host}{port_suffix}{path}"
|
||||
case _:
|
||||
return path
|
||||
|
||||
|
||||
def redact_server_span(span: _ServerSpan | None, scope: Mapping[str, object]) -> None:
|
||||
"""Redact secrets from the request target before the span is exported."""
|
||||
if span is None or not span.is_recording():
|
||||
return
|
||||
raw_path: Final = scope.get("path")
|
||||
path: Final = redact_path(raw_path) if isinstance(raw_path, str) else ""
|
||||
full_url: Final = _url_with_path(scope, path)
|
||||
for key in PATH_ATTRIBUTES:
|
||||
span.set_attribute(key, path)
|
||||
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
|
||||
|
||||
def instrument(self, app: "FastAPI") -> None:
|
||||
if self.provider is None:
|
||||
return
|
||||
from opentelemetry.instrumentation import fastapi # pyright: ignore[reportMissingTypeStubs] # untyped contrib
|
||||
|
||||
fastapi.FastAPIInstrumentor.instrument_app( # pyright: ignore[reportUnknownMemberType] # untyped contrib
|
||||
app,
|
||||
tracer_provider=self.provider,
|
||||
server_request_hook=redact_server_span,
|
||||
excluded_urls=",".join(UNTRACED_PATHS),
|
||||
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:
|
||||
config: Final = OpenTelemetryV2Config.from_env().model_copy(update={"service_name": settings.otel_service_name})
|
||||
return build_tracer_provider(config, exporter=exporter, use_simple_processor=None if exporter is None else True)
|
||||
|
||||
|
||||
def build_telemetry(settings: ControlPlaneSettings, provider: TracerProvider | None = None) -> Telemetry:
|
||||
chosen: Final = (
|
||||
provider if provider is not None else build_provider(settings) if tracing_enabled(settings) 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,
|
||||
)
|
||||
|
|
@ -8,6 +8,7 @@ from types import MappingProxyType
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
|
|
@ -26,6 +27,7 @@ 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
|
||||
|
|
@ -43,6 +45,7 @@ class Container:
|
|||
dataplane_client: DataplaneClient
|
||||
fanout: Fanout
|
||||
audit_log: AuditLog
|
||||
telemetry: Telemetry
|
||||
|
||||
|
||||
class InvalidRolePolicy(ValueError):
|
||||
|
|
@ -111,6 +114,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)
|
||||
|
|
@ -128,6 +132,7 @@ def build_container(
|
|||
policy=policy,
|
||||
)
|
||||
password: Final = settings.ui_password
|
||||
telemetry: Final = build_telemetry(settings, tracer_provider)
|
||||
return Container(
|
||||
settings=settings,
|
||||
policy=policy,
|
||||
|
|
@ -145,7 +150,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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -7,12 +7,14 @@ a management passthrough can tolerate. Retrying is this module's job and only ev
|
|||
import asyncio
|
||||
import random
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Final, Literal, Protocol, TypeAlias
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
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
|
||||
|
||||
|
|
@ -86,6 +88,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,
|
||||
|
|
@ -123,11 +126,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 = self.retry.attempts if method == "GET" else 1
|
||||
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 = self.retry.attempts if method == "GET" else 1
|
||||
outcome: Final = await self._attempt(target, credential.headers, params, json, attempt=0, attempts=attempts)
|
||||
return (
|
||||
outcome
|
||||
if isinstance(outcome, DataplaneUnreachable)
|
||||
|
|
@ -136,36 +143,39 @@ 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:
|
||||
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)
|
||||
else:
|
||||
span.failed(outcome.detail.split(":", 1)[0])
|
||||
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,
|
||||
) -> _RawReply | DataplaneUnreachable:
|
||||
try:
|
||||
response: Final = await self.transport.request(
|
||||
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,
|
||||
)
|
||||
except httpx.HTTPError as error:
|
||||
|
|
|
|||
|
|
@ -109,6 +109,9 @@ class ControlPlaneSettings(BaseSettings):
|
|||
fanout_concurrency: int = 16
|
||||
audit_capacity: int = 10_000
|
||||
api_prefix: str = "/api/v1"
|
||||
otel_enabled: bool | None = None
|
||||
otel_service_name: str = "litellm-control-plane"
|
||||
otel_identity_attributes: bool = True
|
||||
|
||||
@property
|
||||
def oidc(self) -> OidcSettings | None:
|
||||
|
|
|
|||
183
tests/unit/proxy/control_plane/common/test_tracing.py
Normal file
183
tests/unit/proxy/control_plane/common/test_tracing.py
Normal file
|
|
@ -0,0 +1,183 @@
|
|||
from collections.abc import Sequence
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
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.trace import SpanKind, StatusCode, format_trace_id
|
||||
from opentelemetry.trace.span import SpanContext
|
||||
|
||||
from litellm.proxy.control_plane.common.tracing import build_provider, build_telemetry, tracing_enabled
|
||||
from litellm.proxy.control_plane.settings import ControlPlaneSettings
|
||||
from tests.unit.proxy.control_plane.support import MASTER_KEY, 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 test_one_proxy_call_is_one_trace_from_server_span_to_dataplane_header() -> None:
|
||||
exporter = InMemorySpanExporter()
|
||||
client, transport, minter = rig({USER_INFO: [response(200, {"ok": True})]}, build_provider(_settings(), exporter))
|
||||
reply = client.get("/api/v1/dataplanes/eu/proxy/user/info?secret=1", headers=minter.headers("eu"))
|
||||
assert reply.status_code == 200, reply.text
|
||||
spans = 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 = 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 = 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 = InMemorySpanExporter()
|
||||
url = "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 = InMemorySpanExporter()
|
||||
url = "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_names_the_caller_but_not_the_email() -> None:
|
||||
exporter = 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 = 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 = InMemorySpanExporter()
|
||||
outcomes: list[httpx.Response | Exception] = [httpx.ConnectError("boom"), response(503), response(200, {})]
|
||||
client, _, minter = rig({USER_INFO: outcomes}, build_provider(_settings(), exporter))
|
||||
reply = client.get("/api/v1/dataplanes/eu/proxy/user/info", headers=minter.headers("eu"))
|
||||
assert reply.status_code == 200, reply.text
|
||||
attempts = _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_health_is_not_traced() -> None:
|
||||
exporter = InMemorySpanExporter()
|
||||
client, _, _ = rig({}, build_provider(_settings(), exporter))
|
||||
assert client.get("/healthz").status_code == 200
|
||||
assert exporter.get_finished_spans() == ()
|
||||
|
||||
|
||||
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_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,)
|
||||
|
|
@ -8,6 +8,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
|
||||
|
|
@ -81,7 +82,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 = ControlPlaneSettings.model_validate(
|
||||
{"LITELLM_MASTER_KEY": MASTER_KEY, "dataplanes": DATAPLANES, **overrides}
|
||||
|
|
@ -89,6 +92,6 @@ def rig(
|
|||
queues: defaultdict[str, Iterator[httpx.Response | Exception]] = defaultdict(lambda: iter(()))
|
||||
queues.update({url: iter(outcomes) for url, outcomes in routes.items()})
|
||||
transport = RoutedTransport(queues)
|
||||
container = build_container(settings, transport=transport, exchanger=no_exchange)
|
||||
container = build_container(settings, transport=transport, exchanger=no_exchange, tracer_provider=tracer_provider)
|
||||
codec = SessionTokenCodec(MASTER_KEY, settings.session_issuer, timedelta(hours=1), container.policy)
|
||||
return TestClient(create_app(container, Licensed())), transport, Minter(codec, container.policy)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue