はじめに

MLflow 3のトレーシングは、GenAIアプリやエージェントの入出力、中間ステップ、メタデータを記録する強力な可観測性機能です。従来、トレースはMLflowエクスペリメント単位でMLflowコントロールプレーンに保存されていましたが、新たにUnity Catalog(UC)のDeltaテーブルにOpenTelemetry(OTEL)形式で保存できるようになりました(ベータ版)。

この記事では、最新のAPI(MLflow 3.11.0)を使ってトレースをUCに保存する方法を解説します。

従来のトレース保存との違い

従来方式とUC保存方式の違いを整理します。

項目	従来方式(MLflowコントロールプレーン)	UC保存方式(OTEL形式)
保存先	MLflowサービス内	Unity CatalogのDeltaテーブル
アクセス制御	エクスペリメントレベルのACL	UCスキーマ/テーブル権限
トレースID形式	`tr-<UUID>`	URI形式(外部システムとの互換性向上)
保存容量	制限あり	無制限(Deltaテーブル)
SQLクエリ	不可	Databricks SQLウェアハウスで直接クエリ可能
外部ツール連携	MLflow固有	OTEL準拠のサードパーティクライアントから書き込み可能
ビジネスデータとの結合	困難	SQLで結合・分析可能
ガバナンス	エクスペリメント単位	UCのカラムマスキング、行フィルタリング等が適用可能

アーキテクチャ

トレースがDeltaテーブルとして保存されることで、レイクハウス上の他のデータと同様にETLパイプライン、ダッシュボード、Genieによる自然言語分析、PII管理等のガバナンスを適用できるようになります。

前提条件

Unity Catalog対応のワークスペース
ワークスペースのPreviewsページで「OpenTelemetry on Databricks」が有効化済み
CAN USE権限を持つDatabricks SQLウェアハウス
MLflow 3.11.0以降
サポートリージョン: us-east-1, us-east-2, us-west-2, ca-central-1, ap-southeast-1, ap-southeast-2, ap-northeast-1(東京), ap-northeast-2, ap-south-1, eu-central-1, eu-west-1, eu-west-2, sa-east-1

セットアップ

ライブラリのインストール

%pip install "mlflow[databricks]>=3.11.0" --upgrade --force-reinstall
dbutils.library.restartPython()

UCトレース場所を使用してエクスペリメントを作成

最新のAPIではmlflow.set_experimentのtrace_locationパラメータで、エクスペリメント作成時にUCトレース場所を直接バインドします。従来のset_experiment_trace_locationは不要になりました。

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

mlflow.set_tracking_uri("databricks")
os.environ["MLFLOW_TRACING_SQL_WAREHOUSE_ID"] = "<SQL_WAREHOUSE_ID>"

catalog_name = "<UC_CATALOG_NAME>"
schema_name = "<UC_SCHEMA_NAME>"
table_prefix = "<UC_TABLE_PREFIX>"
experiment_name = "/Users/<your_email>/uc_trace_demo"

# スキーマの作成(存在しない場合)
spark.sql(f"CREATE CATALOG IF NOT EXISTS {catalog_name}")
spark.sql(f"CREATE SCHEMA IF NOT EXISTS {catalog_name}.{schema_name}")

# mlflow.set_experimentはupsert操作
experiment = mlflow.set_experiment(
    experiment_name=experiment_name,
    trace_location=UnityCatalog(
        catalog_name=catalog_name,
        schema_name=schema_name,
        table_prefix=table_prefix,
    ),
)
print(f"Experiment ID: {experiment.experiment_id}")
print(f"OTELスパンテーブル: {experiment.trace_location.full_otel_spans_table_name}")

注意: エクスペリメントをUCトレース場所にバインドすると、別のUCトレース場所に再割り当てすることはできません。ただし、複数のエクスペリメントが同じUCトレース場所を共有できます。

テーブルの確認

セットアップコードを実行すると、スキーマ内に4つのテーブルが自動作成されます。

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

display(spark.sql(f"SHOW TABLES IN {catalog_name}.{schema_name}"))

ここでは、2つのビューが含まれています。

トレースの記録

シンプルな関数のトレース

@mlflow.traceデコレータでトレースを記録します。

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

@mlflow.trace
def multiply(x, y):
    return x * y

@mlflow.trace
def calculate(x, y):
    """加算と乗算を行い、結果を返す"""
    sum_result = add(x, y)
    mul_result = multiply(x, y)
    return {"sum": sum_result, "product": mul_result}

result = calculate(3, 5)
print(f"結果: {result}")

注意: トレースをUnity Catalogに保存した場合、ノートブックのセル出力にインラインのトレースビューが表示されません。トレースの確認はMLflowエクスペリメントUIのTracesタブから行ってください。

ちなみに、エクスペリメントには概要タブも追加されていました。

LLMコールのトレース

Foundation Model APIを使用したLLMコールをトレースします。

from mlflow.deployments import get_deploy_client

client = get_deploy_client("databricks")

@mlflow.trace(span_type="CHAIN")
def ask_llm(question: str) -> str:
    """LLMに質問を送り、回答を得る"""
    response = client.predict(
        endpoint="databricks-claude-sonnet-4-6",
        inputs={
            "messages": [
                {"role": "system", "content": "あなたは親切なアシスタントです。簡潔に回答してください。"},
                {"role": "user", "content": question},
            ],
            "max_tokens": 256,
        },
    )
    return response["choices"][0]["message"]["content"]

answer = ask_llm("Databricksにおけるunity catalogの主な利点を3つ教えてください。")
print(answer)

RAGパイプラインのトレース

検索 → コンテキスト構築 → LLM生成の典型的なRAGフローをトレースします。

@mlflow.trace(span_type="RETRIEVER")
def retrieve_documents(query: str) -> list:
    """ダミーのドキュメント検索"""
    docs = [
        {"content": "Unity CatalogはデータとAI資産の統合ガバナンスを提供します。", "score": 0.95},
        {"content": "Delta Lakeはオープンソースのストレージフレームワークです。", "score": 0.82},
        {"content": "MLflowはMLライフサイクル管理のためのプラットフォームです。", "score": 0.78},
    ]
    return docs

@mlflow.trace(span_type="PARSER")
def build_context(documents: list) -> str:
    """検索結果からコンテキストを構築"""
    return "\n".join([doc["content"] for doc in documents])

@mlflow.trace(span_type="CHAIN")
def rag_pipeline(question: str) -> str:
    """RAGパイプライン全体"""
    docs = retrieve_documents(question)
    context = build_context(docs)

    prompt = f"""以下のコンテキストに基づいて質問に回答してください。

コンテキスト:
{context}

質問: {question}
"""
    response = client.predict(
        endpoint="databricks-claude-sonnet-4-6",
        inputs={
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": 512,
        },
    )
    return response["choices"][0]["message"]["content"]

answer = rag_pipeline("Databricksのデータガバナンス機能について教えてください。")
print(answer)

トレースの検索

Python SDKによる検索

mlflow.search_traces()のlocationsパラメータでUCトレース場所を指定します(従来のexperiment_idsは非推奨)。

traces_df = mlflow.search_traces(
    locations=[f"{catalog_name}.{schema_name}.{table_prefix}"],
)
# Arrow変換エラーを回避するため、object型の列を文字列に変換
for col in traces_df.columns:
    if traces_df[col].dtype == "object":
        traces_df[col] = traces_df[col].astype(str)
display(traces_df)

SQLによる直接クエリ

UCのDeltaテーブルに保存されているため、SQLで直接クエリできます。これが従来方式にはなかった大きなメリットです。

SELECT
  trace_id,
  name,
  start_time_unix_nano,
  end_time_unix_nano,
  (end_time_unix_nano - start_time_unix_nano) / 1e6 AS duration_ms,
  status.code AS status_code
FROM <catalog>.<schema>.<table_prefix>_otel_spans
ORDER BY start_time_unix_nano DESC
LIMIT 20

本番運用モニタリングの有効化

UCに保存されたトレースに対して本番運用モニタリングを使用する場合は、エクスペリメントにSQLウェアハウスIDを設定します。

from mlflow.tracing import set_databricks_monitoring_sql_warehouse_id

set_databricks_monitoring_sql_warehouse_id(
    sql_warehouse_id=sql_warehouse_id,
    experiment_id=experiment.experiment_id,
)

注意: この設定は、エクスペリメントにMLflowスコアラーを登録して初めて意味を持ちます。スコアラーを登録すると、モニタリングジョブが本番トレースに対してスコアラークエリを自動実行し、品質評価を継続的に行います。スコアラーを使用しない段階ではこのステップは不要で

制限事項

トレースデータの取り込みは、ワークスペースあたり毎秒200トレース、テーブルあたり毎秒100MBに制限
UCトレース保存時、ノートブックのセル出力にインラインのトレースビューが表示されない場合がある。トレースの確認はMLflowエクスペリメントUIのTracesタブから行う
エクスペリメントはエクスペリメント作成時にのみUCトレース場所にバインド可能(後から再割り当て不可)
2TBを超えるトレースデータではUIパフォーマンスが低下する可能性あり
個別トレースの削除はUI/APIからはできず、SQLでテーブルから直接削除が必要
MLflow MCPサーバーはUCに保存されたトレースとのやり取りをサポートしていない
Knowledge AssistantおよびSupervisor Agentではサポートされていない

まとめ

MLflowトレースをUnity Catalogに保存することで、トレースデータがレイクハウスのファーストクラスのデータになります。従来のMLflowコントロールプレーンへの保存と比較して、UCのガバナンス、SQL分析、ビジネスデータとの結合、スケーラブルなストレージなど、エンタープライズ運用に求められる機能が利用可能になります。

参考リンク

はじめてのDatabricks

はじめてのDatabricks

Databricks無料トライアル

Databricks無料トライアル