はじめに
この記事は 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アプリケーションの提供や、検索エンジンの構築など、様々なサービスに活用できると考えられます。