エージェント向け MLflow Tracing

重要

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

この記事では、 MLflow Tracingと、 AIシステムで生成AIアプリケーションを評価するのに役立つシナリオについて説明します。

ソフトウェア開発では、トレースにはユーザー セッションやリクエスト フローなどのイベントのシーケンスの記録が含まれます。 AI システムのコンテキストでは、トレースは多くの場合、AI システムとのやり取りを指します。 AIシステムのトレースの例としては、プロンプト付きのユーザー メッセージ、ベクトル ルックアップ、および生成AIモデルとのインターフェイスを含む RAG アプリケーションの入力と引数をインストルメント化するものが考えられます。

MLflow Tracingとは何ですか?

MLflowトレーシングを使用すると、生成アプリケーションAIアプリケーションの異なるバージョン間でトレースをログに記録、分析、および比較できます。これにより、生成AI Pythonコードをデバッグし、入力と応答を追跡できます。 これにより、アプリケーションのパフォーマンス低下の原因となっている条件やパラメーターを検出できます。 MLflow Tracing は、 Databricks ツールやインフラストラクチャと緊密に統合されているため、コードを実行するときに、すべてのトレースを Databricks ノートブックや MLflow エクスペリメント UI に保存および表示できます。

LangChain、LlamaIndex、OpenAI、カスタム PyFunc などのライブラリを使用して Databricks で AI システムを開発する場合、MLflow Tracing を使用すると、エージェントの各ステップからのすべてのイベントと中間出力を確認できます。 プロンプト、使用されたモデルとレトリーバー、応答を補強するために取得されたドキュメント、処理にかかる時間、および最終的な出力を簡単に確認できます。 たとえば、モデルが幻覚を見た場合、幻覚につながった各ステップをすばやく調査できます。

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

MLflow トレースには、開発ワークフローの追跡に役立ついくつかの利点があります。 たとえば、次のことができます。

  • 開発中の問題を診断するためのインタラクティブなトレースの視覚化と調査ツール。

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

  • さまざまなフレームワーク、モデル、チャンク サイズ、ソフトウェア開発プラクティスのレイテンシへの影響を調査し、最小限に抑えます。

  • さまざまなモデルによるトークンの使用状況を追跡して、アプリケーションのコストを測定します。

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

MLflow Tracingをインストールする

MLflow Tracing は、MLflow バージョン 2.13.0 以降で利用できます。

%pip install mlflow>=2.13.0 -qqqU
%restart_python

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

開発で MLflow トレースを使用する

MLflow トレースは、パフォーマンスの問題を分析し、エージェントの開発サイクルを加速するのに役立ちます。 次のセクションでは、ノートブックからエージェント開発と MLflow トレースを行うことを前提としています。

注：

ノートブック環境では、MLflow トレースによってエージェントの実行時間に最大で数秒のオーバーヘッドが追加される可能性があります。

注：

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

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

MLflow トレースでは、生成アプリケーションAI トレースを使用してトレースを使用する 3 つの異なる方法が提供されます。 これらの方法の使用例については、「 エージェントにトレースを追加する 」を参照してください。 API リファレンスの詳細については、 MLflow のドキュメントを参照してください。

API

推奨されるユースケース

説明

MLflow 自動ログ

統合型GenAIライブラリの開発

自動ロギングは、LangChain、LlamaIndex、OpenAIなどの一般的なオープンソースフレームワークのトレースを自動的にインストゥルメントします。 ノートブックの先頭に mlflow.<library>.autolog() を追加すると、MLflow はエージェント実行の各ステップのトレースを自動的に記録します。

流暢なAPIs

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

トレースのツリー構造を気にせずに システムをインストルメント化するためのローコードAPIs 。AIMLflow は、Python スタックに基づいて適切な親子ツリー構造 (スパン) を決定します。

MLflowクライアントAPIs

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

MLflowClient 高度なユースケース向けに、よりきめ細かくスレッドセーフな APIs を実装します。 これらの APIs はスパンの親子関係を管理しないため、目的のトレース構造を構築するには、手動で指定する必要があります。 これにはより多くのコードが必要ですが、特にマルチスレッドのユースケースでは、トレースのライフサイクルをより適切に制御できます。

マルチスレッドアプリケーションやコールバックベースのインストゥルメンテーションなど、より詳細な制御が必要なユースケースに推奨されます。

トレースのレビュー

インストゥルメントされたエージェントを実行した後、生成されたトレースをさまざまな方法で確認できます。

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

  • トレースは MLflow エクスペリメントに記録されます。 エクスペリメントページの Traces タブで、過去の痕跡の完全なリストを確認し、それらを検索できます。 エージェントがアクティブ MLflow 実行の下で実行されると、実行ページでトレースを見つけることもできます。

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

制限事項

  • MLflow Tracing は、 Databricksノートブックおよびノートブックジョブ で利用できます。

  • LangChain自動ログ記録はすべてのLangChain予測APIsをサポートしない可能性があります。 MLflowサポートされている の完全なリストについては 、 ドキュメントAPIs を参照してください。