Databricks Model ServingとAWS API Gatewayで作るML API

はじめに

この記事は every Tech Blog Advent Calendar 2024(夏) 13日目の記事です。

こんにちは。株式会社エブリーの開発本部データ&AIチーム(DAI)でデータエンジニアをしている吉田です。

今回は将来的なMLモデルのサービス組み込みに向けた調査の一環として、Databricks Model ServingとAWS API Gatewayを利用してML APIを作成するPoC行ったので、その取り組みについて紹介します。

Databricks Model Serving

Databricks Model Servingは、Databricksが提供するAIモデルをデプロイ・管理するためのサービスです。
Databricks上で構築・学習したモデルや、Databricks Marketplace上で公開されているモデルをデプロイし、REST APIを通じて推論を行うことができます。

https://docs.databricks.com/ja/machine-learning/model-serving/index.html

データ基盤としてDatabricksを利用している現環境では、Databricks Model Servingを利用することで、モデル構築からデプロイまでの一貫した開発フローを提供できると考えました。

AWS API Gateway

AWS API Gatewayは、APIを作成・公開・管理するためのサービスです。

https://aws.amazon.com/jp/api-gateway/

構成

今回のPoCでは、以下を参考に、Databricks Model ServingとAWS API Gatewayを組み合わせてML APIを作成しました。
https://aws.amazon.com/jp/blogs/news/creating-a-machine-learning-powered-rest-api-with-amazon-api-gateway-mapping-templates-and-amazon-sagemaker/

以下のような構成になります。
- Databricks Model Servingでモデルをデプロイ
- Model Servingで提供されるServing EndpointのRest APIをAWS API Gatewayでラップ
- API Gatewayで提供されるエンドポイントにリクエストを送信し、推論結果を取得

Databricks Model Servingが提供するServing EndpointをAWS API Gatewayでラップすることで、AWS API Gatewayが持つAPIの管理や認証、モニタリングなどの便利な機能が利用できます。

実装

モデルの構築

今回はサンプルとして、入力を二倍にして返す単純なモデルを構築します。

モデルの構築

import mlflow


class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input, params=None):
        # in [1,2] -> out [2,4]
        return model_input * 2

モデルの登録

import mlflow
from mlflow.models import infer_signature

import pandas as pd

model = MyModel()
input_data = pd.DataFrame([{"q": 2}, {"q": 10}])
output = model.predict(None, input_data) # 4, 20

signature = infer_signature(input_data, output)

with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        "test_model",
        python_model=MyModel(),
        input_example=input_data,
        signature=signature,
        registered_model_name="test_model",
    )

model_info.model_uri

databricks上でmlflow.pyfunc.log_modelを実行することで、Databricks Model Registryにモデルが登録されます。

モデルのデプロイ

Databricks Model ServingのUIからモデルをデプロイします。
エンドポイント名、利用するモデル、コンピュートリソースなどを設定しデプロイします。
デプロイが完了すると、Serving Endpointが提供されます。

Serving Endpointのテスト

Serving Endpointに対してリクエストを送信し、推論結果を取得します。
REST APIを利用するため、curlコマンドなどでリクエストを送信することができます。
また、Databricks Model ServingのUIや、Databricks SQLからもリクエストを送信することができます。

query

curl \
  -X POST \
  -H "Content-Type: application/json" \
  -d '{"dataframe_split": {"columns": ["q"], "data":[100,20]}}' \
  https://{databricks-host}/serving-endpoints/test-model/invocations | jq .

result

{
  "predictions": [
    {
      "q": 200
    },
    {
      "q": 40
    }
  ]
}

API Gatewayの設定

Databricks Model ServingのServing EndpointをAWS API Gatewayでラップします。
AWS API Gatewayのマッピングテンプレート機能を使い、リクエスト/レスポンスの変換を行います。

統合リクエストのマッピングテンプレート

#set( $queries = $input.params("query") )
{
    "dataframe_split" : {
    "columns" : ["q"],
    "data" : [#foreach( $query in $queries.split(","))$query#if($foreach.hasNext),#end#end]
    }
}

統合レスポンスのマッピングテンプレート

#set($predictions = $input.path("$.predictions"))
{
  "results" : [#foreach( $item in $predictions )$item.q#if($foreach.hasNext),#end#end]
}

APIの呼び出し

API Gatewayのエンドポイントに対してリクエストを送信し、推論結果を取得します。

query

curl "<invoke-url>/test?query=1,2"

result

{
  "results": [2, 4]
}

まとめ

今回はDatabricks Model ServingとAWS API Gatewayを組み合わせてML APIを作成するPoCを行いました。
Databricks Model Servingを利用することで、モデルの構築からデプロイまでの一貫した開発フローを提供できます。
また、AWS API Gatewayを利用することで、APIの管理や認証、モニタリングなどの便利な機能が利用できます。

Databricks Vector Searchなどと組み合わせることで、RAGアプリケーションの提供や、検索エンジンの構築など、様々なサービスに活用できると考えられます。