次の方法で共有

はじめに: GenAI の MLflow トレース (Databricks Notebook)

このページのノートブック バージョンを開く

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

このチュートリアルの最後には、次の内容が含まれます。

  • GenAI アプリ用のリンクされた MLflow 実験を含む Databricks ノートブック
  • MLflow トレースを使用してインストルメント化されたシンプルな GenAI アプリケーション
  • MLflow 実験でそのアプリからのトレース

チュートリアル「クイックスタート」のトレースUI

環境のセットアップ

  1. Databricks ワークスペースに新しいノートブックを作成します。 ノートブックには、GenAI アプリケーションのコンテナーである既定の MLflow 実験があります。 MLflow の実験の詳細については、「 MLflow の概念」セクションを参照してください。

  2. 必要なパッケージをインストールする:

    • mlflow[databricks]:MLflow の最新バージョンを使用して、より多くの機能と改善を得ます。
    • openai: このチュートリアルでは、OpenAI API クライアントを使用して Databricks でホストされるモデルを呼び出します。
%pip install -qq --upgrade "mlflow[databricks]>=3.1.0" openai
dbutils.library.restartPython()

手順 1: トレースを使用してアプリケーションをインストルメント化する

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

まず、 Databricks でホストされる基盤モデルに接続する OpenAI クライアントを作成します。

from databricks.sdk import WorkspaceClient

# 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 でホストされるモデルに接続することもできます。

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

client = openai.OpenAI()

model_name = "gpt-4o-mini"

次に、アプリケーションを定義して実行します。 トレースを使用してアプリをインストルメント化するには、次のものが使用されます。

  • mlflow.openai.autolog(): OpenAI SDK への呼び出しの詳細をキャプチャするための自動インストルメンテーション
  • @mlflow.trace:任意の Python 関数を簡単にトレースできるデコレーター
import mlflow
import os

# 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=200,
    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 以上のライブラリ統合を参照してください。

手順 2: MLflow でトレースを表示する

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

ノートブック トラッキング UI

必要に応じて、MLflow 実験 UI に移動してトレースを表示できます。

  1. 右側のサイドバーで、[実験] アイコン [実験] アイコン をクリックします。
  2. 開いているアイコン [ 新しいウィンドウ] アイコン をクリックします。実験の実行の横にあります。
  3. 生成されたトレースが [ トレース ] タブに表示されます。
  4. トレースをクリックして詳細を表示します。

トレースを理解する

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

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

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

  • 質問された内容
  • 生成された応答
  • 要求にかかった時間
  • 使用されたトークンの数 (コストに影響)

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

次のステップ

ノートブックの例

はじめに: GenAI の MLflow トレース (Databricks Notebook)

ノートブックを入手

  • Last updated on