genesis-flow 1.0.0__py3-none-any.whl

mlflow/tracing/fluent.py

@@ -0,0 +1,1417 @@
+from __future__ import annotations
+
+import contextlib
+import functools
+import importlib
+import inspect
+import json
+import logging
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, Any, Callable, Generator, Literal, Optional, Union
+
+from cachetools import TTLCache
+from opentelemetry import trace as trace_api
+
+from mlflow.entities import NoOpSpan, SpanType, Trace
+from mlflow.entities.span import NO_OP_SPAN_TRACE_ID, LiveSpan, create_mlflow_span
+from mlflow.entities.span_event import SpanEvent
+from mlflow.entities.span_status import SpanStatusCode
+from mlflow.entities.trace_state import TraceState
+from mlflow.entities.trace_status import TraceStatus
+from mlflow.exceptions import MlflowException
+from mlflow.store.tracking import SEARCH_TRACES_DEFAULT_MAX_RESULTS
+from mlflow.tracing import provider
+from mlflow.tracing.client import TracingClient
+from mlflow.tracing.constant import (
+    STREAM_CHUNK_EVENT_NAME_FORMAT,
+    STREAM_CHUNK_EVENT_VALUE_KEY,
+    SpanAttributeKey,
+)
+from mlflow.tracing.provider import is_tracing_enabled, safe_set_span_in_context
+from mlflow.tracing.trace_manager import InMemoryTraceManager
+from mlflow.tracing.utils import (
+    SPANS_COLUMN_NAME,
+    TraceJSONEncoder,
+    capture_function_input_args,
+    encode_span_id,
+    exclude_immutable_tags,
+    get_otel_attribute,
+    set_chat_attributes_special_case,
+)
+from mlflow.tracing.utils.search import extract_span_inputs_outputs, traces_to_df
+from mlflow.tracing.utils.warning import request_id_backward_compatible
+from mlflow.utils import get_results_from_paginated_fn
+from mlflow.utils.annotations import experimental
+
+_logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    import pandas
+
+
+_LAST_ACTIVE_TRACE_ID_GLOBAL = None
+_LAST_ACTIVE_TRACE_ID_THREAD_LOCAL = ContextVar("last_active_trace_id", default=None)
+
+# Cache mapping between evaluation request ID to MLflow backend request ID.
+# This is necessary for evaluation harness to access generated traces during
+# evaluation using the dataset row ID (evaluation request ID).
+_EVAL_REQUEST_ID_TO_TRACE_ID = TTLCache(maxsize=10000, ttl=3600)
+
+
+def trace(
+    func: Optional[Callable[..., Any]] = None,
+    name: Optional[str] = None,
+    span_type: str = SpanType.UNKNOWN,
+    attributes: Optional[dict[str, Any]] = None,
+    output_reducer: Optional[Callable[[list[Any]], Any]] = None,
+) -> Callable[..., Any]:
+    """
+    A decorator that creates a new span for the decorated function.
+
+    When you decorate a function with this :py:func:`@mlflow.trace() <trace>` decorator,
+    a span will be created for the scope of the decorated function. The span will automatically
+    capture the input and output of the function. When it is applied to a method, it doesn't
+    capture the `self` argument. Any exception raised within the function will set the span
+    status to ``ERROR`` and detailed information such as exception message and stacktrace
+    will be recorded to the ``attributes`` field of the span.
+
+    For example, the following code will yield a span with the name ``"my_function"``,
+    capturing the input arguments ``x`` and ``y``, and the output of the function.
+
+    .. code-block:: python
+        :test:
+
+        import mlflow
+
+
+        @mlflow.trace
+        def my_function(x, y):
+            return x + y
+
+    This is equivalent to doing the following using the :py:func:`mlflow.start_span` context
+    manager, but requires less boilerplate code.
+
+    .. code-block:: python
+        :test:
+
+        import mlflow
+
+
+        def my_function(x, y):
+            return x + y
+
+
+        with mlflow.start_span("my_function") as span:
+            x = 1
+            y = 2
+            span.set_inputs({"x": x, "y": y})
+            result = my_function(x, y)
+            span.set_outputs({"output": result})
+
+
+    The @mlflow.trace decorator currently support the following types of functions:
+
+    .. list-table:: Supported Function Types
+        :widths: 20 30
+        :header-rows: 1
+
+        * - Function Type
+          - Supported
+        * - Sync
+          - ✅
+        * - Async
+          - ✅ (>= 2.16.0)
+        * - Generator
+          - ✅ (>= 2.20.2)
+        * - Async Generator
+          - ✅ (>= 2.20.2)
+        * - ClassMethod
+          - ✅ (>= 3.0.0)
+        * - StaticMethod
+          - ✅ (>= 3.0.0)
+
+    For more examples of using the @mlflow.trace decorator, including streaming/async
+    handling, see the `MLflow Tracing documentation <https://www.mlflow.org/docs/latest/tracing/api/manual-instrumentation#decorator>`_.
+
+    .. tip::
+
+        The @mlflow.trace decorator is useful when you want to trace a function defined by
+        yourself. However, you may also want to trace a function in external libraries. In
+        such case, you can use this ``mlflow.trace()`` function to directly wrap the function,
+        instead of using as the decorator. This will create the exact same span as the
+        one created by the decorator i.e. captures information from the function call.
+
+        .. code-block:: python
+            :test:
+
+            import math
+
+            import mlflow
+
+            mlflow.trace(math.factorial)(5)
+
+    Args:
+        func: The function to be decorated. Must **not** be provided when using as a decorator.
+        name: The name of the span. If not provided, the name of the function will be used.
+        span_type: The type of the span. Can be either a string or a
+            :py:class:`SpanType <mlflow.entities.SpanType>` enum value.
+        attributes: A dictionary of attributes to set on the span.
+        output_reducer: A function that reduces the outputs of the generator function into a
+            single value to be set as the span output.
+    """
+
+    def decorator(fn):
+        # Check if the function is a classmethod or staticmethod
+        is_classmethod = isinstance(fn, classmethod)
+        is_staticmethod = isinstance(fn, staticmethod)
+
+        # Extract the original function if it's a descriptor
+        original_fn = fn.__func__ if is_classmethod or is_staticmethod else fn
+
+        # Apply the appropriate wrapper to the original function
+        if inspect.isgeneratorfunction(original_fn) or inspect.isasyncgenfunction(original_fn):
+            wrapped = _wrap_generator(
+                original_fn,
+                name,
+                span_type,
+                attributes,
+                output_reducer,
+            )
+        else:
+            if output_reducer is not None:
+                raise MlflowException.invalid_parameter_value(
+                    "The output_reducer argument is only supported for generator functions."
+                )
+            wrapped = _wrap_function(original_fn, name, span_type, attributes)
+
+        # If the original was a descriptor, wrap the result back as the same type of descriptor
+        if is_classmethod:
+            return classmethod(wrapped)
+        elif is_staticmethod:
+            return staticmethod(wrapped)
+        else:
+            return wrapped
+
+    return decorator(func) if func else decorator
+
+
+def _wrap_function(
+    fn: Callable[..., Any],
+    name: Optional[str] = None,
+    span_type: str = SpanType.UNKNOWN,
+    attributes: Optional[dict[str, Any]] = None,
+) -> Callable[..., Any]:
+    class _WrappingContext:
+        # define the wrapping logic as a coroutine to avoid code duplication
+        # between sync and async cases
+        @staticmethod
+        def _wrapping_logic(fn, args, kwargs):
+            span_name = name or fn.__name__
+
+            with start_span(name=span_name, span_type=span_type, attributes=attributes) as span:
+                span.set_attribute(SpanAttributeKey.FUNCTION_NAME, fn.__name__)
+                inputs = capture_function_input_args(fn, args, kwargs)
+                span.set_inputs(inputs)
+                result = yield  # sync/async function output to be sent here
+                span.set_outputs(result)
+
+                set_chat_attributes_special_case(span, inputs=inputs, outputs=result)
+
+                try:
+                    yield result
+                except GeneratorExit:
+                    # Swallow `GeneratorExit` raised when the generator is closed
+                    pass
+
+        def __init__(self, fn, args, kwargs):
+            self.coro = self._wrapping_logic(fn, args, kwargs)
+
+        def __enter__(self):
+            next(self.coro)
+            return self.coro
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            # Since the function call occurs outside the coroutine,
+            # if an exception occurs, we need to throw it back in, so that
+            # we return control to the coro (in particular, so that the __exit__'s
+            # of start_span and OTel's use_span can execute).
+            if exc_type is not None:
+                self.coro.throw(exc_type, exc_value, traceback)
+            self.coro.close()
+
+    if inspect.iscoroutinefunction(fn):
+
+        async def wrapper(*args, **kwargs):
+            with _WrappingContext(fn, args, kwargs) as wrapping_coro:
+                return wrapping_coro.send(await fn(*args, **kwargs))
+    else:
+
+        def wrapper(*args, **kwargs):
+            with _WrappingContext(fn, args, kwargs) as wrapping_coro:
+                return wrapping_coro.send(fn(*args, **kwargs))
+
+    return _wrap_function_safe(fn, wrapper)
+
+
+def _wrap_generator(
+    fn: Callable[..., Any],
+    name: Optional[str] = None,
+    span_type: str = SpanType.UNKNOWN,
+    attributes: Optional[dict[str, Any]] = None,
+    output_reducer: Optional[Callable[[list[Any]], Any]] = None,
+) -> Callable[..., Any]:
+    """
+    Wrap a generator function to create a span.
+    Generator functions need special handling because of its lazy evaluation nature.
+    Let's say we have a generator function like this:
+    ```
+    @mlflow.trace
+    def generate_stream():
+        # B
+        for i in range(10):
+            # C
+            yield i * 2
+        # E
+
+
+    stream = generate_stream()
+    # A
+    for chunk in stream:
+        # D
+        pass
+    # F
+    ```
+    The execution order is A -> B -> C -> D -> C -> D -> ... -> E -> F.
+    The span should only be "active" at B, C, and E, namely, when the code execution
+    is inside the generator function. Otherwise it will create wrong span tree, or
+    even worse, leak span context and pollute subsequent traces.
+    """
+
+    def _start_stream_span(fn, inputs):
+        try:
+            return start_span_no_context(
+                name=name or fn.__name__,
+                parent_span=get_current_active_span(),
+                span_type=span_type,
+                attributes=attributes,
+                inputs=inputs,
+            )
+        except Exception as e:
+            _logger.debug(f"Failed to start stream span: {e}")
+            return NoOpSpan()
+
+    def _end_stream_span(
+        span: LiveSpan,
+        inputs: Optional[dict[str, Any]] = None,
+        outputs: Optional[list[Any]] = None,
+        output_reducer: Optional[Callable[[list[Any]], Any]] = None,
+        error: Optional[Exception] = None,
+    ):
+        if error:
+            span.add_event(SpanEvent.from_exception(error))
+            span.end(status=SpanStatusCode.ERROR)
+            return
+
+        if output_reducer:
+            try:
+                outputs = output_reducer(outputs)
+            except Exception as e:
+                _logger.debug(f"Failed to reduce outputs from stream: {e}")
+
+        set_chat_attributes_special_case(span, inputs=inputs, outputs=outputs)
+        span.end(outputs=outputs)
+
+    def _record_chunk_event(span: LiveSpan, chunk: Any, chunk_index: int):
+        try:
+            event = SpanEvent(
+                name=STREAM_CHUNK_EVENT_NAME_FORMAT.format(index=chunk_index),
+                # OpenTelemetry SpanEvent only support str-str key-value pairs for attributes
+                attributes={STREAM_CHUNK_EVENT_VALUE_KEY: json.dumps(chunk, cls=TraceJSONEncoder)},
+            )
+            span.add_event(event)
+        except Exception as e:
+            _logger.debug(f"Failing to record chunk event for span {span.name}: {e}")
+
+    if inspect.isgeneratorfunction(fn):
+
+        def wrapper(*args, **kwargs):
+            inputs = capture_function_input_args(fn, args, kwargs)
+            span = _start_stream_span(fn, inputs)
+            generator = fn(*args, **kwargs)
+
+            i = 0
+            outputs = []
+            while True:
+                try:
+                    # NB: Set the span to active only when the generator is running
+                    with safe_set_span_in_context(span):
+                        value = next(generator)
+                except StopIteration:
+                    break
+                except Exception as e:
+                    _end_stream_span(span, error=e)
+                    raise e
+                else:
+                    outputs.append(value)
+                    _record_chunk_event(span, value, i)
+                    yield value
+                    i += 1
+            _end_stream_span(span, inputs, outputs, output_reducer)
+    else:
+
+        async def wrapper(*args, **kwargs):
+            inputs = capture_function_input_args(fn, args, kwargs)
+            span = _start_stream_span(fn, inputs)
+            generator = fn(*args, **kwargs)
+
+            i = 0
+            outputs = []
+            while True:
+                try:
+                    with safe_set_span_in_context(span):
+                        value = await generator.__anext__()
+                except StopAsyncIteration:
+                    break
+                except Exception as e:
+                    _end_stream_span(span, error=e)
+                    raise e
+                else:
+                    outputs.append(value)
+                    _record_chunk_event(span, value, i)
+                    yield value
+                    i += 1
+            _end_stream_span(span, inputs, outputs, output_reducer)
+
+    return _wrap_function_safe(fn, wrapper)
+
+
+def _wrap_function_safe(fn: Callable[..., Any], wrapper: Callable[..., Any]) -> Callable[..., Any]:
+    wrapped = functools.wraps(fn)(wrapper)
+    # Update the signature of the wrapper to match the signature of the original (safely)
+    try:
+        wrapped.__signature__ = inspect.signature(fn)
+    except Exception:
+        pass
+    return wrapped
+
+
+@contextlib.contextmanager
+def start_span(
+    name: str = "span",
+    span_type: Optional[str] = SpanType.UNKNOWN,
+    attributes: Optional[dict[str, Any]] = None,
+) -> Generator[LiveSpan, None, None]:
+    """
+    Context manager to create a new span and start it as the current span in the context.
+
+    This context manager automatically manages the span lifecycle and parent-child relationships.
+    The span will be ended when the context manager exits. Any exception raised within the
+    context manager will set the span status to ``ERROR``, and detailed information such as
+    exception message and stacktrace will be recorded to the ``attributes`` field of the span.
+    New spans can be created within the context manager, then they will be assigned as child
+    spans.
+
+    .. code-block:: python
+        :test:
+
+        import mlflow
+
+        with mlflow.start_span("my_span") as span:
+            x = 1
+            y = 2
+            span.set_inputs({"x": x, "y": y})
+
+            z = x + y
+
+            span.set_outputs(z)
+            span.set_attribute("key", "value")
+            # do something
+
+    When this context manager is used in the top-level scope, i.e. not within another span context,
+    the span will be treated as a root span. The root span doesn't have a parent reference and
+    **the entire trace will be logged when the root span is ended**.
+
+
+    .. tip::
+
+        If you want more explicit control over the trace lifecycle, you can use
+        the `mlflow.start_span_no_context()` API. It provides lower
+        level to start spans and control the parent-child relationships explicitly.
+        However, it is generally recommended to use this context manager as long as it satisfies
+        your requirements, because it requires less boilerplate code and is less error-prone.
+
+    .. note::
+
+        The context manager doesn't propagate the span context across threads by default. see
+        `Multi Threading <https://mlflow.org/docs/latest/tracing/api/manual-instrumentation#multi-threading>`_
+        for how to propagate the span context across threads.
+
+    Args:
+        name: The name of the span.
+        span_type: The type of the span. Can be either a string or
+            a :py:class:`SpanType <mlflow.entities.SpanType>` enum value
+        attributes: A dictionary of attributes to set on the span.
+
+    Returns:
+        Yields an :py:class:`mlflow.entities.Span` that represents the created span.
+    """
+    try:
+        otel_span = provider.start_span_in_context(name)
+
+        # Create a new MLflow span and register it to the in-memory trace manager
+        request_id = get_otel_attribute(otel_span, SpanAttributeKey.REQUEST_ID)
+        mlflow_span = create_mlflow_span(otel_span, request_id, span_type)
+        attributes = dict(attributes) if attributes is not None else {}
+        mlflow_span.set_attributes(attributes)
+        InMemoryTraceManager.get_instance().register_span(mlflow_span)
+
+    except Exception:
+        _logger.debug(f"Failed to start span {name}.", exc_info=True)
+        mlflow_span = NoOpSpan()
+        yield mlflow_span
+        return
+
+    try:
+        # Setting end_on_exit = False to suppress the default span
+        # export and instead invoke MLflow span's end() method.
+        with trace_api.use_span(mlflow_span._span, end_on_exit=False):
+            yield mlflow_span
+    finally:
+        try:
+            mlflow_span.end()
+        except Exception:
+            _logger.debug(f"Failed to end span {mlflow_span.span_id}.", exc_info=True)
+
+
+def start_span_no_context(
+    name: str,
+    span_type: str = SpanType.UNKNOWN,
+    parent_span: Optional[LiveSpan] = None,
+    inputs: Optional[Any] = None,
+    attributes: Optional[dict[str, str]] = None,
+    tags: Optional[dict[str, str]] = None,
+    experiment_id: Optional[str] = None,
+    start_time_ns: Optional[int] = None,
+) -> LiveSpan:
+    """
+    Start a span without attaching it to the global tracing context.
+
+    This is useful when you want to create a span without automatically linking
+    with a parent span and instead manually manage the parent-child relationships.
+
+    The span started with this function must be ended manually using the
+    `end()` method of the span object.
+
+    Args:
+        name: The name of the span.
+        span_type: The type of the span. Can be either a string or
+            a :py:class:`SpanType <mlflow.entities.SpanType>` enum value
+        parent_span: The parent span to link with. If None, the span will be treated as a root span.
+        inputs: The input data for the span.
+        attributes: A dictionary of attributes to set on the span.
+        tags: A dictionary of tags to set on the trace.
+        experiment_id: The experiment ID to associate with the trace. If not provided,
+            the current active experiment will be used.
+        start_time_ns: The start time of the span in nanoseconds. If not provided,
+            the current time will be used.
+
+    Returns:
+        A :py:class:`mlflow.entities.Span` that represents the created span.
+
+    Example:
+        .. code-block:: python
+            :test:
+
+            import mlflow
+
+            root_span = mlflow.start_span_no_context("my_trace")
+
+            # Create a child span
+            child_span = mlflow.start_span_no_context(
+                "child_span",
+                # Manually specify the parent span
+                parent_span=root_span,
+            )
+            # Do something...
+            child_span.end()
+
+            root_span.end()
+
+    """
+    # If parent span is no-op span, the child should also be no-op too
+    if parent_span and parent_span.trace_id == NO_OP_SPAN_TRACE_ID:
+        return NoOpSpan()
+
+    try:
+        # Create new trace and a root span
+        # Once OTel span is created, SpanProcessor.on_start is invoked
+        # TraceInfo is created and logged into backend store inside on_start method
+        otel_span = provider.start_detached_span(
+            name,
+            parent=parent_span._span if parent_span else None,
+            start_time_ns=start_time_ns,
+            experiment_id=experiment_id,
+        )
+
+        if parent_span:
+            trace_id = parent_span.trace_id
+        else:
+            trace_id = get_otel_attribute(otel_span, SpanAttributeKey.REQUEST_ID)
+
+        mlflow_span = create_mlflow_span(otel_span, trace_id, span_type)
+
+        # # If the span is a no-op span i.e. tracing is disabled, do nothing
+        if isinstance(mlflow_span, NoOpSpan):
+            return mlflow_span
+
+        if inputs is not None:
+            mlflow_span.set_inputs(inputs)
+        mlflow_span.set_attributes(attributes or {})
+
+        trace_manager = InMemoryTraceManager.get_instance()
+        if tags := exclude_immutable_tags(tags or {}):
+            # Update trace tags for trace in in-memory trace manager
+            with trace_manager.get_trace(trace_id) as trace:
+                trace.info.tags.update(tags)
+
+        # Register new span in the in-memory trace manager
+        trace_manager.register_span(mlflow_span)
+
+        return mlflow_span
+    except Exception as e:
+        _logger.warning(
+            f"Failed to start span {name}: {e}. For full traceback, set logging level to debug.",
+            exc_info=_logger.isEnabledFor(logging.DEBUG),
+        )
+    return NoOpSpan()
+
+
+@request_id_backward_compatible
+def get_trace(trace_id: str) -> Optional[Trace]:
+    """
+    Get a trace by the given request ID if it exists.
+
+    This function retrieves the trace from the in-memory buffer first, and if it doesn't exist,
+    it fetches the trace from the tracking store. If the trace is not found in the tracking store,
+    it returns None.
+
+    Args:
+        trace_id: The ID of the trace.
+
+
+    .. code-block:: python
+        :test:
+
+        import mlflow
+
+
+        with mlflow.start_span(name="span") as span:
+            span.set_attribute("key", "value")
+
+        trace = mlflow.get_trace(span.trace_id)
+        print(trace)
+
+
+    Returns:
+        A :py:class:`mlflow.entities.Trace` objects with the given request ID.
+    """
+    # Special handling for evaluation request ID.
+    trace_id = _EVAL_REQUEST_ID_TO_TRACE_ID.get(trace_id) or trace_id
+
+    try:
+        return TracingClient().get_trace(trace_id)
+    except MlflowException as e:
+        _logger.warning(
+            f"Failed to get trace from the tracking store: {e}"
+            "For full traceback, set logging level to debug.",
+            exc_info=_logger.isEnabledFor(logging.DEBUG),
+        )
+        return None
+
+
+def search_traces(
+    experiment_ids: Optional[list[str]] = None,
+    filter_string: Optional[str] = None,
+    max_results: Optional[int] = None,
+    order_by: Optional[list[str]] = None,
+    extract_fields: Optional[list[str]] = None,
+    run_id: Optional[str] = None,
+    return_type: Optional[Literal["pandas", "list"]] = None,
+    model_id: Optional[str] = None,
+    sql_warehouse_id: Optional[str] = None,
+    include_spans: bool = True,
+) -> Union["pandas.DataFrame", list[Trace]]:
+    """
+    Return traces that match the given list of search expressions within the experiments.
+
+    .. note::
+
+        If expected number of search results is large, consider using the
+        `MlflowClient.search_traces` API directly to paginate through the results. This
+        function returns all results in memory and may not be suitable for large result sets.
+
+    Args:
+        experiment_ids: List of experiment ids to scope the search. If not provided, the search
+            will be performed across the current active experiment.
+        filter_string: A search filter string.
+        max_results: Maximum number of traces desired. If None, all traces matching the search
+            expressions will be returned.
+        order_by: List of order_by clauses.
+        extract_fields: Specify fields to extract from traces using the format
+            ``"span_name.[inputs|outputs].field_name"`` or ``"span_name.[inputs|outputs]"``.
+
+            .. note::
+
+                This parameter is only supported when the return type is set to "pandas".
+
+            For instance, ``"predict.outputs.result"`` retrieves the output ``"result"`` field from
+            a span named ``"predict"``, while ``"predict.outputs"`` fetches the entire outputs
+            dictionary, including keys ``"result"`` and ``"explanation"``.
+
+            By default, no fields are extracted into the DataFrame columns. When multiple
+            fields are specified, each is extracted as its own column. If an invalid field
+            string is provided, the function silently returns without adding that field's column.
+            The supported fields are limited to ``"inputs"`` and ``"outputs"`` of spans. If the
+            span name or field name contains a dot it must be enclosed in backticks. For example:
+
+            .. code-block:: python
+
+                # span name contains a dot
+                extract_fields = ["`span.name`.inputs.field"]
+
+                # field name contains a dot
+                extract_fields = ["span.inputs.`field.name`"]
+
+                # span name and field name contain a dot
+                extract_fields = ["`span.name`.inputs.`field.name`"]
+
+        run_id: A run id to scope the search. When a trace is created under an active run,
+            it will be associated with the run and you can filter on the run id to retrieve the
+            trace. See the example below for how to filter traces by run id.
+
+        return_type: The type of the return value. The following return types are supported. If
+            the pandas library is installed, the default return type is "pandas". Otherwise, the
+            default return type is "list".
+
+            - `"pandas"`: Returns a Pandas DataFrame containing information about traces
+                where each row represents a single trace and each column represents a field of the
+                trace e.g. trace_id, spans, etc.
+            - `"list"`: Returns a list of :py:class:`Trace <mlflow.entities.Trace>` objects.
+
+        model_id: If specified, search traces associated with the given model ID.
+        sql_warehouse_id: Only used in Databricks. The ID of the SQL warehouse to use for
+            searching traces in inference tables.
+
+        include_spans: If ``True``, include spans in the returned traces. Otherwise, only
+            the trace metadata is returned, e.g., trace ID, start time, end time, etc,
+            without any spans. Default to ``True``.
+
+    Returns:
+        Traces that satisfy the search expressions. Either as a list of
+        :py:class:`Trace <mlflow.entities.Trace>` objects or as a Pandas DataFrame,
+        depending on the value of the `return_type` parameter.
+
+    .. code-block:: python
+        :test:
+        :caption: Search traces with extract_fields
+
+        import mlflow
+
+        with mlflow.start_span(name="span1") as span:
+            span.set_inputs({"a": 1, "b": 2})
+            span.set_outputs({"c": 3, "d": 4})
+
+        mlflow.search_traces(
+            extract_fields=["span1.inputs", "span1.outputs", "span1.outputs.c"],
+            return_type="pandas",
+        )
+
+
+    .. code-block:: python
+        :test:
+        :caption: Search traces with extract_fields and non-dictionary span inputs and outputs
+
+        import mlflow
+
+        with mlflow.start_span(name="non_dict_span") as span:
+            span.set_inputs(["a", "b"])
+            span.set_outputs([1, 2, 3])
+
+        mlflow.search_traces(
+            extract_fields=["non_dict_span.inputs", "non_dict_span.outputs"],
+        )
+
+    .. code-block:: python
+        :test:
+        :caption: Search traces by run ID and return as a list of Trace objects
+
+        import mlflow
+
+
+        @mlflow.trace
+        def traced_func(x):
+            return x + 1
+
+
+        with mlflow.start_run() as run:
+            traced_func(1)
+
+        mlflow.search_traces(run_id=run.info.run_id, return_type="list")
+
+    """
+    from mlflow.tracking.fluent import _get_experiment_id
+
+    # Default to "pandas" only if the pandas library is installed
+    if return_type is None:
+        try:
+            import pandas  # noqa: F401
+
+            return_type = "pandas"
+        except ImportError:
+            return_type = "list"
+
+    if return_type not in ["pandas", "list"]:
+        raise MlflowException.invalid_parameter_value(
+            f"Invalid return type: {return_type}. Return type must be either 'pandas' or 'list'."
+        )
+    elif return_type == "list" and extract_fields:
+        raise MlflowException.invalid_parameter_value(
+            "The `extract_fields` parameter is only supported when return type is set to 'pandas'."
+        )
+    elif return_type == "pandas":
+        # Check if pandas is installed early to avoid unnecessary computation
+        if importlib.util.find_spec("pandas") is None:
+            raise MlflowException(
+                message=(
+                    "The `pandas` library is not installed. Please install `pandas` to use"
+                    " the `return_type='pandas'` option, or set `return_type='list'`."
+                ),
+            )
+
+    if not experiment_ids:
+        if experiment_id := _get_experiment_id():
+            experiment_ids = [experiment_id]
+        else:
+            raise MlflowException(
+                "No active experiment found. Set an experiment using `mlflow.set_experiment`, "
+                "or specify the list of experiment IDs in the `experiment_ids` parameter."
+            )
+
+    def pagination_wrapper_func(number_to_get, next_page_token):
+        return TracingClient().search_traces(
+            experiment_ids=experiment_ids,
+            run_id=run_id,
+            max_results=number_to_get,
+            filter_string=filter_string,
+            order_by=order_by,
+            page_token=next_page_token,
+            model_id=model_id,
+            sql_warehouse_id=sql_warehouse_id,
+            include_spans=include_spans,
+        )
+
+    results = get_results_from_paginated_fn(
+        pagination_wrapper_func,
+        max_results_per_page=SEARCH_TRACES_DEFAULT_MAX_RESULTS,
+        max_results=max_results,
+    )
+
+    if return_type == "pandas":
+        results = traces_to_df(results)
+        if extract_fields:
+            results = extract_span_inputs_outputs(
+                traces=results,
+                fields=extract_fields,
+                col_name=SPANS_COLUMN_NAME,
+            )
+
+    return results
+
+
+def get_current_active_span() -> Optional[LiveSpan]:
+    """
+    Get the current active span in the global context.
+
+    .. attention::
+
+        This only works when the span is created with fluent APIs like `@mlflow.trace` or
+        `with mlflow.start_span`. If a span is created with the
+        `mlflow.start_span_no_context` APIs, it won't be
+        attached to the global context so this function will not return it.
+
+
+    .. code-block:: python
+        :test:
+
+        import mlflow
+
+
+        @mlflow.trace
+        def f():
+            span = mlflow.get_current_active_span()
+            span.set_attribute("key", "value")
+            return 0
+
+
+        f()
+
+    Returns:
+        The current active span if exists, otherwise None.
+    """
+    otel_span = trace_api.get_current_span()
+    # NonRecordingSpan is returned if a tracer is not instantiated.
+    if otel_span is None or isinstance(otel_span, trace_api.NonRecordingSpan):
+        return None
+
+    trace_manager = InMemoryTraceManager.get_instance()
+    request_id = json.loads(otel_span.attributes.get(SpanAttributeKey.REQUEST_ID))
+    return trace_manager.get_span_from_id(request_id, encode_span_id(otel_span.context.span_id))
+
+
+def get_active_trace_id() -> Optional[str]:
+    """
+    Get the active trace ID in the current process.
+
+    This function is thread-safe.
+
+    Example:
+
+    .. code-block:: python
+        :test:
+
+        import mlflow
+
+
+        @mlflow.trace
+        def f():
+            trace_id = mlflow.get_active_trace_id()
+            print(trace_id)
+
+
+        f()
+
+    Returns:
+        The ID of the current active trace if exists, otherwise None.
+    """
+    active_span = get_current_active_span()
+    if active_span:
+        return active_span.trace_id
+    return None
+
+
+def get_last_active_trace_id(thread_local: bool = False) -> Optional[str]:
+    """
+    Get the **LAST** active trace in the same process if exists.
+
+    .. warning::
+
+        This function is not thread-safe by default, returns the last active trace in
+        the same process. If you want to get the last active trace in the current thread,
+        set the `thread_local` parameter to True.
+
+    Args:
+
+        thread_local: If True, returns the last active trace in the current thread. Otherwise,
+            returns the last active trace in the same process. Default is False.
+
+    Returns:
+        The ID of the last active trace if exists, otherwise None.
+
+    .. code-block:: python
+        :test:
+
+        import mlflow
+
+
+        @mlflow.trace
+        def f():
+            pass
+
+
+        f()
+
+        trace_id = mlflow.get_last_active_trace_id()
+
+        # Set a tag on the trace
+        mlflow.set_trace_tag(trace_id, "key", "value")
+
+        # Get the full trace object
+        trace = mlflow.get_trace(trace_id)
+    """
+    return (
+        _LAST_ACTIVE_TRACE_ID_THREAD_LOCAL.get() if thread_local else _LAST_ACTIVE_TRACE_ID_GLOBAL
+    )
+
+
+def _set_last_active_trace_id(trace_id: str):
+    """Internal function to set the last active trace ID."""
+    global _LAST_ACTIVE_TRACE_ID_GLOBAL
+    _LAST_ACTIVE_TRACE_ID_GLOBAL = trace_id
+    _LAST_ACTIVE_TRACE_ID_THREAD_LOCAL.set(trace_id)
+
+
+def update_current_trace(
+    tags: Optional[dict[str, str]] = None,
+    metadata: Optional[dict[str, str]] = None,
+    client_request_id: Optional[str] = None,
+    request_preview: Optional[str] = None,
+    response_preview: Optional[str] = None,
+    state: Optional[Union[TraceState, str]] = None,
+):
+    """
+    Update the current active trace with the given options.
+
+    Args:
+        tags: A dictionary of tags to update the trace with Tags are designed for mutable values,
+            that can be updated after the trace is created via MLflow UI or API.
+        metadata: A dictionary of metadata to update the trace with. Metadata cannot be updated
+            once the trace is logged. It is suitable for recording immutable values like the
+            git hash of the application version that produced the trace.
+        client_request_id: Client supplied request ID to associate with the trace. This is
+            useful for linking the trace back to a specific request in your application or
+            external system. If None, the client request ID is not updated.
+        request_preview: A preview of the request to be shown in the Trace list view in the UI.
+            By default, MLflow will truncate the trace request naively by limiting the length.
+            This parameter allows you to specify a custom preview string.
+        response_preview: A preview of the response to be shown in the Trace list view in the UI.
+            By default, MLflow will truncate the trace response naively by limiting the length.
+            This parameter allows you to specify a custom preview string.
+        state: The state to set on the trace. Can be a TraceState enum value or string.
+            Only "OK" and "ERROR" are allowed. This overrides the overall trace state without
+            affecting the status of the current span.
+
+    Example:
+
+        You can use this function either within a function decorated with ``@mlflow.trace`` or
+        within the scope of the `with mlflow.start_span` context manager. If there is no active
+        trace found, this function will raise an exception.
+
+        Using within a function decorated with `@mlflow.trace`:
+
+        .. code-block:: python
+
+            @mlflow.trace
+            def my_func(x):
+                mlflow.update_current_trace(tags={"fruit": "apple"}, client_request_id="req-12345")
+                return x + 1
+
+        Using within the ``with mlflow.start_span`` context manager:
+
+        .. code-block:: python
+
+            with mlflow.start_span("span"):
+                mlflow.update_current_trace(tags={"fruit": "apple"}, client_request_id="req-12345")
+
+        Updating source information of the trace. These keys are reserved ones and MLflow populate
+        them from environment information by default. You can override them if needed. Please refer
+        to the MLflow Tracing documentation for the full list of reserved metadata keys.
+
+        .. code-block:: python
+
+            mlflow.update_current_trace(
+                metadata={
+                    "mlflow.trace.session": "session-4f855da00427",
+                    "mlflow.trace.user": "user-id-cc156f29bcfb",
+                    "mlflow.source.name": "inference.py",
+                    "mlflow.source.git.commit": "1234567890",
+                    "mlflow.source.git.repoURL": "https://github.com/mlflow/mlflow",
+                },
+            )
+
+        Updating request preview:
+
+        .. code-block:: python
+
+            import mlflow
+            import openai
+
+
+            @mlflow.trace
+            def predict(messages: list[dict]) -> str:
+                # Customize the request preview to show the first and last messages
+                custom_preview = f"{messages[0]['content'][:10]} ... {messages[-1]['content'][:10]}"
+                mlflow.update_current_trace(request_preview=custom_preview)
+
+                # Call the model
+                response = openai.chat.completions.create(
+                    model="o4-mini",
+                    messages=messages,
+                )
+
+                return response.choices[0].message.content
+
+
+            messages = [
+                {"role": "user", "content": "Hi, how are you?"},
+                {"role": "assistant", "content": "I'm good, thank you!"},
+                {"role": "user", "content": "What's your name?"},
+                # ... (long message history)
+                {"role": "assistant", "content": "Bye!"},
+            ]
+            predict(messages)
+
+            # The request preview rendered in the UI will be:
+            #     "Hi, how are you? ... Bye!"
+
+    """
+    active_span = get_current_active_span()
+
+    if not active_span:
+        _logger.warning(
+            "No active trace found. Please create a span using `mlflow.start_span` or "
+            "`@mlflow.trace` before calling `mlflow.update_current_trace`.",
+        )
+        return
+
+    def _warn_non_string_values(d: dict[str, Any], field_name: str):
+        non_string_items = {k: v for k, v in d.items() if not isinstance(v, str)}
+        if non_string_items:
+            _logger.warning(
+                f"Found non-string values in {field_name}. Non-string values in {field_name} will "
+                f"automatically be stringified when the trace is logged. Non-string items: "
+                f"{non_string_items}"
+            )
+
+    _warn_non_string_values(tags or {}, "tags")
+    _warn_non_string_values(metadata or {}, "metadata")
+
+    # Update tags and client request ID for the trace stored in-memory rather than directly
+    # updating the backend store. The in-memory trace will be exported when it is ended.
+    # By doing this, we can avoid unnecessary server requests for each tag update.
+    if request_preview is not None and not isinstance(request_preview, str):
+        raise MlflowException.invalid_parameter_value(
+            "The `request_preview` parameter must be a string."
+        )
+    if response_preview is not None and not isinstance(response_preview, str):
+        raise MlflowException.invalid_parameter_value(
+            "The `response_preview` parameter must be a string."
+        )
+
+    # Update trace info for the trace stored in-memory rather than directly updating the
+    # backend store. The in-memory trace will be exported when it is ended. By doing
+    # this, we can avoid unnecessary server requests for each tag update.
+    request_id = active_span.request_id
+    with InMemoryTraceManager.get_instance().get_trace(request_id) as trace:
+        if request_preview:
+            trace.info.request_preview = request_preview
+        if response_preview:
+            trace.info.response_preview = response_preview
+        if state is not None:
+
+            def _invalid_state_error(value):
+                return MlflowException.invalid_parameter_value(
+                    f"State must be either 'OK' or 'ERROR', but got '{value}'."
+                )
+
+            if state not in (TraceState.OK, TraceState.ERROR):
+                raise _invalid_state_error(state)
+
+            trace.info.state = TraceState(state) if isinstance(state, str) else state
+
+        trace.info.tags.update(tags or {})
+        trace.info.trace_metadata.update(metadata or {})
+        if client_request_id is not None:
+            trace.info.client_request_id = str(client_request_id)
+
+
+@request_id_backward_compatible
+def set_trace_tag(trace_id: str, key: str, value: str):
+    """
+    Set a tag on the trace with the given trace ID.
+
+    The trace can be an active one or the one that has already ended and recorded in the
+    backend. Below is an example of setting a tag on an active trace. You can replace the
+    ``trace_id`` parameter to set a tag on an already ended trace.
+
+    .. code-block:: python
+        :test:
+
+        import mlflow
+
+        with mlflow.start_span(name="span") as span:
+            mlflow.set_trace_tag(span.trace_id, "key", "value")
+
+    Args:
+        trace_id: The ID of the trace to set the tag on.
+        key: The string key of the tag. Must be at most 250 characters long, otherwise
+            it will be truncated when stored.
+        value: The string value of the tag. Must be at most 250 characters long, otherwise
+            it will be truncated when stored.
+    """
+    TracingClient().set_trace_tag(trace_id, key, value)
+
+
+@request_id_backward_compatible
+def delete_trace_tag(trace_id: str, key: str) -> None:
+    """
+    Delete a tag on the trace with the given trace ID.
+
+    The trace can be an active one or the one that has already ended and recorded in the
+    backend. Below is an example of deleting a tag on an active trace. You can replace the
+    ``trace_id`` parameter to delete a tag on an already ended trace.
+
+    .. code-block:: python
+        :test:
+
+        import mlflow
+
+        with mlflow.start_span("my_span") as span:
+            mlflow.set_trace_tag(span.trace_id, "key", "value")
+            mlflow.delete_trace_tag(span.trace_id, "key")
+
+    Args:
+        trace_id: The ID of the trace to delete the tag from.
+        key: The string key of the tag. Must be at most 250 characters long, otherwise
+            it will be truncated when stored.
+    """
+    TracingClient().delete_trace_tag(trace_id, key)
+
+
+@experimental(version="2.17.0")
+def add_trace(trace: Union[Trace, dict[str, Any]], target: Optional[LiveSpan] = None):
+    """
+    Add a completed trace object into another trace.
+
+    This is particularly useful when you call a remote service instrumented by
+    MLflow Tracing. By using this function, you can merge the trace from the remote
+    service into the current active local trace, so that you can see the full
+    trace including what happens inside the remote service call.
+
+    The following example demonstrates how to use this function to merge a trace from a remote
+    service to the current active trace in the function.
+
+    .. code-block:: python
+
+        @mlflow.trace(name="predict")
+        def predict(input):
+            # Call a remote service that returns a trace in the response
+            resp = requests.get("https://your-service-endpoint", ...)
+
+            # Extract the trace from the response
+            trace_json = resp.json().get("trace")
+
+            # Use the remote trace as a part of the current active trace.
+            # It will be merged under the span "predict" and exported together when it is ended.
+            mlflow.add_trace(trace_json)
+
+    If you have a specific target span to merge the trace under, you can pass the target span
+
+    .. code-block:: python
+
+        def predict(input):
+            # Create a local span
+            with mlflow.start_span(name="predict") as span:
+                resp = requests.get("https://your-service-endpoint", ...)
+                trace_json = resp.json().get("trace")
+
+                # Merge the remote trace under the span created above
+                mlflow.add_trace(trace_json, target=span)
+
+    Args:
+        trace: A :py:class:`Trace <mlflow.entities.Trace>` object or a dictionary representation
+            of the trace. The trace **must** be already completed i.e. no further updates should
+            be made to it. Otherwise, this function will raise an exception.
+
+            .. attention:
+
+                The spans in the trace must be ordered in a way that the parent span comes
+                before its children. If the spans are not ordered correctly, this function
+                will raise an exception.
+
+        target: The target span to merge the given trace.
+
+            - If provided, the trace will be merged under the target span.
+            - If not provided, the trace will be merged under the current active span.
+            - If not provided and there is no active span, a new span named "Remote Trace <...>"
+              will be created and the trace will be merged under it.
+    """
+    if not is_tracing_enabled():
+        _logger.debug("Tracing is disabled. Skipping add_trace.")
+        return
+
+    if isinstance(trace, dict):
+        try:
+            trace = Trace.from_dict(trace)
+        except Exception as e:
+            raise MlflowException.invalid_parameter_value(
+                "Failed to load a trace object from the given dictionary. Please ensure the "
+                f"dictionary is in the correct MLflow Trace format. Error: {e}",
+            )
+    elif not isinstance(trace, Trace):
+        raise MlflowException.invalid_parameter_value(
+            f"Invalid trace object: {type(trace)}. Please provide a valid MLflow Trace object "
+            "to use it as a remote trace. You can create a Trace object from its json format by "
+            "using the Trace.from_dict() method."
+        )
+
+    if trace.info.status not in TraceStatus.end_statuses():
+        raise MlflowException.invalid_parameter_value(
+            "The trace must be ended before adding it to another trace. "
+            f"Current status: {trace.info.status}.",
+        )
+
+    if target_span := target or get_current_active_span():
+        _merge_trace(
+            trace=trace,
+            target_trace_id=target_span.trace_id,
+            target_parent_span_id=target_span.span_id,
+        )
+    else:
+        # If there is no target span, create a new root span named "Remote Trace <...>"
+        # and put the remote trace under it. This design aims to keep the trace export
+        # logic simpler and consistent, rather than directly exporting the remote trace.
+        remote_root_span = trace.data.spans[0]
+        span = start_span_no_context(
+            name=f"Remote Trace <{remote_root_span.name}>",
+            inputs=remote_root_span.inputs,
+            attributes={
+                # Exclude request ID attribute not to reuse same request ID
+                k: v
+                for k, v in remote_root_span.attributes.items()
+                if k != SpanAttributeKey.REQUEST_ID
+            },
+            start_time_ns=remote_root_span.start_time_ns,
+        )
+        _merge_trace(
+            trace=trace,
+            target_trace_id=span.trace_id,
+            target_parent_span_id=span.span_id,
+        )
+        span.end(
+            status=trace.info.status,
+            outputs=remote_root_span.outputs,
+            end_time_ns=remote_root_span.end_time_ns,
+        )
+
+
+@experimental(version="2.21.0")
+def log_trace(
+    name: str = "Task",
+    request: Optional[Any] = None,
+    response: Optional[Any] = None,
+    intermediate_outputs: Optional[dict[str, Any]] = None,
+    attributes: Optional[dict[str, Any]] = None,
+    tags: Optional[dict[str, str]] = None,
+    start_time_ms: Optional[int] = None,
+    execution_time_ms: Optional[int] = None,
+) -> str:
+    """
+    Create a trace with a single root span.
+    This API is useful when you want to log an arbitrary (request, response) pair
+    without structured OpenTelemetry spans. The trace is linked to the active experiment.
+
+    Args:
+        name: The name of the trace (and the root span). Default to "Task".
+        request: Input data for the entire trace. This is also set on the root span of the trace.
+        response: Output data for the entire trace. This is also set on the root span of the trace.
+        intermediate_outputs: A dictionary of intermediate outputs produced by the model or agent
+            while handling the request. Keys are the names of the outputs,
+            and values are the outputs themselves. Values must be JSON-serializable.
+        attributes: A dictionary of attributes to set on the root span of the trace.
+        tags: A dictionary of tags to set on the trace.
+        start_time_ms: The start time of the trace in milliseconds since the UNIX epoch.
+            When not specified, current time is used for start and end time of the trace.
+        execution_time_ms: The execution time of the trace in milliseconds since the UNIX epoch.
+
+    Returns:
+        The ID of the logged trace.
+
+    Example:
+
+    .. code-block:: python
+        :test:
+
+        import time
+        import mlflow
+
+        trace_id = mlflow.log_trace(
+            request="Does mlflow support tracing?",
+            response="Yes",
+            intermediate_outputs={
+                "retrieved_documents": ["mlflow documentation"],
+                "system_prompt": ["answer the question with yes or no"],
+            },
+            start_time_ms=int(time.time() * 1000),
+            execution_time_ms=5129,
+        )
+        trace = mlflow.get_trace(trace_id)
+
+        print(trace.data.intermediate_outputs)
+    """
+    if intermediate_outputs:
+        if attributes:
+            attributes.update(SpanAttributeKey.INTERMEDIATE_OUTPUTS, intermediate_outputs)
+        else:
+            attributes = {SpanAttributeKey.INTERMEDIATE_OUTPUTS: intermediate_outputs}
+
+    span = start_span_no_context(
+        name=name,
+        inputs=request,
+        attributes=attributes,
+        tags=tags,
+        start_time_ns=start_time_ms * 1000000 if start_time_ms else None,
+    )
+    span.end(
+        outputs=response,
+        end_time_ns=(start_time_ms + execution_time_ms) * 1000000
+        if start_time_ms and execution_time_ms
+        else None,
+    )
+
+    return span.trace_id
+
+
+def _merge_trace(
+    trace: Trace,
+    target_trace_id: str,
+    target_parent_span_id: str,
+):
+    """
+    Merge the given trace object under an existing trace in the in-memory trace registry.
+
+    Args:
+        trace: The trace object to be merged.
+        target_trace_id: The ID of the parent trace.
+        target_parent_span_id: The parent span ID, under which the child trace should be merged.
+    """
+    trace_manager = InMemoryTraceManager.get_instance()
+
+    # The merged trace should have the same trace ID as the parent trace.
+    with trace_manager.get_trace(target_trace_id) as parent_trace:
+        if not parent_trace:
+            _logger.warning(f"Parent trace with ID {target_trace_id} not found. Skipping merge.")
+            return
+
+        new_trace_id = parent_trace.span_dict[target_parent_span_id]._trace_id
+
+    for span in trace.data.spans:
+        parent_span_id = span.parent_id or target_parent_span_id
+
+        # NB: We clone span one by one in the order it was saved in the original trace. This
+        # works upon the assumption that the parent span always comes before its children.
+        # This is guaranteed in current implementation, but if it changes in the future,
+        # we have to traverse the tree to determine the order.
+        if not trace_manager.get_span_from_id(target_trace_id, parent_span_id):
+            raise MlflowException.invalid_parameter_value(
+                f"Span with ID {parent_span_id} not found. Please make sure the "
+                "spans in the trace are ordered correctly i.e. the parent span comes before "
+                "its children."
+            )
+
+        cloned_span = LiveSpan.from_immutable_span(
+            span=span,
+            parent_span_id=parent_span_id,
+            trace_id=target_trace_id,
+            otel_trace_id=new_trace_id,
+        )
+        trace_manager.register_span(cloned_span)
+
+    # Merge the tags and metadata from the child trace to the parent trace.
+    with trace_manager.get_trace(target_trace_id) as parent_trace:
+        # Order of merging is important to ensure the parent trace's metadata is
+        # not overwritten by the child trace's metadata if they have the same key.
+        parent_trace.info.tags = {**trace.info.tags, **parent_trace.info.tags}
+        parent_trace.info.trace_metadata = {
+            **trace.info.request_metadata,
+            **parent_trace.info.trace_metadata,
+        }