Unity Catalogに保存されているMLflowトレースを照会する

備考

ベータ版

この機能はベータ版です。ワークスペース管理者は、 プレビュー ページからこの機能へのアクセスを制御できます。Databricksのプレビューを管理するを参照してください。

トレースデータをOpenTelemetry形式でUnity Catalogに保存することで、 MLflow Python SDKを使用したり、 Unity Catalogテーブルとビューを使用してDatabricks SQL経由でトレースをクエリしたりできます。

前提条件

トレースをUnity Catalogテーブルに保存し、トレースを生成します。「 MLflow トレースを Unity Catalog に保存する」を参照してください。

MLflow Python SDKを使用してトレースをクエリする

MLflow Python SDK を使用して、トレースオブジェクトを検索および読み込みます。

MLFLOW_TRACING_SQL_WAREHOUSE_ID環境変数を使用して、検索クエリを実行するためのDatabricks SQL ウェアハウスを指定します。
mlflow.search_tracesのlocations引数を使用して、トレースを含む 1 つ以上のMLflow拡張機能またはUnity Catalogスキーマを指定します。
Unity Catalogスキーマの名前、またはUnity CatalogスキーマにリンクされたMLflow拡張機能のIDを指定できます。「セットアップ: Unity Catalogトレース場所を使用してエクスペリメントを作成する」を参照してください。

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

# Specify the name of a catalog and schema containing traces
catalog_name = "<UC_CATALOG>"
schema_name = "<UC_SCHEMA>"
table_prefix = "<UC_TABLE_PREFIX>"

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment(
    experiment_name="...",
    trace_location=UnityCatalog(
        catalog_name=catalog_name,
        schema_name=schema_name,
        table_prefix=table_prefix,
    ),  # optional for existing experiments
)

# Specify the ID of a Databricks SQL warehouse for executing search queries
os.environ["MLFLOW_TRACING_SQL_WAREHOUSE_ID"] = "<SQL_WAREHOUSE_ID>"

traces = mlflow.search_traces(
    filter_string="trace.status = 'OK'",
    order_by=["timestamp_ms DESC"],
    include_spans=False,
)
print(traces)

見つかったトレースを読み込むには:

Python
import os
import mlflow

mlflow.set_tracking_uri("databricks")

# Specify the name of a catalog and schema containing traces
catalog_name = "<UC_CATALOG>"
schema_name = "<UC_SCHEMA>"
table_prefix = "<UC_TABLE_PREFIX>"
# Specify the trace UUID (example: "13ffa97d571048d69d21da12240d5863")
trace_uuid = "<TRACE_UUID>"

# Specify the ID of a Databricks SQL warehouse for executing search queries
os.environ["MLFLOW_TRACING_SQL_WAREHOUSE_ID"] = "<SQL_WAREHOUSE_ID>"

trace = mlflow.get_trace(
    trace_id=f"trace:/{catalog_name}.{schema_name}.{table_prefix}/{trace_uuid}"
)
print(trace)

Databricks SQL を使用してトレースをクエリする

基となるデータはOpenTelemetry準拠のテーブル形式で保存されますが、 MLflowサービスはそれらのテーブルと並行してDatabricks SQLビューを自動的に作成します。これらのビューは、OpenTelemetryデータをMLflow形式に変換します。

トレースデータ量が大きい場合、これらのビューに対するクエリのパフォーマンスが低下する可能性があります。パフォーマンスを維持するために、それらの上にマテリアライズドビューを作成し、マテリアライズドビューを段階的に更新します。最新データで最高のパフォーマンスを得るには、 APIを使用してトレースを照会してください。

Databricksは、基となるテーブルに依存するのではなく、ビューをクエリするかAPIを使用することを推奨しています。なぜなら、テーブルのスキーマは時間の経過とともに変更される可能性があるからです。

`{table_prefix}_trace_unified`

このビューでは、トレースIDごとにグループ化されたすべてのトレースデータを統一的に表示します。各行には、生のスパンデータとトレース情報メタデータが含まれています。メタデータには、 MLflowタグ、メタデータ、評価が含まれます。

スキーマ

Python
trace_id: STRING
client_request_id: STRING
request_time: TIMESTAMP
state: STRING
execution_duration_ms: DECIMAL(30,9)
request: STRING
response: STRING
trace_metadata: MAP<STRING, STRING>
tags: MAP<STRING, STRING>
spans: LIST<STRUCT>
    trace_id: STRING
    span_id: STRING
    trace_state: STRING
    parent_span_id: STRING
    flags: INT
    name: STRING
    kind: STRING
    start_time_unix_nano: BIGINT
    end_time_unix_nano: BIGINT
    attributes: MAP<STRING, STRING>
    dropped_attributes_count: INT
    events: LIST<STRUCT>
        time_unix_nano: BIGINT
        name: STRING
        attributes: MAP<STRING, STRING>
        dropped_attributes_count: INT
    dropped_events_count: INT
    links: LIST<STRUCT>
        trace_id: STRING
        span_id: STRING
        trace_state: STRING
        attributes: MAP<STRING, STRING>
        dropped_attributes_count: INT
        flags: INT
    dropped_links_count: INT
    status: STRUCT
        message: STRING
        code: STRING
    resource: STRUCT
        attributes: MAP<STRING, STRING>
        dropped_attributes_count: INT
    resource_schema_url: STRING
    instrumentation_scope: STRUCT
        name: STRING
        version: STRING
        attributes: MAP<STRING, STRING>
        dropped_attributes_count: INT
    span_schema_url: STRING
assessments: LIST<STRUCT>
    assessment_id: STRING
    trace_id: STRING
    assessment_name: STRING
    source: STRUCT
        source_id: STRING
        source_type: STRING
    create_time: TIMESTAMP
    last_update_time: TIMESTAMP
    expectation: STRUCT
        value: STRING
        serialized_value: STRUCT
            serialization_format: STRING
            value: STRING
            stack_trace: STRING
    feedback: STRUCT
        value: STRING
        error: STRUCT
            error_code: STRING
            error_message: STRING
            stack_trace: STRING
    rationale: STRING
    metadata: MAP<STRING, STRING>
    span_id: STRING
    overrides: STRING
    valid: STRING

`{table_prefix}_trace_metadata`

このビューには、トレース ID 別にグループ化された MLflow タグ、メタデータ、および評価のみが含まれ、MLflow データを取得する場合の統合ビューよりもパフォーマンスが優れています。

スキーマ

Python
trace_id: STRING
client_request_id: STRING
tags: MAP<STRING, STRING>
trace_metadata: MAP<STRING, STRING>
assessments: LIST<STRUCT>
    assessment_id: STRING
    trace_id: STRING
    assessment_name: STRING
    source: STRUCT
        source_id: STRING
        source_type: STRING
    create_time: TIMESTAMP
    last_update_time: TIMESTAMP
    expectation: STRUCT
        value: STRING
        serialized_value: STRUCT
            serialization_format: STRING
            value: STRING
            stack_trace: STRING
    feedback: STRUCT
        value: STRING
        error: STRUCT
            error_code: STRING
            error_message: STRING
            stack_trace: STRING
    rationale: STRING
    metadata: MAP<STRING, STRING>
    span_id: STRING
    overrides: STRING
    valid: STRING

MLflowアノテーションデータ形式

メタデータ、タグ、評価、実行へのリンクなど、MLflow トレースエンティティのデータは、 {table_prefix}_otel_annotationsテーブルに格納されます。各エンティティは型annotation_typeを持つ単一の行として格納され、そのフィールドはトップレベルの列 ( name 、 value 、 comment 、 metadata ) に分割されます。注釈テーブルは追記専用でソフト削除機能があるため、取得時にannotation_idごとに最新の行を取得し（ updated_atで降順に並べ替え）、 deleted_atが設定されている行を除外して重複を削除する必要があります。value列とmetadata列はVARIANT （JSON）です。

この表には以下の列があります。

Python
annotation_id: STRING
target_type: STRING ("TRACE" or "SPAN")
target_id: STRING ("{trace_id}" for TRACE, "{trace_id}:{span_id}" for SPAN)
annotation_type: STRING ("METADATA", "TAG", "FEEDBACK", "EXPECTATION", "RUN_LINK")
name: STRING
value: VARIANT
comment: STRING
metadata: VARIANT
created_at: TIMESTAMP
created_by: STRING
updated_at: TIMESTAMP
updated_by: STRING
deleted_at: TIMESTAMP
deleted_by: STRING

MLflowメタデータ

トレースごとに、これらの行は1つだけ存在します。value列は、トレースのクライアントリクエストID、メタデータマップ、およびリクエスト/レスポンスのプレビューを含むJSON構造体です。

Python
annotation_type: "METADATA"
target_type: "TRACE"
name: "metadata"
value: VARIANT (includes `client_request_id`, `trace_metadata`, `request_preview`, `response_preview`)

MLflowタグ

各タグは別々の行として保存されます。トレースIDとタグキーから決定論的に導出されるannotation_id属性を使用することで、各トレース内で重複を排除できます。

Python
annotation_type: "TAG"
target_type: "TRACE"
name: STRING (the tag key)
value: STRING (the tag value)

MLflow評価

各評価は、その種類に応じてFEEDBACKまたはEXPECTATION行として保存されます。各トレース内で、評価IDに一致するannotation_id属性を使用することで、重複を排除できます。根拠は最上位のcomment列に格納されています。ユーザーが提供した評価メタデータは、内部の MLflow 管理フィールド (キーにmlflow.がプレフィックスとして付いています) と共にmetadata列に格納されますが、ユーザーメタデータを読み取る際には、これらのフィールドは無視してください。

Python
annotation_type: "FEEDBACK" | "EXPECTATION"
target_type: "TRACE"
name: STRING (the assessment name)
value: VARIANT (feedback value, expectation value, or JSON-serialized expectation string)
comment: STRING (the rationale)
metadata: VARIANT (user-supplied assessment metadata)

MLflow実行リンク

トレースとMLflow実行間の各リンクは、個別の行として保存されます。トレースIDと実行IDから決定論的に導出されるannotation_id属性を使用することで、各トレース内で重複を排除できます。

Python
annotation_type: "RUN_LINK"
target_type: "TRACE"
name: "run_link"
value: STRING (the run ID)

クエリパフォーマンスを分析する

クエリの実行速度が遅い原因を診断するには、SQLウェアハウスのクエリ履歴にあるクエリプロファイルを調べます。

Databricksワークスペースの SQLウェアハウス ページに移動します。
SQLウェアハウスを選択し、 [クエリ履歴] タブをクリックします。
ソースとして MLflow が指定されたクエリを検索します。
クエリをクリックすると、そのクエリのプロファイルが表示されます。

クエリプロファイルで、次の点を検査します。

スケジュール時間 : スケジュール時間が長い場合は、ウェアハウスの負荷が高いため、クエリが待機しています。 MLflow UIのドロップダウンメニューを使用して別のSQLウェアハウスに切り替えるか、クライアントで別のウェアハウスを設定してください。
クエリ全体のパフォーマンス : クエリが常に遅い場合は、より大きなSQLウェアを使用し、 trace.timestamp_msの上限と下限を厳しくし、可能な場合は他のフィルタ述語を削除します。

前提条件​

MLflow Python SDKを使用してトレースをクエリする​

Databricks SQL を使用してトレースをクエリする​

{table_prefix}_trace_unified​

スキーマ​

{table_prefix}_trace_metadata​

スキーマ​

MLflowアノテーションデータ形式​

MLflowメタデータ​

MLflowタグ​

MLflow評価​

MLflow実行リンク​

クエリパフォーマンスを分析する​

前提条件

MLflow Python SDKを使用してトレースをクエリする

Databricks SQL を使用してトレースをクエリする

`{table_prefix}_trace_unified`

スキーマ

`{table_prefix}_trace_metadata`

スキーマ

MLflowアノテーションデータ形式

MLflowメタデータ

MLflowタグ

MLflow評価

MLflow実行リンク

クエリパフォーマンスを分析する