/mlflow/tracing/export/mlflow_v3.py

import logging
import threading
from collections import defaultdict
from typing import Sequence

from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.sdk.trace.export import SpanExporter

from mlflow.entities.model_registry import PromptVersion
from mlflow.entities.span import Span
from mlflow.entities.trace import Trace
from mlflow.entities.trace_info import TraceInfo
from mlflow.environment_variables import MLFLOW_ENABLE_ASYNC_TRACE_LOGGING
from mlflow.exceptions import RestException
from mlflow.tracing.client import TracingClient
from mlflow.tracing.constant import SpansLocation, TraceTagKey
from mlflow.tracing.display import get_display_handler
from mlflow.tracing.export.async_export_queue import AsyncTraceExportQueue, Task
from mlflow.tracing.export.utils import try_link_prompts_to_trace
from mlflow.tracing.fluent import _EVAL_REQUEST_ID_TO_TRACE_ID
from mlflow.tracing.trace_manager import InMemoryTraceManager
from mlflow.tracing.utils import (
    add_size_stats_to_trace_metadata,
    encode_span_id,
    get_experiment_id_for_trace,
    maybe_get_request_id,
)
from mlflow.utils.databricks_utils import is_in_databricks_notebook
from mlflow.utils.uri import is_databricks_uri

_logger = logging.getLogger(__name__)


class MlflowV3SpanExporter(SpanExporter):
    """
    An exporter implementation that logs the traces to MLflow Tracking Server
    using the V3 trace schema and API.
    """

    def __init__(self, tracking_uri: str | None = None) -> None:
        self._client = TracingClient(tracking_uri)
        self._is_async_enabled = self._should_enable_async_logging()
        if self._is_async_enabled:
            self._async_queue = AsyncTraceExportQueue()

        # Display handler is no-op when running outside of notebooks.
        self._display_handler = get_display_handler()

        # Tracks whether the store supports span-level logging. Set to False at runtime
        # if log_spans() raises NotImplementedError or returns a 501.
        self._store_supports_log_spans = True

        # Root spans deferred when background thread spans are still running at export time.
        # Keyed by OTel trace ID; popped and exported once all spans in the trace have ended.
        # Protected by _deferred_lock because SimpleSpanProcessor calls export() from the
        # span's own thread, so concurrent calls from multiple threads are possible.
        self._deferred_root_spans: dict[int, ReadableSpan] = {}
        self._deferred_lock = threading.Lock()

    def export(self, spans: Sequence[ReadableSpan]) -> None:
        """
        Export the spans to the destination.

        Args:
            spans: A sequence of OpenTelemetry ReadableSpan objects passed from
                a span processor. All spans (root and non-root) are exported.
        """

        if self._store_supports_log_spans:
            self._export_spans_incrementally(spans)

        self._export_traces(spans)

    def _export_spans_incrementally(self, spans: Sequence[ReadableSpan]) -> None:
        """
        Export spans incrementally as they complete.

        Args:
            spans: Sequence of ReadableSpan objects to export.
            manager: The trace manager instance.
        """
        if is_databricks_uri(self._client.tracking_uri):
            _logger.debug(
                "Databricks tracking server only supports logging spans to UC table, "
                "skipping span exporting."
            )
            return

        mlflow_spans_by_experiment = self._collect_mlflow_spans_for_export(spans)
        for experiment_id, spans_to_log in mlflow_spans_by_experiment.items():
            if self._should_log_async():
                self._async_queue.put(
                    task=Task(
                        handler=self._log_spans,
                        args=(experiment_id, spans_to_log),
                        error_msg="Failed to log spans to the trace server.",
                    )
                )
            else:
                self._log_spans(experiment_id, spans_to_log)

    def _collect_mlflow_spans_for_export(
        self, spans: Sequence[ReadableSpan]
    ) -> dict[str, list[Span]]:
        """
        Collect MLflow spans from ReadableSpans for export, grouped by experiment_id.

        The experiment_id is resolved from the trace info (set during on_start in the
        originating thread) rather than from get_experiment_id_for_trace(), which reads
        thread-local ContextVars that are unavailable in the batch processor's worker thread.

        Args:
            spans: Sequence of ReadableSpan objects.

        Returns:
            Dictionary mapping experiment_id to list of MLflow Span objects.
        """
        manager = InMemoryTraceManager.get_instance()
        spans_by_experiment = defaultdict(list)

        for span in spans:
            mlflow_trace_id = manager.get_mlflow_trace_id_from_otel_id(span.context.trace_id)
            if mlflow_trace_id is None:
                continue
            span_id = encode_span_id(span.context.span_id)
            mlflow_span = manager.get_span_from_id(mlflow_trace_id, span_id)
            if mlflow_span is None:
                continue
            # Get experiment_id from trace info (resolved at on_start time in the
            # originating thread) to survive BatchSpanProcessor thread hops.
            with manager.get_trace(mlflow_trace_id) as trace:
                try:
                    experiment_id = trace.info.experiment_id if trace else None
                except AttributeError:
                    # Remote/distributed traces may have trace_location=None
                    experiment_id = None
            if experiment_id is None:
                experiment_id = get_experiment_id_for_trace(span)
            spans_by_experiment[experiment_id].append(mlflow_span)

        return spans_by_experiment

    def _export_traces(self, spans: Sequence[ReadableSpan]) -> None:
        """
        Export full traces for root spans.

        Args:
            spans: Sequence of ReadableSpan objects.
        """
        manager = InMemoryTraceManager.get_instance()

        # Flush any previously deferred root spans whose background spans have now ended.
        # Copy the keys under the lock, then check has_open_spans() outside the lock to
        # avoid holding _deferred_lock while acquiring InMemoryTraceManager._lock (deadlock risk).
        with self._deferred_lock:
            deferred_ids = list(self._deferred_root_spans.keys())
        for otel_trace_id in deferred_ids:
            if not manager.has_open_spans(otel_trace_id):
                with self._deferred_lock:
                    deferred_span = self._deferred_root_spans.pop(otel_trace_id, None)
                if deferred_span is not None:
                    self._do_export_trace(manager, deferred_span)

        for span in spans:
            if span._parent is not None:
                continue

            # If background-thread child spans are still running, defer the full trace export
            # so that pop_trace is not called until after those spans land in a later batch.
            # This prevents _collect_mlflow_spans_for_export from losing the OTel→MLflow trace
            # ID mapping before those late spans can be logged.
            if manager.has_open_spans(span.context.trace_id):
                with self._deferred_lock:
                    self._deferred_root_spans[span.context.trace_id] = span
                continue

            self._do_export_trace(manager, span)

    def _do_export_trace(self, manager: InMemoryTraceManager, span: ReadableSpan) -> None:
        manager_trace = manager.pop_trace(span.context.trace_id)
        if manager_trace is None:
            _logger.debug(f"Trace for root span {span} not found. Skipping full export.")
            return

        if manager_trace.is_remote_trace and not self._store_supports_log_spans:
            _logger.warning(
                f"Current MLflow server does not support ingesting the span {span.name} "
                "that is created in a remote process. Please upgrade the server version and "
                "use SQL backend to do distributed tracing."
            )
            return

        trace = manager_trace.trace

        # Store mapping from eval request ID to trace ID so that the evaluation
        # harness can access to the trace using mlflow.get_trace(eval_request_id)
        if eval_request_id := trace.info.tags.get(TraceTagKey.EVAL_REQUEST_ID):
            _EVAL_REQUEST_ID_TO_TRACE_ID[eval_request_id] = trace.info.trace_id

        if not maybe_get_request_id(is_evaluate=True):
            self._display_handler.display_traces([trace])

        if self._should_log_async():
            self._async_queue.put(
                task=Task(
                    handler=self._log_trace,
                    args=(trace, manager_trace.prompts),
                    error_msg="Failed to log trace to the trace server.",
                )
            )
        else:
            self._log_trace(trace, prompts=manager_trace.prompts)

    def _log_spans(self, experiment_id: str, spans: list[Span]) -> None:
        """
        Helper method to log spans with error handling.

        Args:
            experiment_id: The experiment ID to log spans to.
            spans: List of spans to log.
        """
        try:
            self._client.log_spans(experiment_id, spans)
        except NotImplementedError:
            # Silently skip if the store doesn't support log_spans. This is expected for stores that
            # don't implement span-level logging, and we don't want to spam warnings for every span.
            self._store_supports_log_spans = False
        except RestException as e:
            # When the FileStore is behind the tracking server, it returns 501 exception.
            # However, the OTLP endpoint returns general HTTP error, not MlflowException, which does
            # not include error_code in the body and handled as a general server side error. Hence,
            # we need to check the message to handle this case.
            if "REST OTLP span logging is not supported" in e.message:
                self._store_supports_log_spans = False
            else:
                _logger.debug(f"Failed to log span to MLflow backend: {e}")
        except Exception as e:
            _logger.debug(f"Failed to log span to MLflow backend: {e}")

    def _log_trace(self, trace: Trace, prompts: Sequence[PromptVersion]) -> None:
        """
        Handles exporting a trace to MLflow using the V3 API and blob storage.
        Steps:
        1. Create the trace in MLflow
        2. Upload the trace data to blob storage using the returned trace info.
        """
        returned_trace_info = None
        try:
            if trace:
                add_size_stats_to_trace_metadata(trace)
                returned_trace_info = self._client.start_trace(trace.info)

                if self._should_log_spans_to_artifacts(returned_trace_info):
                    self._client._upload_trace_data(returned_trace_info, trace.data)
            else:
                _logger.warning("No trace or trace info provided, unable to export")
        except Exception as e:
            _logger.warning(
                f"Failed to send trace to MLflow backend: {e}",
                exc_info=_logger.isEnabledFor(logging.DEBUG),
            )

        # Upload attachments in a separate try-except so trace data still lands
        # even if attachment upload fails. Runs regardless of span storage mode —
        # in TRACKING_STORE mode, spans are in the DB but attachments still go
        # to the artifact repo via the mlflow.artifactLocation tag.
        try:
            if trace and returned_trace_info:
                attachments = {}
                for span in trace.data.spans:
                    attachments.update(span._attachments)
                if attachments:
                    self._client._upload_attachments(returned_trace_info, attachments)
        except Exception as e:
            _logger.warning(
                f"Failed to upload trace attachments: {e}",
                exc_info=_logger.isEnabledFor(logging.DEBUG),
            )

        try:
            # Always run prompt linking asynchronously since (1) prompt linking API calls
            # would otherwise add latency to the export procedure and (2) prompt linking is not
            # critical for trace export (if the prompt fails to link, the user's workflow is
            # minorly affected), so we don't have to await successful linking
            try_link_prompts_to_trace(
                client=self._client,
                trace_id=trace.info.trace_id,
                prompts=prompts,
                synchronous=False,
            )
        except Exception as e:
            _logger.warning(f"Failed to link prompts to trace: {e}")

    def _should_enable_async_logging(self) -> bool:
        if is_in_databricks_notebook():
            # NB: We don't turn on async logging in Databricks notebook by default
            # until we are confident that the async logging is working on the
            # offline workload on Databricks, to derisk the inclusion to the
            # standard image. When it is enabled explicitly via the env var, we
            # will respect that.
            return (
                MLFLOW_ENABLE_ASYNC_TRACE_LOGGING.get()
                if MLFLOW_ENABLE_ASYNC_TRACE_LOGGING.is_set()
                else False
            )

        return MLFLOW_ENABLE_ASYNC_TRACE_LOGGING.get()

    def _should_log_async(self) -> bool:
        # During evaluate or assertion tests, the harness relies on the generated
        # trace objects being immediately available, so log synchronously.
        if maybe_get_request_id(is_evaluate=True):
            return False

        return self._is_async_enabled

    def shutdown(self) -> None:
        # Flush any deferred root spans that are still pending (e.g. if a background
        # span never ended). This prevents them from leaking across exporter lifetimes.
        with self._deferred_lock:
            pending = list(self._deferred_root_spans.items())
            self._deferred_root_spans.clear()
        manager = InMemoryTraceManager.get_instance()
        for _, span in pending:
            self._do_export_trace(manager, span)

    def _should_log_spans_to_artifacts(self, trace_info: TraceInfo) -> bool:
        """
        Whether to log spans to artifacts. Overridden by UC table exporter to False.
        """
        # We only log traces to artifacts when the tracking store doesn't support span logging
        return trace_info.tags.get(TraceTagKey.SPANS_LOCATION) != SpansLocation.TRACKING_STORE.value