メインコンテンツまでスキップ

手動トレース

MLflowの自動トレース では、サポートされているフレームワークの可観測性がすぐに提供されますが、手動トレースでは、生成AI アプリケーションを計測可能にする方法を完全に制御できます。この柔軟性は、詳細なモニタリング機能とデバッグ機能を必要とする本番運用に対応したアプリケーションを構築するために不可欠です。

前提 条件​

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

  • mlflow[databricks]>=3.1 : 生成AI 機能と Databricks 接続を備えたコア MLflow 機能。
  • openai>=1.0.0 : (オプション) カスタムコードがOpenAIと相互作用する場合のみ。必要に応じて、他のSDKと交換します。

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

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

MLflow バージョンの推奨事項

MLflow 2.15.0+ では手動トレース機能を使用できますが、拡張されたトレース機能、洗練されたスパンの種類、改善されたコンテキスト伝達、堅牢なサポートなど、最新の 生成AI 機能を使用するには、 MLflow 3 ( mlflow[databricks]を使用している場合は 3.1 以降) をインストールすることを強くお勧めします 。

ヒント

Databricks ノートブックで実行していますか? MLflow はプレインストールされています。追加の SDK をインストールする必要があるのは、手動でトレースしたコードで SDK を使用している場合のみです。

ローカルで実行していますか? mlflow[databricks] SDK と、コードで呼び出すその他の SDK をインストールする必要があります。

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

手動トレースは、次の必要がある場合に最適です。

カスタムトレース構造

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

ユースケースの例 :RAGパイプライン内の特定のビジネスロジックをトレースし、取得と生成のレイテンシを別々に測定する場合。

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

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

1. ハイレベルな API (推奨)​

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

最適な用途: 最小限のコード変更で関数レベルのトレースを実現

Python
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行で計測可能に
  • 自動例外処理
  • async/generator 関数と連携
  • オートトレースに対応

デコレーター→についてもっと知る

2. 低レベルクライアント API (上級)​

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

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

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

  • カスタムトレースID管理
  • 既存のオブザーバビリティシステムとの統合
  • 複雑なトレース・ライフサイクル要件
  • 非標準のトレースパターン
警告

クライアント API では、次のものを手動で管理する必要があります。

  • 親子関係
  • スパンライフサイクル(開始/終了)
  • 例外処理
  • スレッドセーフ

クライアント APIの詳細 →

APIの比較​

機能

デコレータ

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

クライアント API

自動的な親子関係

Yes

Yes

いいえ - 手動管理

例外処理

自動

自動

手動

オートトレースと連携

Yes

Yes

No

スレッドセーフ

自動

自動

手動

カスタムトレースID

No

No

Yes

どのようなタスクにベストなのか

関数トレース

コードブロックのトレース

高度な制御

一般的なパターン​

オートトレースとの組み合わせ​

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

Python
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

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

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

Python
@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()を使用してこれらのプレビューをカスタマイズできます。

これは、デフォルトのプレビューでは情報が提供されない可能性がある複雑なデータ構造に役立ちます。

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

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

次のステップ​

これらの推奨アクションとチュートリアルで旅を続けてください。

リファレンスガイド​

このガイドで説明されている概念と機能の詳細なドキュメントをご覧ください。

ヒント

ほとんどのユーザーは、高レベルの API (デコレーターとコンテキストマネージャー) から始める必要があります。 これらは、MLflow のエコシステムとの互換性を維持しながら、使いやすさと機能性の最適なバランスを提供します。

に最終更新