エージェント向けの MLflow トレース

[アーティクル]
07/26/2024

重要

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

ソフトウェア開発では、トレースには、ユーザーセッションや要求フローなどのイベントのシーケンスの記録が必要です。 AI システムのコンテキストでは、トレースは多くの場合、AI システムとの間で行われる対話を指します。 AI システムのトレースの例は、プロンプトでのユーザーメッセージ、ベクトル検索、生成 AI モデルとのインターフェイスが含まれた RAG アプリケーションの入力とパラメーターのインストルメント化のように見えることがあります。

MLflow トレースとは

MLflow トレースを使用すると、生成 AI アプリケーションの異なるバージョンにまたがるトレースの記録、分析、比較が可能になります。これにより、生成 AI の Python コードをデバッグしたり、入力や応答を追跡したりできるため、アプリケーションのパフォーマンス低下をもたらす条件またはパラメーターを検出するのに役立ちます。 MLflow トレースは Databricks のツールやインフラストラクチャと密接に統合されているため、コードを実行するときにすべてのトレースを格納し、Databricks ノートブックまたは MLflow 実験 UI に表示できます。

LangChain または PyFunc を使用して Databricks 上で AI システムを開発する場合は、MLflow トレースを使用すると、エージェントの各ステップのすべてのイベントと中間出力を表示できます。プロンプト、どのモデルやリトリバーが使用されたか、応答を強化するためにどのドキュメントが取得されたか、それらにかかった時間、最終的な出力を簡単に表示できます。そのため、モデルが幻覚を起こした場合は、その幻覚をもたらした各ステップをすばやく検査できます。

MLflow トレースを使用する理由

MLflow トレースには、開発ワークフローを追跡するのに役立つ次の利点があります。

開発での問題を診断するための対話型のトレース視覚化および調査ツール。
プロンプトテンプレートとガードレールによって適切な結果が生成されていることを確認します。
さまざまなフレームワーク、モデル、チャンクサイズ、ソフトウェア開発手法が待ち時間に与える影響を調べ、最小限に抑えます。
異なるモデルごとのトークン使用状況を追跡することによってアプリケーションコストを測定します。
ベンチマーク ("ゴールデン") データセットを確立して、異なるバージョンのパフォーマンスを評価します。
オフラインでのレビューと評価のためにトレースを格納します。これには、推論テーブルを使用するように構成されたサービスエンドポイントが必要です。

MLflow トレースをインストールする

MLflow トレースは、MLflow バージョン 2.13.0 以降で使用できます。

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

MLflow トレースを使用する

MLflow トレースは、エージェント開発ワークロードで次の方法で使用できます。

MLflow トレースと LangChain (mlflow.langchain) の統合を使用します。エージェントで mlflow.langchain.autolog を実行し、呼び出し API を実行すると、エージェントの各ステップのトレースを自動的に表示できます。
必要に応じて、Fluent API または MLflow クライアント API を使用して、エージェントの特定の部分にトレースを手動で追加することもできます。

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

MLflow トレースには、トレースを含む生成 AI アプリケーションでトレースを使用するための 3 つの異なる方法が用意されています。「エージェントにトレースを追加する」を参照してください。

API	説明
MLflow の自動ログ記録	(Langchain に推奨) 自動ログ記録は、Langchain などの人気のオープンソースチェーンフレームワークのトレースを自動的にインストルメント化します。
Fluent API	(推奨) トレースのツリー構造を気にすることなく AI システムをインストルメント化するためのローコード API。 Python スタックに基づいて、適切な親子ツリー構造 (スパン) が MLflow によって決定されます。
MLflow クライアント API	MLflowClient によって、高度なユースケース向けのよりきめ細かいスレッドセーフ API が実装されます。これらの API ではスパンの親子関係が管理されないため、目的のトレース構造を構築するには、それを手動で指定する必要があります。これにはより多くのコードが必要ですが、特にマルチスレッドのユースケースでは、トレースのライフサイクルをより適切に制御できます。マルチスレッドアプリケーションやコールバックベースのインストルメンテーションなどの、より細かい制御を必要とするユースケースに推奨されます。

API リファレンスとコード例については、MLflow のドキュメントを参照してください。

トレースを収集するために推論テーブルを有効にする

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

Note

トレースを有効にすると、特に各推論要求のトレースサイズが大きい場合に、エンドポイントの応答速度にオーバーヘッドが発生する可能性があります。 Databricks は、環境とモデルの実装に大きく依存しているため、モデルエンドポイントに対する実際の待機時間への影響に関するサービスレベルアグリーメント (SLA) を保証しません。 Databricks では、運用アプリケーションにデプロイする前に、エンドポイントのパフォーマンスをテストし、トレースのオーバーヘッドに関する分析情報を得ることをお勧めします。

次の表は、さまざまなトレースサイズの推論待機時間への影響を大まかに示しています。

要求あたりのトレースサイズ	待機時間への影響 (ms)
～ 10 KB	～ 1 ms
～ 1 MB	50 ～ 100 ms
10 MB	150 ms ～

制限事項

MLflow トレースは、Databricks ノートブック、ノートブックジョブ、Model Serving で使用できます。
LangChain の自動ログは、すべての LangChain 予測 API をサポートしているわけではありません。サポートされている API の完全な一覧については、MLflow のドキュメントを参照してください。

次の方法で共有