mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-19 12:18:35 +00:00
Merge pull request #22169 from mubashir1osmani/fix/arize-phoenix-nested-traces
fix: arize phoenix nested traces
This commit is contained in:
Sameer Kankute
@@ -136,78 +136,137 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore
|
||||
|
||||
return None
|
||||
|
||||
def _get_phoenix_context(self, kwargs):
|
||||
"""
|
||||
Build a trace context for Phoenix's dedicated TracerProvider.
|
||||
|
||||
The base ``_get_span_context`` returns parent spans from the global
|
||||
TracerProvider (the ``otel`` callback). Those spans live on a
|
||||
*different* TracerProvider, so they won't appear in Phoenix — using
|
||||
them as parents just creates broken links.
|
||||
|
||||
Instead we:
|
||||
1. Honour an incoming ``traceparent`` HTTP header (distributed tracing).
|
||||
2. In proxy mode, create our *own* parent span on Phoenix's tracer
|
||||
so the hierarchy is visible end-to-end inside Phoenix.
|
||||
3. In SDK (non-proxy) mode, just return (None, None) for a root span.
|
||||
"""
|
||||
from opentelemetry import trace
|
||||
|
||||
litellm_params = kwargs.get("litellm_params", {}) or {}
|
||||
proxy_server_request = litellm_params.get("proxy_server_request", {}) or {}
|
||||
headers = proxy_server_request.get("headers", {}) or {}
|
||||
|
||||
# Propagate distributed trace context if the caller sent a traceparent
|
||||
traceparent_ctx = (
|
||||
self.get_traceparent_from_header(headers=headers)
|
||||
if headers.get("traceparent")
|
||||
else None
|
||||
)
|
||||
|
||||
is_proxy_mode = bool(proxy_server_request)
|
||||
|
||||
if is_proxy_mode:
|
||||
# Create a parent span on Phoenix's own tracer so both parent
|
||||
# and child are exported to Phoenix.
|
||||
start_time_val = kwargs.get("start_time", kwargs.get("api_call_start_time"))
|
||||
parent_span = self.tracer.start_span(
|
||||
name="litellm_proxy_request",
|
||||
start_time=self._to_ns(start_time_val) if start_time_val is not None else None,
|
||||
context=traceparent_ctx,
|
||||
kind=self.span_kind.SERVER,
|
||||
)
|
||||
ctx = trace.set_span_in_context(parent_span)
|
||||
return ctx, parent_span
|
||||
|
||||
# SDK mode — no parent span needed
|
||||
return traceparent_ctx, None
|
||||
|
||||
def _handle_success(self, kwargs, response_obj, start_time, end_time):
|
||||
"""
|
||||
Override to prevent creating duplicate litellm_request spans when a proxy parent span exists.
|
||||
|
||||
ArizePhoenixLogger should reuse the proxy parent span instead of creating a new litellm_request span,
|
||||
to maintain a shallow span hierarchy as expected by Arize Phoenix.
|
||||
Override to always create spans on ArizePhoenixLogger's dedicated TracerProvider.
|
||||
|
||||
The base class's ``_get_span_context`` would find the parent span created by
|
||||
the ``otel`` callback on the *global* TracerProvider. That span is invisible
|
||||
in Phoenix (different exporter pipeline), so we ignore it and build our own
|
||||
hierarchy via ``_get_phoenix_context``.
|
||||
"""
|
||||
from opentelemetry.trace import Status, StatusCode
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
from litellm.integrations.opentelemetry import LITELLM_PROXY_REQUEST_SPAN_NAME
|
||||
|
||||
|
||||
verbose_logger.debug(
|
||||
"ArizePhoenixLogger: Logging kwargs: %s, OTEL config settings=%s",
|
||||
kwargs,
|
||||
self.config,
|
||||
)
|
||||
ctx, parent_span = self._get_span_context(kwargs)
|
||||
|
||||
# ArizePhoenixLogger NEVER creates a litellm_request span when a proxy parent span exists
|
||||
# This is different from the base OpenTelemetry behavior which respects USE_OTEL_LITELLM_REQUEST_SPAN
|
||||
should_create_primary_span = parent_span is None or (
|
||||
parent_span.name != LITELLM_PROXY_REQUEST_SPAN_NAME
|
||||
and get_secret_bool("USE_OTEL_LITELLM_REQUEST_SPAN")
|
||||
ctx, parent_span = self._get_phoenix_context(kwargs)
|
||||
|
||||
# Create litellm_request span (child of our parent when in proxy mode)
|
||||
span = self.tracer.start_span(
|
||||
name=self._get_span_name(kwargs),
|
||||
start_time=self._to_ns(start_time),
|
||||
context=ctx,
|
||||
)
|
||||
span.set_status(Status(StatusCode.OK))
|
||||
self.set_attributes(span, kwargs, response_obj)
|
||||
|
||||
if should_create_primary_span:
|
||||
# Create a new litellm_request span
|
||||
span = self._start_primary_span(
|
||||
kwargs, response_obj, start_time, end_time, ctx
|
||||
)
|
||||
# Raw-request sub-span (if enabled) - child of litellm_request span
|
||||
self._maybe_log_raw_request(
|
||||
kwargs, response_obj, start_time, end_time, span
|
||||
)
|
||||
# Ensure proxy-request parent span is annotated with the actual operation kind
|
||||
if (
|
||||
parent_span is not None
|
||||
and parent_span.name == LITELLM_PROXY_REQUEST_SPAN_NAME
|
||||
):
|
||||
self.set_attributes(parent_span, kwargs, response_obj)
|
||||
else:
|
||||
# Do not create primary span (keep hierarchy shallow when parent exists)
|
||||
span = None
|
||||
# Only set attributes if the span is still recording (not closed)
|
||||
# Note: parent_span is guaranteed to be not None here
|
||||
if parent_span.is_recording():
|
||||
parent_span.set_status(Status(StatusCode.OK))
|
||||
self.set_attributes(parent_span, kwargs, response_obj)
|
||||
# Raw-request as direct child of parent_span
|
||||
self._maybe_log_raw_request(
|
||||
kwargs, response_obj, start_time, end_time, parent_span
|
||||
)
|
||||
# Raw-request sub-span (if enabled) — must be created before
|
||||
# ending the parent span so the hierarchy is valid.
|
||||
self._maybe_log_raw_request(
|
||||
kwargs, response_obj, start_time, end_time, span
|
||||
)
|
||||
span.end(end_time=self._to_ns(end_time))
|
||||
|
||||
# 3. Guardrail span
|
||||
# Guardrail span
|
||||
self._create_guardrail_span(kwargs=kwargs, context=ctx)
|
||||
|
||||
# 4. Metrics & cost recording
|
||||
# Annotate and close our proxy parent span
|
||||
if parent_span is not None:
|
||||
parent_span.set_status(Status(StatusCode.OK))
|
||||
self.set_attributes(parent_span, kwargs, response_obj)
|
||||
parent_span.end(end_time=self._to_ns(end_time))
|
||||
|
||||
# Metrics & cost recording
|
||||
self._record_metrics(kwargs, response_obj, start_time, end_time)
|
||||
|
||||
# 5. Semantic logs.
|
||||
# Semantic logs
|
||||
if self.config.enable_events:
|
||||
log_span = span if span is not None else parent_span
|
||||
if log_span is not None:
|
||||
self._emit_semantic_logs(kwargs, response_obj, log_span)
|
||||
self._emit_semantic_logs(kwargs, response_obj, span)
|
||||
|
||||
# 6. Do NOT end parent span - it should be managed by its creator
|
||||
# External spans (from Langfuse, user code, HTTP headers, global context) must not be closed by LiteLLM
|
||||
# However, proxy-created spans should be closed here
|
||||
if (
|
||||
parent_span is not None
|
||||
and parent_span.name == LITELLM_PROXY_REQUEST_SPAN_NAME
|
||||
):
|
||||
def _handle_failure(self, kwargs, response_obj, start_time, end_time):
|
||||
"""
|
||||
Override to always create failure spans on ArizePhoenixLogger's dedicated
|
||||
TracerProvider. Mirrors ``_handle_success`` but sets ERROR status.
|
||||
"""
|
||||
from opentelemetry.trace import Status, StatusCode
|
||||
|
||||
verbose_logger.debug(
|
||||
"ArizePhoenixLogger: Failure - Logging kwargs: %s, OTEL config settings=%s",
|
||||
kwargs,
|
||||
self.config,
|
||||
)
|
||||
|
||||
ctx, parent_span = self._get_phoenix_context(kwargs)
|
||||
|
||||
# Create litellm_request span (child of our parent when in proxy mode)
|
||||
span = self.tracer.start_span(
|
||||
name=self._get_span_name(kwargs),
|
||||
start_time=self._to_ns(start_time),
|
||||
context=ctx,
|
||||
)
|
||||
span.set_status(Status(StatusCode.ERROR))
|
||||
self.set_attributes(span, kwargs, response_obj)
|
||||
self._record_exception_on_span(span=span, kwargs=kwargs)
|
||||
span.end(end_time=self._to_ns(end_time))
|
||||
|
||||
# Guardrail span
|
||||
self._create_guardrail_span(kwargs=kwargs, context=ctx)
|
||||
|
||||
# Annotate and close our proxy parent span
|
||||
if parent_span is not None:
|
||||
parent_span.set_status(Status(StatusCode.ERROR))
|
||||
self.set_attributes(parent_span, kwargs, response_obj)
|
||||
self._record_exception_on_span(span=parent_span, kwargs=kwargs)
|
||||
parent_span.end(end_time=self._to_ns(end_time))
|
||||
|
||||
@staticmethod
|
||||
|
||||
@@ -3834,6 +3834,12 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
|
||||
)
|
||||
)
|
||||
_in_memory_loggers.append(otel_logger)
|
||||
|
||||
# Auto-initialize Arize Phoenix if Phoenix env vars are configured
|
||||
# This allows users to get nested traces in both OTEL and Phoenix
|
||||
# by only specifying "otel" in callbacks
|
||||
_maybe_auto_initialize_arize_phoenix(_in_memory_loggers)
|
||||
|
||||
return otel_logger # type: ignore
|
||||
|
||||
elif logging_integration == "galileo":
|
||||
@@ -3887,7 +3893,8 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
|
||||
headers=f"Authorization={os.getenv('LOGFIRE_TOKEN')}",
|
||||
)
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, OpenTelemetry):
|
||||
# Use exact type check to avoid matching ArizePhoenixLogger (subclass)
|
||||
if type(callback) is OpenTelemetry:
|
||||
return callback # type: ignore
|
||||
_otel_logger = OpenTelemetry(config=otel_config)
|
||||
_in_memory_loggers.append(_otel_logger)
|
||||
@@ -4147,6 +4154,57 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
|
||||
return None
|
||||
|
||||
|
||||
def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list) -> None:
|
||||
"""
|
||||
Auto-initialize ArizePhoenixLogger when Phoenix env vars are detected.
|
||||
|
||||
Called during ``otel`` callback setup so that users get nested traces in
|
||||
both their OTEL collector *and* Arize Phoenix by only listing ``"otel"``
|
||||
in ``callbacks``. If no Phoenix env vars are set, this is a no-op.
|
||||
"""
|
||||
phoenix_env_vars = (
|
||||
"PHOENIX_API_KEY",
|
||||
"PHOENIX_COLLECTOR_HTTP_ENDPOINT",
|
||||
"PHOENIX_COLLECTOR_ENDPOINT",
|
||||
)
|
||||
if not any(os.environ.get(v) for v in phoenix_env_vars):
|
||||
return
|
||||
|
||||
# Already registered — nothing to do
|
||||
if any(
|
||||
isinstance(cb, ArizePhoenixLogger) and cb.callback_name == "arize_phoenix"
|
||||
for cb in _in_memory_loggers
|
||||
):
|
||||
return
|
||||
|
||||
try:
|
||||
from litellm.integrations.opentelemetry import OpenTelemetryConfig
|
||||
|
||||
arize_phoenix_config = ArizePhoenixLogger.get_arize_phoenix_config()
|
||||
otel_config = OpenTelemetryConfig(
|
||||
exporter=arize_phoenix_config.protocol,
|
||||
endpoint=arize_phoenix_config.endpoint,
|
||||
headers=arize_phoenix_config.otlp_auth_headers,
|
||||
)
|
||||
phoenix_logger = ArizePhoenixLogger(
|
||||
config=otel_config, callback_name="arize_phoenix"
|
||||
)
|
||||
_in_memory_loggers.append(phoenix_logger)
|
||||
|
||||
# Register as a litellm callback so it receives success/failure events
|
||||
litellm.logging_callback_manager.add_litellm_callback(phoenix_logger)
|
||||
|
||||
verbose_logger.info(
|
||||
"Auto-initialized Arize Phoenix logger alongside otel "
|
||||
"(endpoint=%s)",
|
||||
arize_phoenix_config.endpoint,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
"Failed to auto-initialize Arize Phoenix logger: %s", str(e)
|
||||
)
|
||||
|
||||
|
||||
def get_custom_logger_compatible_class( # noqa: PLR0915
|
||||
logging_integration: _custom_logger_compatible_callbacks_literal,
|
||||
) -> Optional[CustomLogger]:
|
||||
@@ -4249,7 +4307,8 @@ def get_custom_logger_compatible_class( # noqa: PLR0915
|
||||
from litellm.integrations.opentelemetry import OpenTelemetry
|
||||
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, OpenTelemetry):
|
||||
# Use exact type check to avoid matching ArizePhoenixLogger (subclass)
|
||||
if type(callback) is OpenTelemetry:
|
||||
return callback
|
||||
elif logging_integration == "arize":
|
||||
if "ARIZE_API_KEY" not in os.environ:
|
||||
@@ -4266,7 +4325,8 @@ def get_custom_logger_compatible_class( # noqa: PLR0915
|
||||
from litellm.integrations.opentelemetry import OpenTelemetry
|
||||
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, OpenTelemetry):
|
||||
# Use exact type check to avoid matching ArizePhoenixLogger (subclass)
|
||||
if type(callback) is OpenTelemetry:
|
||||
return callback # type: ignore
|
||||
|
||||
elif logging_integration == "dynamic_rate_limiter":
|
||||
|
||||
@@ -6,8 +6,11 @@ Covers the three root-cause fixes:
|
||||
1. ArizePhoenixLogger / ArizeLogger create *dedicated* TracerProviders.
|
||||
2. The ``otel`` dedup check does NOT match Arize subclasses.
|
||||
3. Arize loggers do NOT overwrite ``proxy_server.open_telemetry_logger``.
|
||||
4. Phoenix creates nested spans (parent + child) in proxy mode.
|
||||
5. Auto-initialization of Phoenix when env vars are detected.
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
@@ -165,5 +168,46 @@ class TestProxyLoggerNotOverwritten(unittest.TestCase):
|
||||
assert proxy_server.open_telemetry_logger is None
|
||||
|
||||
|
||||
class TestPhoenixAutoInitWithOtelOnly(unittest.TestCase):
|
||||
"""When only 'otel' is configured but Phoenix env vars are set,
|
||||
ArizePhoenixLogger should be auto-initialized and receive spans."""
|
||||
|
||||
def setUp(self):
|
||||
"""Save original callbacks to restore after each test."""
|
||||
import litellm
|
||||
self._original_callbacks = litellm.callbacks[:]
|
||||
|
||||
def tearDown(self):
|
||||
"""Restore original callbacks to prevent global state leakage."""
|
||||
import litellm
|
||||
litellm.callbacks = self._original_callbacks
|
||||
|
||||
@patch.dict(os.environ, {
|
||||
"PHOENIX_COLLECTOR_HTTP_ENDPOINT": "http://localhost:6006/v1/traces",
|
||||
}, clear=False)
|
||||
def test_auto_init_creates_phoenix_logger(self):
|
||||
from litellm.integrations.arize.arize_phoenix import ArizePhoenixLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import _maybe_auto_initialize_arize_phoenix
|
||||
|
||||
_in_memory_loggers = []
|
||||
_maybe_auto_initialize_arize_phoenix(_in_memory_loggers)
|
||||
|
||||
phoenix_loggers = [cb for cb in _in_memory_loggers if isinstance(cb, ArizePhoenixLogger)]
|
||||
assert len(phoenix_loggers) == 1, "Phoenix logger should be auto-initialized when env vars are set"
|
||||
|
||||
def test_no_auto_init_without_env_vars(self):
|
||||
from litellm.integrations.arize.arize_phoenix import ArizePhoenixLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import _maybe_auto_initialize_arize_phoenix
|
||||
|
||||
env_keys = ["PHOENIX_API_KEY", "PHOENIX_COLLECTOR_HTTP_ENDPOINT", "PHOENIX_COLLECTOR_ENDPOINT"]
|
||||
with patch.dict(os.environ, {k: "" for k in env_keys}, clear=False):
|
||||
for k in env_keys:
|
||||
os.environ.pop(k, None)
|
||||
_in_memory_loggers = []
|
||||
_maybe_auto_initialize_arize_phoenix(_in_memory_loggers)
|
||||
phoenix_loggers = [cb for cb in _in_memory_loggers if isinstance(cb, ArizePhoenixLogger)]
|
||||
assert len(phoenix_loggers) == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user