MLflow Tracing は、幅広い一般的な Generative AI ライブラリとフレームワークと統合されており、それらすべてに対して 1 行の自動トレース エクスペリエンスを提供します。 これにより、最小限のセットアップで GenAI アプリケーションをすぐに監視できます。
自動トレースは、特定のライブラリまたは SDK の実装に基づいて、LLM 呼び出し、ツールの使用状況、エージェントの対話などのアプリケーションのロジックと中間ステップをキャプチャします。
注
サーバーレス コンピューティング クラスターでは、genAI トレース フレームワークの自動ログ記録は自動的には有効になりません。 トレースする特定の統合に対して適切な mlflow.<library>.autolog() 関数を呼び出して、自動ログ記録を明示的に有効にする必要があります。
自動トレースのしくみ、その前提条件、および手動トレースとの組み合わせの例については、メインの 自動トレース ガイドを参照してください。 以下の簡単な例では、一部の上位統合が強調表示されています。 サポートされている各ライブラリの詳細なガイド(前提条件、高度な例、および特定の動作を含む)は、このセクションのそれぞれのページで入手できます。
一目でわかる上位の統合
最も一般的に使用される統合の一部のクイック スタートの例を次に示します。 タブをクリックすると、基本的な使用例が表示されます。 それぞれの詳細な前提条件とより高度なシナリオについては、専用の統合ページ (タブまたは以下の一覧からリンク) を参照してください。
OpenAI
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
LangChain
import mlflow
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI
# 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"
mlflow.langchain.autolog()
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/langchain-tracing-demo")
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7, max_tokens=1000)
prompt = PromptTemplate.from_template("Tell me a joke about {topic}.")
chain = prompt | llm | StrOutputParser()
chain.invoke({"topic": "artificial intelligence"})
# View trace in MLflow UI
LangGraph
import mlflow
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
# 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"
mlflow.langchain.autolog() # LangGraph uses LangChain's autolog
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/langgraph-tracing-demo")
@tool
def get_weather(city: str):
"""Use this to get weather information."""
return f"It might be cloudy in {city}"
llm = ChatOpenAI(model="gpt-4o-mini")
graph = create_react_agent(llm, [get_weather])
result = graph.invoke({"messages": [("user", "what is the weather in sf?")]})
# View trace in MLflow UI
Anthropic
import mlflow
import anthropic
import os
# 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:
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"
mlflow.anthropic.autolog()
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/anthropic-tracing-demo")
client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))
message = client.messages.create(
model="claude-3-5-sonnet-20241022",
max_tokens=1024,
messages=[{"role": "user", "content": "Hello, Claude"}],
)
# View trace in MLflow UI
DSPy
import mlflow
import dspy
# 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"
mlflow.dspy.autolog()
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/dspy-tracing-demo")
lm = dspy.LM("openai/gpt-4o-mini") # Assumes OPENAI_API_KEY is set
dspy.configure(lm=lm)
class SimpleSignature(dspy.Signature):
input_text: str = dspy.InputField()
output_text: str = dspy.OutputField()
program = dspy.Predict(SimpleSignature)
result = program(input_text="Summarize MLflow Tracing.")
# View trace in MLflow UI
Databricks
import mlflow
import os
from openai import OpenAI # Databricks FMAPI uses OpenAI client
# 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:
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"
mlflow.openai.autolog() # Traces Databricks FMAPI via OpenAI client
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/databricks-fmapi-tracing")
client = OpenAI(
api_key=os.environ.get("DATABRICKS_TOKEN"),
base_url=f"{os.environ.get('DATABRICKS_HOST')}/serving-endpoints"
)
response = client.chat.completions.create(
model="databricks-llama-4-maverick",
messages=[{"role": "user", "content": "Key features of MLflow?"}],
)
# View trace in MLflow UI
Bedrock
import mlflow
import boto3
# 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"
mlflow.bedrock.autolog()
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/bedrock-tracing-demo")
bedrock = boto3.client(
service_name="bedrock-runtime",
region_name="us-east-1" # Replace with your region
)
response = bedrock.converse(
modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
messages=[{"role": "user", "content": "Hello World in one line."}]
)
# View trace in MLflow UI
AutoGen
import mlflow
from autogen import ConversableAgent
import os
# 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:
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"
mlflow.autogen.autolog()
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/autogen-tracing-demo")
config_list = [{"model": "gpt-4o-mini", "api_key": os.environ.get("OPENAI_API_KEY")}]
assistant = ConversableAgent("assistant", llm_config={"config_list": config_list})
user_proxy = ConversableAgent("user_proxy", human_input_mode="NEVER", code_execution_config=False)
user_proxy.initiate_chat(assistant, message="What is 2+2?")
# View trace in MLflow UI
セキュリティで保護された API キー管理
運用環境の場合、Databricks では、AI Gateway または Databricks シークレットを使用して API キーを管理することをお勧めします。 AI Gateway が推奨される方法であり、追加のガバナンス機能を提供します。
Warnung
コードまたはノートブックで API キーを直接コミットしないでください。 機密性の高い資格情報には、常に AI ゲートウェイまたは Databricks シークレットを使用します。
AI ゲートウェイ (推奨)
Databricks では、Gen AI モデルへのアクセスを管理および監視するために 、Mosaic AI Gateway が推奨されています。
AI ゲートウェイで構成された Foundation Model エンドポイントを作成します。
- Databricks ワークスペースで、[サービス]>[新しいエンドポイントの作成] に移動します。
- エンドポイントの種類とプロバイダーを選択します。
- API キーを使用してエンドポイントを構成します。
- エンドポイントの構成中に、 AI ゲートウェイ を有効にし、必要に応じてレート制限、フォールバック、ガードレールを構成します。
- 自動生成されたコードを取得して、エンドポイントのクエリをすばやく開始できます。 Serving> エンドポイント>に移動して、クエリを>。 トレース コードを必ず追加してください。
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)
Databricks シークレット
Databricks シークレットを使用して API キーを管理する:
シークレット スコープを作成し、API キーを格納します。
from databricks.sdk import WorkspaceClient # Set your secret scope and key names secret_scope_name = "llm-secrets" # Choose an appropriate scope name secret_key_name = "api-key" # Choose an appropriate key name # Create the secret scope and store your API key w = WorkspaceClient() w.secrets.create_scope(scope=secret_scope_name) w.secrets.put_secret( scope=secret_scope_name, key=secret_key_name, string_value="your-api-key-here" # Replace with your actual API key )コードでシークレットを取得して使用します。
import mlflow import openai import os # Configure your secret scope and key names secret_scope_name = "llm-secrets" secret_key_name = "api-key" # Retrieve the API key from Databricks secrets os.environ["OPENAI_API_KEY"] = dbutils.secrets.get( scope=secret_scope_name, key=secret_key_name ) # Enable automatic tracing mlflow.openai.autolog() # Use OpenAI client with securely managed API key client = openai.OpenAI() response = client.chat.completions.create( model="gpt-4o-mini", messages=[{"role": "user", "content": "Explain MLflow Tracing"}], max_tokens=100 )
複数の自動トレース統合の有効化
GenAI アプリケーションでは複数のライブラリが組み合わされることが多いため、MLflow Tracing を使用すると、複数の統合に対して自動トレースを同時に有効にして、統一されたトレース エクスペリエンスを実現できます。
たとえば、LangChain と直接 OpenAI トレースの両方を有効にするには、次のようにします。
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 呼び出しの両方のステップを組み合わせた単一のまとまりのあるトレースを生成し、完全なフローを検査できるようにします。 統合の組み合わせのその他の例については、「 自動トレース 」ページを参照してください。
自動トレースの無効化
特定のライブラリの自動トレースは、 mlflow.<library>.autolog(disable=True)を呼び出すことによって無効にすることができます。
すべての自動ログ統合を一度に無効にするには、 mlflow.autolog(disable=True)を使用します。
import mlflow
# Disable for a specific library
mlflow.openai.autolog(disable=True)
# Disable all autologging
mlflow.autolog(disable=True)