star

venv/lib/python3.9/site-packages/mlflow/tracing/__init__.py

@@ -0,0 +1,14 @@
+from mlflow.tracing.display import disable_notebook_display, enable_notebook_display
+from mlflow.tracing.provider import disable, enable, reset, set_destination
+from mlflow.tracing.utils import set_span_chat_messages, set_span_chat_tools
+
+__all__ = [
+    "disable",
+    "enable",
+    "disable_notebook_display",
+    "enable_notebook_display",
+    "set_span_chat_messages",
+    "set_span_chat_tools",
+    "set_destination",
+    "reset",
+]

venv/lib/python3.9/site-packages/mlflow/tracing/artifact_utils.py

@@ -0,0 +1,21 @@
+from mlflow.entities.trace_info import TraceInfo
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INTERNAL_ERROR
+from mlflow.utils.mlflow_tags import MLFLOW_ARTIFACT_LOCATION
+
+TRACE_DATA_FILE_NAME = "traces.json"
+
+
+def get_artifact_uri_for_trace(trace_info: TraceInfo) -> str:
+    """
+    Get the artifact uri for accessing the trace data.
+
+    The artifact root is specified in the trace tags, which is
+    set when logging the trace in the backend.
+    """
+    if MLFLOW_ARTIFACT_LOCATION not in trace_info.tags:
+        raise MlflowException(
+            "Unable to determine trace artifact location.",
+            error_code=INTERNAL_ERROR,
+        )
+    return trace_info.tags[MLFLOW_ARTIFACT_LOCATION]

venv/lib/python3.9/site-packages/mlflow/tracing/assessment.py

@@ -0,0 +1,342 @@
+from typing import Any, Optional, Union
+
+from mlflow.entities.assessment import (
+    Assessment,
+    AssessmentError,
+    AssessmentValueType,
+    Expectation,
+    Feedback,
+    experimental,
+)
+from mlflow.entities.assessment_source import AssessmentSource
+from mlflow.exceptions import MlflowException
+from mlflow.tracking.client import MlflowClient
+
+
+@experimental
+def log_expectation(
+    trace_id: str,
+    name: str,
+    source: Union[str, AssessmentSource],
+    value: AssessmentValueType,
+    metadata: Optional[dict[str, Any]] = None,
+    span_id: Optional[str] = None,
+) -> Assessment:
+    """
+    .. important::
+
+        This API is currently only available for `Databricks Managed MLflow <https://www.databricks.com/product/managed-mlflow>`_.
+
+    Logs an expectation (e.g. ground truth label) to a Trace.
+
+    Args:
+        trace_id: The ID of the trace.
+        name: The name of the expectation assessment e.g., "expected_answer
+        source: The source of the expectation assessment. Must be either an instance of
+                :py:class:`~mlflow.entities.AssessmentSource` or a string that
+                is a valid value in the
+                :py:class:`~mlflow.entities.AssessmentSourceType` enum.
+        value: The value of the expectation. It can be any JSON-serializable value.
+        metadata: Additional metadata for the expectation.
+        span_id: The ID of the span associated with the expectation, if it needs be
+                associated with a specific span in the trace.
+
+    Returns:
+        :py:class:`~mlflow.entities.Assessment`: The created expectation assessment.
+
+    Example:
+
+    The following code annotates a trace with human-provided ground truth.
+
+    .. code-block:: python
+
+        import mlflow
+        from mlflow.entities.assessment import AssessmentSourceType
+
+        mlflow.log_expectation(
+            trace_id="1234",
+            name="expected_answer",
+            value=42,
+            source=AssessmentSourceType.HUMAN,
+        )
+
+    """
+    if value is None:
+        raise MlflowException.invalid_parameter_value("Expectation value cannot be None.")
+
+    return MlflowClient().log_assessment(
+        trace_id=trace_id,
+        name=name,
+        source=_parse_source(source),
+        expectation=Expectation(value) if value is not None else None,
+        metadata=metadata,
+        span_id=span_id,
+    )
+
+
+@experimental
+def update_expectation(
+    trace_id: str,
+    assessment_id: str,
+    name: Optional[str] = None,
+    value: Optional[AssessmentValueType] = None,
+    metadata: Optional[dict[str, Any]] = None,
+) -> Assessment:
+    """
+    .. important::
+
+        This API is currently only available for `Databricks Managed MLflow <https://www.databricks.com/product/managed-mlflow>`_.
+
+    Updates an existing expectation (ground truth) in a Trace.
+
+    Args:
+        trace_id: The ID of the trace.
+        assessment_id: The ID of the expectation assessment to update.
+        name: The updated name of the expectation. Specify only when updating the name.
+        value: The updated value of the expectation. Specify only when updating the value.
+        metadata: Additional metadata for the expectation. Specify only when updating the metadata.
+
+    Returns:
+        :py:class:`~mlflow.entities.Assessment`: The updated feedback assessment.
+
+    Example:
+
+    The following code updates an existing expectation with a new value.
+    To update other fields, provide the corresponding parameters.
+
+    .. code-block:: python
+
+        import mlflow
+        from mlflow.entities.assessment import AssessmentSourceType
+
+        # Create an expectation with value 42.
+        assessment = mlflow.log_expectation(
+            trace_id="1234",
+            name="expected_answer",
+            value=42,
+            source=AssessmentSourceType.HUMAN,
+        )
+
+        # Update the expectation with a new value 43.
+        mlflow.update_expectation(
+            trace_id="1234",
+            assessment_id=assessment.assessment_id,
+            value=43,
+        )
+    """
+    return MlflowClient().update_assessment(
+        assessment_id=assessment_id,
+        trace_id=trace_id,
+        name=name,
+        expectation=Expectation(value) if value is not None else None,
+        metadata=metadata,
+    )
+
+
+@experimental
+def delete_expectation(trace_id: str, assessment_id: str):
+    """
+    .. important::
+
+        This API is currently only available for `Databricks Managed MLflow <https://www.databricks.com/product/managed-mlflow>`_.
+
+    Deletes an expectation associated with a trace.
+
+    Args:
+        trace_id: The ID of the trace.
+        assessment_id: The ID of the expectation assessment to delete.
+    """
+    return MlflowClient().delete_assessment(trace_id=trace_id, assessment_id=assessment_id)
+
+
+@experimental
+def log_feedback(
+    trace_id: str,
+    name: str,
+    source: Union[str, AssessmentSource],
+    value: Optional[AssessmentValueType] = None,
+    error: Optional[AssessmentError] = None,
+    rationale: Optional[str] = None,
+    metadata: Optional[dict[str, Any]] = None,
+    span_id: Optional[str] = None,
+) -> Assessment:
+    """
+    .. important::
+
+        This API is currently only available for `Databricks Managed MLflow <https://www.databricks.com/product/managed-mlflow>`_.
+
+    Logs feedback to a Trace.
+
+    Args:
+        trace_id: The ID of the trace.
+        name: The name of the feedback assessment e.g., "faithfulness"
+        source: The source of the feedback assessment. Must be either an instance of
+                :py:class:`~mlflow.entities.AssessmentSource` or a string that
+                is a valid value in the
+                :py:class:`~mlflow.entities.AssessmentSourceType` enum.
+        value: The value of the feedback.
+        error: An error object representing any issues encountered while computing the
+            feedback, e.g., a timeout error from an LLM judge. Either this or `value`
+            must be provided.
+        rationale: The rationale / justification for the feedback.
+        metadata: Additional metadata for the feedback.
+        span_id: The ID of the span associated with the feedback, if it needs be
+                associated with a specific span in the trace.
+
+    Returns:
+        :py:class:`~mlflow.entities.Assessment`: The created feedback assessment.
+
+    Example:
+
+    The following code annotates a trace with a feedback provided by LLM-as-a-Judge.
+
+    .. code-block:: python
+
+        import mlflow
+        from mlflow.entities.assessment import AssessmentSourceType
+
+        source = AssessmentSource(
+            source_type=Type.LLM_JUDGE,
+            source_id="faithfulness-judge",
+        )
+
+        mlflow.log_feedback(
+            trace_id="1234",
+            name="faithfulness",
+            source=source,
+            value=0.9,
+            rationale="The model is faithful to the input.",
+            metadata={"model": "gpt-4o-mini"},
+        )
+
+    You can also log an error information during the feedback generation process. To do so,
+    provide an instance of :py:class:`~mlflow.entities.AssessmentError` to the `error`
+    parameter, and leave the `value` parameter as `None`.
+
+    .. code-block:: python
+
+        import mlflow
+        from mlflow.entities.assessment import AssessmentError
+
+        source = AssessmentSource(
+            source_type=Type.LLM_JUDGE,
+            source_id="faithfulness-judge",
+        )
+
+        error = AssessmentError(
+            error_code="RATE_LIMIT_EXCEEDED",
+            error_message="Rate limit for the judge exceeded.",
+        )
+
+        mlflow.log_feedback(
+            trace_id="1234",
+            name="faithfulness",
+            source=source,
+            error=error,
+        )
+
+    """
+    if value is None and error is None:
+        raise MlflowException.invalid_parameter_value("Either `value` or `error` must be provided.")
+    return MlflowClient().log_assessment(
+        trace_id=trace_id,
+        name=name,
+        source=_parse_source(source),
+        feedback=Feedback(value, error),
+        rationale=rationale,
+        metadata=metadata,
+        span_id=span_id,
+    )
+
+
+@experimental
+def update_feedback(
+    trace_id: str,
+    assessment_id: str,
+    name: Optional[str] = None,
+    value: Optional[AssessmentValueType] = None,
+    rationale: Optional[str] = None,
+    metadata: Optional[dict[str, Any]] = None,
+) -> Assessment:
+    """
+    .. important::
+
+        This API is currently only available for `Databricks Managed MLflow <https://www.databricks.com/product/managed-mlflow>`_.
+
+    Updates an existing feedback in a Trace.
+
+
+    Args:
+        trace_id: The ID of the trace.
+        assessment_id: The ID of the feedback assessment to update.
+        name: The updated name of the feedback. Specify only when updating the name.
+        value: The updated value of the feedback. Specify only when updating the value.
+        rationale: The updated rationale of the feedback. Specify only when updating the rationale.
+        metadata: Additional metadata for the feedback. Specify only when updating the metadata.
+
+    Returns:
+        :py:class:`~mlflow.entities.Assessment`: The updated feedback assessment.
+
+    Example:
+
+    The following code updates an existing feedback with a new value.
+    To update other fields, provide the corresponding parameters.
+
+    .. code-block:: python
+
+        import mlflow
+        from mlflow.entities.assessment import AssessmentSourceType
+
+        # Create a feedback with value 0.9.
+        assessment = mlflow.log_feedback(
+            trace_id="1234",
+            name="faithfulness",
+            value=0.9,
+            source=AssessmentSourceType.LLM_JUDGE,
+        )
+
+        # Update the feedback with a new value 0.95.
+        mlflow.update_feedback(
+            trace_id="1234",
+            assessment_id=assessment.assessment_id,
+            value=0.95,
+        )
+    """
+    return MlflowClient().update_assessment(
+        trace_id=trace_id,
+        assessment_id=assessment_id,
+        name=name,
+        feedback=Feedback(value) if value is not None else None,
+        rationale=rationale,
+        metadata=metadata,
+    )
+
+
+@experimental
+def delete_feedback(trace_id: str, assessment_id: str):
+    """
+    .. important::
+
+        This API is currently only available for `Databricks Managed MLflow <https://www.databricks.com/product/managed-mlflow>`_.
+
+    Deletes feedback associated with a trace.
+
+    Args:
+        trace_id: The ID of the trace.
+        assessment_id: The ID of the feedback assessment to delete.
+    """
+    return MlflowClient().delete_assessment(trace_id=trace_id, assessment_id=assessment_id)
+
+
+def _parse_source(source: Union[str, AssessmentSource]) -> AssessmentSource:
+    if source is None:
+        raise MlflowException.invalid_parameter_value("`source` must be provided.")
+
+    if isinstance(source, str):
+        return AssessmentSource(source_type=source)
+    elif isinstance(source, AssessmentSource):
+        return source
+
+    raise MlflowException.invalid_parameter_value(
+        "Invalid source type. Must be one of str, AssessmentSource, or AssessmentSourceType."
+    )

venv/lib/python3.9/site-packages/mlflow/tracing/constant.py

@@ -0,0 +1,55 @@
+# NB: These keys are placeholders and subject to change
+class TraceMetadataKey:
+    INPUTS = "mlflow.traceInputs"
+    OUTPUTS = "mlflow.traceOutputs"
+    SOURCE_RUN = "mlflow.sourceRun"
+
+
+class TraceTagKey:
+    TRACE_NAME = "mlflow.traceName"
+    EVAL_REQUEST_ID = "eval.requestId"
+
+
+# A set of reserved attribute keys
+class SpanAttributeKey:
+    EXPERIMENT_ID = "mlflow.experimentId"
+    REQUEST_ID = "mlflow.traceRequestId"
+    INPUTS = "mlflow.spanInputs"
+    OUTPUTS = "mlflow.spanOutputs"
+    SPAN_TYPE = "mlflow.spanType"
+    FUNCTION_NAME = "mlflow.spanFunctionName"
+    START_TIME_NS = "mlflow.spanStartTimeNs"
+    # these attributes are for standardized chat messages and tool definitions
+    # in CHAT_MODEL and LLM spans. they are used for rendering the rich chat
+    # display in the trace UI, as well as downstream consumers of trace data
+    # such as evaluation
+    CHAT_MESSAGES = "mlflow.chat.messages"
+    CHAT_TOOLS = "mlflow.chat.tools"
+    # This attribute is used to populate `intermediate_outputs` property of a trace data
+    # representing intermediate outputs of the trace. This attribute is not empty only on
+    # the root span of a trace created by the `mlflow.log_trace` API. The `intermediate_outputs`
+    # property of the normal trace is generated by the outputs of non-root spans.
+    INTERMEDIATE_OUTPUTS = "mlflow.trace.intermediate_outputs"
+
+
+# All storage backends are guaranteed to support request_metadata key/value up to 250 characters
+MAX_CHARS_IN_TRACE_INFO_METADATA = 250
+# All storage backends are guaranteed to support tag keys up to 250 characters,
+# values up to 4096 characters
+MAX_CHARS_IN_TRACE_INFO_TAGS_KEY = 250
+MAX_CHARS_IN_TRACE_INFO_TAGS_VALUE = 4096
+TRUNCATION_SUFFIX = "..."
+
+# Trace request ID must have the prefix "tr-" appended to the OpenTelemetry trace ID
+TRACE_REQUEST_ID_PREFIX = "tr-"
+
+# Schema version of traces and spans.
+TRACE_SCHEMA_VERSION = 2
+
+# Key for the trace schema version in the trace. This key is also used in
+# Databricks model serving to be careful when modifying it.
+TRACE_SCHEMA_VERSION_KEY = "mlflow.trace_schema.version"
+
+
+STREAM_CHUNK_EVENT_NAME_FORMAT = "mlflow.chunk.item.{index}"
+STREAM_CHUNK_EVENT_VALUE_KEY = "mlflow.chunk.value"

venv/lib/python3.9/site-packages/mlflow/tracing/destination.py

@@ -0,0 +1,79 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import mlflow
+from mlflow.exceptions import MlflowException
+from mlflow.utils.annotations import experimental
+
+
+@experimental
+@dataclass
+class TraceDestination:
+    """A configuration object for specifying the destination of trace data."""
+
+    @property
+    def type(self) -> str:
+        """Type of the destination."""
+        raise NotImplementedError
+
+
+@experimental
+@dataclass
+class MlflowExperiment(TraceDestination):
+    """
+    A destination representing an MLflow experiment.
+
+    By setting this destination in the :py:func:`mlflow.tracing.set_destination` function,
+    MLflow will log traces to the specified experiment.
+
+    Attributes:
+        experiment_id: The ID of the experiment to log traces to. If not specified,
+            the current active experiment will be used.
+        tracking_uri: The tracking URI of the MLflow server to log traces to.
+            If not specified, the current tracking URI will be used.
+    """
+
+    experiment_id: Optional[str] = None
+    tracking_uri: Optional[str] = None
+
+    @property
+    def type(self) -> str:
+        return "experiment"
+
+
+@experimental
+@dataclass
+class Databricks(TraceDestination):
+    """
+    A destination representing a Databricks tracing server.
+
+    By setting this destination in the :py:func:`mlflow.tracing.set_destination` function,
+    MLflow will log traces to the specified experiment.
+
+    If neither experiment_id nor experiment_name is specified, an active experiment
+    when traces are created will be used as the destination.
+    If both are specified, they must refer to the same experiment.
+
+    Attributes:
+        experiment_id: The ID of the experiment to log traces to.
+        experiment_name: The name of the experiment to log traces to.
+    """
+
+    experiment_id: Optional[str] = None
+    experiment_name: Optional[str] = None
+
+    def __post_init__(self):
+        if self.experiment_id is not None:
+            self.experiment_id = str(self.experiment_id)
+
+        if self.experiment_name is not None:
+            experiment_id = mlflow.get_experiment_by_name(self.experiment_name).experiment_id
+            if self.experiment_id is not None and self.experiment_id != experiment_id:
+                raise MlflowException.invalid_parameter_value(
+                    "experiment_id and experiment_name must refer to the same experiment"
+                )
+            self.experiment_id = experiment_id
+
+    @property
+    def type(self) -> str:
+        return "databricks"

venv/lib/python3.9/site-packages/mlflow/tracing/display/__init__.py

@@ -0,0 +1,40 @@
+from mlflow.tracing.display.display_handler import (
+    IPythonTraceDisplayHandler,
+    get_notebook_iframe_html,
+    is_using_tracking_server,
+)
+
+__all__ = [
+    "IPythonTraceDisplayHandler",
+    "get_display_handler",
+    "is_using_tracking_server",
+    "get_notebook_iframe_html",
+]
+
+
+def get_display_handler() -> IPythonTraceDisplayHandler:
+    return IPythonTraceDisplayHandler.get_instance()
+
+
+def disable_notebook_display():
+    """
+    Disables displaying the MLflow Trace UI in notebook output cells.
+    Call :py:func:`mlflow.tracing.enable_notebook_display()` to re-enable display.
+    """
+    IPythonTraceDisplayHandler.disable()
+
+
+def enable_notebook_display():
+    """
+    Enables the MLflow Trace UI in notebook output cells. The display is on
+    by default, and the Trace UI will show up when any of the following operations
+    are executed:
+
+    * On trace completion (i.e. whenever a trace is exported)
+    * When calling the :py:func:`mlflow.search_traces` fluent API
+    * When calling the :py:meth:`mlflow.client.MlflowClient.get_trace`
+      or :py:meth:`mlflow.client.MlflowClient.search_traces` client APIs
+
+    To disable, please call :py:func:`mlflow.tracing.disable_notebook_display()`.
+    """
+    IPythonTraceDisplayHandler.enable()

venv/lib/python3.9/site-packages/mlflow/tracing/display/display_handler.py

@@ -0,0 +1,196 @@
+import html
+import json
+import logging
+from typing import TYPE_CHECKING
+from urllib.parse import urlencode, urljoin
+
+import mlflow
+from mlflow.environment_variables import MLFLOW_MAX_TRACES_TO_DISPLAY_IN_NOTEBOOK
+from mlflow.utils.databricks_utils import is_in_databricks_runtime
+from mlflow.utils.uri import is_http_uri
+
+_logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from mlflow.entities import Trace
+
+
+TRACE_RENDERER_ASSET_PATH = "/static-files/lib/notebook-trace-renderer/index.html"
+
+IFRAME_HTML = """
+<div>
+  <style scoped>
+  button {{
+    border: none;
+    border-radius: 4px;
+    background-color: rgb(34, 114, 180);
+    font-family: -apple-system, "system-ui", "Segoe UI", Roboto, "Helvetica Neue", Arial;
+    font-size: 13px;
+    color: white;
+    margin-top: 8px;
+    margin-bottom: 8px;
+    padding: 8px 16px;
+    cursor: pointer;
+  }}
+  button:hover {{
+    background-color: rgb(66, 153, 224);
+  }}
+  </style>
+  <button
+    onclick="
+        const display = this.nextElementSibling.style.display;
+        const isCollapsed = display === 'none';
+        this.nextElementSibling.style.display = isCollapsed ? null : 'none';
+
+        const verb = isCollapsed ? 'Collapse' : 'Expand';
+        this.innerText = `${{verb}} MLflow Trace`;
+    "
+  >Collapse MLflow Trace</button>
+  <iframe
+    id="trace-renderer"
+    style="width: 100%; height: 500px; border: none; resize: vertical;"
+    src="{src}"
+  />
+</div>
+"""
+
+
+def get_notebook_iframe_html(traces: list["Trace"]):
+    # fetch assets from tracking server
+    uri = urljoin(mlflow.get_tracking_uri(), TRACE_RENDERER_ASSET_PATH)
+    query_string = _get_query_string_for_traces(traces)
+
+    # include mlflow version to invalidate browser cache when mlflow updates
+    src = html.escape(f"{uri}?{query_string}&version={mlflow.__version__}")
+    return IFRAME_HTML.format(src=src)
+
+
+def _serialize_trace_list(traces: list["Trace"]):
+    return json.dumps(
+        # we can't just call trace.to_json() because this
+        # will cause the trace to be serialized twice (once
+        # by to_json and once by json.dumps)
+        [json.loads(trace._serialize_for_mimebundle()) for trace in traces],
+        ensure_ascii=False,
+    )
+
+
+def _get_query_string_for_traces(traces: list["Trace"]):
+    query_params = []
+
+    for trace in traces:
+        query_params.append(("trace_id", trace.info.request_id))
+        query_params.append(("experiment_id", trace.info.experiment_id))
+
+    return urlencode(query_params)
+
+
+def _is_jupyter():
+    try:
+        from IPython import get_ipython
+
+        return get_ipython() is not None
+    except ImportError:
+        return False
+
+
+def is_using_tracking_server():
+    return is_http_uri(mlflow.get_tracking_uri())
+
+
+def is_trace_ui_available():
+    # the notebook display feature only works in
+    # Databricks notebooks, or in Jupyter notebooks
+    # with a tracking server
+    return _is_jupyter() and (is_in_databricks_runtime() or is_using_tracking_server())
+
+
+class IPythonTraceDisplayHandler:
+    _instance = None
+    disabled = False
+
+    @classmethod
+    def get_instance(cls):
+        if cls._instance is None:
+            cls._instance = IPythonTraceDisplayHandler()
+        return cls._instance
+
+    @classmethod
+    def disable(cls):
+        cls.disabled = True
+
+    @classmethod
+    def enable(cls):
+        cls.disabled = False
+        if cls._instance is None:
+            cls._instance = IPythonTraceDisplayHandler()
+
+    def __init__(self):
+        self.traces_to_display = {}
+        if not _is_jupyter():
+            return
+
+        try:
+            from IPython import get_ipython
+
+            # Register a post-run cell display hook to display traces
+            # after the cell has executed. We don't validate that the
+            # user is using a tracking server at this step, because
+            # the user might set it later using mlflow.set_tracking_uri()
+            get_ipython().events.register("post_run_cell", self._display_traces_post_run)
+        except Exception:
+            # swallow exceptions. this function is called as
+            # a side-effect in a few other functions (e.g. log_trace,
+            # get_traces, search_traces), and we don't want to block
+            # the core functionality if the display fails.
+            _logger.debug("Failed to register post-run cell display hook", exc_info=True)
+
+    def _display_traces_post_run(self, result):
+        if self.disabled or not is_trace_ui_available():
+            self.traces_to_display = {}
+            return
+
+        try:
+            from IPython.display import display
+
+            MAX_TRACES_TO_DISPLAY = MLFLOW_MAX_TRACES_TO_DISPLAY_IN_NOTEBOOK.get()
+            traces_to_display = list(self.traces_to_display.values())[:MAX_TRACES_TO_DISPLAY]
+            if len(traces_to_display) == 0:
+                self.traces_to_display = {}
+                return
+
+            display(self.get_mimebundle(traces_to_display), raw=True)
+
+            # reset state
+            self.traces_to_display = {}
+        except Exception:
+            # swallow exceptions. this function is called as
+            # a side-effect in a few other functions (e.g. log_trace,
+            # get_traces, search_traces), and we don't want to block
+            # the core functionality if the display fails.
+            _logger.error("Failed to display traces", exc_info=True)
+            self.traces_to_display = {}
+
+    def get_mimebundle(self, traces: list["Trace"]):
+        if len(traces) == 1:
+            return traces[0]._repr_mimebundle_()
+        else:
+            bundle = {"text/plain": repr(traces)}
+            if is_in_databricks_runtime():
+                bundle["application/databricks.mlflow.trace"] = _serialize_trace_list(traces)
+            else:
+                bundle["text/html"] = get_notebook_iframe_html(traces)
+            return bundle
+
+    def display_traces(self, traces: list["Trace"]):
+        if self.disabled or not is_trace_ui_available():
+            return
+
+        try:
+            if len(traces) == 0:
+                return
+
+            traces_dict = {trace.info.request_id: trace for trace in traces}
+            self.traces_to_display.update(traces_dict)
+        except Exception:
+            _logger.debug("Failed to update traces", exc_info=True)

venv/lib/python3.9/site-packages/mlflow/tracing/export/async_export_queue.py

@@ -0,0 +1,176 @@
+import atexit
+import logging
+import threading
+from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
+from dataclasses import dataclass
+from queue import Empty, Queue
+from queue import Full as queue_Full
+from typing import Callable, Sequence
+
+from mlflow.environment_variables import (
+    MLFLOW_ASYNC_TRACE_LOGGING_MAX_QUEUE_SIZE,
+    MLFLOW_ASYNC_TRACE_LOGGING_MAX_WORKERS,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Task:
+    """A dataclass to represent a simple task."""
+
+    handler: Callable
+    args: Sequence
+    error_msg: str = ""
+
+    def handle(self) -> None:
+        """Handle the task execution. This method must not raise any exception."""
+        try:
+            self.handler(*self.args)
+        except Exception as e:
+            _logger.warning(
+                f"{self.error_msg} Error: {e}.",
+                exc_info=_logger.isEnabledFor(logging.DEBUG),
+            )
+
+
+class AsyncTraceExportQueue:
+    """A queue-based asynchronous tracing export processor."""
+
+    def __init__(self):
+        self._queue: Queue[Task] = Queue(maxsize=MLFLOW_ASYNC_TRACE_LOGGING_MAX_QUEUE_SIZE.get())
+        self._lock = threading.RLock()
+        self._max_workers = MLFLOW_ASYNC_TRACE_LOGGING_MAX_WORKERS.get()
+
+        # Thread event that indicates the queue should stop processing tasks
+        self._stop_event = threading.Event()
+        self._is_active = False
+        self._atexit_callback_registered = False
+
+        self._active_tasks = set()
+
+    def put(self, task: Task):
+        """Put a new task to the queue for processing."""
+        if not self.is_active():
+            self.activate()
+
+        # If stop event is set, wait for the queue to be drained before putting the task
+        if self._stop_event.is_set():
+            self._stop_event.wait()
+
+        try:
+            # Do not block if the queue is full, it will block the main application
+            self._queue.put(task, block=False)
+        except queue_Full:
+            _logger.warning(
+                "Trace export queue is full, trace will be discarded. "
+                "Consider increasing the queue size or number of workers."
+            )
+
+    def _consumer_loop(self) -> None:
+        while not self._stop_event.is_set():
+            self._dispatch_task()
+
+        # Drain remaining tasks when stopping
+        while not self._queue.empty():
+            self._dispatch_task()
+
+    def _dispatch_task(self) -> None:
+        """Dispatch a task from the queue to the worker thread pool."""
+        # NB: Monitor number of active tasks being processed by the workers. If the all
+        #   workers are busy, wait for one of them to finish before draining a new task
+        #   from the queue. This is because ThreadPoolExecutor does not have a built-in
+        #   mechanism to limit the number of pending tasks in the internal queue.
+        #   This ruins the purpose of having a size bound for self._queue, because the
+        #   TPE's internal queue can grow indefinitely and potentially run out of memory.
+        #   Therefore, we should only dispatch a new task when there is a worker available,
+        #   and pend the new tasks in the self._queue which has a size bound.
+        if len(self._active_tasks) >= self._max_workers:
+            _, self._active_tasks = wait(self._active_tasks, return_when=FIRST_COMPLETED)
+
+        try:
+            task = self._queue.get(timeout=1)
+        except Empty:
+            return
+
+        def _handle(task):
+            task.handle()
+            self._queue.task_done()
+
+        try:
+            future = self._worker_threadpool.submit(_handle, task)
+            self._active_tasks.add(future)
+        except Exception as e:
+            # In case it fails to submit the task to the worker thread pool
+            # such as interpreter shutdown, handle the task in this thread
+            _logger.debug(
+                f"Failed to submit task to worker thread pool. Error: {e}",
+                exc_info=True,
+            )
+            _handle(task)
+
+    def activate(self) -> None:
+        """Activate the async queue to accept and handle incoming tasks."""
+        with self._lock:
+            if self._is_active:
+                return
+
+            self._set_up_threads()
+
+            # Callback to ensure remaining tasks are processed before program exit
+            if not self._atexit_callback_registered:
+                atexit.register(self._at_exit_callback)
+                self._atexit_callback_registered = True
+
+            self._is_active = True
+
+    def is_active(self) -> bool:
+        return self._is_active
+
+    def _set_up_threads(self) -> None:
+        """Set up the consumer and worker threads."""
+        with self._lock:
+            self._worker_threadpool = ThreadPoolExecutor(
+                max_workers=self._max_workers,
+                thread_name_prefix="MlflowTraceLoggingWorker",
+            )
+            self._consumer_thread = threading.Thread(
+                target=self._consumer_loop,
+                name="MLflowTraceLoggingConsumer",
+                daemon=True,
+            )
+            self._consumer_thread.start()
+
+    def _at_exit_callback(self) -> None:
+        """Callback function executed when the program is exiting."""
+        try:
+            _logger.info(
+                "Flushing the async trace logging queue before program exit. "
+                "This may take a while..."
+            )
+            self.flush(terminate=True)
+        except Exception as e:
+            _logger.error(f"Error while finishing trace export requests: {e}")
+
+    def flush(self, terminate=False) -> None:
+        """
+        Flush the async logging queue.
+
+        Args:
+            terminate: If True, shut down the logging threads after flushing.
+        """
+        if not self.is_active():
+            return
+
+        self._stop_event.set()
+        self._consumer_thread.join()
+
+        # Wait for all tasks to be processed
+        self._queue.join()
+
+        self._worker_threadpool.shutdown(wait=True)
+        self._is_active = False
+        # Restart threads to listen to incoming requests after flushing, if not terminating
+        if not terminate:
+            self._stop_event.clear()
+            self.activate()

venv/lib/python3.9/site-packages/mlflow/tracing/export/databricks.py

@@ -0,0 +1,94 @@
+import logging
+from typing import Sequence
+
+from google.protobuf.json_format import MessageToDict
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter
+
+from mlflow.entities.trace import Trace
+from mlflow.environment_variables import (
+    MLFLOW_ASYNC_TRACE_LOGGING_RETRY_TIMEOUT,
+    MLFLOW_ENABLE_ASYNC_TRACE_LOGGING,
+)
+from mlflow.protos.databricks_trace_server_pb2 import CreateTrace, DatabricksTracingServerService
+from mlflow.tracing.export.async_export_queue import AsyncTraceExportQueue, Task
+from mlflow.tracing.fluent import _set_last_active_trace_id
+from mlflow.tracing.trace_manager import InMemoryTraceManager
+from mlflow.utils.databricks_utils import get_databricks_host_creds
+from mlflow.utils.rest_utils import (
+    _REST_API_PATH_PREFIX,
+    extract_api_info_for_service,
+    http_request,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+_METHOD_TO_INFO = extract_api_info_for_service(
+    DatabricksTracingServerService, _REST_API_PATH_PREFIX
+)
+
+
+class DatabricksSpanExporter(SpanExporter):
+    """
+    An exporter implementation that logs the traces to Databricks Tracing Server.
+    """
+
+    def __init__(self):
+        self._is_async = MLFLOW_ENABLE_ASYNC_TRACE_LOGGING.get()
+        if self._is_async:
+            _logger.info("MLflow is configured to log traces asynchronously.")
+            self._async_queue = AsyncTraceExportQueue()
+
+    def export(self, spans: Sequence[ReadableSpan]):
+        """
+        Export the spans to the destination.
+
+        Args:
+            spans: A sequence of OpenTelemetry ReadableSpan objects passed from
+                a span processor. Only root spans for each trace should be exported.
+        """
+        for span in spans:
+            if span._parent is not None:
+                _logger.debug("Received a non-root span. Skipping export.")
+                continue
+
+            trace = InMemoryTraceManager.get_instance().pop_trace(span.context.trace_id)
+            if trace is None:
+                _logger.debug(f"Trace for span {span} not found. Skipping export.")
+                continue
+
+            _set_last_active_trace_id(trace.info.request_id)
+
+            if self._is_async:
+                self._async_queue.put(
+                    task=Task(
+                        handler=self._log_trace,
+                        args=(trace,),
+                        error_msg="Failed to log trace to the trace server.",
+                    )
+                )
+            else:
+                self._log_trace(trace)
+
+    def _log_trace(self, trace: Trace):
+        """Create a new Trace record in the Databricks Tracing Server."""
+        request_body = MessageToDict(trace.to_proto(), preserving_proto_field_name=True)
+        endpoint, method = _METHOD_TO_INFO[CreateTrace]
+
+        # NB: Using Databricks SDK's built-in retry logic, which simply retries until the timeout
+        #    is reached, with linearly increasing backoff. Since it doesn't expose additional
+        #    configuration options, we might want to implement our own retry logic in the future.
+        # NB: If async logging is disabled, we don't retry to avoid blocking the application.
+        timeout = MLFLOW_ASYNC_TRACE_LOGGING_RETRY_TIMEOUT.get() if self._is_async else 0
+
+        # Use context manager to ensure the request is closed properly
+        with http_request(
+            host_creds=get_databricks_host_creds(),
+            endpoint=endpoint,
+            method=method,
+            json=request_body,
+            retry_timeout_seconds=timeout,
+        ) as res:
+            if res.status_code != 200:
+                _logger.warning(f"Failed to log trace to the trace server. Response: {res.text}")

venv/lib/python3.9/site-packages/mlflow/tracing/export/databricks_agent_legacy.py

@@ -0,0 +1,53 @@
+import logging
+from typing import Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter
+
+from mlflow.deployments import get_deploy_client
+from mlflow.tracing.destination import TraceDestination
+from mlflow.tracing.trace_manager import InMemoryTraceManager
+
+_logger = logging.getLogger(__name__)
+
+
+class DatabricksAgentSpanExporter(SpanExporter):
+    """
+    An exporter implementation that logs the traces to Databricks Agent Monitoring.
+
+    Args:
+        trace_destination: The destination of the traces.
+
+    TODO: This class should be deprecated in favor of DatabricksSpanExporter, once
+        the Databricks Agent Monitoring is fully migrated to the new tracing server.
+    """
+
+    def __init__(self, trace_destination: TraceDestination):
+        self._databricks_monitor_id = trace_destination.databricks_monitor_id
+        self._trace_manager = InMemoryTraceManager.get_instance()
+        self._deploy_client = get_deploy_client("databricks")
+
+    def export(self, spans: Sequence[ReadableSpan]):
+        """
+        Export the spans to the destination.
+
+        Args:
+            spans: A sequence of OpenTelemetry ReadableSpan objects passed from
+                a span processor. Only root spans for each trace should be exported.
+        """
+        for span in spans:
+            if span._parent is not None:
+                _logger.debug("Received a non-root span. Skipping export.")
+                continue
+
+            trace = self._trace_manager.pop_trace(span.context.trace_id)
+            if trace is None:
+                _logger.debug(f"Trace for span {span} not found. Skipping export.")
+                continue
+
+            # Traces are exported via a serving endpoint that accepts trace JSON as
+            # an input payload, and then will be written to the Inference Table.
+            self._deploy_client.predict(
+                endpoint=self._databricks_monitor_id,
+                inputs={"inputs": [trace.to_json()]},
+            )

venv/lib/python3.9/site-packages/mlflow/tracing/export/inference_table.py

@@ -0,0 +1,74 @@
+import logging
+from typing import Any, Optional, Sequence
+
+from cachetools import TTLCache
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter
+
+from mlflow.environment_variables import (
+    MLFLOW_TRACE_BUFFER_MAX_SIZE,
+    MLFLOW_TRACE_BUFFER_TTL_SECONDS,
+)
+from mlflow.tracing.fluent import _set_last_active_trace_id
+from mlflow.tracing.trace_manager import InMemoryTraceManager
+
+_logger = logging.getLogger(__name__)
+
+
+def pop_trace(request_id: str) -> Optional[dict[str, Any]]:
+    """
+    Pop the completed trace data from the buffer. This method is used in
+    the Databricks model serving so please be careful when modifying it.
+    """
+    return _TRACE_BUFFER.pop(request_id, None)
+
+
+# For Inference Table, we use special TTLCache to store the finished traces
+# so that they can be retrieved by Databricks model serving. The values
+# in the buffer are not Trace dataclass, but rather a dictionary with the schema
+# that is used within Databricks model serving.
+def _initialize_trace_buffer():  # Define as a function for testing purposes
+    return TTLCache(
+        maxsize=MLFLOW_TRACE_BUFFER_MAX_SIZE.get(),
+        ttl=MLFLOW_TRACE_BUFFER_TTL_SECONDS.get(),
+    )
+
+
+_TRACE_BUFFER = _initialize_trace_buffer()
+
+
+class InferenceTableSpanExporter(SpanExporter):
+    """
+    An exporter implementation that logs the traces to Inference Table.
+
+    Currently the Inference Table does not use collector to receive the traces,
+    but rather actively fetches the trace during the prediction process. In the
+    future, we may consider using collector-based approach and this exporter should
+    send the traces instead of storing them in the buffer.
+    """
+
+    def __init__(self):
+        self._trace_manager = InMemoryTraceManager.get_instance()
+
+    def export(self, spans: Sequence[ReadableSpan]):
+        """
+        Export the spans to Inference Table via the TTLCache buffer.
+
+        Args:
+            spans: A sequence of OpenTelemetry ReadableSpan objects passed from
+                a span processor. Only root spans for each trace should be exported.
+        """
+        for span in spans:
+            if span._parent is not None:
+                _logger.debug("Received a non-root span. Skipping export.")
+                continue
+
+            trace = self._trace_manager.pop_trace(span.context.trace_id)
+            if trace is None:
+                _logger.debug(f"Trace for span {span} not found. Skipping export.")
+                continue
+
+            _set_last_active_trace_id(trace.info.request_id)
+
+            # Add the trace to the in-memory buffer so it can be retrieved by upstream
+            _TRACE_BUFFER[trace.info.request_id] = trace.to_dict()

venv/lib/python3.9/site-packages/mlflow/tracing/export/mlflow.py

@@ -0,0 +1,100 @@
+import logging
+from typing import Optional, Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter
+
+from mlflow.entities.trace import Trace
+from mlflow.environment_variables import MLFLOW_ENABLE_ASYNC_LOGGING
+from mlflow.tracing.constant import TraceTagKey
+from mlflow.tracing.display import get_display_handler
+from mlflow.tracing.display.display_handler import IPythonTraceDisplayHandler
+from mlflow.tracing.export.async_export_queue import AsyncTraceExportQueue, Task
+from mlflow.tracing.fluent import _EVAL_REQUEST_ID_TO_TRACE_ID, _set_last_active_trace_id
+from mlflow.tracing.trace_manager import InMemoryTraceManager
+from mlflow.tracing.utils import maybe_get_request_id
+from mlflow.tracking.client import MlflowClient
+
+_logger = logging.getLogger(__name__)
+
+
+class MlflowSpanExporter(SpanExporter):
+    """
+    An exporter implementation that logs the traces to MLflow.
+
+    MLflow backend (will) only support logging the complete trace, not incremental updates
+    for spans, so this exporter is designed to aggregate the spans into traces in memory.
+    Therefore, this only works within a single process application and not intended to work
+    in a distributed environment. For the same reason, this exporter should only be used with
+    SimpleSpanProcessor.
+
+    If we want to support distributed tracing, we should first implement an incremental trace
+    logging in MLflow backend, then we can get rid of the in-memory trace aggregation.
+
+    :meta private:
+    """
+
+    def __init__(
+        self,
+        client: Optional[MlflowClient] = None,
+        display_handler: Optional[IPythonTraceDisplayHandler] = None,
+    ):
+        self._client = client or MlflowClient()
+        self._display_handler = display_handler or get_display_handler()
+        self._trace_manager = InMemoryTraceManager.get_instance()
+        self._async_queue = AsyncTraceExportQueue()
+
+    def export(self, spans: Sequence[ReadableSpan]):
+        """
+        Export the spans to MLflow backend.
+
+        Args:
+            spans: A sequence of OpenTelemetry ReadableSpan objects passed from
+                a span processor. Only root spans for each trace should be exported.
+        """
+        for span in spans:
+            if span._parent is not None:
+                _logger.debug("Received a non-root span. Skipping export.")
+                continue
+
+            trace = self._trace_manager.pop_trace(span.context.trace_id)
+            if trace is None:
+                _logger.debug(f"TraceInfo for span {span} not found. Skipping export.")
+                continue
+
+            _set_last_active_trace_id(trace.info.request_id)
+
+            # 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.request_id
+
+            if not maybe_get_request_id(is_evaluate=True):
+                # Display the trace in the UI if the trace is not generated from within
+                # an MLflow model evaluation context
+                self._display_handler.display_traces([trace])
+
+            self._log_trace(trace)
+
+    def _log_trace(self, trace: Trace):
+        """Log the trace to MLflow backend."""
+        upload_trace_data_task = Task(
+            handler=self._client._upload_trace_data,
+            args=(trace.info, trace.data),
+            error_msg="Failed to log trace to MLflow backend.",
+        )
+
+        upload_ended_trace_info_task = Task(
+            handler=self._client._upload_ended_trace_info,
+            args=(trace.info,),
+            error_msg="Failed to log trace to MLflow backend.",
+        )
+
+        # TODO: Use MLFLOW_ENABLE_ASYNC_TRACE_LOGGING instead and default to async
+        # logging once the async logging implementation becomes stable.
+        if MLFLOW_ENABLE_ASYNC_LOGGING.get():
+            self._async_queue.put(upload_trace_data_task)
+            self._async_queue.put(upload_ended_trace_info_task)
+        else:
+            upload_trace_data_task.handle()
+            upload_ended_trace_info_task.handle()

venv/lib/python3.9/site-packages/mlflow/tracing/processor/databricks.py

@@ -0,0 +1,97 @@
+import json
+import logging
+from typing import Optional
+
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan
+from opentelemetry.sdk.trace import Span as OTelSpan
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter
+
+from mlflow.entities.trace_info import TraceInfo
+from mlflow.entities.trace_status import TraceStatus
+from mlflow.tracing.constant import TRACE_SCHEMA_VERSION, TRACE_SCHEMA_VERSION_KEY, SpanAttributeKey
+from mlflow.tracing.trace_manager import InMemoryTraceManager
+from mlflow.tracing.utils import (
+    deduplicate_span_names_in_place,
+    get_otel_attribute,
+    maybe_get_dependencies_schemas,
+)
+from mlflow.tracking.fluent import _get_experiment_id
+
+_logger = logging.getLogger(__name__)
+
+
+class DatabricksSpanProcessor(SimpleSpanProcessor):
+    """
+    Defines custom hooks to be executed when a span is started or ended (before exporting).
+
+    This process implements simple responsibilities to generate MLflow-style trace
+    object from OpenTelemetry spans and store them in memory.
+    """
+
+    def __init__(
+        self,
+        span_exporter: SpanExporter,
+        experiment_id: Optional[str] = None,
+    ):
+        self.span_exporter = span_exporter
+        self._trace_manager = InMemoryTraceManager.get_instance()
+        self._experiment_id = experiment_id
+
+    def on_start(self, span: OTelSpan, parent_context: Optional[Context] = None):
+        """
+        Handle the start of a span. This method is called when an OpenTelemetry span is started.
+
+        Args:
+            span: An OpenTelemetry Span object that is started.
+            parent_context: The context of the span. Note that this is only passed when the context
+                object is explicitly specified to OpenTelemetry start_span call. If the parent
+                span is obtained from the global context, it won't be passed here so we should not
+                rely on it.
+        """
+
+        request_id = self._create_or_get_request_id(span)
+        span.set_attribute(SpanAttributeKey.REQUEST_ID, json.dumps(request_id))
+
+        tags = {}
+        if dependencies_schema := maybe_get_dependencies_schemas():
+            tags.update(dependencies_schema)
+
+        if span._parent is None:
+            trace_info = TraceInfo(
+                request_id=request_id,
+                experiment_id=self._experiment_id or _get_experiment_id(),
+                timestamp_ms=span.start_time // 1_000_000,  # nanosecond to millisecond
+                execution_time_ms=None,
+                status=TraceStatus.IN_PROGRESS,
+                request_metadata={TRACE_SCHEMA_VERSION_KEY: str(TRACE_SCHEMA_VERSION)},
+                tags=tags,
+            )
+            self._trace_manager.register_trace(span.context.trace_id, trace_info)
+
+    def _create_or_get_request_id(self, span: OTelSpan) -> str:
+        if span._parent is None:
+            return str(span.context.trace_id)  # Use otel-generated trace_id as request_id
+        else:
+            return self._trace_manager.get_request_id_from_trace_id(span.context.trace_id)
+
+    def on_end(self, span: OTelReadableSpan) -> None:
+        """
+        Handle the end of a span. This method is called when an OpenTelemetry span is ended.
+
+        Args:
+            span: An OpenTelemetry ReadableSpan object that is ended.
+        """
+        # Processing the trace only when it is a root span.
+        if span._parent is None:
+            request_id = get_otel_attribute(span, SpanAttributeKey.REQUEST_ID)
+            with self._trace_manager.get_trace(request_id) as trace:
+                if trace is None:
+                    _logger.debug(f"Trace data with request ID {request_id} not found.")
+                    return
+
+                trace.info.execution_time_ms = (span.end_time - span.start_time) // 1_000_000
+                trace.info.status = TraceStatus.from_otel_status(span.status)
+                deduplicate_span_names_in_place(list(trace.span_dict.values()))
+
+        super().on_end(span)

venv/lib/python3.9/site-packages/mlflow/tracing/processor/inference_table.py

@@ -0,0 +1,117 @@
+import json
+import logging
+from typing import Optional
+
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan
+from opentelemetry.sdk.trace import Span as OTelSpan
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter
+
+from mlflow.entities.trace_info import TraceInfo
+from mlflow.entities.trace_status import TraceStatus
+from mlflow.tracing.constant import TRACE_SCHEMA_VERSION, TRACE_SCHEMA_VERSION_KEY, SpanAttributeKey
+from mlflow.tracing.trace_manager import InMemoryTraceManager
+from mlflow.tracing.utils import (
+    deduplicate_span_names_in_place,
+    get_otel_attribute,
+    maybe_get_dependencies_schemas,
+    maybe_get_request_id,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+_HEADER_REQUEST_ID_KEY = "X-Request-Id"
+
+
+# Extracting for testing purposes
+def _get_flask_request():
+    import flask
+
+    if flask.has_request_context():
+        return flask.request
+
+
+class InferenceTableSpanProcessor(SimpleSpanProcessor):
+    """
+    Defines custom hooks to be executed when a span is started or ended (before exporting).
+
+    This processor is used when the tracing destination is Databricks Inference Table.
+    """
+
+    def __init__(self, span_exporter: SpanExporter):
+        self.span_exporter = span_exporter
+        self._trace_manager = InMemoryTraceManager.get_instance()
+
+    def on_start(self, span: OTelSpan, parent_context: Optional[Context] = None):
+        """
+        Handle the start of a span. This method is called when an OpenTelemetry span is started.
+
+        Args:
+            span: An OpenTelemetry Span object that is started.
+            parent_context: The context of the span. Note that this is only passed when the context
+                object is explicitly specified to OpenTelemetry start_span call. If the parent
+                span is obtained from the global context, it won't be passed here so we should not
+                rely on it.
+        """
+        request_id = maybe_get_request_id()
+        if request_id is None:
+            # NB: This is currently used for streaming inference in Databricks Model Serving.
+            # In normal prediction, serving logic pass the request ID using the
+            # `with set_prediction_context` context manager that wraps `model.predict`
+            # call. However, in streaming case, the context manager is not applicable
+            # so we still need to rely on Flask request context (which is set to the
+            # stream response via flask.stream_with_context()
+            if flask_request := _get_flask_request():
+                request_id = flask_request.headers.get(_HEADER_REQUEST_ID_KEY)
+                if not request_id:
+                    _logger.warning(
+                        "Request ID not found in the request headers. Skipping trace processing."
+                    )
+                    return
+            else:
+                _logger.warning(
+                    "Failed to get request ID from the request headers because "
+                    "request context is not available. Skipping trace processing."
+                )
+                return
+
+        span.set_attribute(SpanAttributeKey.REQUEST_ID, json.dumps(request_id))
+        tags = {}
+        if dependencies_schema := maybe_get_dependencies_schemas():
+            tags.update(dependencies_schema)
+
+        if span._parent is None:
+            trace_info = TraceInfo(
+                request_id=request_id,
+                experiment_id=None,
+                timestamp_ms=span.start_time // 1_000_000,  # nanosecond to millisecond
+                execution_time_ms=None,
+                status=TraceStatus.IN_PROGRESS,
+                request_metadata={TRACE_SCHEMA_VERSION_KEY: str(TRACE_SCHEMA_VERSION)},
+                tags=tags,
+            )
+            self._trace_manager.register_trace(span.context.trace_id, trace_info)
+
+    def on_end(self, span: OTelReadableSpan) -> None:
+        """
+        Handle the end of a span. This method is called when an OpenTelemetry span is ended.
+
+        Args:
+            span: An OpenTelemetry ReadableSpan object that is ended.
+        """
+        # Processing the trace only when the root span is found.
+        if span._parent is not None:
+            return
+
+        request_id = get_otel_attribute(span, SpanAttributeKey.REQUEST_ID)
+        with self._trace_manager.get_trace(request_id) as trace:
+            if trace is None:
+                _logger.debug(f"Trace data with request ID {request_id} not found.")
+                return
+
+            trace.info.execution_time_ms = (span.end_time - span.start_time) // 1_000_000
+            trace.info.status = TraceStatus.from_otel_status(span.status)
+            deduplicate_span_names_in_place(list(trace.span_dict.values()))
+
+        super().on_end(span)

venv/lib/python3.9/site-packages/mlflow/tracing/processor/mlflow.py

@@ -0,0 +1,230 @@
+import json
+import logging
+import time
+from typing import Optional
+
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan
+from opentelemetry.sdk.trace import Span as OTelSpan
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter
+
+from mlflow.entities.trace_info import TraceInfo
+from mlflow.entities.trace_status import TraceStatus
+from mlflow.tracing.constant import (
+    MAX_CHARS_IN_TRACE_INFO_METADATA,
+    TRACE_SCHEMA_VERSION,
+    TRACE_SCHEMA_VERSION_KEY,
+    TRUNCATION_SUFFIX,
+    SpanAttributeKey,
+    TraceMetadataKey,
+    TraceTagKey,
+)
+from mlflow.tracing.trace_manager import InMemoryTraceManager, _Trace
+from mlflow.tracing.utils import (
+    deduplicate_span_names_in_place,
+    get_otel_attribute,
+    maybe_get_dependencies_schemas,
+    maybe_get_request_id,
+)
+from mlflow.tracking.client import MlflowClient
+from mlflow.tracking.context.databricks_repo_context import DatabricksRepoRunContext
+from mlflow.tracking.context.git_context import GitRunContext
+from mlflow.tracking.context.registry import resolve_tags
+from mlflow.tracking.default_experiment import DEFAULT_EXPERIMENT_ID
+from mlflow.tracking.fluent import _get_experiment_id
+from mlflow.utils.mlflow_tags import TRACE_RESOLVE_TAGS_ALLOWLIST
+
+_logger = logging.getLogger(__name__)
+
+
+class MlflowSpanProcessor(SimpleSpanProcessor):
+    """
+    Defines custom hooks to be executed when a span is started or ended (before exporting).
+
+    This processor is used when the tracing destination is MLflow Tracking Server.
+    """
+
+    def __init__(
+        self,
+        span_exporter: SpanExporter,
+        client: Optional[MlflowClient] = None,
+        experiment_id: Optional[str] = None,
+    ):
+        self.span_exporter = span_exporter
+        self._client = client or MlflowClient()
+        self._experiment_id = experiment_id
+        self._trace_manager = InMemoryTraceManager.get_instance()
+
+        # We issue a warning when a trace is created under the default experiment.
+        # We only want to issue it once, and typically it can be achieved by using
+        # warnings.warn() with filterwarnings setting. However, the de-duplication does
+        # not work in notebooks (https://github.com/ipython/ipython/issues/11207),
+        # so we instead keep track of the warning issuance state manually.
+        self._issued_default_exp_warning = False
+
+    def on_start(self, span: OTelSpan, parent_context: Optional[Context] = None):
+        """
+        Handle the start of a span. This method is called when an OpenTelemetry span is started.
+
+        Args:
+            span: An OpenTelemetry Span object that is started.
+            parent_context: The context of the span. Note that this is only passed when the context
+                object is explicitly specified to OpenTelemetry start_span call. If the parent span
+                is obtained from the global context, it won't be passed here so we should not rely
+                on it.
+        """
+        request_id = self._trace_manager.get_request_id_from_trace_id(span.context.trace_id)
+
+        if not request_id and span.parent is not None:
+            _logger.debug(
+                "Received a non-root span but the request ID is not found."
+                "The trace has likely been halted due to a timeout expiration."
+            )
+            return
+
+        if not request_id:
+            # If the user started trace/span with fixed start time, this attribute is set
+            start_time_ns = get_otel_attribute(span, SpanAttributeKey.START_TIME_NS)
+
+            trace_info = self._start_trace(span, start_time_ns)
+            self._trace_manager.register_trace(span.context.trace_id, trace_info)
+            request_id = trace_info.request_id
+
+            # NB: This is a workaround to exclude the latency of backend StartTrace API call (within
+            #   _create_trace_info()) from the execution time of the span. The API call takes ~1 sec
+            #   and significantly skews the span duration.
+            if not start_time_ns:
+                span._start_time = time.time_ns()
+
+        span.set_attribute(SpanAttributeKey.REQUEST_ID, json.dumps(request_id))
+
+    def _start_trace(self, span: OTelSpan, start_time_ns: Optional[int]) -> TraceInfo:
+        from mlflow.tracking.fluent import _get_latest_active_run
+
+        metadata = {TRACE_SCHEMA_VERSION_KEY: str(TRACE_SCHEMA_VERSION)}
+
+        # If the span is started within an active MLflow run, we should record it as a trace tag
+        # Note `mlflow.active_run()` can only get thread-local active run,
+        # but tracing routine might be applied to model inference worker threads
+        # in the following cases:
+        #  - langchain model `chain.batch` which uses thread pool to spawn workers.
+        #  - MLflow langchain pyfunc model `predict` which calls `api_request_parallel_processor`.
+        # Therefore, we use `_get_global_active_run()` instead to get the active run from
+        # all threads and set it as the tracing source run.
+        if run := _get_latest_active_run():
+            metadata[TraceMetadataKey.SOURCE_RUN] = run.info.run_id
+
+        experiment_id = self._get_experiment_id_for_trace(span)
+        if experiment_id == DEFAULT_EXPERIMENT_ID and not self._issued_default_exp_warning:
+            _logger.warning(
+                "Creating a trace within the default experiment with id "
+                f"'{DEFAULT_EXPERIMENT_ID}'. It is strongly recommended to not use "
+                "the default experiment to log traces due to ambiguous search results and "
+                "probable performance issues over time due to directory table listing performance "
+                "degradation with high volumes of directories within a specific path. "
+                "To avoid performance and disambiguation issues, set the experiment for "
+                "your environment using `mlflow.set_experiment()` API."
+            )
+            self._issued_default_exp_warning = True
+
+        # Avoid running unnecessary context providers to avoid overhead
+        unfiltered_tags = resolve_tags(ignore=[DatabricksRepoRunContext, GitRunContext])
+        tags = {
+            key: value
+            for key, value in unfiltered_tags.items()
+            if key in TRACE_RESOLVE_TAGS_ALLOWLIST
+        }
+
+        # If the trace is created in the context of MLflow model evaluation, we extract the request
+        # ID from the prediction context. Otherwise, we create a new trace info by calling the
+        # backend API.
+        if request_id := maybe_get_request_id(is_evaluate=True):
+            tags.update({TraceTagKey.EVAL_REQUEST_ID: request_id})
+        if dependencies_schema := maybe_get_dependencies_schemas():
+            tags.update(dependencies_schema)
+        tags.update({TraceTagKey.TRACE_NAME: span.name})
+
+        return self._client._start_tracked_trace(
+            experiment_id=experiment_id,
+            # TODO: This timestamp is not accurate because it is not adjusted to exclude the
+            #   latency of the backend API call. We do this adjustment for span start time
+            #   above, but can't do it for trace start time until the backend API supports
+            #   updating the trace start time.
+            timestamp_ms=(start_time_ns or span.start_time) // 1_000_000,  # ns to ms
+            request_metadata=metadata,
+            tags=tags,
+        )
+
+    def on_end(self, span: OTelReadableSpan) -> None:
+        """
+        Handle the end of a span. This method is called when an OpenTelemetry span is ended.
+
+        Args:
+            span: An OpenTelemetry ReadableSpan object that is ended.
+        """
+        # Processing the trace only when the root span is found.
+        if span._parent is not None:
+            return
+
+        request_id = get_otel_attribute(span, SpanAttributeKey.REQUEST_ID)
+        with self._trace_manager.get_trace(request_id) as trace:
+            if trace is None:
+                _logger.debug(f"Trace data with request ID {request_id} not found.")
+                return
+
+            self._update_trace_info(trace, span)
+            deduplicate_span_names_in_place(list(trace.span_dict.values()))
+
+        super().on_end(span)
+
+    def _get_experiment_id_for_trace(self, span: OTelReadableSpan) -> str:
+        """
+        Determine the experiment ID to associate with the trace.
+
+        The experiment ID can be configured in multiple ways, in order of precedence:
+          1. An experiment ID specified via the span creation API i.e. MlflowClient().start_trace()
+          2. An experiment ID specified via the processor constructor
+          3. An experiment ID of an active run.
+          4. The default experiment ID
+        """
+        from mlflow.tracking.fluent import _get_latest_active_run
+
+        if experiment_id := get_otel_attribute(span, SpanAttributeKey.EXPERIMENT_ID):
+            return experiment_id
+
+        if self._experiment_id:
+            return self._experiment_id
+
+        if run := _get_latest_active_run():
+            return run.info.experiment_id
+
+        return _get_experiment_id()
+
+    def _update_trace_info(self, trace: _Trace, root_span: OTelReadableSpan):
+        """Update the trace info with the final values from the root span."""
+        # The trace/span start time needs adjustment to exclude the latency of
+        # the backend API call. We already adjusted the span start time in the
+        # on_start method, so we reflect the same to the trace start time here.
+        trace.info.timestamp_ms = root_span.start_time // 1_000_000  # nanosecond to millisecond
+        trace.info.execution_time_ms = (root_span.end_time - root_span.start_time) // 1_000_000
+        trace.info.status = TraceStatus.from_otel_status(root_span.status)
+        trace.info.request_metadata.update(
+            {
+                TraceMetadataKey.INPUTS: self._truncate_metadata(
+                    root_span.attributes.get(SpanAttributeKey.INPUTS)
+                ),
+                TraceMetadataKey.OUTPUTS: self._truncate_metadata(
+                    root_span.attributes.get(SpanAttributeKey.OUTPUTS)
+                ),
+            }
+        )
+
+    def _truncate_metadata(self, value: Optional[str]) -> str:
+        """Get truncated value of the attribute if it exceeds the maximum length."""
+        if not value:
+            return ""
+
+        if len(value) > MAX_CHARS_IN_TRACE_INFO_METADATA:
+            trunc_length = MAX_CHARS_IN_TRACE_INFO_METADATA - len(TRUNCATION_SUFFIX)
+            value = value[:trunc_length] + TRUNCATION_SUFFIX
+        return value

venv/lib/python3.9/site-packages/mlflow/tracing/processor/otel.py

@@ -0,0 +1,64 @@
+import json
+import uuid
+from typing import Optional
+
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan
+from opentelemetry.sdk.trace import Span as OTelSpan
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter
+
+from mlflow.entities.trace_info import TraceInfo
+from mlflow.entities.trace_status import TraceStatus
+from mlflow.tracing.constant import TRACE_SCHEMA_VERSION, TRACE_SCHEMA_VERSION_KEY, SpanAttributeKey
+from mlflow.tracing.trace_manager import InMemoryTraceManager
+
+
+class OtelSpanProcessor(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):
+        super().__init__(span_exporter)
+        self.span_exporter = span_exporter
+        self._trace_manager = InMemoryTraceManager.get_instance()
+
+    def on_start(self, span: OTelSpan, parent_context: Optional[Context] = None):
+        """
+        Handle the start of a span. This method is called when an OpenTelemetry span is started.
+
+        Args:
+            span: An OpenTelemetry Span object that is started.
+            parent_context: The context of the span. Note that this is only passed when the context
+                object is explicitly specified to OpenTelemetry start_span call. If the parent
+                span is obtained from the global context, it won't be passed here so we should not
+                rely on it.
+        """
+        # Generate a random request ID and trace info just for the sake of consistency
+        # with other tracing destinations. Doing this makes it much easier to handle
+        # multiple tracing destinations.
+        request_id = uuid.uuid4().hex
+        trace_info = TraceInfo(
+            request_id=request_id,
+            experiment_id=None,
+            timestamp_ms=span.start_time // 1_000_000,  # nanosecond to millisecond
+            execution_time_ms=None,
+            status=TraceStatus.IN_PROGRESS,
+            request_metadata={TRACE_SCHEMA_VERSION_KEY: str(TRACE_SCHEMA_VERSION)},
+            tags={},
+        )
+        span.set_attribute(SpanAttributeKey.REQUEST_ID, json.dumps(request_id))
+
+        self._trace_manager.register_trace(span.context.trace_id, trace_info)
+
+        super().on_start(span, parent_context)
+
+    def on_end(self, span: OTelReadableSpan):
+        # Pops the trace entry from the in-memory trace manager to avoid memory leak
+        if span._parent is None:
+            self._trace_manager.pop_trace(span.context.trace_id)
+
+        super().on_end(span)

venv/lib/python3.9/site-packages/mlflow/tracing/provider.py

@@ -0,0 +1,481 @@
+"""
+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
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Optional
+
+from opentelemetry import context as context_api
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+
+from mlflow.exceptions import MlflowException, MlflowTracingException
+from mlflow.tracing.constant import SpanAttributeKey
+from mlflow.tracing.destination import Databricks, MlflowExperiment, TraceDestination
+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_use_otlp_exporter
+from mlflow.utils.annotations import experimental
+from mlflow.utils.databricks_utils import (
+    is_in_databricks_model_serving_environment,
+    is_mlflow_tracing_enabled_in_model_serving,
+)
+
+if TYPE_CHECKING:
+    from mlflow.entities import Span
+
+# Global tracer provider instance. We manage the tracer provider by ourselves instead of using
+# the global tracer provider provided by OpenTelemetry.
+_MLFLOW_TRACER_PROVIDER = None
+
+# Once() object ensures a function is executed only once in a process.
+# Note that it doesn't work as expected in a distributed environment.
+_MLFLOW_TRACER_PROVIDER_INITIALIZED = Once()
+
+# A trace destination specified by the user via the `set_destination` function.
+# This destination, when set, will take precedence over other configurations.
+_MLFLOW_TRACE_USER_DESTINATION = None
+
+_logger = logging.getLogger(__name__)
+
+
+def start_span_in_context(name: str) -> 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.
+
+    Returns:
+        The newly created OpenTelemetry span.
+    """
+    return _get_tracer(__name__).start_span(name)
+
+
+def start_detached_span(
+    name: str,
+    parent: Optional[trace.Span] = None,
+    experiment_id: Optional[str] = None,
+    start_time_ns: Optional[int] = None,
+) -> Optional[tuple[str, trace.Span]]:
+    """
+    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) if parent else None
+    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)
+    return tracer.start_span(name, context=context, attributes=attributes, start_time=start_time_ns)
+
+
+@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)
+    token = context_api.attach(context)
+    return token  # noqa: RET504
+
+
+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.
+    """
+    context_api.detach(token)
+
+
+@experimental
+def set_destination(destination: TraceDestination):
+    """
+    Set a custom span destination 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.
+
+    To reset the destination, call the :py:func:`mlflow.tracing.reset()` function.
+
+    Args:
+        destination: A ``TraceDestination`` object that specifies the destination of the trace data.
+
+    Example:
+
+        .. code-block:: python
+
+            import mlflow
+            from mlflow.tracing.destination import MlflowExperiment
+
+            # Setting the destination to an MLflow experiment with ID "123"
+            mlflow.tracing.set_destination(MlflowExperiment(experiment_id="123"))
+
+            # Reset the destination (to an active experiment as default)
+            mlflow.tracing.reset()
+    """
+    if not isinstance(destination, TraceDestination):
+        raise MlflowException.invalid_parameter_value(
+            f"Invalid destination type: {type(destination)}. "
+            "The destination must be an instance of TraceDestination."
+        )
+
+    # The destination needs to be persisted because the tracer setup can be re-initialized
+    # e.g. when the tracing is disabled and re-enabled, or tracking URI is changed, etc.
+    global _MLFLOW_TRACE_USER_DESTINATION
+    _MLFLOW_TRACE_USER_DESTINATION = destination
+
+    _setup_tracer_provider()
+
+
+def _get_tracer(module_name: str):
+    """
+    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.
+    """
+    # Initiate tracer provider only once in the application lifecycle
+    _MLFLOW_TRACER_PROVIDER_INITIALIZED.do_once(_setup_tracer_provider)
+    return _MLFLOW_TRACER_PROVIDER.get_tracer(module_name)
+
+
+def _get_trace_exporter():
+    """
+    Get the exporter instance that is used by the current tracer provider.
+    """
+    if _MLFLOW_TRACER_PROVIDER:
+        processors = _MLFLOW_TRACER_PROVIDER._active_span_processor._span_processors
+        # There should be only one processor used for MLflow tracing
+        processor = processors[0]
+        return processor.span_exporter
+
+
+def _setup_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 _MLFLOW_TRACER_PROVIDER_INITIALIZED flag accordingly.
+    """
+    global _MLFLOW_TRACER_PROVIDER
+
+    if disabled:
+        _MLFLOW_TRACER_PROVIDER = trace.NoOpTracerProvider()
+        return
+
+    # 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.
+    if _MLFLOW_TRACE_USER_DESTINATION is not None:
+        if isinstance(_MLFLOW_TRACE_USER_DESTINATION, MlflowExperiment):
+            from mlflow import MlflowClient
+            from mlflow.tracing.export.mlflow import MlflowSpanExporter
+            from mlflow.tracing.processor.mlflow import MlflowSpanProcessor
+
+            client = MlflowClient(tracking_uri=_MLFLOW_TRACE_USER_DESTINATION.tracking_uri)
+            exporter = MlflowSpanExporter(client)
+            processor = MlflowSpanProcessor(
+                exporter, client, _MLFLOW_TRACE_USER_DESTINATION.experiment_id
+            )
+
+        elif isinstance(_MLFLOW_TRACE_USER_DESTINATION, Databricks):
+            from mlflow.tracing.export.databricks import DatabricksSpanExporter
+            from mlflow.tracing.processor.databricks import DatabricksSpanProcessor
+
+            exporter = DatabricksSpanExporter()
+            processor = DatabricksSpanProcessor(
+                span_exporter=exporter, experiment_id=_MLFLOW_TRACE_USER_DESTINATION.experiment_id
+            )
+
+        # TODO: Remove this branch once we fully migrate to the new tracing server
+        else:
+            from mlflow.tracing.export.databricks_agent_legacy import DatabricksAgentSpanExporter
+            from mlflow.tracing.processor.databricks import DatabricksSpanProcessor
+
+            exporter = DatabricksAgentSpanExporter(_MLFLOW_TRACE_USER_DESTINATION)
+            processor = DatabricksSpanProcessor(span_exporter=exporter, experiment_id=None)
+
+    elif should_use_otlp_exporter():
+        # Export to OpenTelemetry Collector when configured
+        from mlflow.tracing.processor.otel import OtelSpanProcessor
+
+        exporter = get_otlp_exporter()
+        processor = OtelSpanProcessor(exporter)
+
+    elif is_in_databricks_model_serving_environment():
+        # Export to Inference Table when running in Databricks Model Serving
+        if not is_mlflow_tracing_enabled_in_model_serving():
+            _MLFLOW_TRACER_PROVIDER = trace.NoOpTracerProvider()
+            return
+
+        from mlflow.tracing.export.inference_table import InferenceTableSpanExporter
+        from mlflow.tracing.processor.inference_table import InferenceTableSpanProcessor
+
+        exporter = InferenceTableSpanExporter()
+        processor = InferenceTableSpanProcessor(exporter)
+
+    else:
+        # Default to MLflow Tracking Server
+        from mlflow.tracing.export.mlflow import MlflowSpanExporter
+        from mlflow.tracing.processor.mlflow import MlflowSpanProcessor
+
+        exporter = MlflowSpanExporter()
+        processor = MlflowSpanProcessor(exporter)
+
+    tracer_provider = TracerProvider()
+    tracer_provider.add_span_processor(processor)
+    _MLFLOW_TRACER_PROVIDER = tracer_provider
+
+    from mlflow.tracing.utils.warning import suppress_warning
+
+    # 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")
+
+
+@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()) == 1
+
+        # Disable tracing
+        mlflow.tracing.disable()
+        f()
+        assert len(mlflow.search_traces()) == 1
+
+    """
+    if not is_tracing_enabled():
+        return
+
+    _setup_tracer_provider(disabled=True)
+    _MLFLOW_TRACER_PROVIDER_INITIALIZED.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()) == 1
+
+        # Disable tracing
+        mlflow.tracing.disable()
+        f()
+        assert len(mlflow.search_traces()) == 1
+
+        # Re-enable tracing
+        mlflow.tracing.enable()
+        f()
+        assert len(mlflow.search_traces()) == 2
+
+    """
+    if is_tracing_enabled() and _MLFLOW_TRACER_PROVIDER_INITIALIZED.done:
+        _logger.info("Tracing is already enabled")
+        return
+
+    _setup_tracer_provider()
+    _MLFLOW_TRACER_PROVIDER_INITIALIZED.done = True
+
+
+def trace_disabled(f):
+    """
+    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, **kwargs):
+        is_func_called = False
+        result = None
+        try:
+            if is_tracing_enabled():
+                disable()
+                try:
+                    is_func_called, result = True, f(*args, **kwargs)
+                finally:
+                    enable()
+            else:
+                is_func_called, result = True, f(*args, **kwargs)
+        # We should only catch the exception from disable() and enable()
+        # and let other exceptions propagate.
+        except MlflowTracingException as e:
+            _logger.warning(
+                f"An error occurred while disabling or re-enabling 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),
+            )
+            # If the exception is raised before the original function
+            # is called, we should call the original function
+            if not is_func_called:
+                result = f(*args, **kwargs)
+
+        return result
+
+    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.
+    _setup_tracer_provider(disabled=True)
+    # Flip _MLFLOW_TRACE_PROVIDER_INITIALIZED flag to False so that
+    # the next tracing operation will re-initialize the provider.
+    _MLFLOW_TRACER_PROVIDER_INITIALIZED.done = False
+
+    # Reset the custom destination set by the user
+    global _MLFLOW_TRACE_USER_DESTINATION
+    _MLFLOW_TRACE_USER_DESTINATION = None
+
+
+@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 _MLFLOW_TRACER_PROVIDER_INITIALIZED.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)

venv/lib/python3.9/site-packages/mlflow/tracing/trace_manager.py

@@ -0,0 +1,187 @@
+import contextlib
+import logging
+import threading
+from dataclasses import dataclass, field
+from typing import Generator, Optional
+
+from mlflow.entities import LiveSpan, Trace, TraceData, TraceInfo
+from mlflow.environment_variables import MLFLOW_TRACE_TIMEOUT_SECONDS
+from mlflow.tracing.constant import SpanAttributeKey
+from mlflow.tracing.utils.timeout import get_trace_cache_with_timeout
+
+_logger = logging.getLogger(__name__)
+
+
+# Internal representation to keep the state of a trace.
+# Dict[str, Span] is used instead of TraceData to allow access by span_id.
+@dataclass
+class _Trace:
+    info: TraceInfo
+    span_dict: dict[str, LiveSpan] = field(default_factory=dict)
+
+    def to_mlflow_trace(self) -> Trace:
+        trace_data = TraceData()
+        for span in self.span_dict.values():
+            # Convert LiveSpan, mutable objects, into immutable Span objects before persisting.
+            trace_data.spans.append(span.to_immutable_span())
+            if span.parent_id is None:
+                # Accessing the OTel span directly get serialized value directly.
+                trace_data.request = span._span.attributes.get(SpanAttributeKey.INPUTS)
+                trace_data.response = span._span.attributes.get(SpanAttributeKey.OUTPUTS)
+        return Trace(self.info, trace_data)
+
+    def get_root_span(self) -> Optional[LiveSpan]:
+        for span in self.span_dict.values():
+            if span.parent_id is None:
+                return span
+        return None
+
+
+class InMemoryTraceManager:
+    """
+    Manage spans and traces created by the tracing system in memory.
+    """
+
+    _instance_lock = threading.Lock()
+    _instance = None
+
+    @classmethod
+    def get_instance(cls):
+        if cls._instance is None:
+            with cls._instance_lock:
+                if cls._instance is None:
+                    cls._instance = InMemoryTraceManager()
+        return cls._instance
+
+    def __init__(self):
+        # In-memory cache to store request_id -> _Trace mapping.
+        self._traces = get_trace_cache_with_timeout()
+
+        # Store mapping between OpenTelemetry trace ID and MLflow request ID
+        self._trace_id_to_request_id: dict[int, str] = {}
+        self._lock = threading.Lock()  # Lock for _traces
+
+    def register_trace(self, trace_id: int, trace_info: TraceInfo):
+        """
+        Register a new trace info object to the in-memory trace registry.
+
+        Args:
+            trace_id: The trace ID for the new trace.
+            trace_info: The trace info object to be stored.
+        """
+        # Check for a new timeout setting whenever a new trace is created.
+        self._check_timeout_update()
+        with self._lock:
+            self._traces[trace_info.request_id] = _Trace(trace_info)
+            self._trace_id_to_request_id[trace_id] = trace_info.request_id
+
+    def update_trace_info(self, trace_info: TraceInfo):
+        """
+        Update the trace info object in the in-memory trace registry.
+
+        Args:
+            trace_info: The updated trace info object to be stored.
+        """
+        with self._lock:
+            if trace_info.request_id not in self._traces:
+                _logger.debug(f"Trace data with request ID {trace_info.request_id} not found.")
+                return
+            self._traces[trace_info.request_id].info = trace_info
+
+    def register_span(self, span: LiveSpan):
+        """
+        Store the given span in the in-memory trace data.
+
+        Args:
+            span: The span to be stored.
+        """
+        if not isinstance(span, LiveSpan):
+            _logger.debug(f"Invalid span object {type(span)} is passed. Skipping.")
+            return
+
+        with self._lock:
+            trace_data_dict = self._traces[span.request_id].span_dict
+            trace_data_dict[span.span_id] = span
+
+    @contextlib.contextmanager
+    def get_trace(self, request_id: str) -> Generator[Optional[_Trace], None, None]:
+        """
+        Yield the trace info for the given request_id.
+        This is designed to be used as a context manager to ensure the trace info is accessed
+        with the lock held.
+        """
+        with self._lock:
+            yield self._traces.get(request_id)
+
+    def get_span_from_id(self, request_id: str, span_id: str) -> Optional[LiveSpan]:
+        """
+        Get a span object for the given request_id and span_id.
+        """
+        with self._lock:
+            trace = self._traces.get(request_id)
+
+        return trace.span_dict.get(span_id) if trace else None
+
+    def get_root_span_id(self, request_id) -> Optional[str]:
+        """
+        Get the root span ID for the given trace ID.
+        """
+        with self._lock:
+            trace = self._traces.get(request_id)
+
+        if trace:
+            for span in trace.span_dict.values():
+                if span.parent_id is None:
+                    return span.span_id
+
+        return None
+
+    def get_request_id_from_trace_id(self, trace_id: int) -> Optional[str]:
+        """
+        Get the request ID for the given trace ID.
+        """
+        return self._trace_id_to_request_id.get(trace_id)
+
+    def set_request_metadata(self, request_id: str, key: str, value: str):
+        """
+        Set the request metadata for the given request ID.
+        """
+        with self.get_trace(request_id) as trace:
+            if trace:
+                trace.info.request_metadata[key] = value
+
+    def pop_trace(self, trace_id: int) -> Optional[Trace]:
+        """
+        Pop the trace data for the given id and return it as a ready-to-publish Trace object.
+        """
+        with self._lock:
+            request_id = self._trace_id_to_request_id.pop(trace_id, None)
+            trace = self._traces.pop(request_id, None)
+        return trace.to_mlflow_trace() if trace else None
+
+    def _check_timeout_update(self):
+        """
+        TTL/Timeout may be updated by users after initial cache creation. This method checks
+        for the update and create a new cache instance with the updated timeout.
+        """
+        new_timeout = MLFLOW_TRACE_TIMEOUT_SECONDS.get()
+        if new_timeout != getattr(self._traces, "timeout", None):
+            if len(self._traces) > 0:
+                _logger.warning(
+                    f"The timeout of the trace buffer has been updated to {new_timeout} seconds. "
+                    "This operation discards all in-progress traces at the moment. Please make "
+                    "sure to update the timeout when there are no in-progress traces."
+                )
+
+            with self._lock:
+                # We need to check here again in case this method runs in parallel
+                if new_timeout != getattr(self._traces, "timeout", None):
+                    self._traces = get_trace_cache_with_timeout()
+
+    @classmethod
+    def reset(self):
+        """Clear all the aggregated trace data. This should only be used for testing."""
+        if self._instance:
+            with self._instance._lock:
+                self._instance._traces.clear()
+            self._instance = None

venv/lib/python3.9/site-packages/mlflow/tracing/utils/__init__.py

@@ -0,0 +1,451 @@
+# TODO: Split this file into multiple files and move under utils directory.
+from __future__ import annotations
+
+import inspect
+import json
+import logging
+import uuid
+from collections import Counter
+from dataclasses import asdict, is_dataclass
+from functools import lru_cache
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+from opentelemetry import trace as trace_api
+from packaging.version import Version
+
+import mlflow
+from mlflow.entities.span_status import SpanStatusCode
+from mlflow.exceptions import BAD_REQUEST, MlflowTracingException
+from mlflow.tracing.constant import SpanAttributeKey
+from mlflow.utils.mlflow_tags import IMMUTABLE_TAGS
+
+_logger = logging.getLogger(__name__)
+
+SPANS_COLUMN_NAME = "spans"
+
+if TYPE_CHECKING:
+    from mlflow.client import MlflowClient
+    from mlflow.entities import LiveSpan
+    from mlflow.types.chat import ChatMessage, ChatTool
+
+
+def capture_function_input_args(func, args, kwargs) -> Optional[dict[str, Any]]:
+    try:
+        # Avoid capturing `self`
+        func_signature = inspect.signature(func)
+        bound_arguments = func_signature.bind(*args, **kwargs)
+        bound_arguments.apply_defaults()
+
+        # Remove `self` from bound arguments if it exists
+        if bound_arguments.arguments.get("self"):
+            del bound_arguments.arguments["self"]
+
+        return bound_arguments.arguments
+    except Exception:
+        _logger.warning(f"Failed to capture inputs for function {func.__name__}.")
+        return None
+
+
+class TraceJSONEncoder(json.JSONEncoder):
+    """
+    Custom JSON encoder for serializing non-OpenTelemetry compatible objects in a trace or span.
+
+    Trace may contain types that require custom serialization logic, such as Pydantic models,
+    non-JSON-serializable types, etc.
+    """
+
+    def default(self, obj):
+        try:
+            import langchain
+
+            # LangChain < 0.3.0 does some trick to support Pydantic 1.x and 2.x, so checking
+            # type with installed Pydantic version might not work for some models.
+            # https://github.com/langchain-ai/langchain/blob/b66a4f48fa5656871c3e849f7e1790dfb5a4c56b/libs/core/langchain_core/pydantic_v1/__init__.py#L7
+            if Version(langchain.__version__) < Version("0.3.0"):
+                from langchain_core.pydantic_v1 import BaseModel as LangChainBaseModel
+
+                if isinstance(obj, LangChainBaseModel):
+                    return obj.dict()
+        except ImportError:
+            pass
+
+        try:
+            import pydantic
+
+            if isinstance(obj, pydantic.BaseModel):
+                # NB: Pydantic 2.0+ has a different API for model serialization
+                if Version(pydantic.VERSION) >= Version("2.0"):
+                    return obj.model_dump()
+                else:
+                    return obj.dict()
+        except ImportError:
+            pass
+
+        # Some dataclass object defines __str__ method that doesn't return the full object
+        # representation, so we use dict representation instead.
+        # E.g. https://github.com/run-llama/llama_index/blob/29ece9b058f6b9a1cf29bc723ed4aa3a39879ad5/llama-index-core/llama_index/core/chat_engine/types.py#L63-L64
+        if is_dataclass(obj):
+            try:
+                return asdict(obj)
+            except TypeError:
+                pass
+
+        # Some object has dangerous side effect in __str__ method, so we use class name instead.
+        if not self._is_safe_to_encode_str(obj):
+            return type(obj)
+
+        try:
+            return super().default(obj)
+        except TypeError:
+            return str(obj)
+
+    def _is_safe_to_encode_str(self, obj) -> bool:
+        """Check if it's safe to encode the object as a string."""
+        try:
+            # These Llama Index objects are not safe to encode as string, because their __str__
+            # method consumes the stream and make it unusable.
+            # E.g. https://github.com/run-llama/llama_index/blob/54f2da61ba8a573284ab8336f2b2810d948c3877/llama-index-core/llama_index/core/base/response/schema.py#L120-L127
+            from llama_index.core.base.response.schema import (
+                AsyncStreamingResponse,
+                StreamingResponse,
+            )
+            from llama_index.core.chat_engine.types import StreamingAgentChatResponse
+
+            if isinstance(
+                obj, (AsyncStreamingResponse, StreamingResponse, StreamingAgentChatResponse)
+            ):
+                return False
+        except ImportError:
+            pass
+
+        return True
+
+
+@lru_cache(maxsize=1)
+def encode_span_id(span_id: int) -> str:
+    """
+    Encode the given integer span ID to a 16-byte hex string.
+    # https://github.com/open-telemetry/opentelemetry-python/blob/9398f26ecad09e02ad044859334cd4c75299c3cd/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py#L507-L508
+    # NB: We don't add '0x' prefix to the hex string here for simpler parsing in backend.
+    #   Some backend (e.g. Databricks) disallow this prefix.
+    """
+    return trace_api.format_span_id(span_id)
+
+
+@lru_cache(maxsize=1)
+def encode_trace_id(trace_id: int) -> str:
+    """
+    Encode the given integer trace ID to a 32-byte hex string.
+    """
+    return trace_api.format_trace_id(trace_id)
+
+
+def decode_id(span_or_trace_id: str) -> int:
+    """
+    Decode the given hex string span or trace ID to an integer.
+    """
+    return int(span_or_trace_id, 16)
+
+
+def build_otel_context(trace_id: int, span_id: int) -> trace_api.SpanContext:
+    """
+    Build an OpenTelemetry SpanContext object from the given trace and span IDs.
+    """
+    return trace_api.SpanContext(
+        trace_id=trace_id,
+        span_id=span_id,
+        # NB: This flag is OpenTelemetry's concept to indicate whether the context is
+        # propagated from remote parent or not. We don't support distributed tracing
+        # yet so always set it to False.
+        is_remote=False,
+    )
+
+
+def deduplicate_span_names_in_place(spans: list[LiveSpan]):
+    """
+    Deduplicate span names in the trace data by appending an index number to the span name.
+
+    This is only applied when there are multiple spans with the same name. The span names
+    are modified in place to avoid unnecessary copying.
+
+    E.g.
+        ["red", "red"] -> ["red_1", "red_2"]
+        ["red", "red", "blue"] -> ["red_1", "red_2", "blue"]
+
+    Args:
+        spans: A list of spans to deduplicate.
+    """
+    span_name_counter = Counter(span.name for span in spans)
+    # Apply renaming only for duplicated spans
+    span_name_counter = {name: 1 for name, count in span_name_counter.items() if count > 1}
+    # Add index to the duplicated span names
+    for span in spans:
+        if count := span_name_counter.get(span.name):
+            span_name_counter[span.name] += 1
+            span._span._name = f"{span.name}_{count}"
+
+
+def get_otel_attribute(span: trace_api.Span, key: str) -> Optional[str]:
+    """
+    Get the attribute value from the OpenTelemetry span in a decoded format.
+
+    Args:
+        span: The OpenTelemetry span object.
+        key: The key of the attribute to retrieve.
+
+    Returns:
+        The attribute value as decoded string. If the attribute is not found or cannot
+        be parsed, return None.
+    """
+    try:
+        return json.loads(span.attributes.get(key))
+    except Exception:
+        _logger.debug(f"Failed to get attribute {key} with from span {span}.", exc_info=True)
+
+
+def _try_get_prediction_context():
+    # NB: Tracing is enabled in mlflow-skinny, but the pyfunc module cannot be imported as it
+    #     relies on numpy, which is not installed in skinny.
+    try:
+        from mlflow.pyfunc.context import get_prediction_context
+    except ImportError:
+        return
+
+    return get_prediction_context()
+
+
+def maybe_get_request_id(is_evaluate=False) -> Optional[str]:
+    """Get the request ID if the current prediction is as a part of MLflow model evaluation."""
+    context = _try_get_prediction_context()
+    if not context or (is_evaluate and not context.is_evaluate):
+        return None
+
+    if not context.request_id and is_evaluate:
+        raise MlflowTracingException(
+            f"Missing request_id for context {context}. "
+            "request_id can't be None when is_evaluate=True.",
+            error_code=BAD_REQUEST,
+        )
+
+    return context.request_id
+
+
+def maybe_get_dependencies_schemas() -> Optional[dict]:
+    context = _try_get_prediction_context()
+    if context:
+        return context.dependencies_schemas
+
+
+def exclude_immutable_tags(tags: dict[str, str]) -> dict[str, str]:
+    """Exclude immutable tags e.g. "mlflow.user" from the given tags."""
+    return {k: v for k, v in tags.items() if k not in IMMUTABLE_TAGS}
+
+
+def generate_request_id() -> str:
+    return uuid.uuid4().hex
+
+
+def construct_full_inputs(func, *args, **kwargs) -> dict[str, Any]:
+    """
+    Construct the full input arguments dictionary for the given function,
+    including positional and keyword arguments.
+    """
+    signature = inspect.signature(func)
+    # this does not create copy. So values should not be mutated directly
+    arguments = signature.bind_partial(*args, **kwargs).arguments
+
+    if "self" in arguments:
+        arguments.pop("self")
+
+    return arguments
+
+
+def set_span_chat_messages(
+    span: LiveSpan,
+    messages: Union[dict, ChatMessage],
+    append=False,
+):
+    """
+    Set the `mlflow.chat.messages` attribute on the specified span. This
+    attribute is used in the UI, and also by downstream applications that
+    consume trace data, such as MLflow evaluate.
+
+    Args:
+        span: The LiveSpan to add the attribute to
+        messages: A list of standardized chat messages (refer to the
+                 `spec <../llms/tracing/tracing-schema.html#chat-completion-spans>`_
+                 for details)
+        append: If True, the messages will be appended to the existing messages. Otherwise,
+                the attribute will be overwritten entirely. Default is False.
+                This is useful when you want to record messages incrementally, e.g., log
+                input messages first, and then log output messages later.
+
+    Example:
+
+    .. code-block:: python
+        :test:
+
+        import mlflow
+        from mlflow.tracing import set_span_chat_messages
+
+
+        @mlflow.trace
+        def f():
+            messages = [{"role": "user", "content": "hello"}]
+            span = mlflow.get_current_active_span()
+            set_span_chat_messages(span, messages)
+            return 0
+
+
+        f()
+    """
+    from mlflow.types.chat import ChatMessage
+
+    sanitized_messages = []
+    for message in messages:
+        if isinstance(message, dict):
+            ChatMessage.validate_compat(message)
+            sanitized_messages.append(message)
+        elif isinstance(message, ChatMessage):
+            # NB: ChatMessage is used for both request and response messages. In OpenAI's API spec,
+            #   some fields are only present in either the request or response (e.g., tool_call_id).
+            #   Those fields should not be recorded unless set explicitly, so we set
+            #   exclude_unset=True here to avoid recording unset fields.
+            sanitized_messages.append(message.model_dump_compat(exclude_unset=True))
+
+    if append:
+        existing_messages = span.get_attribute(SpanAttributeKey.CHAT_MESSAGES) or []
+        sanitized_messages = existing_messages + sanitized_messages
+
+    span.set_attribute(SpanAttributeKey.CHAT_MESSAGES, sanitized_messages)
+
+
+def set_span_chat_tools(span: LiveSpan, tools: list[ChatTool]):
+    """
+    Set the `mlflow.chat.tools` attribute on the specified span. This
+    attribute is used in the UI, and also by downstream applications that
+    consume trace data, such as MLflow evaluate.
+
+    Args:
+        span: The LiveSpan to add the attribute to
+        tools: A list of standardized chat tool definitions (refer to the
+              `spec <../llms/tracing/tracing-schema.html#chat-completion-spans>`_
+              for details)
+
+    Example:
+
+    .. code-block:: python
+        :test:
+
+        import mlflow
+        from mlflow.tracing import set_span_chat_tools
+
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "add",
+                    "description": "Add two numbers",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "a": {"type": "number"},
+                            "b": {"type": "number"},
+                        },
+                        "required": ["a", "b"],
+                    },
+                },
+            }
+        ]
+
+
+        @mlflow.trace
+        def f():
+            span = mlflow.get_current_active_span()
+            set_span_chat_tools(span, tools)
+            return 0
+
+
+        f()
+    """
+    from mlflow.types.chat import ChatTool
+
+    if not isinstance(tools, list):
+        raise MlflowTracingException(
+            f"Invalid tools type {type(tools)}. Expected a list of ChatTool.",
+            error_code=BAD_REQUEST,
+        )
+
+    sanitized_tools = []
+    for tool in tools:
+        if isinstance(tool, dict):
+            ChatTool.validate_compat(tool)
+            sanitized_tools.append(tool)
+        elif isinstance(tool, ChatTool):
+            sanitized_tools.append(tool.model_dump_compat(exclude_unset=True))
+
+    span.set_attribute(SpanAttributeKey.CHAT_TOOLS, sanitized_tools)
+
+
+def start_client_span_or_trace(
+    client: MlflowClient,
+    name: str,
+    span_type: str,
+    parent_span: Optional[LiveSpan] = None,
+    inputs: Optional[dict[str, Any]] = None,
+    attributes: Optional[dict[str, Any]] = None,
+    start_time_ns: Optional[int] = None,
+) -> LiveSpan:
+    """
+    An utility to start a span or trace using MlflowClient based on the current active span.
+    """
+    if parent_span := parent_span or mlflow.get_current_active_span():
+        return client.start_span(
+            name=name,
+            request_id=parent_span.request_id,
+            parent_id=parent_span.span_id,
+            span_type=span_type,
+            inputs=inputs,
+            attributes=attributes,
+            start_time_ns=start_time_ns,
+        )
+    else:
+        return client.start_trace(
+            name=name,
+            span_type=span_type,
+            inputs=inputs,
+            attributes=attributes,
+            start_time_ns=start_time_ns,
+        )
+
+
+def end_client_span_or_trace(
+    client: MlflowClient,
+    span: LiveSpan,
+    outputs: Optional[dict[str, Any]] = None,
+    attributes: Optional[dict[str, Any]] = None,
+    status: str = SpanStatusCode.OK,
+    end_time_ns: Optional[int] = None,
+) -> LiveSpan:
+    """
+    An utility to end a span or trace using MlflowClient based on the current active span.
+    """
+    if span.parent_id is not None:
+        return client.end_span(
+            request_id=span.request_id,
+            span_id=span.span_id,
+            outputs=outputs,
+            attributes=attributes,
+            status=status,
+            end_time_ns=end_time_ns,
+        )
+    else:
+        span.set_status(status)
+        span.set_outputs(outputs)
+        return client.end_trace(
+            request_id=span.request_id,
+            outputs=outputs,
+            attributes=attributes,
+            status=status,
+            end_time_ns=end_time_ns,
+        )

venv/lib/python3.9/site-packages/mlflow/tracing/utils/exception.py

@@ -0,0 +1,21 @@
+import functools
+
+from mlflow.exceptions import MlflowTracingException
+
+
+def raise_as_trace_exception(f):
+    """
+    A decorator to make sure that the decorated function only raises MlflowTracingException.
+
+    Any exceptions are caught and translated to MlflowTracingException before exiting the function.
+    This is helpful for upstream functions to handle tracing related exceptions properly.
+    """
+
+    @functools.wraps(f)
+    def wrapper(*args, **kwargs):
+        try:
+            return f(*args, **kwargs)
+        except Exception as e:
+            raise MlflowTracingException(e) from e
+
+    return wrapper

venv/lib/python3.9/site-packages/mlflow/tracing/utils/once.py

@@ -0,0 +1,35 @@
+# Customized from https://github.com/open-telemetry/opentelemetry-python/blob/754fc36a408dd45e86d4a0f820f84e692f14b4c1/opentelemetry-api/src/opentelemetry/util/_once.py
+from threading import Lock
+from typing import Callable
+
+
+class Once:
+    """Execute a function exactly once and block all callers until the function returns"""
+
+    def __init__(self) -> None:
+        self.__lock = Lock()
+        self.__done = False
+
+    @property
+    def done(self):
+        with self.__lock:
+            return self.__done
+
+    @done.setter
+    def done(self, value):
+        with self.__lock:
+            self.__done = value
+
+    def do_once(self, func: Callable[[], None]):
+        """
+        Execute ``func`` if it hasn't been executed or return.
+        Will block until ``func`` has been called by one thread.
+        """
+        if self.__done:
+            return
+
+        with self.__lock:
+            if not self.__done:
+                func()
+                self.__done = True
+                return

venv/lib/python3.9/site-packages/mlflow/tracing/utils/otlp.py

@@ -0,0 +1,63 @@
+import os
+from typing import Optional
+
+from opentelemetry.sdk.trace.export import SpanExporter
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import RESOURCE_DOES_NOT_EXIST
+
+
+def should_use_otlp_exporter() -> bool:
+    return _get_otlp_endpoint() is not None
+
+
+def get_otlp_exporter() -> SpanExporter:
+    """
+    Get the OTLP exporter based on the configured protocol.
+    """
+    endpoint = _get_otlp_endpoint()
+    protocol = _get_otlp_protocol()
+    if protocol == "grpc":
+        try:
+            from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+        except ImportError:
+            raise MlflowException(
+                "gRPC OTLP exporter is not available. Please install the required dependency by "
+                "running `pip install opentelemetry-exporter-otlp-proto-grpc`.",
+                error_code=RESOURCE_DOES_NOT_EXIST,
+            )
+
+        return OTLPSpanExporter(endpoint=endpoint)
+    elif protocol == "http/protobuf":
+        try:
+            from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+        except ImportError as e:
+            raise MlflowException(
+                "HTTP OTLP exporter is not available. Please install the required dependency by "
+                "running `pip install opentelemetry-exporter-otlp-proto-http`.",
+                error_code=RESOURCE_DOES_NOT_EXIST,
+            ) from e
+
+        return OTLPSpanExporter(endpoint=endpoint)
+    else:
+        raise MlflowException.invalid_parameter_value(
+            f"Unsupported OTLP protocol '{protocol}' is configured. Please set "
+            "the protocol to either 'grpc' or 'http/protobuf'."
+        )
+
+
+def _get_otlp_endpoint() -> Optional[str]:
+    """
+    Get the OTLP endpoint from the environment variables.
+    Ref: https://opentelemetry.io/docs/languages/sdk-configuration/otlp-exporter/#endpoint-configuration
+    """
+    # Use `or` instead of default value to do lazy eval
+    return os.environ.get("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT") or os.environ.get(
+        "OTEL_EXPORTER_OTLP_ENDPOINT"
+    )
+
+
+def _get_otlp_protocol() -> str:
+    return os.environ.get("OTEL_EXPORTER_OTLP_TRACES_PROTOCOL") or os.environ.get(
+        "OTEL_EXPORTER_OTLP_PROTOCOL", "grpc"
+    )

venv/lib/python3.9/site-packages/mlflow/tracing/utils/search.py

@@ -0,0 +1,292 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any, Literal, NamedTuple, Optional, Union
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+
+SPANS_COLUMN_NAME = "spans"
+
+if TYPE_CHECKING:
+    import pandas
+
+    import mlflow.entities
+    from mlflow.entities import Trace
+
+
+def traces_to_df(traces: list[Trace]) -> "pandas.DataFrame":
+    """
+    Convert a list of MLflow Traces to a pandas DataFrame with one column called "traces"
+    containing string representations of each Trace.
+    """
+    import pandas as pd
+
+    from mlflow.entities.trace import Trace  # import here to avoid circular import
+
+    rows = [trace.to_pandas_dataframe_row() for trace in traces]
+    return pd.DataFrame.from_records(data=rows, columns=Trace.pandas_dataframe_columns())
+
+
+def extract_span_inputs_outputs(
+    traces: Union[list["mlflow.entities.Trace"], "pandas.DataFrame"],
+    fields: list[str],
+    col_name: Optional[str] = None,
+) -> "pandas.DataFrame":
+    """
+    Extracts the specified input and output fields from the spans contained in the specified traces.
+
+    Args:
+        traces: A list of :py:class:`mlflow.entities.Trace` or a pandas DataFrame containing traces.
+        fields: A list of field strings of the form 'span_name.[inputs|outputs]' or
+            'span_name.[inputs|outputs].field_name'.
+        col_name: The name of the column in the traces DataFrame containing the spans. If `traces`
+            is a list of MLflow Traces, this argument should not be provided.
+    """
+    try:
+        import pandas as pd
+    except ImportError as e:
+        raise MlflowException(
+            message=(
+                "The `pandas` library is not installed. Please install `pandas` to use the"
+                f"`mlflow.tracing.extract` function. Error: {e}"
+            ),
+        )
+
+    parsed_fields = _parse_fields(fields)
+
+    if isinstance(traces, list):
+        if col_name is not None:
+            raise MlflowException(
+                message=(
+                    "If `traces` is a list of MLflow Traces, `col_name` should not be provided."
+                ),
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        traces = traces_to_df(traces)
+        col_name = SPANS_COLUMN_NAME
+
+    if isinstance(traces, pd.DataFrame):
+        return _extract_from_traces_pandas_df(df=traces, col_name=col_name, fields=parsed_fields)
+
+    raise MlflowException(
+        message=(
+            "`traces` must be a list of MLflow Traces or a pandas DataFrame. Got: {type(traces)}"
+        ),
+        error_code=INVALID_PARAMETER_VALUE,
+    )
+
+
+class _PeekableIterator:
+    """
+    Wraps an iterator and allows peeking at the next element without consuming it.
+    """
+
+    def __init__(self, it):
+        self.it = iter(it)
+        self._next = None
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self._next is not None:
+            next_value = self._next
+            self._next = None
+            return next_value
+        return next(self.it)
+
+    def peek(self):
+        if self._next is None:
+            try:
+                self._next = next(self.it)
+            except StopIteration:
+                return None
+        return self._next
+
+
+class _ParsedField(NamedTuple):
+    """
+    Represents a parsed field from a string of the form 'span_name.[inputs|outputs]' or
+    'span_name.[inputs|outputs].field_name'.
+    """
+
+    span_name: str
+    field_type: Literal["inputs", "outputs"]
+    field_name: Optional[str]
+
+    def __str__(self) -> str:
+        return (
+            f"{self.span_name}.{self.field_type}.{self.field_name}"
+            if self.field_name is not None
+            else f"{self.span_name}.{self.field_type}"
+        )
+
+
+_BACKTICK = "`"
+
+
+class _FieldParser:
+    def __init__(self, field: str) -> None:
+        self.field = field
+        self.chars = _PeekableIterator(field)
+
+    def peek(self) -> str:
+        return self.chars.peek()
+
+    def next(self) -> str:
+        return next(self.chars)
+
+    def has_next(self) -> bool:
+        return self.peek() is not None
+
+    def consume_until_char_or_end(self, stop_char: Optional[str] = None) -> str:
+        """
+        Consume characters until the specified character is encountered or the end of the
+        string. If char is None, consume until the end of the string.
+        """
+        consumed = ""
+        while (c := self.peek()) and c != stop_char:
+            consumed += self.next()
+        return consumed
+
+    def _parse_span_name(self) -> str:
+        if self.peek() == _BACKTICK:
+            self.next()
+            span_name = self.consume_until_char_or_end(_BACKTICK)
+            if self.peek() != _BACKTICK:
+                raise MlflowException.invalid_parameter_value(
+                    f"Expected closing backtick: {self.field!r}"
+                )
+            self.next()
+        else:
+            span_name = self.consume_until_char_or_end(".")
+
+        if self.peek() != ".":
+            raise MlflowException.invalid_parameter_value(
+                f"Expected dot after span name: {self.field!r}"
+            )
+        self.next()
+        return span_name
+
+    def _parse_field_type(self) -> str:
+        field_type = self.consume_until_char_or_end(".")
+        if field_type not in ("inputs", "outputs"):
+            raise MlflowException.invalid_parameter_value(
+                f"Invalid field type: {field_type!r}. Expected 'inputs' or 'outputs'."
+            )
+
+        if self.has_next():
+            self.next()  # Consume the dot
+        return field_type
+
+    def _parse_field_name(self) -> str:
+        if self.peek() == _BACKTICK:
+            self.next()
+            field_name = self.consume_until_char_or_end(_BACKTICK)
+            if self.peek() != _BACKTICK:
+                raise MlflowException.invalid_parameter_value(
+                    f"Expected closing backtick: {self.field!r}"
+                )
+            self.next()
+
+            # There should be no more characters after the closing backtick
+            if self.has_next():
+                raise MlflowException.invalid_parameter_value(
+                    f"Unexpected characters after closing backtick: {self.field!r}"
+                )
+
+        else:
+            field_name = self.consume_until_char_or_end()
+
+        return field_name
+
+    def parse(self) -> _ParsedField:
+        span_name = self._parse_span_name()
+        field_type = self._parse_field_type()
+        field_name = self._parse_field_name() if self.has_next() else None
+        return _ParsedField(span_name=span_name, field_type=field_type, field_name=field_name)
+
+
+def _parse_fields(fields: list[str]) -> list[_ParsedField]:
+    """
+    Parses the specified field strings of the form 'span_name.[inputs|outputs]' or
+    'span_name.[inputs|outputs].field_name' into _ParsedField objects.
+    """
+    return [_FieldParser(field).parse() for field in fields]
+
+
+def _extract_from_traces_pandas_df(
+    df: "pandas.DataFrame", col_name: str, fields: list[_ParsedField]
+) -> "pandas.DataFrame":
+    """
+    Extracts the specified fields from the spans contained in the specified column of the
+    specified traces DataFrame.
+    """
+
+    from mlflow.entities import Span
+
+    if col_name not in df.columns:
+        raise MlflowException(
+            message=(
+                f"Column '{col_name}' not found in traces DataFrame."
+                f" Available columns: {df.columns}"
+            ),
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+    new_columns: dict[str, list[Any]] = defaultdict(list)
+    for _, row in df.iterrows():
+        spans_dict: dict[str, list[Span]] = defaultdict(list)
+        for span in _extract_spans_from_row(row[col_name]):
+            spans_dict[span.name].append(span)
+
+        for field in fields:
+            matching_spans = spans_dict.get(field.span_name, [])
+            matching_value = _find_matching_value(field, matching_spans)
+            new_columns[str(field)].append(matching_value)
+
+    df_with_new_fields = df.copy()
+    for field in fields:
+        df_with_new_fields[str(field)] = new_columns[str(field)]
+
+    return df_with_new_fields
+
+
+def _find_matching_value(field: _ParsedField, spans: list["mlflow.entities.Span"]) -> Optional[Any]:
+    """
+    Find the value of the field in the list of spans. If the field is not found, return None.
+    """
+    for span in spans:
+        span_inputs_or_outputs = getattr(span, field.field_type)
+        if (
+            isinstance(span_inputs_or_outputs, dict)
+            and field.field_name is not None
+            and field.field_name in span_inputs_or_outputs
+        ):
+            return span_inputs_or_outputs.get(field.field_name)
+        elif field.field_name is None:
+            return span_inputs_or_outputs
+
+
+def _extract_spans_from_row(
+    row_content: Optional[list[dict[str, Any]]],
+) -> list["mlflow.entities.Span"]:
+    """
+    Parses and extracts MLflow Spans from the row content of a traces pandas DataFrame.
+    """
+    from mlflow.entities import Span
+
+    if row_content is None:
+        return []
+
+    try:
+        return [Span.from_dict(span_dict) for span_dict in row_content]
+    except Exception as e:
+        raise MlflowException(
+            message=(
+                f"Failed to extract spans from traces DataFrame row content: {row_content}."
+                f" Error: {e}"
+            ),
+            error_code=INVALID_PARAMETER_VALUE,
+        ) from e

venv/lib/python3.9/site-packages/mlflow/tracing/utils/timeout.py

@@ -0,0 +1,250 @@
+import atexit
+import logging
+import threading
+import time
+from collections import OrderedDict
+
+from cachetools import Cache, TTLCache
+
+from mlflow.entities.span_event import SpanEvent
+from mlflow.entities.span_status import SpanStatusCode
+from mlflow.environment_variables import (
+    MLFLOW_TRACE_BUFFER_MAX_SIZE,
+    MLFLOW_TRACE_BUFFER_TTL_SECONDS,
+    MLFLOW_TRACE_TIMEOUT_CHECK_INTERVAL_SECONDS,
+    MLFLOW_TRACE_TIMEOUT_SECONDS,
+)
+from mlflow.exceptions import MlflowTracingException
+
+_logger = logging.getLogger(__name__)
+
+_TRACE_EXPIRATION_MSG = (
+    "Trace {request_id} is timed out after {ttl} seconds. The operation may be stuck or "
+    "taking too long to complete. To increase the timeout, set the environment variable "
+    "MLFLOW_TRACE_TIMEOUT_SECONDS to a larger value."
+)
+
+
+def get_trace_cache_with_timeout() -> Cache:
+    """
+    Return a cache object that stores traces in-memory while they are in-progress.
+
+    If the timeout is specified, this returns a customized cache that logs the
+    expired traces to the backend. Otherwise, this returns a regular cache.
+    """
+
+    if timeout := MLFLOW_TRACE_TIMEOUT_SECONDS.get():
+        return MlflowTraceTimeoutCache(
+            timeout=timeout,
+            maxsize=MLFLOW_TRACE_BUFFER_MAX_SIZE.get(),
+        )
+
+    # NB: Ideally we should return the vanilla Cache object only with maxsize.
+    # But we used TTLCache before introducing the timeout feature (that does not
+    # monitor timeout periodically nor log the expired traces). To keep the
+    # backward compatibility, we return TTLCache.
+    return TTLCache(
+        ttl=MLFLOW_TRACE_BUFFER_TTL_SECONDS.get(),
+        maxsize=MLFLOW_TRACE_BUFFER_MAX_SIZE.get(),
+    )
+
+
+class _TimedCache(Cache):
+    """
+    This code is ported from cachetools library to avoid depending on the private class.
+    https://github.com/tkem/cachetools/blob/d44c98407030d2e91cbe82c3997be042d9c2f0de/src/cachetools/__init__.py#L376
+    """
+
+    class _Timer:
+        def __init__(self, timer):
+            self.__timer = timer
+            self.__nesting = 0
+
+        def __call__(self):
+            if self.__nesting == 0:
+                return self.__timer()
+            else:
+                return self.__time
+
+        def __enter__(self):
+            if self.__nesting == 0:
+                self.__time = time = self.__timer()
+            else:
+                time = self.__time
+            self.__nesting += 1
+            return time
+
+        def __exit__(self, *exc):
+            self.__nesting -= 1
+
+        def __reduce__(self):
+            return _TimedCache._Timer, (self.__timer,)
+
+        def __getattr__(self, name):
+            return getattr(self.__timer, name)
+
+    def __init__(self, maxsize, timer=time.monotonic, getsizeof=None):
+        Cache.__init__(self, maxsize, getsizeof)
+        self.__timer = _TimedCache._Timer(timer)
+
+    def __repr__(self, cache_repr=Cache.__repr__):
+        with self.__timer as time:
+            self.expire(time)
+            return cache_repr(self)
+
+    def __len__(self, cache_len=Cache.__len__):
+        with self.__timer as time:
+            self.expire(time)
+            return cache_len(self)
+
+    @property
+    def currsize(self):
+        with self.__timer as time:
+            self.expire(time)
+            return super().currsize
+
+    @property
+    def timer(self):
+        """The timer function used by the cache."""
+        return self.__timer
+
+    def clear(self):
+        with self.__timer as time:
+            self.expire(time)
+            Cache.clear(self)
+
+    def get(self, *args, **kwargs):
+        with self.__timer:
+            return Cache.get(self, *args, **kwargs)
+
+    def pop(self, *args, **kwargs):
+        with self.__timer:
+            return Cache.pop(self, *args, **kwargs)
+
+    def setdefault(self, *args, **kwargs):
+        with self.__timer:
+            return Cache.setdefault(self, *args, **kwargs)
+
+
+class MlflowTraceTimeoutCache(_TimedCache):
+    """
+    A different implementation of cachetools.TTLCache that logs the expired traces to the backend.
+
+    NB: Do not use this class outside a singleton context. This class is not thread-safe.
+    """
+
+    def __init__(self, timeout: int, maxsize: int):
+        super().__init__(maxsize=maxsize)
+        self._timeout = timeout
+
+        # Set up the linked list ordered by expiration time
+        self._root = TTLCache._Link()
+        self._root.prev = self._root
+        self._root.next = self._root
+        self._links = OrderedDict()
+
+        self._start_expire_check_loop()
+
+    @property
+    def timeout(self) -> int:
+        # Timeout should not be changed after the cache is created
+        # because the linked list will not be updated accordingly.
+        return self._timeout
+
+    def __setitem__(self, key, value, cache_setitem=Cache.__setitem__):
+        """Set the item in the cache, and also in the linked list if it is a new key"""
+        with self.timer as time:
+            cache_setitem(self, key, value)
+
+        if key not in self._links:
+            # Add the new item to the tail of the linked list
+            # Inspired by https://github.com/tkem/cachetools/blob/d44c98407030d2e91cbe82c3997be042d9c2f0de/src/cachetools/__init__.py#L432
+            tail = self._root.prev
+            link = TTLCache._Link(key)
+            link.expires = time + self._timeout
+            link.next = self._root
+            link.prev = tail
+            tail.next = link
+            self._root.prev = link
+            self._links[key] = link
+
+    def __delitem__(self, key, cache_delitem=Cache.__delitem__):
+        """Delete the item from the cache and the linked list."""
+        cache_delitem(self, key)
+        link = self._links.pop(key)
+        link.unlink()
+
+    def _start_expire_check_loop(self):
+        # Close the daemon thread when the main thread exits
+        atexit.register(self.clear)
+
+        self._expire_checker_thread = threading.Thread(
+            target=self._expire_check_loop, daemon=True, name="TTLCacheExpireLoop"
+        )
+        self._expire_checker_stop_event = threading.Event()
+        self._expire_checker_thread.start()
+
+    def _expire_check_loop(self):
+        while not self._expire_checker_stop_event.is_set():
+            try:
+                self.expire()
+            except Exception as e:
+                _logger.debug(f"Failed to expire traces: {e}")
+                # If an error is raised from the expiration method, stop running the loop.
+                # Otherwise, the expire task might get heavier and heavier due to the
+                # increasing number of expired items.
+                break
+
+            time.sleep(MLFLOW_TRACE_TIMEOUT_CHECK_INTERVAL_SECONDS.get())
+
+    def expire(self, time=None):
+        """
+        Trigger the expiration of traces that have exceeded the timeout.
+
+        Args:
+            time: Unused. Only for compatibility with the parent class.
+        """
+        expired = self._get_expired_traces()
+
+        # End the expired traces and set the status to ERROR in background thread
+        for request_id in expired:
+            trace = self[request_id]
+            if root_span := trace.get_root_span():
+                try:
+                    root_span.set_status(SpanStatusCode.ERROR)
+                    msg = _TRACE_EXPIRATION_MSG.format(request_id=request_id, ttl=self._timeout)
+                    exception_event = SpanEvent.from_exception(MlflowTracingException(msg))
+                    root_span.add_event(exception_event)
+                    root_span.end()  # Calling end() triggers span export
+                    _logger.info(msg + " You can find the aborted trace in the MLflow UI.")
+                except Exception as e:
+                    _logger.debug(f"Failed to export an expired trace {request_id}: {e}")
+
+                # NB: root_span.end() should pop the trace from the cache. But we need to
+                # double-check it because it may not happens due to some errors.
+                if request_id in self:
+                    del self[request_id]
+
+    def _get_expired_traces(self) -> list[str]:
+        """
+        Find all expired traces and return their request IDs.
+
+        The linked list is ordered by expiration time, so we can traverse the list from the head
+        and return early whenever we find a trace that has not expired yet.
+        """
+        time = self.timer()
+        curr = self._root.next
+
+        if curr.expires and time < curr.expires:
+            return []
+
+        expired = []
+        while curr is not self._root and not (time < curr.expires):
+            expired.append(curr.key)
+            curr = curr.next
+        return expired
+
+    def clear(self):
+        super().clear()
+        self._expire_checker_stop_event.set()
+        self._expire_checker_thread.join()

venv/lib/python3.9/site-packages/mlflow/tracing/utils/token.py

@@ -0,0 +1,19 @@
+import contextvars
+from dataclasses import dataclass
+from typing import Optional
+
+from mlflow.entities import LiveSpan
+
+
+@dataclass
+class SpanWithToken:
+    """
+    A utility container to hold an MLflow span and its corresponding OpenTelemetry token.
+
+    The token is a special object that is generated when setting a span as active within
+    the Open Telemetry span context. This token is required when inactivate the span i.e.
+    detaching the span from the context.
+    """
+
+    span: LiveSpan
+    token: Optional[contextvars.Token] = None

venv/lib/python3.9/site-packages/mlflow/tracing/utils/warning.py

@@ -0,0 +1,45 @@
+import importlib
+import logging
+
+_logger = logging.getLogger(__name__)
+
+
+class LogDemotionFilter(logging.Filter):
+    def __init__(self, module: str, message: str):
+        super().__init__()
+        self.module = module
+        self.message = message
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        if record.name == self.module and self.message in record.getMessage():
+            record.levelno = logging.DEBUG  # Change the log level to DEBUG
+            record.levelname = "DEBUG"
+
+            # Check the log level for the logger is debug or not
+            logger = logging.getLogger(self.module)
+            return logger.isEnabledFor(logging.DEBUG)
+        return True
+
+    def __eq__(self, other):
+        if isinstance(other, LogDemotionFilter):
+            return self.module == other.module and self.message == other.message
+        return False
+
+
+def suppress_warning(module: str, message: str):
+    """
+    Convert 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.
+
+    Args:
+        module: The module name of the logger that raises the warning.
+        message: The (part of) message in the log that needs to be demoted to DEBUG level
+    """
+    try:
+        logger = getattr(importlib.import_module(module), "logger", None)
+        log_filter = LogDemotionFilter(module, message)
+        if logger and not any(f == log_filter for f in logger.filters):
+            logger.addFilter(log_filter)
+    except Exception as e:
+        _logger.debug(f"Failed to suppress the warning for {module}", exc_info=e)
+        raise