エージェント向け 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 Tracing は、 Databricksノートブックおよびノートブックジョブ で利用できます。

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