Salesforce Data Cloud の Amazon SageMaker モデルの呼び出し機能を検証する

こんにちは。エンジニアの山下です。

今回は Salesforce Data Cloud と Amazon SageMaker の連携機能について書きたいと思います。

やること

今回は Data Cloud と SageMaker の連携機能を検証するわけですが、実はこの検証に SageMaker は使用しません。本記事のタイトルにも SageMaker のモデルを呼び出すとは一言も書いていなかったりします。

要するにこういうことです。今回は SageMaker のエンドポイントを装ったダミーのサービスを立ち上げて Data Cloud に接続します。そしてダミーのサービスから適当な結果を返して Data Cloud がちゃんとそれを処理してくれることを確認します。つまり SageMaker のエンドポイントをスタブ化するということですね。

このような検証方法を取るのにはそれなりの理由があるのですが、その理由については後ほど説明したいと思います。

Data Cloud の設定

そもそも SageMaker との連携のために何が必要なのかがわからなければ始まらないので、まずは Data Cloud の設定画面から見ていきます。

Data Cloud で SageMaker のモデルを呼び出すには、Einstein Studio で Predictive Model としてモデルを登録する必要があります。以下がその画面です。

SageMaker と接続するための専用の項目がありますね。こちらを選択すると以下の画面に遷移します。

Data Cloud と SageMaker は API 経由で連携するので、ここではその API エンドポイントを設定する必要があります。今回はスタブを使うので Secret Key などの項目は適当で大丈夫です。設定しても特に使用しません。

上記の設定項目の下にはリクエスト・レスポンスの形式を選択するドロップダウンリストがあり、それぞれのフォーマットを JSON か CSV から選択できます。

この画面にはフォーマットのサンプルも載っています。リクエスト・レスポンスのパラメータについては後ほど必要になったタイミングで詳述します。

最後にモデルの入力・出力パラメータを定義します。

アイコンが右端からはみ出していて不穏だ ... 。

この画面の Outputs の JSON Key はノーヒントで特殊な形式の入力を要求してくるので、割と戸惑うポイントです。

公式ドキュメント には $.predictions のような形で入力せよとだけ書いてあり、かつ先ほどあった出力フォーマットのサンプルには以下のような記述があるので、これらから察してくださいということのようです。

{
  "predictions": [0.66, 0.55]
}

配列の下にオブジェクトを置くことも可能で、例えば $.predictions.value とした場合、以下のような JSON から出力を読み取れます。

{
  "predictions": [
    { "value": 0.66 },
    { "value": 0.55 },
    :
  ]
}

意外と高機能ですね。

ちなみに CSV の場合は JSON Key の代わりに Position という項目の入力を要求されます。こちらは各行の何列目の要素が該当する出力項目になるかを指定する形になります。

SageMaker エンドポイントの要件

先の設定画面から、Data Cloud と SageMaker が API エンドポイントを通じて連携を行うことがわかりました。しかし、疑り深い方は「この API エンドポイントを通じて SageMaker に固有の特殊なやり取りが交わされるのではないか」と思うかもしれません。

結論から述べると、そんなことは全くありません。ここで使用する API エンドポイントは一般的な Web サーバが提供する普通の API です。

そもそも SageMaker では、モデルのデプロイを行った場合、裏で機械学習モデルを搭載した Docker コンテナがホストされる仕組みになっているのですが、このコンテナの実体は単なる Web サーバにすぎません。実際、AWS の 公式ドキュメント には以下のように買いてあります。

To receive inference requests, the container must have a web server listening on port 8080 and must accept POST requests to the /invocations and /ping endpoints.

驚くほど単純で逆に不安になりますね ... 。

上記の説明には推論用エンドポイントを /invocations というパスにせよという指定がありますが、Data Cloud は URL で指定されたエンドポイントを愚直に呼び出してくれるので、今回はこの縛りすら特に守る必要はありません。

要するに普通の Web API を用意して Data Cloud に設定すればそれで連携可能なので、API Gateway + Lambda などでサクッと用意していきます。*1

SageMaker エンドポイントのスタブの準備

ここでリクエスト・レスポンスのフォーマットを確認しておきます。Data Cloud から発行されるリクエストの JSON 形式での入力は以下の形式になります。

{
  "instances": [
    { "features": [<param1>, <param2>, ...] },
    { "features": [<param1>, <param2>, ...] },
    :
  ]
}

Data Cloud では基本的に DMO = Data Model Object の１レコードを入力としてモデルの推論を行うことを前提としています。従って、上記の instances はこれらの DMO の集合に対応し、instances の各要素の features の中に入力パラメータ群が格納されます。

他方、レスポンスは先に Data Cloud に設定した Outputs に適合していれば何でも OK です。今回は以下の形式にしました。

{
  "predictions": [
    { "value": 0.66 },
    { "value": 0.55 },
    :
  ]
}

というわけで実装します。

export const handler = async (event) => {
  const predictions = event['instances'].map(_ => {
    return {"value": 0.5};
  });
  return {
    predictions: predictions,
  };
};

スタブなので返却する推論結果は 0.5 で固定しています。

推論の実行

SageMaker と連携して作成した Predictive Model は以下の箇所で利用可能です。

• Flow
• Prediction Job

ちなみに Data Cloud の Batch Data Transform でも Einstein Studio のモデルによる推論が利用可能ですが、こちらは Einstein Studio で作成したモデルに利用が制限されているようで、SageMaker と連携して作成した Predictive Model は利用できないようです。

今回は Prediction Job を利用して推論を実行します。Einstein Studio の Integrations タブから Job を作成することができます。

先にも述べた通り、Einstein Studio のモデルは DMO のレコードが入力になることが前提となっています。そのため、Job を作成する際には、どの DMO レコードのどのフィールドをモデルへの入力として利用するのかを選択します。

他、Job の実行方式や Job の実行結果を格納する DMO の名前の設定などを行います。こちらは特に難しい点はないので割愛します。

Job 作成直後はステータスが Inactive になっているので、忘れずに Activate しておきます。数分程度でステータスが Active に変わるので、後は Run を実行するだけです。実行後、特に問題がなければ Job の Last Run Status が Success に変わります。

また Run の実行前に、スタブのコードに console.log(event) を書き添えておくと、以下のようなログが CloudWatch Logs で確認できます。

2024-07-17T04:00:00.000Z 8d20289a-ada8-4f5e-a7e6-81106fe38731    INFO    {
  instances: [
    { features: [Array] }, { features: [Array] },
    { features: [Array] }, { features: [Array] },
    :
  ]
}

というわけで、Data Cloud とスタブの連携ができました。簡単でしたね。

なぜスタブを使ったのか

SageMaker のエンドポイントにスタブを利用したため、検証としてはやや遠回りな印象を与えたかもしれませんが、実は今回の手順が Data Cloud と SageMaker を連携させるための必須の手順になります。

というのは、何を隠そう、現時点では Data Cloud と SageMaker を直接接続する手段がないためです。

本投稿の冒頭で Data Cloud の設定画面を掲載しましたが、よくよく見ると、API エンドポイントの設定画面にある認証関連項目は API Key のみとなっています。しかし、SageMaker のエンドポイントでは AWS Signature V4 による署名が要求され、この署名のためには少なくとも２つ以上のヘッダー要素の指定が必要となります。この認証項目の不整合により、Data Cloud の現在の枠組みで SageMaker エンドポイントの認証を通すことが不可能になっています。

従って、Data Cloud と SageMaker を連携させるには、AWS Lambda をはじめとする仲介役が必ず必要になります。この制約により、SageMaker エンドポイントの本質が単なる Web サーバであることを理解していない開発者はそもそも連携を始めることすらできない、というのが Data Cloud の現状です。

しかし、逆にスタブの呼び出しさえできれば、その実体は Lambda なので、そこから SageMaker のモデルを呼び出すことはもちろん、他のクラウドベンダーのモデルも呼び出し放題です。それでいいのか。

そんなわけで、今回はスタブを使った Data Cloud と SageMaker の連携機能の検証方法についてお送りしました。

なお API Gateway + Lambda + SageMaker の構成でエンドポイントを構成する方法については Amazon にブログ記事があるので、そちらを参照するのがよさそうですね。

aws.amazon.com

まとめ

今回の内容をまとめると、

• SageMaker のエンドポイントの本質は単なる Web サーバである
• Data Cloud と SageMaker は認証項目の不整合で直接連携はできない
• Data Cloud と SageMaker を連携させるには Lambda などの仲介役が必要

といった感じになりそうです。

Data Cloud と SageMaker の連携が直接できないのは歯痒いですが、SageMaker の裏側の仕組みについての知見が深まり、個人的にはそれなりに収穫のある検証でした。

Data Cloud はまだ新しいサービスなので、今後の改良で SageMaker と直接連携ができるようになることを期待したいですね。

以上です。最後までお読みいただき、ありがとうございました。

*1:今回はレスポンスに JSON を使用するので Lambda を使用していますが、レスポンスを CSV で返す場合は Lambda の採用が問題になることがあります。実は Lambda で JavaScript や Python を使った場合、レスポンスが JSON 形式であるという暗黙の仮定がなされ、CSV をうまく返却することができません。従って、この問題が発生しない Java や C# などの言語を Lambda で使うか、そもそも Lambda を使わずに EC2 などで API サーバを立てる必要があります。