Databricksから外部LLMへアクセスするにはMLFlow Deployments APIを使いましょう、というお話
はじめに
LLMのサービングインフラとしてDatabricksを使用する場合、以下の２つが代表的な活用パターンだと思います。
- Databricks上で直接LLMを動かす
- ご自身のカスタムモデルなどをRESTエンドポイントとしてDatabricks上にデプロイし、DatabricksのサーバーレスGPUサービングを利用して運用する
- Databricksを外部のSaaS型LLMへの仲介サーバーとして使用する
- OpenAI GPT-4などの外部LLMへのアクセスゲートウェイやハブとして、Databricks上に中継用のエンドポイントを立ち上げ運用する。小さいCPUインスタンスでも十分。
このうち、「1. Databricks上で直接LLMを動かす」については、さらに２つのオプションがあります。それぞれの詳細はドキュメントをご覧ください。
ただ、今回の本題は「2. Databricksを外部のSaaS型LLMへの仲介サーバーとして使用する」に関してです。まずは、これまではどのように実装していたかをご説明します。
これまで＆これから
MLFlow ２.4まではMLFlow Model Registryへの登録が必要
これまで外部のLLMにDatabricksを経由してアクセスしようと思うとこちらのブログのように、
- MLFlowに所望の外部モデルを登録（モデル本体ではなく、モデルへのアクセス情報（API KEYなど））
- 登録された外部モデルを使用してサービングエンドポイントを作成
- RESTでエンドポイントにアクセス
というな手順を踏む必要がありました。
特に最初の手順「1. MLFlowに所望のモデルを登録」が必要なことで、なんとなく直感性に欠け、煩雑さが増していた印象があります。
MLFlow ２.5からはAI Gatewayが登場
その課題を含め、LLMサービングをより簡易化するためにMLFlow 2.5.0から登場したのがAI Gatewayです。
上記アナウンスにもある通り、Databricks上でもマネージドサービスとしてAI GatewayがPrivate Previewという形で提供されてきました。
いずれにしてもAI Gatewayを使用することで、外部モデルへのエンドポイントをRouteとして非常に簡単に登録でき、クライアントからもゲートウェイの提供する統一的なAPIでアクセスが可能なため、LLM更新や切り替えに伴うアプリケーションのメンテナンスコストの最小化なども期待されました。さらに外部LLMへのアクセス用Credentialなどをゲートウェイ内で一元管理できることから、セキュリティ面での期待もありました。
MLFlow ２.9からはDeployments APIに移行
ただ、そんなAI Gatewayですが、MLFlow 2.9.0からは非推奨となり、代わりにMLflow Deployments APIへの移行が推奨されています。
従いまして、これまでDatabricks上でPrivate PreviewのAI Gatewayを使用されてきた方、または、それ以外の環境でOSSのMLFlow AI Gatewayを使用されてきた方はこの機会にMLflow Deployments APIへの移行を御検討いただくのが良いかと思います。
なお、Deployments APIは、2024年１月４日時点で以下の外部LLMに対応しています。
- openai
- Azure OpenAI
- anthropic
- cohere
- aws-bedrock
- palm
- ai21labs
- databricks-model-serving（※これはつまり、Databricks上で自身または他者がサービングしているモデルを外部モデルとして定義することができるということです）
AI GatewayからDeployments APIへの移行
ではここからは、AI GatewayからDeployments APIへの移行について、具体的なコードを用いながら説明をしていこうと思います。
AI Gatewayを用いた実装例
まずはAI Gatewayで外部LLMのエンドポイントを作る場合です。
ここでは例としてOpenAIのGPT-3.5をCompletion APIで使用するケースを見ていきます。
管理者側
AI Gatewayの特徴は各LLMへのアクセス情報をRouteというオブジェクトとしてゲートウェイ側に登録することです。また、登録後のRouteを管理するためのAPIも複数用意されています。
import mlflow.gateway
mlflow.gateway.set_gateway_uri("databricks")
route_name = "my_llm"
# Routeを新規に作る
mlflow.gateway.create_route(
name=route_name,
route_type="llm/v1/completions",
model={
"name": "gpt-3.5-turbo-instruct",
"provider": "openai",
"openai_config": {
"openai_api_key": dbutils.secrets.get(scope=SECRETS_NAME, key=KEY_NAME),
}
}
)
# 既存のRouteの一覧を取得する
mlflow.gateway.search_routes()
# 既存の特定のRouteへの参照を取得する
mlflow.gateway.get_route(route_name)
# 既存のRouteを削除する
mlflow.gateway.delete_route(route_name)
クライアント側
クライアントのコードも非常にシンプルです。AI GatewayのAPIを使用することで、バックエンドのLLMの種別に関わらず統一されたクライアントアプリケーションコードを書くことができます。
import mlflow.gateway
mlflow.gateway.set_gateway_uri("databricks")
route_name = "my_llm"
response = mlflow.gateway.query(
route=route_name,
data={
"prompt": "MLOpsは何ですか？",
"temperature": 0.3,
"max_tokens": 512,
}
)
print(response)
Deployments APIを用いた実装例
では、ここからDeployments APIベースのコードに変更していきます。
MLFlowのバージョンアップ
まず、これはMLFlow 2.9.0からの機能なので、Databricks Runtimeにプリインストールされているバージョンがそれよりも古い場合には以下の通りにバージョンアップを実施します。
%pip install mlflow[genai]>=2.9.0
dbutils.library.restartPython()
管理者側
管理者側のコードで一番の違いは、AI GatewayでRouteと呼ばれていたものが、Endpointという名前で統一されたことくらいでしょうか？それ以外にもメソッド名の変更など細かな違いはあれど、大きくは同じような雰囲気で書けると思います。
import mlflow.deployments
client = mlflow.deployments.get_deploy_client("databricks")
endpoint_name = "my_llm"
# Endpointを新規に作る
client.create_endpoint(
name=endpoint_name,
config={
"served_entities": [{
"external_model": {
"name": "gpt-3.5-turbo-instruct",
"provider": "openai",
"task": "llm/v1/completions",
"openai_config": {
"openai_api_key": "{{secrets/SECRETS_NAME/KEY_NAME}}"
}
}
}]
}
)
# 既存のEndpointの一覧を取得する
endpoint_list = client.list_endpoints()
# 既存の特定のEndpointへの参照を取得する
endpoint = client.get_endpoint(endpoint_name)
# 既存のEndpointを削除する
deleted_endpoint = client.delete_endpoint(endpoint_name)
クライアント側
クライアント側のコードもRouteからEndpointに変わったこと以外は、細かな名称の変更くらいでそこまで大きなコード修正は無いように見えます。
import mlflow.deployments
client = mlflow.deployments.get_deploy_client("databricks")
endpoint_name = "my_llm"
response = client.predict(
endpoint=endpoint_name,
inputs={
"prompt": "MLOpsは何ですか？",
"temperature": 0.3,
"max_tokens": 512,
}
)
print(response)
Deployments APIにのみ存在する機能は「Update」
前述の通り、AI Gatewayには登録されているRouteを更新するAPIが存在していませんでした。したがって、更新処理を実施したい場合は一度当該Routeをゲートウェイから削除した上で、新たな設定情報で新規作成する必要がありました。
しかし、Deployments Serverは既存のEndpointのUpdateが可能です。例えば以下のように使用できます。
例1：OpenAIのGPT-３.5をdavinci-002に変更する
updated_endpoint = client.update_endpoint(
endpoint=endpoint_name,
config={
"served_entities": [{
"external_model": {
"name": "davinci-002",
"provider": "openai",
"task": "llm/v1/completions",
"openai_config": {
"openai_api_key": "{{secrets/SECRETS_NAME/KEY_NAME}}"
}
}
}]
},
)
print(updated_endpoint)
この場合、taskはv1/completionsのままなので、クライアント側のコードは以下の通り先ほどと同様のままで正常に動作します。
import mlflow.deployments
client = mlflow.deployments.get_deploy_client("databricks")
endpoint_name = "my_llm"
response = client.predict(
endpoint=endpoint_name,
inputs={
"prompt": "MLOpsは何ですか？",
"temperature": 0.3,
"max_tokens": 512,
}
)
print(response)
例2：OpenAIのGPT-３.5をCompletion APIからChat Completion APIに変更
updated_endpoint = client.update_endpoint(
endpoint=endpoint_name,
config={
"served_entities": [{
"external_model": {
"name": "gpt-3.5-turbo",
"provider": "openai",
"task": "llm/v1/chat",
"openai_config": {
"openai_api_key": "{{secrets/SECRETS_NAME/KEY_NAME}}"
}
}
}]
},
)
print(updated_endpoint)
この場合はtaskがv1/chat、つまり、Chat Completion APIに変更されたので、クライアント側のコードも以下のように書き直してあげる必要があります。
import mlflow.deployments
client = mlflow.deployments.get_deploy_client("databricks")
endpoint_name = "openai-chat-endpoint-hiouchiy"
data = dict(
messages=[
{"role": "system", "content": "あなたはAIのことならなんでも知っている世界的に著名な博士です。"},
{"role": "user", "content": "MLOpsは何ですか？"}
],
temperature=.3,
max_tokens=512
)
response = client.predict(endpoint=endpoint_name, inputs=data)
print(response)
まとめ
MLFlow AI GatewayからMLFlow Deployments Serverへの移行方法を具体的なコードとともに記載しました。AI Gatewayが非推奨となったことを知った時、個人的には多少動揺がありましたが、AIの進化と同様に周辺ソフトウェアも日進月歩で進化しておりますので、やむを得ないものと理解しています。基本的にはより使いやすく、より便利になっていくはずなので、ぜひこの変化を楽しんでいきましょう。こちらのブログでも極力タイムリーに最新情報をお届けしてまいります。
Discussion