はじめに: 生成AI向けMLflow Tracing (Databricksノートブック)

このクイックスタートは、開発環境としてノートブックを使用する場合に、生成AI アプリをMLflow Tracing Databricksと統合するのに役立ちます。ローカル IDE を使用している場合は、代わりに IDE クイックスタートを使用してください。

達成できること

このチュートリアルを終了すると、次のことができるようになります。

生成AIアプリ用のMLflowエクスペリメントがリンクされたDatabricksノートブック
MLflow Tracing で計測可能になったシンプルな生成AI アプリケーション
MLflowエクスペリメント内のそのアプリからのトレース

trace

前提条件

Databricksノートブック : Databricksワークスペースに新しいノートブックを作成します。ノートブックには、GenAI アプリケーションのコンテナーである安全MLflowエクスペリメントがあります。 MLflowエクスペリメントの詳細については、MLflowコンセプトのセクションをご覧ください。

ステップ 1: MLflowを更新する

最高の生成AI エクスペリエンスを得るには、MLflow を最新バージョンに更新してください。

Python
%pip install --upgrade "mlflow[databricks]>=3.1" openai
dbutils.library.restartPython()

ステップ 2: アプリケーションを計測可能にする

以下のコードスニペットは、LLM を使用して文テンプレートを完成させる単純な生成AI アプリを定義します。

OpenAI クライアントを初期化して、Databricks がホストする LLM または OpenAI がホストする LLM に接続します。

Databricks-hosted LLMs
OpenAI SDK

以下のコードスニペットでは、OpenAI クライアントを使用して、Databricks がホストする LLM に接続します。アプリはAnthropicの Claude Sonnet LLM使用しますが、利用可能な基盤モデルから選択することもできます。

Python
import mlflow
from databricks.sdk import WorkspaceClient

# Enable MLflow's autologging to instrument your application with Tracing
mlflow.openai.autolog()

# Create an OpenAI client that is connected to Databricks-hosted LLMs
w = WorkspaceClient()
client = w.serving_endpoints.get_open_ai_client()

model_name = "databricks-claude-sonnet-4"

ネイティブの OpenAI SDK を使用して、OpenAI でホストされるモデルに接続します。以下のコードスニペットでは gpt-4o-miniを使用していますが、使用可能な OpenAI モデルから選択できます。

Python
import mlflow
import openai

# Ensure your OPENAI_API_KEY is set in your environment
# os.environ["OPENAI_API_KEY"] = "<YOUR_API_KEY>" # Uncomment and set if not globally configured

# Enable MLflow's autologging to instrument your application with Tracing
mlflow.openai.autolog()

client = openai.OpenAI()

model_name = "gpt-4o-mini"

アプリケーションを定義して実行します。

@mlflow.traceデコレータを使用すると、任意の Python 関数を簡単にトレースでき、 OpenAI 自動インストルメンテーションと組み合わせて、OpenAI SDK への呼び出しの詳細をキャプチャできます。

Python
import mlflow
import os

# Ensure your OPENAI_API_KEY is set in your environment
# os.environ["OPENAI_API_KEY"] = "<YOUR_API_KEY>" # Uncomment and set if not globally configured

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

# Set up MLflow tracking to Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/openai-tracing-demo")

# Use the trace decorator to capture the application's entry point
@mlflow.trace
def my_app(input: str):
  # This call is automatically instrumented by `mlflow.openai.autolog()`
  response = client.chat.completions.create(
    model=model_name,
    temperature=0.1,
    max_tokens=100,
    messages=[
      {
        "role": "system",
        "content": "You are a helpful assistant.",
      },
      {
        "role": "user",
        "content": input,
      },
    ]
  )
  return response.choices[0].message.content

result = my_app(input="What is MLflow?")
print(result)

アプリにトレースを追加する方法の詳細については、トレースインストルメンテーションガイドと20 を超えるライブラリ統合を参照してください。

ステップ 3: MLflowでトレースを表示する

トレースはノートブックのセルの下に表示されます。

ノートブックのトレース UI

オプションとして、 MLflowエクスペリメントUIにアクセスしてトレースを表示することもできます。

エクスペリメントアイコンをクリックします右側のサイドバーにあります。
開くアイコンをクリックします体験実行の隣にあります。
生成されたトレースが [トレース] タブに表示されます。
トレースをクリックして、その詳細を表示します。

トレースを理解する

作成したトレースには、次の情報が表示されます。

ルートスパン : my_app(...)関数への入力を表します
- 子スパン : OpenAI完了リクエストを表す
属性 : モデル名、トークン数、タイミング情報などのメタデータが含まれます
入力 : モデルに送信されたメッセージ
出力 : モデルから受信した応答

この単純なトレースは、次のようなアプリケーションの動作に関する貴重な知見をすでに提供しています。

質問されたこと
どのようなレスポンスが生成されたか
リクエストにかかった時間
使用されたトークンの数 (コストに影響)

RAGシステムやマルチステップエージェントのようなより複雑なアプリケーションの場合、 MLflow Tracing は各コンポーネントとステップの内部動作を明らかにすることで、さらに多くの価値を提供します。

ガイドと参考資料

このガイドの概念と機能の詳細については、以下を参照してください。

MLflow Tracingガイド- MLflow Tracingについてさらに詳しく知りたい場合は、ここから始めてください。
MLflow Tracing統合- 自動トレース統合を備えた 20 以上のライブラリ
トレースの概念 - MLflow Tracingの基本を理解する

達成できること​

前提 条件​

ステップ 1: MLflowを更新する​

ステップ 2: アプリケーションを計測可能にする​

ステップ 3: MLflowでトレースを表示する​

トレースを理解する​

ガイドと参考資料​