👻

mlflow 3でGenAIへの対応がされたので使ってみた

に公開

今回は、mlflow 3で対応された生成AIに対する機能について実際につかてみたので共有します。

mlflowとは?

MLflowは、機械学習の実践者やチームが機械学習プロセスの複雑な処理を行えるよう支援するために特別に構築されたオープンソースプラットフォームです。MLflowは機械学習プロジェクトのライフサイクル全体に焦点を当て、各フェーズの管理、追跡、再現性を保証します。MLモデルを作るためには適切に実験管理をすることで最適なモデルをデプロイするための仕組みが必要であり、mlflowはそのようなユースケースに最も効力を発揮します。

https://mlflow.org/docs/latest/ml/

mlflow 3では何ができる?

MLflow 3では、機械学習モデルやAIアプリケーション、エージェント向けにクラス最高の実験追跡をはじめとして可観測性、パフォーマンス評価を提供してくれます。MLflow 3を使用すると、以下のことを実現できます。

  • 開発ノートブック内のインタラクティブなクエリから本番環境のバッチ処理やリアルタイムサービングのデプロイメントまで、あらゆる環境におけるモデル・プロンプト・エージェント、AIアプリケーションのパフォーマンスを一元的に追跡・分析
  • MLflowのトレース機能と評価機能を活用した強化されたパフォーマンス比較エクスペリエンスにより、本番環境に最適なモデル、プロンプト、エージェント、AI アプリケーションを選択

https://mlflow.org/docs/latest/genai/mlflow-3/

以下、新しく実装された機能について公式ページから抜粋します。

生成AI向けのトラッキングの改善

  • モデル成果物だけでなく、生成AIアプリケーションに特化したバージョン管理メカニズムを導入
  • 新しいLoggedModelエンティティはメタデータハブとして機能し、各概念アプリケーションバージョンを、特定の外部コードや設定、トレースや評価実行といった他のMLflowエンティティにリンク
  • 新しいバージョン管理メカニズムは、従来のMLモデルやディープラーニングチェックポイントにもシームレスに機能します。

包括的なパフォーマンス追跡と観測可能性

  • 強化されたモデル追跡機能により、モデル、実行、トレース、プロンプト、評価指標間の包括的なリネージを提供
  • 新しいモデル中心の設計により、異なる開発環境や本番環境からのトレースと指標をグループ化し、モデルバージョン間の詳細な比較を実現

プロダクショングレードの生成AI評価

  • 評価および監視機能は、GenAI アプリケーションの品質をライフサイクル全体にわたって体系的に測定、改善、維持に寄与
  • 開発から本番環境まで、同じ品質スコアラーを使用することで、コストとレイテンシを管理しながらアプリケーションが正確で信頼性の高いレスポンスを提供

ヒューマンインザループのフィードバック

  • モデル予測における人間によるアノテーションとフィードバックを追跡できるようになり、人間参加型の評価サイクルを効率化
  • データサイエンティスト、ドメインエキスパート、そして関係者が協力してモデルの品質を効率的に向上できる、協調的な環境を実現

最先端のプロンプト最適化

  • MLflowプロンプトレジストリには、最先端の研究に基づいて構築されたプロンプト最適化機能を追加
  • 評価フィードバックとラベル付きデータセットを使用してプロンプトを自動的に改善
  • バージョン管理、追跡、そして体系的なプロンプトエンジニアリングワークフローが対応

ウェブサイトとドキュメントの刷新

  • 生成AI開発と従来の機械学習ワークフローという2つの主要なユーザージャーニーをサポートするために全面的に再設計
  • 新しい構造では、生成AI機能(LLM、プロンプトエンジニアリング、トレースなど)と、実験追跡、モデルレジストリ、デプロイ、評価といった従来の機械学習機能の専用セクションを提供

実際に使ってみる

こちらの公式ページのサンプルを動かしてみます。

https://mlflow.org/docs/latest/genai/mlflow-3/genai-agent/

環境構築

uvを利用して以下のように環境作りました。

uv init mlflow3_test -p 3.12
cd mlflow3_test
uv add mlflow openai

プロンプトのバージョン管理

まずはプロンプトのバージョン管理を試します。以下のサンプルを使います。

prompt.py
import mlflow

# define a prompt template
prompt_template = """\
You are an expert AI assistant. Answer the user's question with clarity, accuracy, and conciseness.

## Question:
{{question}}

## Guidelines:
- Keep responses factual and to the point.
- If relevant, provide examples or step-by-step instructions.
- If the question is ambiguous, clarify before answering.

Respond below:
"""

# register the prompt
prompt = mlflow.genai.register_prompt(
    name="ai_assistant_prompt",
    template=prompt_template,
    commit_message="Initial version of AI assistant",
)

mlflow.genai.register_promptを用いることでプロンプトを登録することができます。nameにプロンプトの名前を、templateにそのテンプレートを登録することで管理できます。同じnameで登録すると、実行するたびにバージョンがインクリメントされます。

コードを実行するとmlflowにプロンプトが登録されます。

uv run prompt.py

実行した状態でmlflowのUIにアクセスすると以下のようにバージョンが管理されていることが確認できます。Promptの画面からアクセスすると指定した名前(今回ではai_assistant_prompt)という名前で登録されていることが確認できます。

uv run mlflow ui

chainの実装

次はlangchainを使ったchainを使って実装してみます。コードは以下になります。ユーザが入力した国の首都を取得する部分と、その首都での観光場所を取得するに段階に分けています。

chain.py
import mlflow
from langchain_openai import OpenAI
import os
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough


# set the active model for linking traces
mlflow.set_active_model(name="langchain_model")

# Enable autologging so that interactive traces from the client are automatically linked to a LoggedModel
mlflow.langchain.autolog()


# 国名から首都を調べるチェーン
capital_prompt = PromptTemplate.from_template(
    "{country}の首都を教えてください。首都名のみを回答してください。"
)

# 首都から観光スポットを探すチェーン
tourism_prompt = PromptTemplate.from_template(
    "{capital}の有名な観光スポットを3つ教えてください。"
)

# LLMとパーサーの設定
llm = OpenAI(temperature=0)
output_parser = StrOutputParser()

# チェーンの作成
capital_chain = capital_prompt | llm | output_parser
tourism_chain = tourism_prompt | llm | output_parser

# 2つのチェーンを組み合わせ
full_chain = (
    {"capital": capital_chain, "country": RunnablePassthrough()}
    | RunnablePassthrough.assign(tourism=lambda x: tourism_chain.invoke({"capital": x["capital"]}))
)

# 実行例
if __name__ == "__main__":
    country_name = "日本"
    result = full_chain.invoke({"country": country_name})
    
    print(f"国: {result['country']}")
    print(f"首都: {result['capital']}")
    print(f"観光スポット:\n{result['tourism']}")

これを実行すると、mlflowのモデルのトラッキングに記録されます。トレースについても呼び出し順や実行時間、入出力まで一貫して記録できていることが確認できます。

uv run chain.py
uv run mlflow ui

まとめ

今回はmlflow 3の生成AIに関する機能を使ってみました。mlflow上でも生成AIの実験結果を確認できるようになり、広範囲なMLモデルの管理ができるようになってとてもツールとしてレベルが上がったと思います。MLOpsなどでぜひ皆さんも使ってみてください。

Discussion