エージェント向け MLflow Tracing

重要

この機能はパブリックプレビュー段階です。

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

MLflow Tracingとは何ですか?

MLflow トレースは 、ジェネレーション AI アプリケーションの実行に関する詳細情報をキャプチャします。 トレースでは、要求の各中間ステップに関連付けられた入力、出力、メタデータがログに記録されるため、バグや予期しない動作の原因を特定できます。 たとえば、モデルが幻覚を見た場合、幻覚につながった各ステップをすばやく調査できます。

MLflow Tracing は Databricks ツールおよびインフラストラクチャと統合されているため、 Databricks ノートブックまたは MLflow エクスペリメント UI にトレースを保存および表示できます。

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

MLflow Tracing を使用する理由は何ですか?

MLflow トレースには、いくつかの利点があります。

  • インタラクティブなトレースの視覚化を確認し、調査ツールを使用して問題を診断します。

  • プロンプトテンプレートとガードレールが妥当な結果を生成することを確認します。

  • さまざまなフレームワーク、モデル、チャンク サイズのレイテンシを分析します。

  • 異なるモデル間でのトークンの使用を測定することにより、アプリケーションのコストを見積もります。

  • ベンチマークの「ゴールデン」データセットを確立して、さまざまなバージョンのパフォーマンスを評価します。

  • 本番運用モデルエンドポイントからのトレースを保存して、問題をデバッグし、オフラインでのレビューと評価を実行します。

エージェントにトレースを追加する

MLflow Tracing では、生成アプリケーションAI アプリケーションにトレースを追加するための 3 つの方法がサポートされています。 API リファレンスの詳細については、 MLflow のドキュメントを参照してください。

API

推奨されるユースケース

説明

MLflow 自動ログ

統合されたGenAIライブラリによる開発

自動ロギングは、LangChain、LlamaIndex、OpenAI などのサポートされているオープンソースフレームワークのトレースを自動的にログに記録します。

流暢なAPIs

Pyfuncを使用したカスタムエージェント

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

MLflowクライアントAPIs

マルチスレッドなどの高度なユースケース

MLflowClient 高度なユースケース向けに、きめ細かくスレッドセーフな APIs を提供します。 スパンの親子関係は手動で管理する必要があります。 これにより、特にマルチスレッドのユースケースで、トレースのライフサイクルをより適切に制御できます。

MLflow Tracingをインストールする

MLflow トレースは、MLflow バージョン 2.13.0 以降で使用でき、<DBR< 15.4 LTS ML 以降にプレインストールされています。 必要に応じて、次のコードを使用して MLflow をインストールします。

%pip install mlflow>=2.13.0 -qqqU
%restart_python

または、互換性のある MLflow バージョンを含む最新バージョンの databricks-agentsをインストールすることもできます。

%pip install databricks-agents

autologgingを使用してエージェントにトレースを追加する

LangChain や OpenAI などの GenAI ライブラリがトレースをサポートしている場合は、コードに mlflow.<library>.autolog() を追加して自動ログを有効にします。 例えば：

mlflow.langchain.autolog()

注：

Databricks Runtime 15.4 LTS ML以降、MLflowトレースはノートブック内のデフォルトによって有効になります。たとえば、LangChain を使用してトレースを無効にするには、ノートブックで mlflow.langchain.autolog(log_traces=False) を実行します。

MLflow では、トレース自動ログ記録用の追加のライブラリがサポートされています。 統合ライブラリの完全な一覧については、 MLflow トレースのドキュメントを参照してください。

Fluent APIs を使用して、エージェントにトレースを手動で追加する

の FluentAPIs では、コードの実行フローに基づいてトレース階層が自動的に作成されます。MLflow

関数を修飾する

@mlflow.trace デコレーターを使用して、デコレートされた関数のスコープの範囲を作成します。

MLflow Span オブジェクトは、トレース ステップを整理します。Spans は、ワークフロー内の個々の操作またはステップ (API 呼び出しやベクトル ストア クエリなど) に関する情報をキャプチャします。

スパンは、関数が呼び出されたときに開始され、関数が戻ったときに終了します。 MLflow は、関数の入力と出力、および関数から発生した例外を記録します。

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

@mlflow.trace(name="agent", span_type="TYPE", attributes={"key": "value"})
def my_function(x, y):
    return x + y

トレースコンテキストマネージャーの使用

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

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

外部関数をラップする

外部ライブラリ関数をトレースするには、関数を mlflow.traceでラップします。

from sklearn.metrics import accuracy_score

y_pred = [0, 2, 1, 3]
y_true = [0, 1, 2, 3]

traced_accuracy_score = mlflow.trace(accuracy_score)
traced_accuracy_score(y_true, y_pred)

### Fluent API example

The following example shows how to use the Fluent APIs `mlflow.trace` and `mlflow.start_span` to trace the `quickstart-agent`:

```python
import mlflow
from mlflow.deployments import get_deploy_client

class QAChain(mlflow.pyfunc.PythonModel):
    def __init__(self):
        self.client = get_deploy_client("databricks")

    @mlflow.trace(name="quickstart-agent")
    def predict(self, model_input, system_prompt, params):
        messages = [
                {
                    "role": "system",
                    "content": system_prompt,
                },
                {
                    "role": "user",
                    "content":  model_input[0]["query"]
                }
          ]

        traced_predict = mlflow.trace(self.client.predict)
        output = traced_predict(
            endpoint=params["model_name"],
            inputs={
                "temperature": params["temperature"],
                "max_tokens": params["max_tokens"],
                "messages": messages,
            },
        )

        with mlflow.start_span(name="_final_answer") as span:
          # Initiate another span generation
            span.set_inputs({"query": model_input[0]["query"]})

            answer = output["choices"][0]["message"]["content"]

            span.set_outputs({"generated_text": answer})
            # Attributes computed at runtime can be set using the set_attributes() method.
            span.set_attributes({
              "model_name": params["model_name"],
                        "prompt_tokens": output["usage"]["prompt_tokens"],
                        "completion_tokens": output["usage"]["completion_tokens"],
                        "total_tokens": output["usage"]["total_tokens"]
                    })
              return answer

トレースを追加した後、関数を実行します。 次に、前のセクションの predict() 関数の例を続けます。 トレースは、呼び出しメソッド ( predict()) を実行すると自動的に表示されます。

SYSTEM_PROMPT = """
You are an assistant for Databricks users. You answer Python, coding, SQL, data engineering, spark, data science, DW and platform, API, or infrastructure administration questions related to Databricks. If the question is unrelated to one of these topics, kindly decline to answer. If you don't know the answer, say that you don't know; don't try to make up an answer. Keep the answer as concise as possible. Use the following pieces of context to answer the question at the end:
"""

model = QAChain()

prediction = model.predict(
  [
      {"query": "What is in MLflow 5.0"},
  ],
  SYSTEM_PROMPT,
  {
    # Using Databricks Foundation Model for easier testing, feel free to replace it.
    "model_name": "databricks-dbrx-instruct",
    "temperature": 0.1,
    "max_tokens": 1000,
  }
)

MLflowクライアントAPIs

MlflowClient トレースの開始と終了、スパンの管理、およびスパン フィールドの設定を行うための、きめ細かくスレッド セーフな APIs を公開します。 これにより、トレースのライフサイクルと構造を完全に制御できます。 これらの APIs は、マルチスレッド アプリケーションやコールバックなど、Fluent APIs が要件に対して不十分な場合に便利です。

MLflow クライアントを使用して完全なトレースを作成する手順を次に示します。

  1. MLflowClient のインスタンスを作成するには、 client = MlflowClient()します。

  2. client.start_trace()メソッドを使用してトレースを開始します。これにより、トレース コンテキストが開始され、絶対ルート スパンが開始され、ルート スパン オブジェクトが返されます。 このメソッドは、 start_span() API の前に実行する必要があります。

    1. トレースの属性、入力、および出力を client.start_trace()で設定します。

    注：

    Fluent APIsには、start_trace() メソッドに相当するものはありません。 これは、Fluent APIs がトレース コンテキストを自動的に初期化し、管理された状態に基づいてルート スパンであるかどうかを判断するためです。

  3. start_trace() API はスパンを返します。 リクエスト ID、トレースの一意の識別子 ( trace_idとも呼ばれる)、および返されたスパンの ID を span.request_id と span.span_idを使用して取得します。

  4. client.start_span(request_id, parent_id=span_id)を使用して子スパンを開始し、スパンの属性、入力、および出力を設定します。

    1. この方法では、 request_id と parent_id でスパンをトレース階層の正しい位置に関連付ける必要があります。 別のスパンオブジェクトを返します。

  5. 子スパンを終了するには、 client.end_span(request_id, span_id)を呼び出します。

  6. 作成する子スパンに対して、手順 3 から 5 を繰り返します。

  7. すべての子スパンが終了したら、 client.end_trace(request_id) を呼び出してトレースを閉じ、記録します。

from mlflow.client import MlflowClient

mlflow_client = MlflowClient()

root_span = mlflow_client.start_trace(
  name="simple-rag-agent",
  inputs={
          "query": "Demo",
          "model_name": "DBRX",
          "temperature": 0,
          "max_tokens": 200
         }
  )

request_id = root_span.request_id

# Retrieve documents that are similar to the query
similarity_search_input = dict(query_text="demo", num_results=3)

span_ss = mlflow_client.start_span(
      "search",
      # Specify request_id and parent_id to create the span at the right position in the trace
        request_id=request_id,
        parent_id=root_span.span_id,
        inputs=similarity_search_input
  )
retrieved = ["Test Result"]

# You must explicitly end the span
mlflow_client.end_span(request_id, span_id=span_ss.span_id, outputs=retrieved)

root_span.end_trace(request_id, outputs={"output": retrieved})

トレースのレビュー

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

  • トレースの視覚化は、セル出力にインラインでレンダリングされます。

  • トレースは MLflow エクスペリメントに記録されます。 エクスペリメントページの「トレース」タブで、履歴トレースの完全なリストを確認および検索できます。エージェントをアクティブな MLflow 実行の下で実行すると、 実行 ページにトレースが表示されます。

  • search_traces() API を使用してプログラムでトレースを取得します。

MLflowトレースを本番運用で利用する

MLflow Tracingは Mosaic AI Model Servingとも統合されており、問題を効率的にデバッグし、パフォーマンスを監視し、オフライン評価用のゴールデンデータセットを作成することができます。 サービス エンドポイントで MLflow トレースが有効になっている場合、トレースは response 列の下の推論テーブルに記録されます。

サービス エンドポイントで MLflow トレースを有効にするには、エンドポイント構成で ENABLE_MLFLOW_TRACING 環境変数を Trueに設定する必要があります。 カスタム環境変数を使用してエンドポイントをデプロイする方法については、「 プレーンテキスト環境変数の追加」を参照してください。 deploy() API を使用してエージェントをデプロイした場合、トレースは自動的に推論テーブルに記録されます。「アプリケーションを生成するためのエージェントのデプロイ」AI。

注：

推論テーブルへのトレースの書き込みは非同期的に行われるため、開発中にノートブック環境と同じオーバーヘッドが追加されることはありません。 ただし、エンドポイントの応答速度にオーバーヘッドが発生する可能性があります (特に、各推論要求のトレース サイズが大きい場合)。 Databricks は、環境とモデルの実装に大きく依存するため、モデル エンドポイントに対する実際の待機時間の影響について、サービス レベル アグリーメント (SLA) を保証するものではありません。 Databricks 、本番運用アプリケーションにデプロイする前に、エンドポイントのパフォーマンスをテストし、トレースのオーバーヘッドについて理解を深めることをお勧めします。

次の表は、さまざまなトレース サイズに対する推論レイテンシーへの影響を大まかに示しています。

要求あたりのトレース サイズ

待機時間への影響 (ミリ秒)

~10キロバイト

~ 1 ミリ秒

~ 1メガバイト

50~100ミリ秒

10メガバイト

150ミリ秒~

制限事項

  • MLflow Tracing は、 Databricksダッシュボード、ダッシュボード ジョブ、およびモデルサーバー で利用できます。

LangChain 自動ログは、 LangChain すべての予測 APIsをサポートしているわけではありません。 サポートされている APIsの完全なリストについては、 MLflow のドキュメントを参照してください。