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

MLflow Tracing 統合

MLflow Tracing は、さまざまな一般的な生成AI ライブラリやフレームワークと統合されており、それらすべてに対して 1 行の自動トレース エクスペリエンスを提供します。 これにより、最小限のセットアップで生成AIアプリケーションの可観測性を即座に得ることができます。

自動トレースは、特定のライブラリまたは SDK の実装に基づいて、アプリケーションのロジックと中間ステップ (LLM 呼び出し、ツールの使用、エージェントの対話など) をキャプチャします。

自動トレースの仕組み、前提条件、および手動トレースとの組み合わせ例については、メインの自動 トレース ガイドを参照してください。以下の簡単な例は、いくつかの主要な統合を示しています。サポートされている各ライブラリの詳細なガイドは、前提条件、高度な例、および特定の動作をカバーしており、このセクションのそれぞれのページで入手できます。

トップの統合の概要​

ここでは、最も一般的に使用される統合のクイックスタートの例を示します。タブをクリックすると、基本的な使用例が表示されます。それぞれの詳細な前提条件とより高度なシナリオについては、専用の統合ページ (以下のタブまたは一覧からリンクされています) を参照してください。

Python
import mlflow
import openai

# If running this code outside of a Databricks notebook (e.g., locally),
# uncomment and set the following environment variables to point to your Databricks workspace:
# import os
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"

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

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

openai_client = openai.OpenAI()

messages = [
    {
        "role": "user",
        "content": "What is the capital of France?",
    }
]

response = openai_client.chat.completions.create(
    model="gpt-4o-mini",
    messages=messages,
    temperature=0.1,
    max_tokens=100,
)
# View trace in MLflow UI

OpenAI統合ガイド

安全なAPIキー管理​

本番運用環境の場合、Databricks AIGateway または シークレットを使用してDatabricks APIキーを管理することをお勧めします。AI Gateway は推奨される方法であり、追加のガバナンス機能を提供します。

警告

API キーをコードやノートブックに直接コミットしないでください。機密性の高い資格情報には、常に AI Gateway または Databricks のシークレットを使用してください。

DatabricksMosaic AI Gatewayでは、Gen モデルの管理とモニタリングへのアクセスにAI を推奨しています。

AI Gateway で構成された基盤モデル エンドポイントを作成します。

  1. Databricks ワークスペースで、「 サービス提供 」> 「新しいエンドポイントの作成 」に移動します。
  2. エンドポイントのタイプとプロバイダーを選択します。
  3. API キーを使用してエンドポイントを構成します。
  4. エンドポイントの設定時に、 AI Gateway を有効にし、必要に応じてレート制限、フォールバック、ガードレールを設定します。
  5. 自動生成されたコードを取得して、エンドポイントのクエリをすばやく開始できます。エンドポイント> Serving > Query > に移動します 。 トレース コードを必ず追加してください。
Python
import mlflow
from openai import OpenAI
import os

# How to get your Databricks token: https://docs.databricks.com/en/dev-tools/auth/pat.html
# DATABRICKS_TOKEN = os.environ.get('DATABRICKS_TOKEN')
# Alternatively in a Databricks notebook you can use this:
DATABRICKS_TOKEN = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()

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

# Set up MLflow tracking (if running outside Databricks)
# If running in a Databricks notebook, these are not needed.
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/my-genai-app")

client = OpenAI(
  api_key=DATABRICKS_TOKEN,
  base_url="<YOUR_HOST_URL>/serving-endpoints"
)

chat_completion = client.chat.completions.create(
  messages=[
  {
    "role": "system",
    "content": "You are an AI assistant"
  },
  {
    "role": "user",
    "content": "What is MLflow?"
  }
  ],
  model="<YOUR_ENDPOINT_NAME>",
  max_tokens=256
)

print(chat_completion.choices[0].message.content)

複数の自動トレース統合の有効化​

生成AIアプリケーションは複数のライブラリを組み合わせることが多いため、 MLflow Tracing では複数の統合に対して同時に自動トレースを有効にし、統一されたトレースエクスペリエンスを提供できます。

たとえば、LangChain と直接 OpenAI トレースの両方を有効にするには、次のようにします。

Python
import mlflow

# Enable MLflow Tracing for both LangChain and OpenAI
mlflow.langchain.autolog()
mlflow.openai.autolog()

# Your code using both LangChain and OpenAI directly...
# ... an example can be found on the Automatic Tracing page ...

MLflow は、LangChain と直接 OpenAI LLM 呼び出しの両方のステップを組み合わせた 1 つのまとまりのあるトレースを生成するため、フロー全体を検査できます。統合の組み合わせのその他の例については、 自動トレース のページを参照してください。

自動トレースの無効化​

特定のライブラリの自動トレースを無効にするには、 mlflow.<library>.autolog(disable=True).すべての自動ログ統合を一度に無効にするには、 mlflow.autolog(disable=True)を使用します。

Python
import mlflow

# Disable for a specific library
mlflow.openai.autolog(disable=True)

# Disable all autologging
mlflow.autolog(disable=True)
に最終更新