株式会社エブリーの技術ブログ

全469件

2024/06/13

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を利用することで、モデル構築からデプロイまでの一貫した開発フローを提供できると考えました。 DatabricksでのML開発フロー 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 models モデルのデプロイ Databricks Model ServingのUIからモデルをデプロイします。エンドポイント名、利用するモデル、コンピュートリソースなどを設定しデプロイします。デプロイが完了すると、Serving Endpointが提供されます。 databricks model serving 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のマッピングテンプレート機能を使い、リクエスト/レスポンスの変換を行います。 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アプリケーションの提供や、検索エンジンの構築など、様々なサービスに活用できると考えられます。

株式会社エブリー の技術ブログ

コンテンツ

株式会社エブリーの技術ブログ