wehub-resource-sync
1911 lines
72 KiB
Python
1911 lines
72 KiB
Python
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import functools
|
|
import importlib
|
|
import inspect
|
|
import json
|
|
import logging
|
|
import os
|
|
import warnings
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from contextvars import ContextVar
|
|
from typing import TYPE_CHECKING, Any, Callable, Generator, Literal, ParamSpec, TypeVar, overload
|
|
|
|
from cachetools import TTLCache
|
|
from opentelemetry import trace as trace_api
|
|
|
|
from mlflow.entities import Link, NoOpSpan, Session, SpanLogLevel, 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_location import TraceLocationBase
|
|
from mlflow.entities.trace_state import TraceState
|
|
from mlflow.entities.trace_status import TraceStatus
|
|
from mlflow.environment_variables import MLFLOW_SEARCH_TRACES_MAX_THREADS
|
|
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,
|
|
TraceMetadataKey,
|
|
)
|
|
from mlflow.tracing.context import _USER_TRACE_CONTEXT
|
|
from mlflow.tracing.provider import (
|
|
_get_trace_exporter,
|
|
get_current_otel_span,
|
|
is_tracing_enabled,
|
|
safe_set_span_in_context,
|
|
with_active_span,
|
|
)
|
|
from mlflow.tracing.sampling import _SAMPLING_RATIO_OVERRIDE
|
|
from mlflow.tracing.trace_manager import InMemoryTraceManager
|
|
from mlflow.tracing.utils import (
|
|
TraceJSONEncoder,
|
|
capture_function_input_args,
|
|
encode_span_id,
|
|
exclude_immutable_tags,
|
|
get_otel_attribute,
|
|
)
|
|
from mlflow.tracing.utils.search import traces_to_df
|
|
from mlflow.utils import get_results_from_paginated_fn
|
|
from mlflow.utils.annotations import deprecated, deprecated_parameter, experimental
|
|
from mlflow.utils.validation import _validate_list_param
|
|
|
|
_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)
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def _set_sampling_ratio_override(sampling_ratio_override: float | None):
|
|
"""Context manager to set the sampling ratio override for the OTel sampler."""
|
|
if sampling_ratio_override is not None:
|
|
token = _SAMPLING_RATIO_OVERRIDE.set(sampling_ratio_override)
|
|
try:
|
|
yield
|
|
finally:
|
|
_SAMPLING_RATIO_OVERRIDE.reset(token)
|
|
else:
|
|
yield
|
|
|
|
|
|
# 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)
|
|
|
|
_P = ParamSpec("_P")
|
|
_R = TypeVar("_R")
|
|
|
|
|
|
@overload
|
|
def trace(
|
|
func: Callable[_P, _R],
|
|
name: str | None = None,
|
|
span_type: str = SpanType.UNKNOWN,
|
|
attributes: dict[str, Any] | None = None,
|
|
output_reducer: Callable[[list[Any]], Any] | None = None,
|
|
trace_destination: TraceLocationBase | None = None,
|
|
sampling_ratio_override: float | None = None,
|
|
log_level: SpanLogLevel | str | None = None,
|
|
links: list[Link] | None = None,
|
|
) -> Callable[_P, _R]: ...
|
|
|
|
|
|
@overload
|
|
def trace(
|
|
func: None = None,
|
|
name: str | None = None,
|
|
span_type: str = SpanType.UNKNOWN,
|
|
attributes: dict[str, Any] | None = None,
|
|
output_reducer: Callable[[list[Any]], Any] | None = None,
|
|
trace_destination: TraceLocationBase | None = None,
|
|
sampling_ratio_override: float | None = None,
|
|
log_level: SpanLogLevel | str | None = None,
|
|
links: list[Link] | None = None,
|
|
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: ...
|
|
|
|
|
|
def trace(
|
|
func: Callable[..., Any] | None = None,
|
|
name: str | None = None,
|
|
span_type: str = SpanType.UNKNOWN,
|
|
attributes: dict[str, Any] | None = None,
|
|
output_reducer: Callable[[list[Any]], Any] | None = None,
|
|
trace_destination: TraceLocationBase | None = None,
|
|
sampling_ratio_override: float | None = None,
|
|
log_level: SpanLogLevel | str | None = None,
|
|
links: list[Link] | None = 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.
|
|
trace_destination: The destination to log the trace to, such as MLflow Experiment. If
|
|
not provided, the destination will be an active MLflow experiment or an destination
|
|
set by the :py:func:`mlflow.tracing.set_destination` function. This parameter
|
|
should only be used for root span and setting this for non-root spans will be
|
|
ignored with a warning.
|
|
sampling_ratio_override: The sampling ratio override for this specific function. Must be
|
|
between 0.0 and 1.0. If provided, this overrides the global
|
|
``MLFLOW_TRACE_SAMPLING_RATIO`` setting. A value of 1.0 means all traces are sampled,
|
|
0.5 means 50% are sampled, and 0.0 means no traces are sampled. If not provided (None),
|
|
the global sampling ratio is used. Note: This only applies to root spans; nested calls
|
|
always trace if the parent is traced.
|
|
log_level: Optional severity level to attach to the span. Accepts a
|
|
:py:class:`SpanLogLevel <mlflow.entities.SpanLogLevel>` or its name
|
|
(e.g. ``"INFO"``, ``"DEBUG"``). If not provided, the span level is
|
|
resolved from the span type at end time.
|
|
links: A list of :py:class:`Link <mlflow.entities.Link>` objects to associate with the span.
|
|
"""
|
|
|
|
# Validate sampling_ratio_override
|
|
if sampling_ratio_override is not None and not (0.0 <= sampling_ratio_override <= 1.0):
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"sampling_ratio_override must be between 0.0 and 1.0, got {sampling_ratio_override}"
|
|
)
|
|
|
|
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,
|
|
trace_destination,
|
|
sampling_ratio_override,
|
|
log_level,
|
|
links,
|
|
)
|
|
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,
|
|
trace_destination,
|
|
sampling_ratio_override,
|
|
log_level,
|
|
links,
|
|
)
|
|
|
|
# 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: str | None = None,
|
|
span_type: str = SpanType.UNKNOWN,
|
|
attributes: dict[str, Any] | None = None,
|
|
trace_destination: TraceLocationBase | None = None,
|
|
sampling_ratio_override: float | None = None,
|
|
log_level: SpanLogLevel | str | None = None,
|
|
links: list[Link] | None = 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,
|
|
trace_destination=trace_destination,
|
|
log_level=log_level,
|
|
links=links,
|
|
) 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)
|
|
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:
|
|
if exc_value is None:
|
|
exc_value = exc_type()
|
|
self.coro.throw(exc_value)
|
|
self.coro.close()
|
|
|
|
if inspect.iscoroutinefunction(fn):
|
|
|
|
async def wrapper(*args, **kwargs):
|
|
with (
|
|
_set_sampling_ratio_override(sampling_ratio_override),
|
|
_WrappingContext(fn, args, kwargs) as wrapping_coro,
|
|
):
|
|
return wrapping_coro.send(await fn(*args, **kwargs))
|
|
|
|
else:
|
|
|
|
def wrapper(*args, **kwargs):
|
|
with (
|
|
_set_sampling_ratio_override(sampling_ratio_override),
|
|
_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: str | None = None,
|
|
span_type: str = SpanType.UNKNOWN,
|
|
attributes: dict[str, Any] | None = None,
|
|
output_reducer: Callable[[list[Any]], Any] | None = None,
|
|
trace_destination: TraceLocationBase | None = None,
|
|
sampling_ratio_override: float | None = None,
|
|
log_level: SpanLogLevel | str | None = None,
|
|
links: list[Link] | None = 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,
|
|
experiment_id=getattr(trace_destination, "experiment_id", None),
|
|
log_level=log_level,
|
|
links=links,
|
|
)
|
|
except Exception as e:
|
|
_logger.debug(f"Failed to start stream span: {e}")
|
|
return NoOpSpan()
|
|
|
|
def _end_stream_span(
|
|
span: LiveSpan,
|
|
inputs: dict[str, Any] | None = None,
|
|
outputs: list[Any] | None = None,
|
|
output_reducer: Callable[[list[Any]], Any] | None = None,
|
|
error: Exception | None = 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}")
|
|
|
|
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):
|
|
with _set_sampling_ratio_override(sampling_ratio_override):
|
|
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):
|
|
with _set_sampling_ratio_override(sampling_ratio_override):
|
|
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
|
|
# Add unique marker for MLflow trace detection
|
|
wrapped.__mlflow_traced__ = True
|
|
return wrapped
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def start_span(
|
|
name: str = "span",
|
|
span_type: str | None = SpanType.UNKNOWN,
|
|
attributes: dict[str, Any] | None = None,
|
|
trace_destination: TraceLocationBase | None = None,
|
|
log_level: SpanLogLevel | str | None = None,
|
|
run_id: str | None = None,
|
|
links: list[Link] | None = 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.
|
|
trace_destination: The destination to log the trace to, such as MLflow Experiment. If
|
|
not provided, the destination will be an active MLflow experiment or an destination
|
|
set by the :py:func:`mlflow.tracing.set_destination` function. This parameter
|
|
should only be used for root span and setting this for non-root spans will be
|
|
ignored with a warning.
|
|
log_level: Optional severity level to attach to the span. Accepts a
|
|
:py:class:`SpanLogLevel <mlflow.entities.SpanLogLevel>` or its name
|
|
(e.g. ``"INFO"``, ``"DEBUG"``). If not provided, the span level is
|
|
resolved from the span type at end time.
|
|
run_id: The ID of the MLflow run to associate with the trace. This parameter is
|
|
only applied when creating a root span. If provided without an explicit
|
|
`trace_destination`, the trace will be logged to the run's experiment. If an
|
|
active MLflow run is already set via `mlflow.start_run()`, this parameter takes
|
|
precedence over the active run.
|
|
links: A list of :py:class:`Link <mlflow.entities.Link>` objects to associate with
|
|
the span.
|
|
|
|
Returns:
|
|
Yields an :py:class:`mlflow.entities.Span` that represents the created span.
|
|
"""
|
|
# If tracing is disabled via context(enabled=False), return NoOpSpan
|
|
config = _USER_TRACE_CONTEXT.get()
|
|
if config is not None and config.enabled is False:
|
|
yield NoOpSpan()
|
|
return
|
|
|
|
try:
|
|
experiment_id = getattr(trace_destination, "experiment_id", None)
|
|
|
|
otel_span = provider.start_span_in_context(
|
|
name,
|
|
experiment_id=experiment_id,
|
|
)
|
|
|
|
# If the span was dropped by the sampler (e.g., due to sampling ratio),
|
|
# still propagate the OTel context so that child spans inherit the same
|
|
# trace ID and are also consistently dropped.
|
|
if not otel_span.is_recording():
|
|
noop_span = NoOpSpan(otel_span=otel_span)
|
|
else:
|
|
noop_span = None
|
|
|
|
# Create a new MLflow span and register it to the in-memory trace manager
|
|
request_id = get_otel_attribute(otel_span, SpanAttributeKey.REQUEST_ID)
|
|
|
|
# SpanProcessor should have already registered the span in the in-memory trace manager
|
|
trace_manager = InMemoryTraceManager.get_instance()
|
|
mlflow_span = trace_manager.get_span_from_id(
|
|
request_id, encode_span_id(otel_span.context.span_id)
|
|
)
|
|
mlflow_span.set_span_type(span_type)
|
|
attributes = dict(attributes) if attributes is not None else {}
|
|
mlflow_span.set_attributes(attributes)
|
|
if log_level is not None:
|
|
mlflow_span.set_log_level(log_level)
|
|
|
|
if run_id is not None:
|
|
if mlflow_span.parent_id is not None:
|
|
_logger.warning(
|
|
"The `run_id` parameter can only be used for root spans, but the span "
|
|
f"`{name}` is not a root span. The specified value `{run_id}` "
|
|
"will be ignored."
|
|
)
|
|
else:
|
|
with trace_manager.get_trace(request_id) as trace:
|
|
trace.info.trace_metadata[TraceMetadataKey.SOURCE_RUN] = run_id
|
|
|
|
except Exception:
|
|
_logger.debug(f"Failed to start span {name}.", exc_info=True)
|
|
noop_span = NoOpSpan()
|
|
|
|
if noop_span is None:
|
|
for link in links or []:
|
|
try:
|
|
mlflow_span.add_link(link)
|
|
except MlflowException:
|
|
_logger.warning("Skipping invalid link: %s", link)
|
|
|
|
# Yield NoOp spans outside the try/except block so that exceptions thrown
|
|
# back into the generator (via contextmanager's throw()) propagate correctly
|
|
# instead of being caught by the broad "except Exception" above.
|
|
if noop_span is not None:
|
|
with safe_set_span_in_context(noop_span):
|
|
yield noop_span
|
|
return
|
|
|
|
try:
|
|
with with_active_span(mlflow_span):
|
|
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: LiveSpan | None = None,
|
|
inputs: Any | None = None,
|
|
attributes: dict[str, str] | None = None,
|
|
tags: dict[str, str] | None = None,
|
|
metadata: dict[str, str] | None = None,
|
|
experiment_id: str | None = None,
|
|
start_time_ns: int | None = None,
|
|
log_level: SpanLogLevel | str | None = None,
|
|
links: list[Link] | None = 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.
|
|
metadata: A dictionary of metadata 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.
|
|
log_level: Optional severity level to attach to the span. Accepts a
|
|
:py:class:`SpanLogLevel <mlflow.entities.SpanLogLevel>` or its name
|
|
(e.g. ``"INFO"``, ``"DEBUG"``). If not provided, the span level is
|
|
resolved from the span type at end time.
|
|
links: A list of :py:class:`Link <mlflow.entities.Link>` objects to associate with
|
|
the span.
|
|
|
|
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 tracing is disabled via context(enabled=False), return NoOpSpan
|
|
config = _USER_TRACE_CONTEXT.get()
|
|
if config is not None and config.enabled is False:
|
|
return NoOpSpan()
|
|
|
|
# If parent span is no-op span, the child should also be no-op. Preserve the
|
|
# parent OTel span when present so descendants keep inheriting the dropped
|
|
# trace context instead of starting a fresh root trace.
|
|
if parent_span and parent_span.trace_id == NO_OP_SPAN_TRACE_ID:
|
|
return NoOpSpan(otel_span=parent_span._span)
|
|
|
|
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 the span was dropped by the sampler, return a NoOpSpan that
|
|
# preserves the OTel span's context so that safe_set_span_in_context
|
|
# propagates the correct trace ID to child spans.
|
|
if not otel_span.is_recording():
|
|
return NoOpSpan(otel_span=otel_span)
|
|
|
|
if parent_span:
|
|
trace_id = parent_span.trace_id
|
|
else:
|
|
trace_id = get_otel_attribute(otel_span, SpanAttributeKey.REQUEST_ID)
|
|
|
|
# SpanProcessor should have already registered the span in the in-memory trace manager
|
|
trace_manager = InMemoryTraceManager.get_instance()
|
|
mlflow_span = trace_manager.get_span_from_id(
|
|
trace_id, encode_span_id(otel_span.context.span_id)
|
|
)
|
|
mlflow_span.set_span_type(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 {})
|
|
if log_level is not None:
|
|
mlflow_span.set_log_level(log_level)
|
|
|
|
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)
|
|
|
|
if metadata:
|
|
with trace_manager.get_trace(trace_id) as trace:
|
|
trace.info.trace_metadata.update(metadata)
|
|
|
|
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()
|
|
|
|
for link in links or []:
|
|
try:
|
|
mlflow_span.add_link(link)
|
|
except MlflowException:
|
|
_logger.warning("Skipping invalid link: %s", link)
|
|
|
|
return mlflow_span
|
|
|
|
|
|
@deprecated_parameter("request_id", "trace_id")
|
|
def get_trace(trace_id: str, silent: bool = False, flush: bool = False) -> Trace | None:
|
|
"""
|
|
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.
|
|
silent: If True, suppress the warning message when the trace is not found. The API will
|
|
return None without any warning. Default to False.
|
|
flush: If True and the trace is not found, flush any pending async trace writes and
|
|
retry. Useful in tests or scripts where async logging may not have completed.
|
|
Default to False.
|
|
|
|
.. 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
|
|
|
|
exc: MlflowException | None = None
|
|
try:
|
|
return TracingClient().get_trace(trace_id)
|
|
except MlflowException as e:
|
|
exc = e
|
|
|
|
if flush:
|
|
_flush_pending_async_trace_writes()
|
|
exc = None
|
|
try:
|
|
return TracingClient().get_trace(trace_id)
|
|
except MlflowException as e:
|
|
exc = e
|
|
|
|
if not silent:
|
|
hint = (
|
|
" If using async trace logging, pass flush=True to wait for pending writes."
|
|
if not flush
|
|
else ""
|
|
)
|
|
_logger.warning(
|
|
f"Failed to get trace from the tracking store: {exc}.{hint} "
|
|
"For full traceback, set logging level to debug.",
|
|
exc_info=_logger.isEnabledFor(logging.DEBUG),
|
|
)
|
|
else:
|
|
_logger.debug(f"Failed to get trace from the tracking store: {exc}.", exc_info=True)
|
|
return None
|
|
|
|
|
|
def _flush_pending_async_trace_writes(terminate: bool = False) -> None:
|
|
"""Flush all pending async trace writes through the BSP and exporter queues.
|
|
|
|
Two-layer flush:
|
|
1. flush_all_batch_processors() drains BSP-registered processors and their exporters.
|
|
2. The direct exporter flush handles the no-BSP path (MLFLOW_USE_BATCH_SPAN_PROCESSOR=false),
|
|
where the processor is not in the registry but the exporter still has an async queue.
|
|
|
|
Args:
|
|
terminate: If True, shut down background threads after flushing. Used in test teardown
|
|
to prevent thread leaks between tests.
|
|
"""
|
|
# Lazy import to avoid circular dependency:
|
|
# base_mlflow imports _set_last_active_trace_id from this module.
|
|
from mlflow.tracing.processor.base_mlflow import flush_all_batch_processors
|
|
|
|
try:
|
|
flush_all_batch_processors(terminate=terminate)
|
|
except Exception:
|
|
_logger.debug("Failed to flush batch processors.", exc_info=True)
|
|
try:
|
|
if trace_exporter := _get_trace_exporter():
|
|
if hasattr(trace_exporter, "_async_queue"):
|
|
trace_exporter._async_queue.flush(terminate=terminate)
|
|
except Exception:
|
|
_logger.debug("Failed to flush trace exporter async queue.", exc_info=True)
|
|
|
|
|
|
def _get_search_locations(locations: list[str] | None) -> list[str]:
|
|
from mlflow.tracking.fluent import _get_experiment_id
|
|
|
|
if locations:
|
|
return locations
|
|
|
|
if experiment_id := _get_experiment_id():
|
|
return [experiment_id]
|
|
|
|
raise MlflowException(
|
|
"No active experiment found. Set an experiment using `mlflow.set_experiment`, "
|
|
"or specify the list of experiment IDs in the `locations` parameter."
|
|
)
|
|
|
|
|
|
@deprecated_parameter("experiment_ids", "locations")
|
|
def search_traces(
|
|
experiment_ids: list[str] | None = None,
|
|
filter_string: str | None = None,
|
|
max_results: int | None = None,
|
|
order_by: list[str] | None = None,
|
|
extract_fields: list[str] | None = None,
|
|
run_id: str | None = None,
|
|
return_type: Literal["pandas", "list"] | None = None,
|
|
model_id: str | None = None,
|
|
sql_warehouse_id: str | None = None,
|
|
include_spans: bool = True,
|
|
locations: list[str] | None = None,
|
|
flush: bool = False,
|
|
) -> "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.
|
|
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:
|
|
.. deprecated:: 3.6.0
|
|
This parameter is deprecated and will be removed in a future version.
|
|
|
|
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: DEPRECATED. Use the `MLFLOW_TRACING_SQL_WAREHOUSE_ID` environment
|
|
variable instead. The ID of the SQL warehouse to use for
|
|
searching traces in inference tables or UC tables. Only used in Databricks.
|
|
|
|
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``.
|
|
|
|
locations: A list of locations to search over. To search over experiments, provide
|
|
a list of experiment IDs. To search over UC tables on databricks, provide
|
|
a list of locations in the format
|
|
`<catalog_name>.<schema_name>[.<table_prefix>]`.
|
|
If not provided, the search will be performed across the current active experiment.
|
|
|
|
flush: If ``True``, flush any pending async trace writes before searching. Useful
|
|
in tests or scripts to ensure all traces are visible. Default to ``False``.
|
|
|
|
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")
|
|
|
|
"""
|
|
|
|
if sql_warehouse_id is not None:
|
|
warnings.warn(
|
|
"The `sql_warehouse_id` parameter is deprecated. Please use the "
|
|
"`MLFLOW_TRACING_SQL_WAREHOUSE_ID` environment variable instead.",
|
|
category=FutureWarning,
|
|
)
|
|
os.environ["MLFLOW_TRACING_SQL_WAREHOUSE_ID"] = sql_warehouse_id
|
|
|
|
if extract_fields is not None:
|
|
warnings.warn(
|
|
"The `extract_fields` parameter is deprecated and will be removed in a future version.",
|
|
category=FutureWarning,
|
|
stacklevel=2,
|
|
)
|
|
|
|
# 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'`."
|
|
),
|
|
)
|
|
|
|
_validate_list_param("locations", locations, allow_none=True)
|
|
|
|
if flush:
|
|
_flush_pending_async_trace_writes()
|
|
|
|
if not experiment_ids and not locations:
|
|
_logger.debug("Searching traces in the current active experiment")
|
|
locations = _get_search_locations(locations)
|
|
|
|
if (
|
|
locations
|
|
and any("." in loc for loc in locations)
|
|
and (not filter_string or "trace.timestamp_ms" not in filter_string.lower())
|
|
):
|
|
warnings.warn(
|
|
"Searching traces without a time range constraint on UC table locations can be slow "
|
|
"and expensive. Consider adding a `trace.timestamp_ms` filter to your `filter_string` "
|
|
"to limit the scan, e.g. filter_string=\"trace.timestamp_ms > '2024-01-01'\".",
|
|
category=UserWarning,
|
|
stacklevel=2,
|
|
)
|
|
|
|
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,
|
|
include_spans=include_spans,
|
|
locations=locations,
|
|
)
|
|
|
|
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, extract_fields=extract_fields)
|
|
|
|
return results
|
|
|
|
|
|
@experimental(version="3.10.0")
|
|
def search_sessions(
|
|
max_results: int = 100,
|
|
run_id: str | None = None,
|
|
model_id: str | None = None,
|
|
include_spans: bool = True,
|
|
locations: list[str] | None = None,
|
|
) -> list[Session]:
|
|
"""
|
|
Return complete sessions that match the given search criteria.
|
|
|
|
A session is a collection of traces that share the same session ID, typically representing
|
|
a multi-turn conversation or a series of related interactions. This API retrieves complete
|
|
sessions by first identifying unique session IDs from traces, then fetching all traces
|
|
belonging to each session in parallel.
|
|
|
|
Args:
|
|
max_results: Maximum number of sessions to return. Default is 100.
|
|
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
|
|
traces.
|
|
model_id: If specified, search traces associated with the given model ID.
|
|
include_spans: If ``True``, include spans in the returned traces. Otherwise, only
|
|
the trace metadata is returned. Default is ``True``.
|
|
locations: A list of locations to search over. To search over experiments, provide
|
|
a list of experiment IDs. To search over UC tables on databricks, provide
|
|
a list of locations in the format
|
|
`<catalog_name>.<schema_name>[.<table_prefix>]`.
|
|
If not provided, the search will be performed across the current active experiment.
|
|
|
|
Returns:
|
|
A list of :py:class:`Session <mlflow.entities.Session>` objects, where each session
|
|
contains :py:class:`Trace <mlflow.entities.Trace>` objects that share the same
|
|
session ID. Sessions are ordered by the timestamp of their first trace (most recent first).
|
|
Each Session object provides convenient access via ``session.id`` and supports
|
|
iteration with ``for trace in session``.
|
|
|
|
.. code-block:: python
|
|
:caption: Basic usage - search sessions in an experiment
|
|
|
|
import mlflow
|
|
|
|
# Get all sessions from the current experiment
|
|
sessions = mlflow.search_sessions()
|
|
|
|
# Each session provides convenient access to ID and traces
|
|
for session in sessions:
|
|
print(f"Session {session.id} has {len(session)} traces")
|
|
for trace in session:
|
|
print(f" Trace: {trace.info.trace_id}")
|
|
|
|
.. code-block:: python
|
|
:caption: Use sessions with evaluation
|
|
|
|
import mlflow
|
|
|
|
# Get sessions for evaluation
|
|
sessions = mlflow.search_sessions(
|
|
locations=["experiment_id"],
|
|
max_results=50,
|
|
)
|
|
|
|
# Flatten sessions for evaluation if needed
|
|
all_traces = [trace for session in sessions for trace in session]
|
|
mlflow.genai.evaluate(data=all_traces)
|
|
"""
|
|
_validate_list_param("locations", locations, allow_none=True)
|
|
|
|
if not locations:
|
|
_logger.debug("Searching sessions in the current active experiment")
|
|
locations = _get_search_locations(locations)
|
|
|
|
session_id_key = TraceMetadataKey.TRACE_SESSION
|
|
|
|
# Step 1: Page through traces to collect unique session IDs (up to max_results)
|
|
seen_session_ids: set[str] = set()
|
|
session_ids: list[str] = []
|
|
page_token: str | None = None
|
|
|
|
while len(session_ids) < max_results:
|
|
traces: list[Trace] = TracingClient().search_traces(
|
|
max_results=SEARCH_TRACES_DEFAULT_MAX_RESULTS,
|
|
order_by=["timestamp DESC"],
|
|
page_token=page_token,
|
|
include_spans=False,
|
|
run_id=run_id,
|
|
model_id=model_id,
|
|
locations=locations,
|
|
)
|
|
|
|
for trace in traces:
|
|
session_id = trace.info.request_metadata.get(session_id_key)
|
|
if session_id and session_id not in seen_session_ids:
|
|
seen_session_ids.add(session_id)
|
|
session_ids.append(session_id)
|
|
if len(session_ids) >= max_results:
|
|
break
|
|
|
|
page_token = traces.token if hasattr(traces, "token") else None
|
|
if not page_token:
|
|
break
|
|
|
|
if not session_ids:
|
|
return []
|
|
|
|
# Step 2: Fetch complete traces for each session in parallel
|
|
def fetch_session_traces(session_id: str) -> list[Trace]:
|
|
session_filter = f"metadata.`{session_id_key}` = '{session_id}'"
|
|
|
|
return search_traces(
|
|
filter_string=session_filter,
|
|
max_results=None, # Get all traces in the session
|
|
order_by=["timestamp ASC"], # Order by time within session
|
|
include_spans=include_spans,
|
|
run_id=run_id,
|
|
return_type="list",
|
|
model_id=model_id,
|
|
locations=locations,
|
|
)
|
|
|
|
max_workers = min(len(session_ids), MLFLOW_SEARCH_TRACES_MAX_THREADS.get())
|
|
with ThreadPoolExecutor(
|
|
max_workers=max_workers,
|
|
thread_name_prefix="search_sessions",
|
|
) as executor:
|
|
session_results = list(executor.map(fetch_session_traces, session_ids))
|
|
|
|
# Filter out empty sessions, wrap in Session objects, and preserve order
|
|
sessions: list[Session] = [Session(s) for s in session_results if s]
|
|
|
|
return sessions
|
|
|
|
|
|
def get_current_active_span() -> LiveSpan | None:
|
|
"""
|
|
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 = get_current_otel_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 = otel_span.attributes.get(SpanAttributeKey.REQUEST_ID)
|
|
|
|
# Span is not registered in the in-memory trace manager, meaning that the current active
|
|
# span is not created by MLflow, but rather by other OpenTelemetry sdk. Return a span object
|
|
# that wraps the otel span.
|
|
if not request_id:
|
|
return create_mlflow_span(otel_span, otel_span.context.trace_id)
|
|
|
|
request_id = json.loads(request_id)
|
|
return trace_manager.get_span_from_id(request_id, encode_span_id(otel_span.context.span_id))
|
|
|
|
|
|
def get_active_trace_id() -> str | None:
|
|
"""
|
|
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.
|
|
"""
|
|
if active_span := get_current_active_span():
|
|
return active_span.trace_id
|
|
return None
|
|
|
|
|
|
def get_last_active_trace_id(thread_local: bool = False) -> str | None:
|
|
"""
|
|
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: dict[str, str] | None = None,
|
|
metadata: dict[str, str] | None = None,
|
|
client_request_id: str | None = None,
|
|
request_preview: str | None = None,
|
|
response_preview: str | None = None,
|
|
state: TraceState | str | None = None,
|
|
model_id: str | None = None,
|
|
session_id: str | None = None,
|
|
user: str | None = 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.
|
|
model_id: The ID of the model to associate with the trace. If not set, the active
|
|
model ID is associated with the trace.
|
|
session_id: Session ID to associate with the trace. Stored as metadata under the
|
|
``mlflow.trace.session`` key.
|
|
user: User identifier to associate with the trace. Stored as metadata under the
|
|
``mlflow.trace.user`` key.
|
|
|
|
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 user, session, and source information of the trace:
|
|
|
|
.. code-block:: python
|
|
|
|
mlflow.update_current_trace(
|
|
session_id="session-4f855da00427",
|
|
user="user-id-cc156f29bcfb",
|
|
metadata={
|
|
"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):
|
|
if non_string_items := {k: v for k, v in d.items() if not isinstance(v, str)}:
|
|
_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}"
|
|
)
|
|
|
|
tags = tags or {}
|
|
metadata = dict(metadata) if metadata else {}
|
|
|
|
if session_id is not None:
|
|
metadata[TraceMetadataKey.TRACE_SESSION] = session_id
|
|
if user is not None:
|
|
metadata[TraceMetadataKey.TRACE_USER] = user
|
|
if model_id:
|
|
metadata[TraceMetadataKey.MODEL_ID] = model_id
|
|
|
|
_warn_non_string_values(tags, "tags")
|
|
_warn_non_string_values(metadata, "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)
|
|
trace.info.trace_metadata.update(metadata)
|
|
if client_request_id is not None:
|
|
trace.info.client_request_id = str(client_request_id)
|
|
|
|
|
|
@deprecated_parameter("request_id", "trace_id")
|
|
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)
|
|
|
|
|
|
@deprecated_parameter("request_id", "trace_id", version="3.0.0")
|
|
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)
|
|
|
|
|
|
def add_trace(trace: Trace | dict[str, Any], target: LiveSpan | None = 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
|
|
},
|
|
# ensure this span has a smaller start time than the remote trace
|
|
# so when it's loaded the order is correct when sorting by start time
|
|
# TODO: deprecate this function once we fully support OTel traces
|
|
start_time_ns=remote_root_span.start_time_ns - 1,
|
|
)
|
|
_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,
|
|
)
|
|
|
|
|
|
# TODO: remove this function in 3.7.0
|
|
@deprecated(since="3.6.0")
|
|
def log_trace(
|
|
name: str = "Task",
|
|
request: Any | None = None,
|
|
response: Any | None = None,
|
|
intermediate_outputs: dict[str, Any] | None = None,
|
|
attributes: dict[str, Any] | None = None,
|
|
tags: dict[str, str] | None = None,
|
|
start_time_ms: int | None = None,
|
|
execution_time_ms: int | None = 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, flush=True)
|
|
|
|
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)
|
|
# end the cloned span to ensure it's processed by the exporter
|
|
cloned_span.end(end_time_ns=span.end_time_ns)
|
|
|
|
# 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,
|
|
}
|