star

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