databricksのModel ServingにカスタムAgentをデプロイする時の手順について。

こんにちは、CCCMKホールディングスAIエンジニアの三浦です。

日中はとても暑いのですが、朝早い時間に外に出るとほんの少し涼しさを感じることがあります。8月も後半に入ったんだなぁと改めて思いました。

以前databricksのModel Servingを使ってRAG AgentをAPIで提供する方法についてご紹介しました。

techblog.cccmkhd.co.jp

この時はdatabricksのPlaygroundで作ったAgentをModel Servingにデプロイしたのですが、最近LangGraphで独自に作ったAgentをModel Servingにデプロイする方法を調べていました。色々と詰まるところがあったのですが、無事にModel Servingにデプロイし、APIで利用できるようにするところまで進めることが出来ました。

今回の記事ではLangGraphで開発したAgentをModel Servingにデプロイするまでの手順と、進めるにあたってポイントになった点をまとめていきたいと思います。また、前回は米国リージョンで検証をしましたが、Model ServingがAzureの日本リージョンでも利用できるようになったため、今回は日本リージョンで検証を行っています。

作業の流れ

AgentをdatabricksのModel Servingで稼働させるためにはいくつかのステップを踏む必要があります。①MLflowにAgentを記録する、②Unity CatalogのModelに登録する、③Model Servingにデプロイする、です。

この中で特に気を付ける必要があるのが①のステップです。①がうまく出来ていないと③のデプロイのステップで失敗してしまうことが時々ありました。

①MLflowにAgentを記録する

最初にMLflowにAgentを記録する必要があります。以下のドキュメントに記載されている、コードベースの記録という方法を取りました。

docs.databricks.com

まず以下のようなAgentを定義したPythonのスクリプトを用意します(wikipedia_agent.pyという名前で保存しました。)。このAgentはLangGraph-Supervisorを使って開発したMulti-Agentで、ユーザーが「プログラミング言語の歴史についてまとめてレポートにしてください。」といった指示を与えると、レポートを作るために必要な情報をwikipediaから検索し、それらをまとめてレポートを作ります。どんな情報を検索したらよいのか、検索計画も自身で立てることが出来るようになっています。

from databricks_langchain.chat_models import ChatDatabricks
from langchain_community.tools import WikipediaQueryRun
from langchain_community.utilities import WikipediaAPIWrapper
from langchain_core.prompts import PromptTemplate
from langgraph.prebuilt import create_react_agent
from langgraph_supervisor import create_supervisor
import mlflow

from prompts import (planner_prompt, wikipedia_prompt, reporter_prompt, supervisor_prompt)

# llm
llm = ChatDatabricks(model="gpt-4o")

# tool
wiki_tool = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper(lang="ja"))

planner = create_react_agent(
    name="planner",
    model=llm,
    prompt=planner_prompt,
    tools=[]
)

wikipedia_searcher = create_react_agent(
    name="wikipedia_searcher",
    model=llm,
    prompt=wikipedia_prompt,
    tools=[wiki_tool]
)

reporter = create_react_agent(
    name="reporter",
    model=llm,
    prompt=reporter_prompt,
    tools=[]
)

mlflow.langchain.autolog()

supervisor = create_supervisor(
    model=llm,
    prompt=supervisor_prompt,
    agents=[planner, wikipedia_searcher, reporter]
).compile()

mlflow.models.set_model(supervisor)

MLflowのコードベースの記録をする場合には、記録する対象のModel(Agent)をmlflow.models.set_modelで指定する必要があります。また、各Agentに与えるシステムプロンプトは別のPythonスクリプトにまとめました(prompts.pyという名前で保存しました。)。

planner_prompt = 'あなたユーザーの指示を達成するために必要になる情報をどのようにwikipediaから検索したらよいか計画を立案する役割を持ったAgentです。'
wikipedia_prompt = 'あなたは与えられた検索クエリでwikipediaから情報検索し、返す役割を持ったAgentです。'
reporter_prompt = 'あなたはこれまでに収集した情報をまとめ、ユーザーへの回答を生成する役割を持ったAgentです。'
supervisor_prompt = """
あなたはユーザーからの命令を達成するため、複数のAgentにタスクを渡す熟練の管理者です。あなたの管理下のAgentの役割を以下にまとめます。

- planner: ユーザーの指示を達成するために必要になる情報をどのようにwikipediaから検索したらよいか計画を立案する
- wikipedia_searcher: 与えられた検索クエリでwikipediaから情報検索し、返す
- reporter: これまでに収集した情報をまとめ、ユーザーへの回答を生成する

あなたは管理下のAgentに対し1つずつタスクを与え、ユーザーの指示に答えて下さい。"""

これらのスクリプトと同階層でノートブックを新規作成しMLflowに記録する処理を実装していくのですが、2点注意する点があります。

MLflowのバージョンは3.2以降を使う

mlflow.langchain.log_model()を使ってAgentを記録する時にMLflowのバージョンが3.1.xだと"AttributeError('module langchain.chains has no attribute base')"というエラーが出て失敗してしまう現象が発生しました。

この現象は8月にリリースされたMLflow3.2.0では解消されているため、3.2以降をインストールして作業をする必要があります。

Databricks runtimeは17.0MLを使う

最初、Agentを記録する際にDatabricks runtimeのバージョンは16.4 ML(Scala 2.12)を使用していたのですが、Model Servingにデプロイする時のコンテナイメージの構築時にライブラリの依存関係のエラーが発生しました。おそらくAgentを記録する際に使用した環境に含まれていたPythonのライブラリのバージョンに起因するものと考え、Databricks runtimeのバージョンを1つ最新の17.0MLにしてみたところ、エラーが発生せずにデプロイが完了しました。

準備が整っていればAgentの記録はmlflow.langchain.log_modelを呼び出すだけで行うことが出来ます。lc_modelにはAgentを定義したPythonスクリプトを、code_pathsにはAgentが参照するモジュールを指定します。databricksの他のリソースをAgentからアクセスさせる場合にはresourcesに指定します。また、input_exampleの指定は必須です。

# mlflowにModelとして記録(log_model)する
from mlflow.models.resources import DatabricksServingEndpoint

# 入力例は必須
input_example = {
    "messages":[
        {"role":"user","content":"プログラミング言語の歴史についてまとめてレポートにしてください。"}
    ]
}

with mlflow.start_run(run_name="wikipedia_agent"):
    logged_agent_info = mlflow.langchain.log_model(
        name="wikipedia_agent",
        lc_model="./wikipedia_agent.py",
        input_example=input_example,
        code_paths=["./prompts.py"],
        resources=[
            DatabricksServingEndpoint(
                endpoint_name="gpt-4o"
            )
        ],
    )

②Unity CatalogのModelに登録する

先ほどのAgentを記録するノートブックで記録が完了すると、記録されたAgentのURLが表示されます。そちらにアクセスし、遷移したページの右上に表示されている"Register model"をクリックします。

クリックすると、Agentを記録するUnity CatalogのModel名の入力を求められるので、記入し、"Register"を実行します。

しばらくすると、Unity Catalogの対象のSchema配下のModelsに登録したAgentが表示されていることが確認できます。

③Model Servingにデプロイする

Unity Catalogに記録されたModelsのページから、"Serve this model"をクリックするとModel Servingへのデプロイ画面に進むことが出来ます。

デプロイする際に"Inference tables"を有効にするかどうかを設定することが出来ます。Inference tablesはModel ServingにデプロイされたAgentへの入力と出力を自動的に記録してくれる便利な機能です。なのですが、これを有効にするのに少し詰まってしまいました。おそらく次のドキュメントに記載されているレガシーなInference tablesの設定とAI GatewayのInference tablesの設定を混同してしまったからだと思います。

learn.microsoft.com

私がInference tablesを有効に出来たときの手順をご紹介します。

最初にServed entitiesにある"Enable tracing"のチェックはオフにします。