wehub-resource-sync
106 lines
4.5 KiB
Python
106 lines
4.5 KiB
Python
import logging
|
|
|
|
from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan
|
|
from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter
|
|
|
|
from mlflow.entities.span import create_mlflow_span
|
|
from mlflow.entities.trace_info import TraceInfo, TraceLocation, TraceState
|
|
from mlflow.environment_variables import (
|
|
MLFLOW_ENABLE_OTEL_GENAI_SEMCONV,
|
|
MLFLOW_TRACE_ENABLE_OTLP_DUAL_EXPORT,
|
|
)
|
|
from mlflow.tracing.constant import TRACE_SCHEMA_VERSION, TRACE_SCHEMA_VERSION_KEY, SpanAttributeKey
|
|
from mlflow.tracing.processor.otel_metrics_mixin import OtelMetricsMixin
|
|
from mlflow.tracing.trace_manager import InMemoryTraceManager
|
|
from mlflow.tracing.utils import generate_trace_id_v3
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
class OtelSpanProcessor(OtelMetricsMixin, BatchSpanProcessor):
|
|
"""
|
|
SpanProcessor implementation to export MLflow traces to a OpenTelemetry collector.
|
|
|
|
Extending OpenTelemetry BatchSpanProcessor to add some custom hooks to be executed when a span
|
|
is started or ended (before exporting).
|
|
"""
|
|
|
|
def __init__(self, span_exporter: SpanExporter, export_metrics: bool) -> None:
|
|
super().__init__(span_exporter)
|
|
self._export_metrics = export_metrics
|
|
# In opentelemetry-sdk 1.34.0, the `span_exporter` field was removed from the
|
|
# `BatchSpanProcessor` class.
|
|
# https://github.com/open-telemetry/opentelemetry-python/issues/4616
|
|
#
|
|
# The `span_exporter` field was restored as a property in 1.34.1
|
|
# https://github.com/open-telemetry/opentelemetry-python/pull/4621
|
|
#
|
|
# We use a try-except block to maintain compatibility with both versions.
|
|
try:
|
|
self.span_exporter = span_exporter
|
|
except AttributeError:
|
|
pass
|
|
|
|
# Only register traces with trace manager when NOT in dual export mode
|
|
# In dual export mode, MLflow span processors handle trace registration
|
|
self._should_register_traces = not MLFLOW_TRACE_ENABLE_OTLP_DUAL_EXPORT.get()
|
|
if self._should_register_traces:
|
|
self._trace_manager = InMemoryTraceManager.get_instance()
|
|
|
|
def on_start(self, span: OTelReadableSpan, parent_context=None):
|
|
if self._should_register_traces:
|
|
if not span.parent:
|
|
trace_info = self._create_trace_info(span)
|
|
trace_id = trace_info.trace_id
|
|
self._trace_manager.register_trace(span.context.trace_id, trace_info)
|
|
else:
|
|
trace_id = self._trace_manager.get_mlflow_trace_id_from_otel_id(
|
|
span.context.trace_id
|
|
)
|
|
self._trace_manager.register_span(create_mlflow_span(span, trace_id))
|
|
|
|
super().on_start(span, parent_context)
|
|
|
|
def on_end(self, span: OTelReadableSpan):
|
|
if self._export_metrics:
|
|
self.record_metrics_for_span(span)
|
|
|
|
if self._should_register_traces and not span.parent:
|
|
manager_trace = self._trace_manager.pop_trace(span.context.trace_id)
|
|
# Emit user-defined tags as individual span attributes so they survive OTLP export
|
|
# and can be restored to trace tag rows on the server side.
|
|
if (
|
|
manager_trace is not None
|
|
and manager_trace.trace.info.tags
|
|
and span._attributes is not None
|
|
):
|
|
for key, value in manager_trace.trace.info.tags.items():
|
|
if not key.startswith("mlflow."):
|
|
span._attributes[SpanAttributeKey.TRACE_TAG_PREFIX + key] = value
|
|
|
|
if MLFLOW_ENABLE_OTEL_GENAI_SEMCONV.get():
|
|
span = self._translate_span(span)
|
|
|
|
super().on_end(span)
|
|
|
|
def _translate_span(self, span: OTelReadableSpan) -> OTelReadableSpan:
|
|
try:
|
|
from mlflow.tracing.export.genai_semconv.translator import translate_span_to_genai
|
|
|
|
return translate_span_to_genai(span)
|
|
except Exception:
|
|
_logger.debug("Failed to translate span to GenAI semconv", exc_info=True)
|
|
return span
|
|
|
|
def _create_trace_info(self, span: OTelReadableSpan) -> TraceInfo:
|
|
"""Create a TraceInfo object from an OpenTelemetry span."""
|
|
return TraceInfo(
|
|
trace_id=generate_trace_id_v3(span),
|
|
trace_location=TraceLocation.from_experiment_id(None),
|
|
request_time=span.start_time // 1_000_000, # nanosecond to millisecond
|
|
execution_duration=None,
|
|
state=TraceState.IN_PROGRESS,
|
|
trace_metadata={TRACE_SCHEMA_VERSION_KEY: str(TRACE_SCHEMA_VERSION)},
|
|
tags={},
|
|
)
|