Skip to main content
Last updated on

Store MLflow traces in Unity Catalog

Beta

This feature is in Beta. Workspace admins can control access to this feature from the Previews page. See Manage Databricks previews.

Databricks supports storing MLflow traces in Unity Catalog tables using an OpenTelemetry-compatible format (OTel). By default, MLflow stores traces organized by experiments in the MLflow control plane service. However, storing traces in Unity Catalog using OTel format provides the following benefits:

  • Access control is managed through Unity Catalog schema and table permissions rather than experiment-level ACLs. Users with access to the Unity Catalog tables can view all traces stored in those tables, regardless of which experiment the traces belong to.
  • Trace IDs use URI format instead of the tr-<UUID> format, improving compatibility with external systems.
  • Store unlimited traces in Delta tables, enabling long-term retention and analysis of trace data.
  • Query trace data directly using SQL through a Databricks SQL warehouse, enabling advanced analytics and custom reporting.
  • OTel format ensures compatibility with other OpenTelemetry clients and tools.

Prerequisites​

  • A Unity Catalog-enabled workspace.

  • Ensure the "OpenTelemetry on Databricks" preview is enabled, along with "Variant Shredding for Optimized Read Performance on Semi-Structured Data." See Manage Databricks previews.

  • Permissions to create catalogs and schemas in Unity Catalog.

  • A Databricks SQL warehouse with CAN USE permissions. Save the warehouse ID for later reference.

  • While this feature is in Beta, your workspace must be in one of the following regions:

    • us-east1

  • MLflow Python library version 3.11 or later installed in your environment:

    Bash
    pip install mlflow[databricks]>=3.11.0 --upgrade --force-reinstall

Setup: Create an experiment with a Unity Catalog trace location​

Run the following code to create and bind an experiment to a Unity Catalog trace location:

Python
# Example values for the placeholders below:
# MLFLOW_TRACING_SQL_WAREHOUSE_ID: "abc123def456" (found in SQL warehouse URL)
# experiment_name: "/Users/user@company.com/traces"
# catalog_name: "main" or "my_catalog"
# schema_name: "mlflow_traces" or "production_traces"
# table_prefix: "my_otel"

import os
import mlflow
from mlflow.entities.trace_location import UnityCatalog

mlflow.set_tracking_uri("databricks")

# Specify the ID of a SQL warehouse you have access to.
os.environ["MLFLOW_TRACING_SQL_WAREHOUSE_ID"] = "<SQL_WAREHOUSE_ID>"
# Specify the name of the MLflow Experiment to use for viewing traces in the UI.
experiment_name = "<MLFLOW_EXPERIMENT_NAME>"
# Specify the name of the Catalog to use for storing traces.
catalog_name = "<UC_CATALOG_NAME>"
# Specify the name of the Schema to use for storing traces.
schema_name = "<UC_SCHEMA_NAME>"
# Specify the name of the prefix appended to every table storing trace data.
table_prefix = "<UC_TABLE_PREFIX>"

# mlflow.set_experiment is an upsert operation
experiment = mlflow.set_experiment(
    experiment_name=experiment_name,
    trace_location=UnityCatalog(
        catalog_name=catalog_name,
        schema_name=schema_name,
        table_prefix=table_prefix,  # defaults to experiment id if not provided
    ),
)

print(f"Experiment ID: {experiment.experiment_id}")
print(experiment.trace_location.full_otel_spans_table_name)

You can also use mlflow.create_experiment with the same trace_location parameter. Unlike set_experiment, create_experiment does not set the active experiment, so you must call set_experiment afterward in order to ensure that traces are routed to the correct location:

Python
experiment_id = mlflow.create_experiment(
    name=experiment_name,
    trace_location=UnityCatalog(
        catalog_name=catalog_name,
        schema_name=schema_name,
        table_prefix=table_prefix,
    ),
)

# trace_location is optional here since
# the experiment is already bound to the UC trace location above.
mlflow.set_experiment(experiment_id=experiment_id)

Once you bind an experiment to a UC trace location, you cannot reassign the experiment to a different UC trace location. However, multiple experiments can share the same UC trace location.

Verify tables​

After running the setup code, four new Unity Catalog tables appear in the schema in the Catalog Explorer UI:

  • <table_prefix>_otel_annotations
  • <table_prefix>_otel_logs
  • <table_prefix>_otel_metrics
  • <table_prefix>_otel_spans

Grant permissions​

A Databricks user or service principal needs the following Unity Catalog privileges to write or read MLflow traces from the Unity Catalog tables:

  1. USE_CATALOG on the catalog.
  2. USE_SCHEMA on the schema.
  3. MODIFY and SELECT on each of the <table_prefix>_<type> tables.
note

ALL_PRIVILEGES is not sufficient for accessing Unity Catalog trace tables. You must explicitly grant MODIFY and SELECT.

Log traces to the Unity Catalog tables​

After creating the tables, you can write traces to them from various sources by specifying the trace location. How you do this depends on the source of the traces.

The Unity Catalog trace location can be specified using the mlflow.set_experiment Python API.

Python
import mlflow

from mlflow.entities.trace_location import UnityCatalog

mlflow.set_tracking_uri("databricks")

# Specify the catalog, schema, and table prefix to use for storing Traces
catalog_name = "<UC_CATALOG_NAME>"
schema_name = "<UC_SCHEMA_NAME>"
table_prefix = "<UC_TABLE_PREFIX>"

# For existing experiments, it is not necessary to specify `trace_location`. MLflow
# retrieves the UC trace location bound to the experiment and routes traces to
# that location.
mlflow.set_experiment(
    experiment_name="...",
    trace_location=UnityCatalog(
        catalog_name=catalog_name,
        schema_name=schema_name,
        table_prefix=table_prefix,
    ),  # optional for existing experiments
)

# Create and ingest an example trace using the `@mlflow.trace` decorator
@mlflow.trace
def test(x):
    return x + 1

test(100)

View traces in the UI​

View traces stored in OTel format the same way you view other traces:

  1. In your Workspace, go to Experiments.

  2. Find the experiment where your traces are logged. For example, the experiment set by mlflow.set_experiment("/Shared/my-genai-app-traces").

  3. Click the Traces tab to see a list of all traces logged to that experiment.

    Trace List View

  4. If you stored your traces in a Unity Catalog table, Databricks retrieves traces using an SQL warehouse. Select a SQL warehouse from the drop-down menu.

For more information on using the UI to search for traces, see View traces in the Databricks MLflow UI.

Enable production monitoring​

To use production monitoring with traces stored in Unity Catalog, you must configure a SQL warehouse ID for the experiment. The monitoring job requires this configuration to execute scorer queries against Unity Catalog tables.

Set the SQL warehouse ID using set_databricks_monitoring_sql_warehouse_id():

Python
from mlflow.tracing import set_databricks_monitoring_sql_warehouse_id

# Set the SQL warehouse ID for monitoring
set_databricks_monitoring_sql_warehouse_id(
    sql_warehouse_id="<SQL_WAREHOUSE_ID>",
    experiment_id="<EXPERIMENT_ID>"  # Optional, uses active experiment if not specified
)

Alternatively, you can set the MLFLOW_TRACING_SQL_WAREHOUSE_ID environment variable before starting monitoring.

If you skip this step, monitoring jobs fail with an error indicating that the mlflow.monitoring.sqlWarehouseId experiment tag is missing.

To configure monitoring for Unity Catalog traces, you need the following workspace-level permissions:

  • CAN USE permission on the SQL warehouse
  • CAN EDIT permission on the MLflow experiment
  • Permission on the monitoring job (automatically granted when you register the first scorer)

The monitoring job runs under the identity of the user who first registered a scorer on the experiment. This user's permissions determine what the monitoring job can access.

Limitations​

  • Trace ingestion is limited to 200 traces per second per workspace and 100 MB per second per table.

  • An experiment can only be bound to a Unity Catalog trace location at experiment creation time.

  • Traces stored in Unity Catalog are not supported with Knowledge Assistant or Supervisor Agent.

  • Deleting individual traces is not supported for traces stored in Unity Catalog. To remove traces, you must delete rows directly from the underlying Unity Catalog tables using SQL. This differs from experiment traces, which can be deleted using the MLflow UI or API.

  • MLflow MCP server does not support interacting with traces stored in Unity Catalog.

Next steps​

On this page