DatabricksでLangChainのMLflow連携を試してみる

Posted at 2025-01-09

こちらのノートブックをDatabricksで動かします。Databricksの環境に合わせて一部修正しています。ランタイムはDBR 16.1 MLを使います。

MLflow

MLflow は、機械学習ライフサイクル全体でワークフローとアーティファクトを管理するための多用途でオープンソースのプラットフォームです。多くの人気のあるMLライブラリと組み込みの統合があり、任意のライブラリ、アルゴリズム、またはデプロイメントツールと一緒に使用できます。拡張性を念頭に設計されており、新しいワークフロー、ライブラリ、およびツールをサポートするプラグインを作成できます。

LangChain統合のコンテキストでは、MLflowは次の機能を提供します：

エクスペリメント追跡: MLflowは、モデル、コード、プロンプト、メトリクスなど、LangChainエクスペリメントのアーティファクトを追跡および保存します。
依存関係管理: MLflowはモデルの依存関係を自動的に記録し、開発環境と本番環境の一貫性を確保します。
モデル評価: MLflowはLangChainアプリケーションの評価のためのネイティブ機能を提供します。
トレース: MLflowは、LangChainのチェーン、エージェント、リトリーバー、その他のコンポーネントを通じたデータフローを視覚的にトレースすることを可能にします。

注: トレース機能はMLflowバージョン2.14.0以降でのみ利用可能です。

このノートブックでは、MLflowを使用してLangChainエクスペリメントを追跡する方法を示します。この機能に関する詳細情報や、LangChainとMLflowを使用したチュートリアルや例を探索するには、MLflowのLangChain統合に関するドキュメントを参照してください。

セットアップ

MLflow Pythonパッケージをインストールします：

以下のライブラリは後半で使います。

%pip install google-search-results numexpr

MLflowをアップグレードしないと後のオートロギングなどが正常に動きません。バージョン2.19.0で動作を確認できました。

%pip install mlflow -qU

この例では、OpenAI LLMを使用します。以下のコマンドをスキップして、別のLLMを使用しても構いません。

%pip install langchain-openai langchain_community -qU
%restart_python

Databricksでは、最初からMLflowトラッキングサーバーが動いているのでURIの指定は不要です。

import os

# MLflowトラッキングサーバーが稼働している場合はMLflowトラッキングURIを設定
#os.environ["MLFLOW_TRACKING_URI"] = ""
os.environ["OPENAI_API_KEY"] = dbutils.secrets.get("demo-token-takaaki.yayoi", "openai_api_key")

まず、モデルとアーティファクトを追跡するために専用のMLflowクスペリメントを作成しましょう。このステップをスキップしてデフォルトのクスペリメントを使用することもできますが、ランとアーティファクトを別々のエクスペリメントに整理して、散らかりを避け、クリーンで構造化されたワークフローを維持することを強くお勧めします。

エクスペリメントはフルパスで指定します。

import mlflow

mlflow.set_experiment("/Workspace/Users/takaaki.yayoi@databricks.com/20250108_langchain_summarization/LangChain MLflow Integration")

概要

次の方法のいずれかを使用して、MLflowをLangChainアプリケーションに統合します：

自動ロギング: mlflow.langchain.autolog()コマンドを使用してシームレスなトラッキングを有効にします。これは、LangChain MLflow統合を活用するための最初の推奨オプションです。
手動ロギング: MLflow APIを使用してLangChainのチェーンとエージェントをログに記録し、実験で追跡する内容を細かく制御します。
カスタムコールバック: チェーンを呼び出す際にMLflowコールバックを手動で渡し、特定の呼び出しを追跡するなど、ワークロードの半自動カスタマイズを可能にします。

シナリオ 1: MLFlow 自動ロギング

自動ロギングを開始するには、単に mlflow.langchain.autolog() を呼び出します。この例では、log_models パラメータを True に設定しています。これにより、チェーンの定義とその依存ライブラリがMLflowモデルとして記録され、包括的なトラッキング体験が提供されます。

import mlflow

mlflow.langchain.autolog(
    # これらは自動的にログに記録される情報を制御するためのオプション設定です (デフォルト: False)
    # 引数の完全なリストについては、https://mlflow.org/docs/latest/llms/langchain/autologging.html#id1 を参照してください
    log_models=True,
    log_input_examples=True,
    log_model_signatures=True,
)

翻訳チェーンの定義

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model_name="gpt-4o-mini", temperature=0)
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "あなたは {input_language} を {output_language} に翻訳する有能なアシスタントです。",
        ),
        ("human", "{input}"),
    ]
)
parser = StrOutputParser()

chain = prompt | llm | parser

翻訳チェーンの呼び出し

MLflowがモデル、トレース、アーティファクトをトラッキングサーバーにログするため、通常より数秒長くかかる場合があります。

test_input = {
    "input_language": "English",
    "output_language": "Japanese",
    "input": "I love programming.",
}

chain.invoke(test_input)

上のセルを実行することで、トレースおよびモデルのロギングが行われます。

MLflow Tracking UIを探索して、どの情報がログされているかを詳しく理解しましょう。

トレース - 実験の「トレース」タブに移動し、最初の行のリクエストIDリンクをクリックします。表示されるトレースツリーは、チェーン呼び出しのコールスタックを視覚化し、チェーン内の各コンポーネントがどのように実行されるかについての深い洞察を提供します。
MLflowモデル - log_model=Trueを設定したため、MLflowはチェーン定義を追跡するために自動的にMLflowランを作成します。最新のランページに移動し、「アーティファクト」タブを開くと、依存関係、入力例、モデル署名など、MLflowモデルとしてログされたファイルアーティファクトが一覧表示されます。

ロギングされたチェーンを呼び出す

次に、モデルを再度ロードし、同じ予測を再現できることを確認して、一貫性と信頼性を確保しましょう。

モデルをロードする方法は2つあります。

mlflow.langchain.load_model(MODEL_URI) - これはモデルを元のLangChainオブジェクトとしてロードします。
mlflow.pyfunc.load_model(MODEL_URI) - これはモデルをPythonModelラッパー内でロードし、スキーマの強制などの追加ロジックを含むpredict() APIで予測ロジックをカプセル化します。

# MLflow UIに表示されるRun IDをYOUR_RUN_IDに置き換えてください
loaded_model = mlflow.langchain.load_model("runs:/YOUR_RUN_ID/model")
loaded_model.invoke(test_input)

オートロギングの設定

mlflow.langchain.autolog() 関数は、MLflow にログ記録されるアーティファクトを細かく制御するためのいくつかのパラメーターを提供します。利用可能な設定の包括的なリストについては、最新の MLflow LangChain Autologging Documentation を参照してください。

シナリオ 2: コードからエージェントを手動でログ記録する

前提条件

この例では、検索エンジンAPIである SerpAPI をツールとして使用し、エージェントがGoogle検索結果を取得します。LangChainは SerpAPI とネイティブに統合されており、わずか1行のコードでエージェントのツールを構成できます。

始めるには:

必要なPythonパッケージをpipでインストールします: pip install google-search-results numexpr
SerpAPIの公式ウェブサイトでアカウントを作成し、APIキーを取得します。
環境変数にAPIキーを設定します: os.environ["SERPAPI_API_KEY"] = "YOUR_API_KEY"

エージェントの定義

この例では、エージェントの定義を コードとして ログに記録し、Pythonオブジェクトを直接フィードしてシリアライズ形式で保存するのではなく、このアプローチを取ります。このアプローチにはいくつかの利点があります:

シリアライズ不要: モデルをコードとして保存することで、シリアライズの必要がなくなります。これは、ネイティブにサポートされていないコンポーネントを扱う際に問題となるシリアライズを回避できるためです。また、異なる環境でモデルをデシリアライズする際の互換性問題のリスクも排除します。
透明性の向上: 保存されたコードファイルを確認することで、モデルが何をするのかについて貴重な洞察を得ることができます。これは、モデルの動作がロードされるまで不透明であり、リモートコード実行などのセキュリティリスクを潜在的に露呈する可能性があるpickleのようなシリアライズ形式とは対照的です。

まず、エージェントインスタンスを定義する別の .py ファイルを作成します。

時間を節約するために、次のセルを実行してエージェント定義コードを含むPythonファイル agent.py を生成できます。実際の開発シナリオでは、別のノートブックや手作業で作成したPythonスクリプトで定義します。

os.environ["SERPAPI_API_KEY"] = "SerpAPIのAPIキー"

script_content = """
from langchain.agents import AgentType, initialize_agent, load_tools
from langchain_openai import ChatOpenAI
import mlflow

llm = ChatOpenAI(model_name="gpt-4o-mini", temperature=0)
tools = load_tools(["serpapi", "llm-math"], llm=llm)
agent = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION)

# 重要: 記録するインスタンスを登録するために set_model() をコールします
mlflow.models.set_model(agent)
"""

with open("agent.py", "w") as f:
    f.write(script_content)

上のセルを実行すると、ノートブックと同じパスにagent.pyファイルが作成されます。

エージェントのログ

元のノートブックに戻り、次のセルを実行して agent.py ファイルに定義したエージェントをログに記録します。

question = "F1レーシングカーで月まで運転するのにどれくらいかかりますか？"

with mlflow.start_run(run_name="search-math-agent") as run:
    info = mlflow.langchain.log_model(
        lc_model="agent.py",  # エージェント定義への相対コードパスを指定
        artifact_path="model",
        input_example=question,
    )

print("エージェントがMLflowに正常にログされました！")

エージェントがMLflowに正常にログされました！

エージェントが記録されていることを確認できます。

ログされたエージェントの呼び出し

エージェントを再度読み込み、呼び出します。モデルを読み込む方法は2つあります。

# デフォルト設定で自動ロギングをオンにして、エージェント呼び出しのトレースを確認できるようにします。
mlflow.langchain.autolog()

# モデルを再読み込み
agent = mlflow.pyfunc.load_model(info.model_uri)

# 呼び出し
agent.predict(question)

エクスペリメントの 「Traces」 タブに移動し、最初の行のリクエストIDリンクをクリックします。トレースは、エージェントが単一の予測呼び出し内で複数のタスクをどのように操作するかを視覚化します：

質問に答えるために必要なサブタスクを決定する。
F1レーシングカーの速度を検索する。
地球から月までの距離を検索する。
LLMを使用して除算を計算する。

と書いてありますが、上の画面にあるようにノートブック上から直接トレースを確認できます。

シナリオ 3. MLflow コールバックの使用

MLflow コールバック は、LangChain アプリケーションを MLflow で追跡するための半自動化された方法を提供します。利用可能な主なコールバックは2つあります：

MlflowLangchainTracer: 主にトレースを生成するために使用され、mlflow >= 2.14.0 で利用可能です。
MLflowCallbackHandler: メトリクスとアーティファクトを MLflow トラッキングサーバーにログします。

MlflowLangchainTracer

チェーンまたはエージェントが MlflowLangchainTracer コールバックで呼び出されると、MLflow はコールスタックのトレースを自動的に生成し、それを MLflow トラッキングサーバーにログします。結果は mlflow.langchain.autolog() と全く同じですが、特定の呼び出しのみをトレースしたい場合に特に便利です。オートロギングは同じノートブック/スクリプト内のすべての呼び出しに適用されます。

from mlflow.langchain.langchain_tracer import MlflowLangchainTracer

mlflow_tracer = MlflowLangchainTracer()

# この呼び出しはトレースを生成します
chain.invoke(test_input, config={"callbacks": [mlflow_tracer]})

# この呼び出しはトレースを生成しません
chain.invoke(test_input)

Databricksの場合デフォルトでトレースがオンになっているので、2つトレースが記録されました。

コールバックの渡し方

LangChain では、コールバックインスタンスを渡す方法が2つあります: (1) リクエスト時のコールバック - invoke メソッドに渡すか、with_config() でバインドする (2) コンストラクタのコールバック - チェーンのコンストラクタに設定する。MlflowLangchainTracer をコールバックとして使用する場合、リクエスト時のコールバックを使用する必要があります。コンストラクタに設定すると、コールバックはトップレベルのオブジェクトにのみ適用され、子コンポーネントに伝播されないため、不完全なトレースになります。この動作の詳細については、コールバックのドキュメントを参照してください。

# OK
chain.invoke(test_input, config={"callbacks": [mlflow_tracer]})
chain.with_config(callbacks=[mlflow_tracer])
# NG
chain = TheNameOfSomeChain(callbacks=[mlflow_tracer])

サポートされているメソッド

MlflowLangchainTracer は、Runnable Interfaces の次の呼び出しメソッドをサポートしています。

標準インターフェース: invoke, stream, batch
非同期インターフェース: astream, ainvoke, abatch, astream_log, astream_events

他のメソッドは完全な互換性が保証されていません。

MlflowCallbackHandler

MlflowCallbackHandler は、LangChain Community コードベースに存在するコールバックハンドラーです。

このコールバックはチェーン/エージェントの呼び出しに渡すことができますが、flush_tracker() メソッドを呼び出して明示的に終了させる必要があります。

コールバックを使用してチェーンが呼び出されると、次のアクションが実行されます:

新しい MLflow ランを作成するか、アクティブな MLflow 実験内で利用可能なアクティブなランを取得します。
LLM 呼び出しの回数、トークン使用量、その他の関連メトリクスなどのメトリクスを記録します。チェーン/エージェントに LLM 呼び出しが含まれており、spacy ライブラリがインストールされている場合、flesch_kincaid_grade などのテキストの複雑さメトリクスを記録します。
JSON ファイルとして内部ステップを記録します（これはトレースのレガシーバージョンです）。
チェーンの入力と出力を Pandas Dataframe として記録します。
チェーン/エージェントインスタンスで flush_tracker() メソッドを呼び出し、チェーン/エージェントを MLflow モデルとして記録します。

チェーンをシンプルなものに作り直します。

from langchain_community.callbacks import MlflowCallbackHandler

mlflow_callback = MlflowCallbackHandler()

prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "あなたは有能なアシスタントです。",
        ),
        ("human", "{input}"),
    ]
)
parser = StrOutputParser()

chain = prompt | llm | parser

chain.invoke({"input": "What is LangChain callback?"}, config={"callbacks": [mlflow_callback]})

mlflow_callback.flush_tracker()

こちらを実行すると、トレースが表示されることに加え、MLflowランとして上述の情報が記録されます。

参考文献

この機能の詳細や、MLflow を使用した LangChain のチュートリアルや例については、LangChain 統合のための MLflow ドキュメントを参照してください。

MLflow は、LangChain 統合のためのいくつかのチュートリアルや例も提供しています:

はじめてのDatabricks

Databricks無料トライアル

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up