MLflow 3では、生成AI向けの機能のサポートが強化された (MLflow for GenAI)。実際に触ってみたので、以下の機能について簡単に整理する。
- Tracing (Observability)
- Prompt Registry
Tracing (Observability)
アプリケーション内の処理フローをエンドツーエンドで監視できる機能。特に、複雑な多段階処理で構成されることの多いLLMアプリケーション（典型的にはRAGやAgentなど）において、モデルの入出力や、トークン消費量といったメトリクスを追跡するのに役立つ。
従来のMLflow Trackingがパラメータやメトリクスのロギングに主眼を置いていたのに対し、Tracingでは処理経路そのものの監視・追跡に重点が置かれている。
コンセプト
Tracingでは、監視する対象としてTraceとSpanという2つの概念が存在する。
Trace
LLMアプリケーションにおける実行単位であり、一連の処理フロー全体に対応する概念。各TraceにはIDや実行ステータス、タグなどのメタデータが関連付けられて管理される。
Span
Traceを構成する個々の処理単位に対応する概念。典型的には、LLMのAPIコールやツールの実行などが該当する。Spanのタイプとしては以下のようなものが組み込みでサポートされている。必要に応じて自作のSpanタイプを定義することも可能。
| Spanタイプ | 説明 |
|---|---|
CHAT_MODEL |
チャットモデルとのインタラクション |
TOOL |
（主にAgentによる）ツールの実行 |
RETRIEVER |
ベクトルDBへの文書検索 |
PARSER |
LLMレスポンスのパース処理 |
MLflowにおけるTraceとSpanの設計は、オブザーバビリティの標準仕様であるOpenTelemetryに基づいている。そのため、DatadogなどOpenTelemetryベースのツールとも互換性がある。
MLflow UI
UIでは以下の情報が可視化される。
- Trace全体やSpanごとの実行時間
- Trace全体やSpanごとの入出力（I/O）
これにより、LLMアプリケーションのフロー全体を通じて、個々の処理単位の挙動を詳細に監視でき、ボトルネックの特定と性能改善に役立つ。
使い方
トレーシングの方法には、コード一行で導入可能なAutomatic Tracingと、柔軟にカスタマイズ可能なManual Tracingの2通りがある。
Automatic Tracing
さまざまなツールと統合されており、mlflow.<framework-name>.autologを呼び出すだけで簡単にトレーシングが可能。以下はOpenAI APIの例。LangChainやLiteLLMなど、他のフレームワークにも対応している。
import mlflow
from dotenv import load_dotenv
from openai import OpenAI
load_dotenv()
mlflow.openai.autolog()
client = OpenAI()
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "What is MLflow?"}],
max_tokens=100,
)
Manual Tracing
処理単位に@mlflow.traceを付けることで、手動でSpanを記録できる。トレーシングの粒度を自分で制御したい場合に有効。前述のAutomatic Tracingと併用することも可能。
import time
import mlflow
from dotenv import load_dotenv
from mlflow.entities import SpanType
from openai import OpenAI
load_dotenv()
@mlflow.trace(name="RAG Pipeline", span_type=SpanType.CHAIN)
def answer_question(question: str) -> str:
context = retrieve_context(question)
client = OpenAI()
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{"role": "system", "content": f"Context: {context}"},
{"role": "user", "content": question},
],
max_tokens=100,
)
return response.choices[0].message.content
@mlflow.trace(span_type=SpanType.RETRIEVER)
def retrieve_context(question: str) -> str:
time.sleep(3) # dummy
return f"Relevant context for: {question}"
result = answer_question("What is MLflow Tracing?")
（補足）Human Feedbackとの統合
Databricks Managed版では、MLflow UIを通して各Traceに対してHuman Feedbackを紐づけて登録できる機能も提供されている。LLMの出力は多様で、定量的な評価が困難なケースも多い。そこで、UIを通してユーザーからのフィードバックを取得し、それを活用することで継続的な改善を図ることができる。
（画像は公式サイトから引用）
リリースノートによると、2025年7月現在この機能はDatabricks Managed版のみで利用可能。OSS版への提供は今後予定されているとのこと。
Prompt Registry
プロンプトエンジニアリングは、自然言語という自由度の高いフォーマットを扱うがゆえに、変更履歴の追跡や管理が困難で属人化しやすい。Prompt Registryはプロンプト管理を支援する以下のような機能をサポートしている。
- バージョン管理：プロンプトはバージョンごとに管理され、コミットメッセージと共に変更の意図も記録できる。
- メタデータ管理：作成者や用途などの情報をタグとして登録可能。検索やフィルタリングに活用できる。
- プロンプトの比較：バージョン間の差分をUIで視覚的に確認できる。どこが変更されたのかが一目で分かる。
これらの機能により、プロンプトの変更や改善を見通しよく行うことができ、継続的な改善サイクルを安心して回せるようになる。
使い方
プロンプトの作成 / 編集
register_promptを使ってプロンプトを保存できる。同名のプロンプトが既に存在する場合はバージョンが自動でインクリメントされる。なお、このようにPython SDKを利用した方法のほか、MLflow UI上からの作成 / 編集も可能。
import mlflow
template = """\
以下の与えられた文章を要約してください：
{{ sentences }}
"""
prompt = mlflow.genai.register_prompt(
name="summarization-prompt",
template=template,
commit_message="Initial commit",
tags={
"author": "author@example.com",
"task": "summarization",
},
)
プロンプトの呼び出し
登録したプロンプトは、load_promptを使って呼び出せる。
import mlflow
prompt = mlflow.genai.load_prompt("prompts:/summarization-prompt/1")
おわりに
MLflow for GenAIでは、LLMOpsに関する多くの機能がサポートされていることがわかった。LLMアプリケーション開発やプロンプトエンジニアリングは複雑化しやすく、管理も煩雑になりがちなため、このような機能を活用することで管理コストを下げ、開発に集中できる環境を整えることが重要だと感じた。
この記事では触れなかったものの、LLMの評価機能やプロンプトの最適化、モデルのサービング機能なども提供されているので、今後必要に応じて試してみたい。