エージェント向けの MLflow トレース
重要
この機能はパブリック プレビュー段階にあります。
この記事では、MLflow トレースと、それが AI システムでの生成 AI アプリケーションの評価に役立つシナリオについて説明します。
ソフトウェア開発では、トレースには、ユーザー セッションや要求フローなどのイベントのシーケンスの記録が必要です。 AI システムのコンテキストでは、トレースは多くの場合、AI システムとの間で行われる対話を指します。 AI システムのトレースの例は、プロンプトでのユーザー メッセージ、ベクトル検索、生成 AI モデルとのインターフェイスが含まれた RAG アプリケーションの入力とパラメーターのインストルメント化のように見えることがあります。
MLflow トレースとは
MLflow トレースを使用すると、生成 AI アプリケーションの異なるバージョンにまたがるトレースの記録、分析、比較が可能になります。 これにより、生成型 AI Python コードをデバッグし、入力と応答を追跡できます。 そうすることで、アプリケーションのパフォーマンス低下に寄与する条件やパラメーターを検出するのに役立ちます。 MLflow トレースは Databricks のツールやインフラストラクチャと密接に統合されているため、コードを実行するときにすべてのトレースを格納し、Databricks ノートブックまたは MLflow 実験 UI に表示できます。
LangChain、LlamaIndex、OpenAI、カスタム PyFunc などのライブラリを使用して Databricks 上で AI システムを開発する場合は、MLflow トレースを使用すると、エージェントの各ステップのすべてのイベントと中間出力を表示できます。 プロンプト、どのモデルやリトリバーが使用されたか、応答を強化するためにどのドキュメントが取得されたか、それらにかかった時間、最終的な出力を簡単に表示できます。 たとえば、モデルが幻覚を起こした場合は、その幻覚をもたらした各手順をすばやく検査できます。
MLflow トレースを使用する理由
MLflow トレースには、開発ワークフローを追跡するのに役立ついくつかの利点があります。 たとえば、次のようなことができます。
- 開発での問題を診断するために対話型のトレース視覚化を確認し、調査ツールを使用します。
- プロンプト テンプレートとガードレールによって適切な結果が生成されていることを確認します。
- さまざまなフレームワーク、モデル、チャンク サイズ、ソフトウェア開発手法が待ち時間に与える影響を調べ、最小限に抑えます。
- 異なるモデルごとのトークン使用状況を追跡することによってアプリケーション コストを測定します。
- ベンチマーク ("ゴールデン") データセットを確立して、異なるバージョンのパフォーマンスを評価します。
- 生産モデル エンドポイントからのトレースを格納して問題をデバッグし、オフラインのレビューと評価を実行します。
MLflow トレースをインストールする
MLflow トレースは、MLflow バージョン 2.13.0 以降で使用できます。
%pip install mlflow>=2.13.0 -qqqU
%restart_python
あるいは、%pip install databricks-agents を実行して、互換性のある MLflow バージョンを含む最新バージョンの databricks-agents をインストールできます。
開発中に MLflow トレースを使用する
MLflow トレースは、パフォーマンスの問題を分析し、エージェント開発サイクルを加速するのに役立ちます。 以降のセクションでは、ノートブックからエージェント開発と MLflow トレースを実行していることを前提としています。
Note
ノートブック環境では、MLflow トレースがエージェントの実行時に最大数秒のオーバーヘッドを合計する場合があります。 これは主に、MLflow 実験へのトレースのログ記録の待機時間に起因します。 運用モデル エンドポイントでは、MLflow トレースのパフォーマンスへの影響が大幅に小さくなります。 実稼働での「MLflow トレースの使用」を参照してください。
Note
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 はエージェント実行の各ステップのトレースを自動的に記録します。 |
| Fluent API | Pyfunc を使用したカスタム エージェント | トレースのツリー構造を気にすることなく AI システムをインストルメント化するためのロー コード API。 Python スタックに基づいて、適切な親子ツリー構造 (スパン) が MLflow によって決定されます。 |
| MLflow クライアント API | マルチスレッドなどの高度なユース ケース | MLflowClient は高度なユース ケース向けのよりきめ細かいスレッドセーフ API を実装しています。 これらの API ではスパンの親子関係が管理されないため、目的のトレース構造を構築するには、それを手動で指定する必要があります。 これにはより多くのコードが必要ですが、特にマルチスレッドのユース ケースでは、トレースのライフサイクルをより適切に制御できます。マルチスレッド アプリケーションやコールバック ベースのインストルメンテーションなどの、より細かい制御を必要とするユース ケースに推奨されます。 |
トレースを確認
インストルメント化されたエージェントを実行した後、生成されたトレースをさまざまな方法で確認できます。
- トレース視覚エフェクトは、セル出力にインラインでレンダリングされます。
- トレースは MLflow 実験に記録されます。 履歴トレースの完全な一覧を確認し、実験ページの
Tracesで検索してください。 エージェントがアクティブな MLflow 実行で実行されている場合は、[実行] ページでトレースを見つけることもできます。 - search_traces() API を使用してプログラムでトレースを取得します。
本番環境で MLflow トレースを使用する
また、MLflow Tracing は Mosaic AI モデル Serving と統合されているため、問題を効率的にデバッグし、パフォーマンスを監視し、オフライン評価用のゴールデン データセットを作成できます。 MLflow トレースがサービス エンドポイントで有効になっている場合、トレースは response 列の下の 推論テーブル に記録されます。
サービス エンドポイントの 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 のドキュメントを参照してください。