mlflow.tracing

Attention

The mlflow.tracing namespace only contains a few utility functions fo managing traces. The main entry point for MLflow Tracing is Tracing Fluent APIs defined directly under the mlflow namespace, or the low-level Tracing Client APIs

mlflow.tracing.disable()

Disable tracing.

Note

This function sets up OpenTelemetry to use NoOpTracerProvider and effectively disables all tracing operations.

Example:

import mlflow


@mlflow.trace
def f():
    return 0


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

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

mlflow.tracing.disable_notebook_display()

Disables displaying the MLflow Trace UI in notebook output cells. Call mlflow.tracing.enable_notebook_display() to re-enable display.

mlflow.tracing.enable()

Enable tracing.

Example:

import mlflow


@mlflow.trace
def f():
    return 0


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

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

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

mlflow.tracing.enable_notebook_display()

Enables the MLflow Trace UI in notebook output cells. The display is on by default, and the Trace UI will show up when any of the following operations are executed:

• On trace completion (i.e. whenever a trace is exported)

• When calling the mlflow.search_traces() fluent API

• When calling the mlflow.client.MlflowClient.get_trace() or mlflow.client.MlflowClient.search_traces() client APIs

To disable, please call mlflow.tracing.disable_notebook_display().

mlflow.tracing.set_span_chat_messages(span: LiveSpan, messages: list[RequestMessage])

Set the mlflow.chat.messages attribute on the specified span. This attribute is used in the UI, and also by downstream applications that consume trace data, such as MLflow evaluate.

Parameters

• span – The LiveSpan to add the attribute to

• messages – A list of standardized chat messages (refer to the spec for details)

Example:

import mlflow
from mlflow.tracing import set_span_chat_messages


@mlflow.trace
def f():
    messages = [{"role": "user", "content": "hello"}]
    span = mlflow.get_current_active_span()
    set_span_chat_messages(span, messages)
    return 0


f()

mlflow.tracing.set_span_chat_tools(span: LiveSpan, tools: list[ChatTool])

Set the mlflow.chat.tools attribute on the specified span. This attribute is used in the UI, and also by downstream applications that consume trace data, such as MLflow evaluate.

Parameters

• span – The LiveSpan to add the attribute to

• tools –

  A list of standardized chat tool definitions (refer to the spec for details)

Example:

import mlflow
from mlflow.tracing import set_span_chat_tools

tools = [
    {
        "type": "function",
        "function": {
            "name": "add",
            "description": "Add two numbers",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {"type": "number"},
                    "b": {"type": "number"},
                },
                "required": ["a", "b"],
            },
        },
    }
]


@mlflow.trace
def f():
    span = mlflow.get_current_active_span()
    set_span_chat_tools(span, tools)
    return 0


f()