mlflow--mlflow/mlflow/tracing/provider.py

"""
This module provides a set of functions to manage the global tracer provider for MLflow tracing.

Every tracing operation in MLflow *MUST* be managed through this module, instead of directly
using the OpenTelemetry APIs. This is because MLflow needs to control the initialization of the
tracer provider and ensure that it won't interfere with the other external libraries that might
use OpenTelemetry e.g. PromptFlow, Snowpark.
"""

import contextvars
import functools
import json
import logging
import os
import random
import threading
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any, Callable, ParamSpec, TypeVar

from opentelemetry import context as context_api
from opentelemetry import trace
from opentelemetry.context.contextvars_context import ContextVarsRuntimeContext
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import SpanProcessor, TracerProvider
from opentelemetry.sdk.trace.id_generator import IdGenerator
from opentelemetry.sdk.trace.sampling import ParentBased

import mlflow
from mlflow.entities.trace_location import (
    MlflowExperimentLocation,
    TraceLocationBase,
    UCSchemaLocation,
    UnityCatalog,
)
from mlflow.environment_variables import (
    MLFLOW_TRACE_ENABLE_OTLP_DUAL_EXPORT,
    MLFLOW_TRACE_SAMPLING_RATIO,
    MLFLOW_TRACE_USE_ISOLATED_RANDOM_ID_GENERATOR,
    MLFLOW_USE_BATCH_SPAN_PROCESSOR,
    MLFLOW_USE_DEFAULT_TRACER_PROVIDER,
)
from mlflow.exceptions import MlflowException, MlflowTracingException
from mlflow.tracing.config import reset_config
from mlflow.tracing.constant import SpanAttributeKey
from mlflow.tracing.destination import TraceDestination, UserTraceDestinationRegistry
from mlflow.tracing.sampling import _MlflowSampler
from mlflow.tracing.utils.exception import raise_as_trace_exception
from mlflow.tracing.utils.once import Once
from mlflow.tracing.utils.otlp import (
    get_otlp_exporter,
    should_export_otlp_metrics,
    should_use_otlp_exporter,
)
from mlflow.tracing.utils.warning import suppress_warning
from mlflow.utils.databricks_utils import (
    is_in_databricks_model_serving_environment,
    is_mlflow_tracing_enabled_in_model_serving,
)
from mlflow.utils.uri import is_databricks_uri

if TYPE_CHECKING:
    from mlflow.entities import Span


P = ParamSpec("P")
R = TypeVar("R")

# A trace destination specified by the user via the `set_destination` function.
_MLFLOW_TRACE_USER_DESTINATION = UserTraceDestinationRegistry()


_logger = logging.getLogger(__name__)


_private_random_generators = set()
if hasattr(os, "register_at_fork"):
    # Re-seed the private random instances in forked child processes to prevent them
    # from generating the same ID sequence as the parent process.

    def _reseed():
        # use `tuple` to create a snapshot of the set for iterating,
        # to avoid concurrency issue.
        for r in tuple(_private_random_generators):
            r.seed()

    os.register_at_fork(after_in_child=_reseed)


class _IsolatedRandomIdGenerator(IdGenerator):
    """
    An OTel IdGenerator that uses a private ``random.Random`` instance, isolated from
    Python's global ``random`` module state.

    Unlike the default OTel ``RandomIdGenerator``, this is immune to ``random.seed()``
    calls in user code, preventing duplicate trace/span IDs across re-runs.

    Enable via ``MLFLOW_TRACE_USE_ISOLATED_RANDOM_ID_GENERATOR=true``.
    """

    def __init__(self) -> None:
        self._random = random.Random()
        _private_random_generators.add(self._random)

    def __del__(self):
        _private_random_generators.remove(self._random)

    def generate_span_id(self) -> int:
        span_id = self._random.getrandbits(64)
        while span_id == trace.INVALID_SPAN_ID:
            span_id = self._random.getrandbits(64)
        return span_id

    def generate_trace_id(self) -> int:
        trace_id = self._random.getrandbits(128)
        while trace_id == trace.INVALID_TRACE_ID:
            trace_id = self._random.getrandbits(128)
        return trace_id


class _TracerProviderWrapper:
    """
    A facade for the tracer provider.
    MLflow uses different tracer providers depending on the MLFLOW_USE_DEFAULT_TRACER_PROVIDER
    environment variable setting.
    1. Use an isolated tracer provider instance managed by MLflow. This is the default behavior such
       that MLflow does not break an environment where MLflow and OpenTelemetry SDK are used in
       different purposes.
    2. Use the global OpenTelemetry tracer provider singleton and traces created by MLflow and
        OpenTelemetry SDK will be exported to the same destination.
    """

    def __init__(self):
        self._isolated_tracer_provider = None
        self._isolated_tracer_provider_once = Once()
        # Separate once flag for global provider mode. We use MLflow's own flag instead of
        # OTel's _TRACER_PROVIDER_SET_ONCE so that MLflow can initialize its span processors
        # even when an external library (e.g., auto-instrumentation) has already set up
        # the global tracer provider before MLflow initializes.
        self._global_provider_init_once = Once()

    @property
    def once(self) -> Once:
        if MLFLOW_USE_DEFAULT_TRACER_PROVIDER.get():
            return self._isolated_tracer_provider_once
        return self._global_provider_init_once

    def get(self) -> TracerProvider:
        if MLFLOW_USE_DEFAULT_TRACER_PROVIDER.get():
            return self._isolated_tracer_provider
        return trace.get_tracer_provider()

    def _retire_current_batch_processors(self) -> None:
        """Flush and shut down the outgoing provider's MLflow batch processors.

        Dropping a provider without this leaks its BatchSpanProcessor daemon
        thread (#24209). Retires both MLflow-owned kinds that carry one:
        ``BaseMlflowSpanProcessor`` (async MLflow export) and ``OtelSpanProcessor``
        (OTLP export). External processors on a shared global provider are left
        untouched.
        """
        from mlflow.tracing.processor.base_mlflow import (
            BaseMlflowSpanProcessor,
            retire_batch_processor,
        )
        from mlflow.tracing.processor.otel import OtelSpanProcessor

        current = self.get()
        if not isinstance(current, TracerProvider):
            return
        active = getattr(current, "_active_span_processor", None)
        for processor in list(getattr(active, "_span_processors", ())):
            if isinstance(processor, BaseMlflowSpanProcessor):
                retire_batch_processor(processor)
            elif isinstance(processor, OtelSpanProcessor):
                # Flush before shutdown so queued spans are exported, not dropped.
                try:
                    processor.force_flush()
                    processor.shutdown()
                except Exception:
                    _logger.debug(f"Failed to retire OtelSpanProcessor {processor}", exc_info=True)

    def set(self, tracer_provider: TracerProvider):
        self._retire_current_batch_processors()
        if MLFLOW_USE_DEFAULT_TRACER_PROVIDER.get():
            self._isolated_tracer_provider = tracer_provider
        else:
            # Bypass the once flag otherwise the update will be ignored.
            # We check the once flag inside `get_or_init_tracer`. For other cases, trace provider
            # should be forcibly updated.
            trace._TRACER_PROVIDER = tracer_provider

    def get_or_init_tracer(self, module_name: str) -> trace.Tracer:
        self.once.do_once(_initialize_tracer_provider)
        return self.get().get_tracer(module_name)

    def reset(self):
        self._retire_current_batch_processors()
        if MLFLOW_USE_DEFAULT_TRACER_PROVIDER.get():
            self._isolated_tracer_provider = None
            self._isolated_tracer_provider_once._done = False
        else:
            trace._TRACER_PROVIDER = None
            self._global_provider_init_once._done = False

    def _swap_raw(self, tracer_provider: TracerProvider) -> TracerProvider:
        """Install ``tracer_provider`` and return the previous one, without retiring it.

        Unlike ``set()``, the outgoing provider is left running so ``trace_disabled``
        can restore it after hiding it behind a NoOp, without rebuilding its
        BatchSpanProcessor thread. Contract: a concurrent global mutation
        (``enable()``/``set_destination()`` from another thread) during the window
        can be overwritten by the restore; mutating global tracing state from
        multiple threads at once is unsupported.
        """
        if MLFLOW_USE_DEFAULT_TRACER_PROVIDER.get():
            old = self._isolated_tracer_provider
            self._isolated_tracer_provider = tracer_provider
            return old
        old = trace._TRACER_PROVIDER
        trace._TRACER_PROVIDER = tracer_provider
        return old


provider = _TracerProviderWrapper()
mlflow_runtime_context = ContextVarsRuntimeContext()

# Serialize trace_disabled's swap/restore and count active frames so nested and
# concurrent calls swap only on the outermost entry and restore on the outermost
# exit; otherwise overlapping frames can leave a NoOp installed permanently (#24209).
_trace_disabled_lock = threading.Lock()
_trace_disabled_state: dict[str, Any] = {"depth": 0, "saved_provider": None, "saved_once": None}


def get_context_api():
    """
    Get the context API to use for the current tracer provider mode.
    """
    return mlflow_runtime_context if MLFLOW_USE_DEFAULT_TRACER_PROVIDER.get() else context_api


def get_current_context() -> context_api.Context | None:
    """
    Get the current context.
    If the current tracer provider mode is unified mode (mlflow+otel), return None
    to let downstream otel APIs to fetch it from otel runtime context.
    """
    return (
        mlflow_runtime_context.get_current() if MLFLOW_USE_DEFAULT_TRACER_PROVIDER.get() else None
    )


def get_current_otel_span() -> trace.Span | None:
    """
    Get the current otel span from the current context.
    """
    return trace.get_current_span(context=get_current_context())


def start_span_in_context(name: str, experiment_id: str | None = None) -> trace.Span:
    """
    Start a new OpenTelemetry span in the current context.

    Note that this function doesn't set the started span as the active span in the context. To do
    that, the upstream also need to call `use_span()` function in the OpenTelemetry trace APIs.

    Args:
        name: The name of the span.
        experiment_id: The ID of the experiment to log the span to. If not specified, the span will
            be logged to the active experiment or explicitly set trace destination.

    Returns:
        The newly created OpenTelemetry span.
    """
    attributes = {}
    if experiment_id:
        attributes[SpanAttributeKey.EXPERIMENT_ID] = json.dumps(experiment_id)
    span = _get_tracer(__name__).start_span(
        name,
        attributes=attributes,
        context=get_current_context(),
    )

    if experiment_id and getattr(span, "_parent", None):
        _logger.warning(
            "The `experiment_id` parameter can only be used for root spans, but the span "
            f"`{name}` is not a root span. The specified value `{experiment_id}` will be ignored."
        )
        span._span.attributes.pop(SpanAttributeKey.EXPERIMENT_ID, None)
    return span


@contextmanager
def with_active_span(span: "Span"):
    """
    A context manager that sets the given MLflow span as the active span in the current context.

    A fork of OpenTelemetry's `use_span` context manager, but use MLflow's `set_span_in_context` and
    `detach_span_from_context` functions to set and detach the span from the context, in order to
    switch the context depending on the `MLFLOW_USE_DEFAULT_TRACER_PROVIDER` environment variable.
    """
    try:
        token = set_span_in_context(span)
        try:
            yield
        finally:
            detach_span_from_context(token)
    except Exception as exc:
        span.record_exception(exc)
        raise


def start_detached_span(
    name: str,
    parent: trace.Span | None = None,
    experiment_id: str | None = None,
    start_time_ns: int | None = None,
) -> tuple[str, trace.Span] | None:
    """
    Start a new OpenTelemetry span that is not part of the current trace context, but with the
    explicit parent span ID if provided.

    Args:
        name: The name of the span.
        parent: The parent OpenTelemetry span. If not provided, the span will be created as a root
                span.
        experiment_id: The ID of the experiment. This is used to associate the span with a specific
            experiment in MLflow.
        start_time_ns: The start time of the span in nanoseconds.
            If not provided, the current timestamp is used.

    Returns:
        The newly created OpenTelemetry span.
    """
    tracer = _get_tracer(__name__)
    context = (
        trace.set_span_in_context(parent, context=get_current_context())
        if parent
        else get_current_context()
    )
    attributes = {}

    # Set start time and experiment to attribute so we can pass it to the span processor
    if start_time_ns:
        attributes[SpanAttributeKey.START_TIME_NS] = json.dumps(start_time_ns)
    if experiment_id:
        attributes[SpanAttributeKey.EXPERIMENT_ID] = json.dumps(experiment_id)
    span = tracer.start_span(
        name,
        context=context,
        attributes=attributes,
        start_time=start_time_ns,
    )

    if experiment_id and getattr(span, "_parent", None):
        _logger.warning(
            "The `experiment_id` parameter can only be used for root spans, but the span "
            f"`{name}` is not a root span. The specified value `{experiment_id}` will be ignored."
        )
        span._span.attributes.pop(SpanAttributeKey.EXPERIMENT_ID, None)
    return span


@contextmanager
def safe_set_span_in_context(span: "Span"):
    """
    A context manager that sets the given OpenTelemetry span as the active span in the current
    context.

    Args:
        span: An MLflow span object to set as the active span.

    Example:

    .. code-block:: python

        import mlflow


        with mlflow.start_span("my_span") as span:
            span.set_attribute("my_key", "my_value")

        # The span is automatically detached from the context when the context manager exits.
    """
    token = set_span_in_context(span)
    try:
        yield
    finally:
        detach_span_from_context(token)


def set_span_in_context(span: "Span") -> contextvars.Token:
    """
    Set the given OpenTelemetry span as the active span in the current context.

    Args:
        span: An MLflow span object to set as the active span.

    Returns:
        A token object that will be required when detaching the span from the context.
    """
    context = trace.set_span_in_context(span._span, context=get_current_context())
    if MLFLOW_USE_DEFAULT_TRACER_PROVIDER.get():
        # When using the default tracer provider, attach to MLflow's runtime context so that span
        # will not get mixed with the native OpenTelemetry runtime context.
        token = mlflow_runtime_context.attach(context)
    else:
        token = context_api.attach(context)
    return token


def detach_span_from_context(token: contextvars.Token):
    """
    Remove the active span from the current context.

    Args:
        token: The token returned by `_set_span_to_active` function.
    """
    if MLFLOW_USE_DEFAULT_TRACER_PROVIDER.get():
        mlflow_runtime_context.detach(token)
    else:
        context_api.detach(token)


def set_destination(destination: TraceLocationBase, *, context_local: bool = False):
    """
    Set a custom span location to which MLflow will export the traces.

    A destination specified by this function will take precedence over
    other configurations, such as tracking URI, OTLP environment variables.

    Args:
        destination: A trace location object that specifies the location of the trace data.
            Currently, the following locations are supported:

            - :py:class:`~mlflow.entities.trace_location.MlflowExperimentLocation`: Logs traces to
                an MLflow experiment.

        context_local: If False (default), the destination is set globally. If True, the destination
            is isolated per async task or thread, providing isolation in concurrent applications.

    Example:

        **Logging traces to MLflow Experiment:**

        .. code-block:: python

            from mlflow.entities.trace_location import MlflowExperimentLocation

            mlflow.tracing.set_destination(MlflowExperimentLocation(experiment_id="123"))

        Note: This has the same effect as setting the active MLflow experiment via the
        ``MLFLOW_EXPERIMENT_ID`` environment variable or the ``mlflow.set_experiment`` API,
        but with narrower scope.

        **Isolate the destination between async tasks or threads:**

        .. code-block:: python

            from mlflow.tracing.destination import Databricks

            mlflow.tracing.set_destination(
                MlflowExperimentLocation(experiment_id="123"),
                context_local=True,
            )

        The destination set with the ``context_local`` flag will only be effective within the
        current async task or thread. This is particularly useful when you want to send traces
        to different destinations from a multi-tenant application.

        ** Reset the destination:**

        .. code-block:: python

            mlflow.tracing.reset()

    """
    if isinstance(destination, TraceDestination):
        # NB: Deprecation warnings are issued in the constructor of the destination classes
        # so we don't need to issue a warning here.
        destination = destination.to_location()

    if not isinstance(destination, TraceLocationBase):
        raise MlflowException.invalid_parameter_value(
            f"Invalid destination type: {type(destination)}. "
            "The destination must be an instance of TraceLocation."
        )

    if isinstance(destination, UnityCatalog):
        raise MlflowException.invalid_parameter_value(
            "UnityCatalog table-prefix destinations are not supported by "
            "`mlflow.tracing.set_destination`. Use `set_experiment` with a "
            "UnityCatalog location instead."
        )

    if isinstance(destination, UCSchemaLocation):
        _logger.warning(
            "Passing `UCSchemaLocation` to `mlflow.tracing.set_destination` is deprecated "
            "and will be removed in a future MLflow version. Use `set_experiment` with a "
            "UnityCatalog location instead. See "
            "https://docs.databricks.com/aws/en/mlflow3/genai/tracing/trace-unity-catalog"
        )
        if mlflow.get_tracking_uri() is None or not mlflow.get_tracking_uri().startswith(
            "databricks"
        ):
            mlflow.set_tracking_uri("databricks")
            _logger.info(
                "Automatically setting the tracking URI to `databricks` "
                "because the tracing destination is set to Databricks."
            )

    _MLFLOW_TRACE_USER_DESTINATION.set(destination, context_local=context_local)
    _initialize_tracer_provider()


def _get_tracer(module_name: str) -> trace.Tracer:
    """
    Get a tracer instance for the given module name.

    If the tracer provider is not initialized, this function will initialize the tracer provider.
    Other simultaneous calls to this function will block until the initialization is done.
    """
    return provider.get_or_init_tracer(module_name)


def _get_span_processor():
    """
    Get the MLflow span processor instance from the current tracer provider.

    When multiple span processors are registered (e.g., OTLP + MLflow in dual-export mode),
    this scans for a BaseMlflowSpanProcessor explicitly rather than assuming index 0.
    Returns None if no MLflow span processor is found or the provider is not SDK-based.
    """
    # Inline import to avoid circular dependency:
    # provider -> base_mlflow -> fluent -> entities
    from mlflow.tracing.processor.base_mlflow import BaseMlflowSpanProcessor

    tracer_provider = provider.get()
    if not isinstance(tracer_provider, TracerProvider):
        return None

    if active_span_processor := getattr(tracer_provider, "_active_span_processor", None):
        processors = getattr(active_span_processor, "_span_processors", None)
    else:
        processors = None
    if not processors:
        return None

    return next(
        (p for p in processors if isinstance(p, BaseMlflowSpanProcessor)),
        processors[0],
    )


def _get_trace_exporter():
    """
    Get the exporter instance from the MLflow span processor in the current tracer provider.
    """
    if processor := _get_span_processor():
        return getattr(processor, "span_exporter", None)
    return None


def _initialize_tracer_provider(disabled=False):
    """
    Instantiate a tracer provider and set it as the global tracer provider.

    Note that this function ALWAYS updates the global tracer provider, regardless of the current
    state. It is the caller's responsibility to ensure that the tracer provider is initialized
    only once, and update the _INITIALIZED flag accordingly.
    """
    processors = _get_span_processors(disabled=disabled)
    if not processors:
        provider.set(trace.NoOpTracerProvider())
        return

    # Demote the "Failed to detach context" log raised by the OpenTelemetry logger to DEBUG
    # level so that it does not show up in the user's console. This warning may indicate
    # some incorrect context handling, but in many cases just false positive that does not
    # cause any issue in the generated trace.
    # Note that we need to apply it permanently rather than just the scope of prediction call,
    # because the exception can happen for streaming case, where the error log might be
    # generated when the iterator is consumed and we don't know when it will happen.
    suppress_warning("opentelemetry.context", "Failed to detach context")

    # These Otel warnings are occasionally raised in a valid case, e.g. we timeout a trace
    # but some spans are still active. We suppress them because they are not actionable.
    suppress_warning("opentelemetry.sdk.trace", "Setting attribute on ended span")
    suppress_warning("opentelemetry.sdk.trace", "Calling end() on an ended span")

    # When using the global tracer provider mode (MLFLOW_USE_DEFAULT_TRACER_PROVIDER=false),
    # check if a TracerProvider is already set (e.g., set with `set_tracer_provider` API for
    # Uvicorn, FastAPI, etc.). If so, add MLflow's span processors to the existing provider
    # instead of replacing it. This enables unified tracing where both external and MLflow spans
    # are captured in the same trace.
    if not MLFLOW_USE_DEFAULT_TRACER_PROVIDER.get():
        existing_provider = trace.get_tracer_provider()
        if isinstance(existing_provider, TracerProvider):
            _logger.debug(
                "An existing TracerProvider is detected. Adding MLflow's span processors "
                "to the existing provider for unified tracing."
            )
            try:
                existing_processors = set(existing_provider._active_span_processor._span_processors)
                for processor in processors:
                    if not any(isinstance(p, type(processor)) for p in existing_processors):
                        existing_provider.add_span_processor(processor)
                return
            except Exception as e:
                _logger.debug(
                    f"An error occurred while adding span processors to the existing provider: {e}."
                    " Overriding the existing provider with MLflow trace provider.",
                    exc_info=True,
                )

    # NB: If otel resource env vars are set explicitly, don't create an empty resource
    # so that they are propagated to otel spans.
    otel_service_name = os.environ.get("OTEL_SERVICE_NAME")
    otel_resource_attributes = os.environ.get("OTEL_RESOURCE_ATTRIBUTES")
    resource = None
    sdk_attributes = {
        "telemetry.sdk.language": "python",
        "telemetry.sdk.name": "mlflow",
        "telemetry.sdk.version": mlflow.__version__,
    }
    if not otel_service_name and not otel_resource_attributes:
        # Setting an empty resource to avoid triggering resource aggregation, which causes
        # an issue in LiteLLM tracing: https://github.com/mlflow/mlflow/issues/16296
        # Add telemetry resource: https://opentelemetry.io/docs/specs/semconv/resource/#telemetry-sdk
        resource = Resource(sdk_attributes)
    else:
        # Update the env var to let otel initialize the resource with MLflow's SDK attributes
        attributes = {**_parse_otel_resource_attributes(otel_resource_attributes), **sdk_attributes}
        os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ",".join([
            f"{k}={v}" for k, v in attributes.items()
        ])

    if MLFLOW_TRACE_USE_ISOLATED_RANDOM_ID_GENERATOR.get():
        tracer_provider = TracerProvider(
            resource=resource,
            sampler=_get_trace_sampler(),
            id_generator=_IsolatedRandomIdGenerator(),
        )
    else:
        tracer_provider = TracerProvider(
            resource=resource,
            sampler=_get_trace_sampler(),
        )
    for processor in processors:
        tracer_provider.add_span_processor(processor)

    provider.set(tracer_provider)


def _parse_otel_resource_attributes(otel_resource_attributes: str | None) -> dict[str, Any]:
    """
    Parse the otel resource attributes from a comma separated key-value pairs string.
    """
    attributes = {}
    if not otel_resource_attributes:
        return attributes

    try:
        for item in otel_resource_attributes.split(","):
            key, value = item.split("=", maxsplit=1)
            attributes[key.strip()] = value.strip()
    except Exception:
        _logger.debug("Failed to parse otel resource attributes, skipping", exc_info=True)
        return {}
    return attributes


def _get_trace_sampler() -> ParentBased | None:
    """
    Get the sampler configuration based on environment variable.

    Returns a ``ParentBased`` sampler that wraps ``_MlflowSampler`` so that:
    - Root spans are sampled using the configured ratio (global or per-function override).
    - Child spans with a sampled parent are always sampled.
    - Child spans with a non-sampled parent are always dropped.

    Returns:
        ParentBased sampler wrapping _MlflowSampler, or None for default sampling.
    """
    sampling_ratio = MLFLOW_TRACE_SAMPLING_RATIO.get()
    if sampling_ratio is not None:
        if not (0.0 <= sampling_ratio <= 1.0):
            _logger.warning(
                f"{MLFLOW_TRACE_SAMPLING_RATIO} must be between 0.0 and 1.0, got {sampling_ratio}. "
                "Ignoring the invalid value and using default sampling (1.0)."
            )
            return None
        return ParentBased(root=_MlflowSampler(sampling_ratio))
    return None


def _resolve_experiment_uc_location() -> UnityCatalog | None:
    from mlflow.tracking._tracking_service.utils import _get_store
    from mlflow.tracking.fluent import _get_experiment_id

    tracking_uri = mlflow.get_tracking_uri()
    if not tracking_uri or not is_databricks_uri(tracking_uri):
        return None

    try:
        experiment_id = _get_experiment_id()
        if not experiment_id:
            return None

        experiment = _get_store().get_experiment(experiment_id)
        if not experiment:
            return None

        return experiment.trace_location
    except Exception:
        _logger.debug(
            "Failed to auto-resolve UC location for active experiment",
            exc_info=True,
        )
        return None


def _get_span_processors(disabled: bool = False) -> list[SpanProcessor]:
    """
    Get the list of span processors based on configuration.

    Args:
        disabled: If True, returns an empty list of processors because tracing is disabled.

    Returns:
        List of span processors to be added to the TracerProvider.
    """
    if disabled:
        return []

    processors = []

    # TODO: Update this logic to pluggable registry where
    #  1. Partners can implement span processor/exporter and destination class.
    #  2. They can register their implementation to the registry via entry points.
    #  3. MLflow will pick the implementation based on given destination id.
    trace_destination = _MLFLOW_TRACE_USER_DESTINATION.get()

    # If no explicit destination is set, check whether the active experiment is
    # linked to a UC table-prefix location. This could happen if the experiment is set
    # via another mechanism (e.g. env vars) rather than `mlflow.set_experiment`.
    if trace_destination is None:
        if trace_destination := _resolve_experiment_uc_location():
            _MLFLOW_TRACE_USER_DESTINATION.set(trace_destination)

    if trace_destination:
        # In PrPr, users must set the destination to a Unity Catalog location to export traces.
        if isinstance(trace_destination, (UCSchemaLocation, UnityCatalog)):
            from mlflow.tracing.export.uc_table import DatabricksUCTableSpanExporter
            from mlflow.tracing.processor.uc_table import DatabricksUCTableSpanProcessor

            exporter = DatabricksUCTableSpanExporter(tracking_uri=mlflow.get_tracking_uri())
            processor = DatabricksUCTableSpanProcessor(span_exporter=exporter)
            processors.append(processor)
            _logger.debug("Added DatabricksUCTableSpanProcessor based on trace destination")

        elif isinstance(trace_destination, MlflowExperimentLocation):
            if is_in_databricks_model_serving_environment():
                _logger.info(
                    "Traces will be sent to the destination set by "
                    "`mlflow.tracing.set_destination` API. To enable saving traces to "
                    "both MLflow experiment and inference table, remove this API call "
                    "from your model and set `MLFLOW_EXPERIMENT_ID` env var instead."
                )
            processor = _get_mlflow_span_processor(tracking_uri=mlflow.get_tracking_uri())
            processors.append(processor)
            _logger.debug("Added MlflowSpanProcessor based on trace destination")

        # Trace destination has highest precedence; definitely ignore defaults
        # Ignore OTLP (unless dual export is set and we should use OTLP)
        if not (should_use_otlp_exporter() and MLFLOW_TRACE_ENABLE_OTLP_DUAL_EXPORT.get()):
            _logger.debug("Rely on trace destination for tracing")
            return processors

    # If no explicit trace destination OR we passed the dual exporter check, honor OTLP
    # configuration.
    if should_use_otlp_exporter():
        from mlflow.tracing.processor.otel import OtelSpanProcessor

        exporter = get_otlp_exporter()
        otel_processor = OtelSpanProcessor(
            span_exporter=exporter,
            # Only export metrics from the Otel processor if dual export is not enabled.
            # Otherwise, both Otel and MLflow processors will export metrics, causing
            # duplication.
            export_metrics=should_export_otlp_metrics()
            and not MLFLOW_TRACE_ENABLE_OTLP_DUAL_EXPORT.get(),
        )
        processors.append(otel_processor)
        _logger.debug("Added OtelSpanProcessor")

        # We have now added the OTLP processor.
        # If dual export is not set, return.
        # If dual export is set AND we have already added a set_destination processor, return.
        # If dual export is set but no set_destination processor added, skip return and go
        # to default processing to catch the default processor, if present.
        if (not MLFLOW_TRACE_ENABLE_OTLP_DUAL_EXPORT.get()) or any(
            not isinstance(p, OtelSpanProcessor) for p in processors
        ):
            return processors
    # Finally, default MLflow-based processors (inference table in serving, else tracking URI).
    if is_in_databricks_model_serving_environment():
        if not is_mlflow_tracing_enabled_in_model_serving():
            # May still return an OTLP processor from above, or nothing at all.
            _logger.debug("Mlflow tracing is not enabled in model serving")
            return processors

        from mlflow.tracing.export.inference_table import InferenceTableSpanExporter
        from mlflow.tracing.processor.inference_table import InferenceTableSpanProcessor

        exporter = InferenceTableSpanExporter()
        processor = InferenceTableSpanProcessor(exporter)
        processors.append(processor)
        _logger.debug("Added InferenceTableSpanProcessor")
    else:
        processor = _get_mlflow_span_processor(tracking_uri=mlflow.get_tracking_uri())
        processors.append(processor)
        _logger.debug("Added MLflow span processors")

    return processors


def _get_mlflow_span_processor(tracking_uri: str):
    """
    Get the MLflow span processor instance that is used by the current tracer provider.
    """
    # Inline imports to avoid circular dependency:
    # provider -> base_mlflow -> fluent -> entities
    from mlflow.tracing.export.mlflow_v3 import MlflowV3SpanExporter
    from mlflow.tracing.processor.mlflow_v3 import MlflowV3SpanProcessor

    exporter = MlflowV3SpanExporter(tracking_uri=tracking_uri)
    return MlflowV3SpanProcessor(
        span_exporter=exporter,
        export_metrics=should_export_otlp_metrics(),
        use_batch_processor=MLFLOW_USE_BATCH_SPAN_PROCESSOR.get() and exporter._is_async_enabled,
    )


@raise_as_trace_exception
def disable():
    """
    Disable tracing.

    .. note::

        This function sets up `OpenTelemetry` to use
        `NoOpTracerProvider <https://github.com/open-telemetry/opentelemetry-python/blob/4febd337b019ea013ccaab74893bd9883eb59000/opentelemetry-api/src/opentelemetry/trace/__init__.py#L222>`_
        and effectively disables all tracing operations.

    Example:

    .. code-block:: python
        :test:

        import mlflow


        @mlflow.trace
        def f():
            return 0


        # Tracing is enabled by default
        f()
        assert len(mlflow.search_traces(flush=True)) == 1

        # Disable tracing
        mlflow.tracing.disable()
        f()
        assert len(mlflow.search_traces(flush=True)) == 1

    """
    if not is_tracing_enabled():
        return

    _initialize_tracer_provider(disabled=True)
    provider.once._done = True


@raise_as_trace_exception
def enable():
    """
    Enable tracing.

    Example:

    .. code-block:: python
        :test:

        import mlflow


        @mlflow.trace
        def f():
            return 0


        # Tracing is enabled by default
        f()
        assert len(mlflow.search_traces(flush=True)) == 1

        # Disable tracing
        mlflow.tracing.disable()
        f()
        assert len(mlflow.search_traces(flush=True)) == 1

        # Re-enable tracing
        mlflow.tracing.enable()
        f()
        assert len(mlflow.search_traces(flush=True)) == 2

    """
    if is_tracing_enabled() and provider.once._done:
        _logger.info("Tracing is already enabled")
        return

    _initialize_tracer_provider()
    provider.once._done = True


def trace_disabled(f: Callable[P, R]) -> Callable[P, R]:
    """
    A decorator that temporarily disables tracing for the duration of the decorated function.

    .. code-block:: python

        @trace_disabled
        def f():
            with mlflow.start_span("my_span") as span:
                span.set_attribute("my_key", "my_value")

            return


        # This function will not generate any trace
        f()

    :meta private:
    """

    @functools.wraps(f)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        # Hide the real provider behind a NoOp for the call and restore the same
        # object after, instead of disable()+enable() which rebuilt (and leaked)
        # a BatchSpanProcessor thread every call (#24209).
        entered = False
        try:
            with _trace_disabled_lock:
                if _trace_disabled_state["depth"] == 0:
                    if not is_tracing_enabled():
                        return f(*args, **kwargs)
                    # Pin `once` on so a lazy get_or_init_tracer() in f() can't
                    # re-initialize over the NoOp.
                    _trace_disabled_state["saved_once"] = provider.once._done
                    _trace_disabled_state["saved_provider"] = provider._swap_raw(
                        trace.NoOpTracerProvider()
                    )
                    provider.once._done = True
                _trace_disabled_state["depth"] += 1
                entered = True
        except MlflowTracingException as e:
            _logger.warning(
                f"An error occurred while disabling tracing: {e} The original "
                "function will still be executed, but the tracing state may not be "
                "as expected. For full traceback, set logging level to debug.",
                exc_info=_logger.isEnabledFor(logging.DEBUG),
            )
            return f(*args, **kwargs)

        try:
            return f(*args, **kwargs)
        finally:
            if entered:
                with _trace_disabled_lock:
                    _trace_disabled_state["depth"] -= 1
                    if _trace_disabled_state["depth"] == 0:
                        provider._swap_raw(_trace_disabled_state["saved_provider"])
                        provider.once._done = _trace_disabled_state["saved_once"]
                        _trace_disabled_state["saved_provider"] = None
                        _trace_disabled_state["saved_once"] = None

    return wrapper


def reset():
    """
    Reset the flags that indicates whether the MLflow tracer provider has been initialized.
    This ensures that the tracer provider is re-initialized when next tracing
    operation is performed.
    """
    # Set NoOp tracer provider to reset the global tracer to the initial state.
    _initialize_tracer_provider(disabled=True)
    # Flip the "once" flag to False so that
    # the next tracing operation will re-initialize the provider.
    provider.reset()

    # Reset the custom destination set by the user
    _MLFLOW_TRACE_USER_DESTINATION.reset()

    # Reset the tracing configuration to defaults
    reset_config()


@raise_as_trace_exception
def is_tracing_enabled() -> bool:
    """
    Check if tracing is enabled based on whether the global tracer
    is instantiated or not.

    Trace is considered as "enabled" if the followings
    1. The default state (before any tracing operation)
    2. The tracer is not either ProxyTracer or NoOpTracer
    """
    if not provider.once._done:
        return True

    tracer = _get_tracer(__name__)
    # Occasionally ProxyTracer instance wraps the actual tracer
    if isinstance(tracer, trace.ProxyTracer):
        tracer = tracer._tracer
    return not isinstance(tracer, trace.NoOpTracer)