はじめに
Rehab for JAPAN AI/MLOpsエンジニアの松本です。
弊社のプロダクトでLLMを活用するにあたり、LLMのinput/outputを監視できるようにする必要があり、MLflowを試してみたので簡単にご紹介します。
ターゲット
- 無料で使えるLLMのロギングツールについて知りたい方
- MLflow初心者
- PythonとDockerを扱える方
（※LLMのAPIを呼び出す処理の書き方については本記事では扱いません）
前提知識・背景
MLflowとは？
MLflowは機械学習ライフサイクル管理のためのオープンソースプラットフォームです。
実験管理、コードの再現性確保、モデルのデプロイと管理を容易にするもので、元来、機械学習プロジェクトで活用されてきたツールですが、近年ではLLM向けの機能も続々と開発されています。
※LLMに関わる機能の大半はまだ"Experimental"となっており、今後のアップデートで仕様が変わったり無くなったりする可能性がある旨、公式ドキュメントにも記載されていますのでご注意ください。
何故MLflow？
LLMOpsツールといえば、LangSmithやWeights & Biasesなどが有名ですが、いずれも無料で利用できるのは個人利用までで、チーム利用・商用になると有償になる為、まずは固定費のかからないオープンソースから試すことにしました。
準備作業
1. 環境構築
公式Dockerイメージを使うのが最も楽です。
FROM ghcr.io/mlflow/mlflow:v2.17.0
RUN apt-get update -y && apt-get install -y git
※最新バージョンはこちらで確認できますので、適宜変更してください
※gitをインストールする訳：
v2.17.0では処理を動かす度に"Failed to import Git" というwarningが出て、こちらのissueを参考にgitをインストールしたところ解決した為です。
2. MLflowサーバーを立ち上げる
作成したコンテナ内で以下のコマンドを実行します。
ポート番号は指定しなければデフォルトで5000番が使われますが、他が良ければ--portオプションで変更可能です。
mlflow server --host 0.0.0.0 --port 5000
実行したらブラウザで http://localhost:5000 を開けばUIが表示されます。
MLflowのロギング機能紹介
◾️ 前提：MLflowのコードを書く際に毎回必要な設定
import mlflow
# サーバーを指定
MLFLOW_SERVER_URI = "http://localhost:5000"
mlflow.set_tracking_uri(uri=MLFLOW_SERVER_URI)
# ひとまとめに管理したい単位で名前をつけます。例えばモデル名や機能名など。
mlflow.set_experiment("Tracing Demo")
set_experimentで指定した名前はサイドバーに表示され、UI上で結果を確認する際の識別子となっています。
◾️ ロギング用コード紹介
🔷 自動ロギングについて
LangChainやOpenAIなど、一部のライブラリに関しては、以下のように一行書くだけで自動ロギングできるようですが、2024年10月現在、Geminiには対応していないようだったので、今回は別の方法を試しました。
mlflow.openai.autolog()
🔷 デコレータをつける方法
関数にデコレータをつけるだけで、その関数の引数と戻り値をrequest/responseという形でロギングできます。これならGeminiや独自モデル含め、どんなモデルにも使えるので便利です。
@mlflow.trace
def call_llm(prompt):
# LLMにリクエストを投げる処理（関数の中身は自由）
response = send_request_to_llm(prompt)
# 必要な内容を取り出すなどの後処理をして返す（戻り値は複数に分ける事も可能）
return response
処理実行後、前述のサイドバーでexperiment名をクリックしたら、"Traces"タブをクリックすると以下のような一覧表が表示されます。
🔷 attributesについて
実行時に使用した条件などのメタデータをセットできます。dict型で渡せば中身は自由です。
MODEL_NAME = "gemini-1.5-flash-001"
@mlflow.trace(attributes={"model": MODEL_NAME, "okimochi": "o(^-^)o"})
def call_llm(prompt):
# 以下省略
一覧上には表示されませんが、各リクエストIDをクリックすると表示される実行ごとの詳細画面で確認できます。
◾️ MLflowのデータ管理構造
作業ディレクトリにmlartifactsとmlrunsという2つのディレクトリが作られ、その中にデータが格納されていきます。traceデコレータで記録されるrequest/responseデータはmlartifactsの中にjson形式で格納されます。
.
├── mlartifacts/ # モデルやデータ等の成果物が格納される場所
│ └── 1234567890/ # experiment_id (set_experimentで指定した名前ごと)
│ └── traces/ # Tracesデータはこのディレクトリ名にまとめられる
│ └── 234567890123456/ # 実行ごとに付与される識別番号
│ └── traces.json # request/responseデータ保存場所
│
└── mlruns/ # 実行ごとのパラメータ・メタデータが格納される場所
└── 1234567890/ # 以下はexperiment_id/実行ID/の構成
└── 234567890123456/
└── （以下省略）
まとめ・所感
今回、MLflowを使って以下の事が確認できました。
- 環境構築の簡便さ
- ロギング処理実装の容易さ
- ロギングできる内容の自由度
- LLMの入出力に関わるUI仕様
所感として、すでにLLMにリクエストを投げる処理を実装済みであっても、さほど手間をかけずに導入できるのではないかと思いました。
MLflowには他にも様々な機能があり、プロンプトエンジニアリング用のUIもあるようなので、今後試してみたいと思っています。
最後までお読み頂き、ありがとうございました🙇
Discussion