MLflow Tracingによるエージェントの可観測性

この記事では、DatabricksのMLflow Tracing を使用して、生成AIアプリケーションに可観測性を追加する方法について説明します。

MLflow Tracingとは

MLflow Tracingは、開発からデプロイまで、生成AIアプリケーションのエンドツーエンドの監視を提供します。トレーシングは、 Databricksのジェネレーション AI ツールセットと完全に統合されており、開発および本番運用のライフサイクル全体にわたる詳細な知見を捉えることができます。

インライントレースは、生成AIアプリの各ステップの詳細な情報をキャプチャします

以下は、生成AIアプリケーションでのトレースの主なユースケースです。

デバッグの効率化 : トレースにより、ジェネレーション AI アプリケーションの各ステップが可視化され、問題の診断と解決が容易になります。
オフライン評価 : トレースは、エージェント評価のための貴重なデータを生成するため、エージェントの品質を長期的に測定し、改善することができます。
本番運用モニタリング : トレースにより、エージェントの動作と詳細な実行ステップを可視化し、本番運用におけるエージェントのパフォーマンスを監視および最適化できます。
監査ログ : MLflow Tracingは、エージェントのアクションと決定の包括的な監査ログを生成します。これは、コンプライアンスを確保し、予期しない問題が発生した場合のデバッグをサポートするために不可欠です。

必要条件

MLflow Tracingは、MLflow バージョン 2.13.0 以降で使用できます。 Databricks では、最新の機能と改善点にアクセスするために、最新バージョンの MLflow をインストールすることをお勧めします。

Bash
%pip install mlflow>=2.13.0 -qqqU
%restart_python

自動トレース

MLflow の自動ログ記録を使用すると、コードに 1 行 (. mlflow.<library>.autolog()) を追加することで、エージェントをすばやくインストルメント化できます。

MLflow では、最も一般的なエージェント作成ライブラリの自動ログ記録がサポートされています。各オーサリングライブラリの詳細については、MLflow自動ログのドキュメントを参照してください。

ライブラリ	自動ログのバージョンのサポート	自動ロギングコマンド
LangChain	0.1.0 ~最新	`mlflow.langchain.autolog()`
Langgraph	0.1.1 ~ 最新	`mlflow.langgraph.autolog()`
OpenAI	1.0.0 ~ 最新	`mlflow.openai.autolog()`
LlamaIndex	0.10.44 ~ 最新	`mlflow.llamaindex.autolog()`
DSPy	2.5.17 ~ 最新	`mlflow.dspy.autolog()`
Amazon Bedrock	1.33.0 ~ 最新 (boto3)	`mlflow.bedrock.autolog()`
Anthropic	0.30.0 ~ 最新	`mlflow.anthropic.autolog()`
AutoGen	0.2.36 ~ 0.2.40	`mlflow.autogen.autolog()`
Google Gemini	1.0.0 ~ 最新	`mlflow.gemini.autolog()`
CrewAI	0.80.0 ~ 最新	`mlflow.crewai.autolog()`
LiteLLM	1.52.9 ~ 最新	`mlflow.litellm.autolog()`
Groq	0.13.0 ~ 最新	`mlflow.groq.autolog()`
Mistral	1.0.0 ~ 最新	`mlflow.mistral.autolog()`

自動ログを無効にする

自動ロギングトレースは、 Databricks Runtime 15.4 ML 以降の次のライブラリのデフォルトによって有効になります。

LangChain
Langgraph
OpenAI
LlamaIndex

これらのライブラリの自動ロギングトレースを無効にするには、ノートブックで次のコマンドを実行します。

Python
`mlflow.<library>.autolog(log_traces=False)`

トレースを手動で追加する

自動ロギングはエージェントをインストゥルメントする便利な方法を提供しますが、エージェントをより詳細にインストゥルメントしたり、自動ロギングでキャプチャされないトレースを追加したりすることもできます。このような場合は、 MLflow Tracing APIs を使用して、トレースを手動で追加します。

MLflow Tracing APIsは、トレースのツリー構造の管理を気にせずにトレースを追加できるローコードAPIsです。MLflow は、Python スタックを使用して、適切な親子スパン関係を自動的に決定します。

自動ログ記録と手動トレースの組み合わせ

手動トレース APIs は、自動ログ記録と共に使用できます。 MLflow は、自動ログ記録と手動トレースによって作成されたスパンを組み合わせて、エージェントの実行の完全なトレースを作成します。自動ログ記録と手動トレースの組み合わせの例については、「 MLflow Tracingを使用したツール呼び出しエージェントのインストルメント化」を参照してください。

`@mlflow.trace`デコレータを使用したトレース関数

コードを手動でインストゥルメントする最も簡単な方法は、 @mlflow.trace デコレータで関数をデコレートすることです。 MLflow トレースデコレーターは、デコレートされた関数のスコープを使用して "スパン" を作成し、トレース内の実行単位を表し、トレース視覚化に 1 行として表示されます。スパンは、関数の入力と出力、レイテンシ、および関数から発生した例外をキャプチャします。

たとえば、次のコードでは、入力引数 x と y と出力をキャプチャする my_function という名前の範囲を作成します。

Python
import mlflow

@mlflow.trace
def add(x: int, y: int) -> int:
  return x + y

また、スパン名、スパンタイプをカスタマイズしたり、スパンにカスタム属性を追加したりすることもできます。

Python
from mlflow.entities import SpanType

@mlflow.trace(
  # By default, the function name is used as the span name. You can override it with the `name` parameter.
  name="my_add_function",
  # Specify the span type using the `span_type` parameter.
  span_type=SpanType.TOOL,
  # Add custom attributes to the span using the `attributes` parameter. By default, MLflow only captures input and output.
  attributes={"key": "value"}
)
def add(x: int, y: int) -> int:
  return x + y

コンテキストマネージャーを使用して任意のコードブロックをトレース

関数だけでなく、任意のコードブロックの範囲を作成するには、コードブロックをラップするコンテキストマネージャーとして mlflow.start_span() を使用します。スパンは、コンテキストが入力されたときに開始され、コンテキストが終了したときに終了します。 span の入力と出力は、コンテキストマネージャーによって生成される span オブジェクトの setter メソッドを使用して手動で提供する必要があります。詳細については、MLflowドキュメント - コンテキストハンドラーを参照してください。

Python
with mlflow.start_span(name="my_span") as span:
  span.set_inputs({"x": x, "y": y})
  result = x + y
  span.set_outputs(result)
  span.set_attribute("key", "value")

下位レベルのトレース・ライブラリ

MLflow には、トレースツリー構造を明示的に制御するための低レベルの APIs も用意されています。 MLflow のドキュメント - 手動インストルメンテーションを参照してください。

トレースの例: 自動ログ記録と手動トレースの組み合わせ

次の例では、OpenAI の自動ログ記録と手動トレースを組み合わせて、ツール呼び出しエージェントを完全にインストゥルメントします。

Python
import json
from openai import OpenAI
import mlflow
from mlflow.entities import SpanType

client = OpenAI()

# Enable OpenAI autologging to capture LLM API calls
# (*Not necessary if you are using the Databricks Runtime 15.4 ML and above, where OpenAI autologging is enabled by default)
mlflow.openai.autolog()

# Define the tool function. Decorate it with `@mlflow.trace` to create a span for its execution.
@mlflow.trace(span_type=SpanType.TOOL)
def get_weather(city: str) -> str:
  if city == "Tokyo":
    return "sunny"
  elif city == "Paris":
    return "rainy"
  return "unknown"


tools = [
  {
    "type": "function",
    "function": {
      "name": "get_weather",
      "parameters": {
        "type": "object",
        "properties": {"city": {"type": "string"}},
      },
    },
  }
]

_tool_functions = {"get_weather": get_weather}

# Define a simple tool-calling agent
@mlflow.trace(span_type=SpanType.AGENT)
def run_tool_agent(question: str):
  messages = [{"role": "user", "content": question}]

  # Invoke the model with the given question and available tools
  response = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=messages,
    tools=tools,
  )
  ai_msg = response.choices[0].message
  messages.append(ai_msg)

  # If the model requests tool calls, invoke the function(s) with the specified arguments
  if tool_calls := ai_msg.tool_calls:
    for tool_call in tool_calls:
      function_name = tool_call.function.name
      if tool_func := _tool_functions.get(function_name):
        args = json.loads(tool_call.function.arguments)
        tool_result = tool_func(**args)
      else:
        raise RuntimeError("An invalid tool is returned from the assistant!")

      messages.append(
        {
          "role": "tool",
          "tool_call_id": tool_call.id,
          "content": tool_result,
        }
      )

    # Send the tool results to the model and get a new response
    response = client.chat.completions.create(
      model="gpt-4o-mini", messages=messages
    )

  return response.choices[0].message.content

# Run the tool calling agent
question = "What's the weather like in Paris today?"
answer = run_tool_agent(question)

タグによるトレースの注釈付け

MLflow トレースタグは、会話 ID、ユーザー ID、Git コミットハッシュなどのカスタムメタデータをトレースに追加できるキーと値のペアです。タグは MLflow UI に表示され、トレースをフィルター処理して検索します。

タグは、 MLflow APIs または MLflow UI を使用して、進行中または完了したトレースに設定できます。次の例は、進行中のトレースにタグを追加する方法を示しています mlflow.update_current_trace() API を使用します。

Python
@mlflow.trace
def my_func(x):
    mlflow.update_current_trace(tags={"fruit": "apple"})
    return x + 1

トレースのタグ付けと、トレースのフィルター処理と検索に使用する方法の詳細については、 MLflow のドキュメント - トレースタグの設定に関するページを参照してください。

トレースの確認

エージェントの実行後にトレースを確認するには、次のいずれかのオプションを使用します。

インライン視覚化 : Databricks ノートブックでは、トレースはセル出力にインラインでレンダリングされます。
MLflowエクスペリメント :Databricksで、 エクスペリメント >エクスペリメントの選択> トレース に移動して、エクスペリメントのすべてのトレースを表示および検索します。
MLflow 実行 : エージェントがアクティブな MLflow 実行の下で実行されると、MLflow UI の実行ページにトレースが表示されます。
エージェント評価 UI : Mosaic AI エージェント評価では、評価結果の [詳細なトレースビューを表示 ] をクリックして、各エージェント実行のトレースを確認できます。
トレース検索 API : プログラムでトレースを取得するには、トレース検索 API を使用します。

トレースを使用したエージェントの評価

トレースデータは、エージェントを評価するための貴重なリソースとして機能します。モデルの実行に関する詳細な情報をキャプチャすることで、MLflow Tracingはオフライン評価に役立ちます。トレースデータを使用して、ゴールデンデータセットに対するエージェントのパフォーマンスを評価し、問題を特定し、エージェントのパフォーマンスを向上させることができます。

Python
%pip install -U mlflow databricks-agents
%restart_python

Python
import mlflow

# Get the recent 50 successful traces from the experiment
traces = mlflow.search_traces(
    max_results=50,
    filter_string="status = 'OK'",
)

traces.drop_duplicates("request", inplace=True) # Drop duplicate requests.
traces["trace"] = traces["trace"].apply(lambda x: x.to_json()) # Convert the trace to JSON format.

# Evaluate the agent with the trace data
mlflow.evaluate(data=traces, model_type="databricks-agent")

エージェント評価の詳細については、「評価を実行して結果を表示する」を参照してください。

推論テーブルを使用してデプロイされたエージェントを監視する

エージェントが Mosaic AI Model Serving にデプロイされた後、推論テーブルを使用してエージェントを監視できます。推論テーブルには、レビューアプリからのリクエスト、レスポンス、エージェントトレース、エージェントフィードバックの詳細なログが含まれています。この情報により、問題のデバッグ、パフォーマンスの監視、オフライン評価用のゴールデンデータセットの作成が可能になります。

エージェントのデプロイの推論テーブルを有効にするには、「 AI エージェントの推論テーブルを有効にする」を参照してください。

オンライントレースのクエリ

ノートブックを使用して推論テーブルをクエリし、結果を分析します。

トレースを視覚化するには、 display(<the request logs table>) を実行し、検査する行を選択します。

Python
# Query the inference table
df = spark.sql("SELECT * FROM <catalog.schema.my-inference-table-name>")
display(df)

エージェントの監視

「ジェネレーション AI アプリを監視する方法」を参照してください。

トレースのオーバーヘッドレイテンシ

トレースは、パフォーマンスへの影響を最小限に抑えるために非同期的に書き込まれます。ただし、トレースは、特に各推論要求のトレースサイズが大きい場合に、エンドポイントの応答速度に遅延を追加します。 Databricks は、エンドポイントをテストして、レイテンシの影響のトレースを理解してから、本番運用にデプロイすることをお勧めします。

次の表は、トレースサイズによる待機時間の影響の概算を示しています。

要求あたりのトレースサイズ	応答速度の遅延 (ミリ秒) への影響
~10キロバイト	~1ミリ秒
~1メガバイト	50 ~ 100 ミリ秒
10MB	150ミリ秒~

トラブルシューティング

トラブルシューティングと一般的な質問については、 MLflow のドキュメント: トレースハウツーガイドと MLflow のドキュメントの FAQを参照してください。

MLflow Tracingとは​

必要条件​

自動トレース​

自動ログを無効にする​

トレースを手動で追加する​

自動ログ記録と手動トレースの組み合わせ​

@mlflow.traceデコレータを使用したトレース関数​

コンテキストマネージャーを使用して任意のコードブロックをトレース​

下位レベルのトレース・ライブラリ​

トレースの例: 自動ログ記録と手動トレースの組み合わせ​

タグによるトレースの注釈付け​

トレースの確認​

トレースを使用したエージェントの評価​

推論テーブルを使用してデプロイされたエージェントを監視する​

オンライン トレースのクエリ​

エージェントの監視​

トレースのオーバーヘッドレイテンシ​

トラブルシューティング​