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.configure(span_processors: list[typing.Callable[[ForwardRef('LiveSpan')], NoneType]] | None = None) → mlflow.tracing.config.TracingConfigContext[source]

Note

Experimental: This function may change or be removed in a future release without warning.

Configure MLflow tracing. Can be used as function or context manager.

Only updates explicitly provided arguments, leaving others unchanged.

Parameters

span_processors – List of functions to process spans before export. This is helpful for filtering/masking particular attributes from the span to prevent sensitive data from being logged or for reducing the size of the span. Each function must accept a single argument of type LiveSpan and should not return any value. When multiple functions are provided, they are applied sequentially in the order they are provided.

Returns

Context manager for temporary configuration changes.: When used as a function, the configuration changes persist. When used as a context manager, changes are reverted on exit.

Return type

TracingConfigContext

Examples

def pii_filter(span):
    """Example PII filter that masks sensitive data in span attributes."""
    # Mask sensitive inputs
    if inputs := span.inputs:
        for key, value in inputs.items():
            if "password" in key.lower() or "token" in key.lower():
                span.set_inputs({**inputs, key: "[REDACTED]"})

    # Mask sensitive outputs
    if outputs := span.outputs:
        if isinstance(outputs, dict):
            for key in outputs:
                if "secret" in key.lower():
                    outputs[key] = "[REDACTED]"
            span.set_outputs(outputs)

    # Mask sensitive attributes
    for attr_key in list(span.attributes.keys()):
        if "api_key" in attr_key.lower():
            span.set_attribute(attr_key, "[REDACTED]")

# Permanent configuration change
mlflow.tracing.configure(span_processors=[pii_filter])

# Temporary configuration change
with mlflow.tracing.configure(span_processors=[pii_filter]):
    # PII filtering enabled only in this block
    pass

mlflow.tracing.disable()[source]

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()[source]: Disables displaying the MLflow Trace UI in notebook output cells. Call mlflow.tracing.enable_notebook_display() to re-enable display.

mlflow.tracing.enable()[source]

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()[source]

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.reset()[source]: Reset the flags that indicates whether the MLflow tracer provider has been initialized. This ensures that the tracer provider is re-initialized when next tracing operation is performed.

mlflow.tracing.set_databricks_monitoring_sql_warehouse_id(sql_warehouse_id: str, experiment_id: Optional[str] = None) → None[source]

Note

Experimental: This function may change or be removed in a future release without warning.

Set the SQL warehouse ID used for Databricks production monitoring on traces logged to the given MLflow experiment. This only has an effect for experiments with UC schema as trace location.

Parameters

sql_warehouse_id – The SQL warehouse ID to use for monitoring.
experiment_id – The MLflow experiment ID. If not provided, the current active experiment will be used.

mlflow.tracing.set_destination(destination: mlflow.entities.trace_location.TraceLocationBase, *, context_local: bool = False)[source]

Set a custom span location to which MLflow will export the traces.

A destination specified by this function will take precedence over other configurations, such as tracking URI, OTLP environment variables.

Parameters

destination –
A trace location object that specifies the location of the trace data. Currently, the following locations are supported:
- MlflowExperimentLocation: Logs traces to
  an MLflow experiment.
- UCSchemaLocation: Logs traces to a
  Databricks Unity Catalog schema. Only available in Databricks.
context_local – If False (default), the destination is set globally. If True, the destination is isolated per async task or thread, providing isolation in concurrent applications.

Example

Logging traces to MLflow Experiment:

from mlflow.entities.trace_location import MlflowExperimentLocation

mlflow.tracing.set_destination(MlflowExperimentLocation(experiment_id="123"))

Note: This has the same effect as setting the active MLflow experiment via the MLFLOW_EXPERIMENT_ID environment variable or the mlflow.set_experiment API, but with narrower scope.

Logging traces to Databricks Unity Catalog:

from mlflow.entities.trace_location import UCSchemaLocation

mlflow.tracing.set_destination(
    UCSchemaLocation(catalog_name="catalog", schema_name="schema")
)

Isolate the destination between async tasks or threads:

from mlflow.tracing.destination import Databricks

mlflow.tracing.set_destination(
    MlflowExperimentLocation(experiment_id="123"),
    context_local=True,
)

The destination set with the context_local flag will only be effective within the current async task or thread. This is particularly useful when you want to send traces to different destinations from a multi-tenant application.

** Reset the destination:**

mlflow.tracing.reset()

mlflow.tracing.set_experiment_trace_location(location: UCSchemaLocation, experiment_id: Optional[str] = None, sql_warehouse_id: Optional[str] = None) → UCSchemaLocation[source]

Note

Experimental: This function may change or be removed in a future release without warning.

Configure the storage location for traces of an experiment.

Unity Catalog tables for storing trace data will be created in the specified schema. When tracing is enabled, all traces for the specified experiment will be stored in the provided Unity Catalog schema.

Note

If the experiment is already linked to a storage location, this will raise an error. Use mlflow.tracing.unset_experiment_trace_location to remove the existing storage location first and then set a new one.

Parameters

location – The storage location for experiment traces in Unity Catalog.
experiment_id – The MLflow experiment ID to set the storage location for. If not specified, the current active experiment will be used.
sql_warehouse_id – SQL warehouse ID for creating views and querying. If not specified, uses the value from MLFLOW_TRACING_SQL_WAREHOUSE_ID, fallback to the default SQL warehouse if the environment variable is not set.

Returns

The UCSchemaLocation object representing the configured storage location, including the table names of the spans and logs tables.

Example

import mlflow
from mlflow.entities import UCSchemaLocation

location = UCSchemaLocation(catalog_name="my_catalog", schema_name="my_schema")

result = mlflow.tracing.set_experiment_trace_location(
    location=location,
    experiment_id="12345",
)
print(result.full_otel_spans_table_name)  # my_catalog.my_schema.otel_spans_table


@mlflow.trace
def add(x):
    return x + 1


add(1)  # this writes the trace to the storage location set above

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

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()

mlflow.tracing.unset_experiment_trace_location(location: UCSchemaLocation, experiment_id: Optional[str] = None) → None[source]

Note

Experimental: This function may change or be removed in a future release without warning.

Unset the storage location for traces of an experiment.

This function removes the experiment storage location configuration, including the view and the experiment tag.

Parameters

location – The storage location to unset.
experiment_id – The MLflow experiment ID to unset the storage location for. If not provided, the current active experiment will be used.

Example

import mlflow
from mlflow.entities import UCSchemaLocation

mlflow.tracing.unset_experiment_trace_location(
    location=UCSchemaLocation(catalog_name="my_catalog", schema_name="my_schema"),
    experiment_id="12345",
)