手動トレース

MLflow の 自動トレース は、サポートされているフレームワークの即時の監視機能を提供しますが、手動トレースを使用すると、GenAI アプリケーションのインストルメント化方法を完全に制御できます。 この柔軟性は、詳細な監視とデバッグ機能を必要とする運用対応アプリケーションを構築するために不可欠です。

[前提条件]

MLflow 3

このガイドには、次のパッケージが必要です。

• mlflow[databricks]>=3.1: GenAI 機能と Databricks 接続を備えたコア MLflow 機能。
• openai>=1.0.0: (省略可能) カスタム コードが OpenAI と対話する場合にのみ、必要に応じて他の SDK に置き換えます。

基本的な要件をインストールします。

%pip install --upgrade "mlflow[databricks]>=3.1"
# %pip install --upgrade openai>=1.0.0 # Install if needed

MLflow 2.x

このガイドには、次のパッケージが必要です。

• mlflow[databricks]>=2.15.0,<3.0.0: Databricks 接続を使用した MLflow のコア機能。
• openai>=1.0.0: (省略可能) カスタム コードが OpenAI と対話する場合のみ。

基本的な要件をインストールします。

%pip install --upgrade "mlflow[databricks]>=2.15.0,<3.0.0"
# pip install --upgrade openai>=1.0.0 # Install if needed

手動トレースを使用する場合

手動トレースは、必要に応じ適切な選択肢です。

Fine-Grained コントロール

カスタム トレース構造

• トレースするコードのどの部分を正確に定義するか
• スパンのカスタム階層を作成する
• 範囲の境界と関係を制御する

ユース ケースの例: 取得と生成の待機時間を個別に測定する RAG パイプライン内の特定のビジネス ロジックをトレースする。

カスタム フレームワーク

サポートされていないライブラリ

• 独自のフレームワークまたは内部フレームワークをインストルメント化する
• カスタム LLM ラッパーにトレースを追加する
• 公式統合の前に新しいライブラリをサポートする

ユース ケースの例: 会社の内部 LLM ゲートウェイまたはカスタム エージェント フレームワークにトレースを追加する。

高度なシナリオ

複雑なワークフロー

• マルチスレッドまたは非同期操作
• カスタム集計を使用した応答のストリーミング
• 複雑なネストされた操作
• カスタム トレースメタデータと属性

ユース ケースの例: エージェントがカスタム ビジネス ロジックを使用して複雑なワークフローを実行するマルチエージェント システムのトレース。

手動トレースのアプローチ

MLflow には、手動トレース用の 3 つの抽象化レベルが用意されており、それぞれ異なるユース ケースに適しています。

1. High-Level API (推奨)

高度な API は、最小限のコード変更でトレースを追加する直感的な方法を提供します。 トレースライフサイクル、例外追跡、親子関係を自動的に処理します。

デコレーター

最適: コード変更を最小限に抑えた関数レベルのトレース

import mlflow
from mlflow.entities import SpanType

@mlflow.trace(span_type=SpanType.CHAIN)
def process_request(query: str) -> str:
    # Your code here - automatically traced!
    result = generate_response(query)
    return result

@mlflow.trace(span_type=SpanType.LLM)
def generate_response(query: str) -> str:
    # Nested function - parent-child relationship handled automatically
    return llm.invoke(query)

主な利点:

• 1 行インストルメンテーション
• 例外の自動処理
• 非同期/ジェネレーター関数で動作します
• 自動トレースと互換性がある

デコレーター →の詳細を確認する

コンテキスト マネージャー

最適な方法: コード ブロックと複雑なワークフローのトレース

import mlflow

with mlflow.start_span(name="data_processing") as span:
    # Set inputs at the start
    span.set_inputs({"query": query, "filters": filters})

    # Your processing logic
    data = fetch_data(query, filters)
    processed = transform_data(data)

    # Set outputs before exiting
    span.set_outputs({"count": len(processed), "status": "success"})

主な利点:

• 柔軟なスパン境界
• 動的な入出力設定
• スパン ライフサイクルを細かく制御する
• 関数以外のコード ブロックに最適

コンテキスト マネージャーの→の詳細を確認する

2. Low-Level クライアント API (詳細)

トレース ライフサイクルを完全に制御する必要があるシナリオでは、クライアント API は MLflow のトレース バックエンドに直接アクセスできます。

from mlflow import MlflowClient

client = MlflowClient()

# Start a trace
root_span = client.start_trace("complex_workflow")

# Create child spans with explicit parent relationships
child_span = client.start_span(
    name="data_retrieval",
    request_id=root_span.request_id,
    parent_id=root_span.span_id,
    inputs={"query": query}
)

# End spans explicitly
client.end_span(
    request_id=child_span.request_id,
    span_id=child_span.span_id,
    outputs={"documents": documents}
)

# End the trace
client.end_trace(request_id=root_span.request_id)

クライアント API を使用する場合:

• カスタム トレース ID 管理
• 既存の可観測性システムとの統合
• 複雑なトレース ライフサイクル要件
• 標準以外のトレース パターン

クライアント API →の詳細を確認する

API の比較

一般的なパターン

自動トレースとの組み合わせ

手動トレースは、MLflow の自動トレース機能とシームレスに統合されます。

import mlflow
import openai

# Enable auto-tracing for OpenAI
mlflow.openai.autolog()

@mlflow.trace(span_type="CHAIN")
def rag_pipeline(query: str):
    # Manual span for retrieval
    with mlflow.start_span(name="retrieval") as span:
        docs = retrieve_documents(query)
        span.set_outputs({"doc_count": len(docs)})

    # Auto-traced OpenAI call
    response = openai.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": f"Answer based on: {docs}\n\nQuery: {query}"}]
    )

    return response.choices[0].message.content

複雑なワークフロー トレース

複数のステップを含む複雑なワークフローの場合は、入れ子になったスパンを使用して詳細な実行フローをキャプチャします。

@mlflow.trace(name="data_pipeline")
def process_data_pipeline(data_source: str):
    # Extract phase
    with mlflow.start_span(name="extract") as extract_span:
        raw_data = extract_from_source(data_source)
        extract_span.set_outputs({"record_count": len(raw_data)})

    # Transform phase
    with mlflow.start_span(name="transform") as transform_span:
        transformed = apply_transformations(raw_data)
        transform_span.set_outputs({"transformed_count": len(transformed)})

    # Load phase
    with mlflow.start_span(name="load") as load_span:
        result = load_to_destination(transformed)
        load_span.set_outputs({"status": "success"})

    return result

UI での要求と応答のプレビューのカスタマイズ

MLflow UI には、[トレース] タブに Request 列と Response 列があり、トレース全体の入力と出力のプレビューが表示されます。 既定では、これらは切り捨てられます。 手動トレース (特にトレースのルート スパンを作成する @mlflow.trace デコレーターまたはコンテキスト マネージャー) を使用する場合は、 mlflow.update_current_trace()を使用してこれらのプレビューをカスタマイズできます。

これは、既定のプレビューでは情報が得られない可能性がある複雑なデータ構造に役立ちます。

import mlflow
import openai # Assuming openai is used, replace if not

# This example assumes you have an OpenAI client initialized and API key set up.
# client = openai.OpenAI()

@mlflow.trace
def predict(messages: list[dict]) -> str:
    # Customize the request preview to show the first and last messages
    custom_preview = f'{messages[0]["content"][:10]} ... {messages[-1]["content"][:10]}'
    mlflow.update_current_trace(request_preview=custom_preview)

    # Call the model
    # response = openai.chat.completions.create(
    #     model="gpt-4o-mini",
    #     messages=messages,
    # )
    # return response.choices[0].message.content
    return f"Response based on {len(messages)} messages."

messages = [
    {"role": "user", "content": "Hi, how are you?"},
    {"role": "assistant", "content": "I'm good, thank you!"},
    {"role": "user", "content": "What's your name?"},
    # ... (long message history)
    {"role": "assistant", "content": "Bye!"},
]
predict(messages)

これにより、特定のデータ構造に対してより有益になるようにプレビューを調整できます。

次のステップ

これらの推奨されるアクションとチュートリアルを使用して、体験を続けます。

• デコレーターと Fluent API - 高度な API を使用するほとんどのユース ケースでここから始める
• Low-Level クライアント API - フル コントロールを必要とする高度なシナリオについて説明します
• アプリのデバッグと監視 - 手動でトレースしたアプリをデバッグに使用する

リファレンス ガイド

このガイドで説明されている概念と機能の詳細なドキュメントを確認します。

• トレース データ モデル - スパンとトレース構造を理解する
• トレースの概念 - トレースの基礎について学習する
• 自動トレース - 手動トレースと自動トレースの組み合わせ