StepFunctionsを使ってSageMakerエンドポイントのデプロイを実行する

こんにちは、CTO室AI推進部@ken11です。
みなさんカナリアは好きですか?
僕は大好きです。
華麗なパスさばき、高速ドリブル、迫力のシュート、カナリア軍団ブラジル代表のプレーはいつ見ても美しい…

ええ、もちろんここはエンジニアブログなのでサッカーの話をしに来たわけではありません。サッカーの話がしたい人は恵比寿に来てもらえれば夜な夜なCLの試合を見守る僕とサッカー談義を繰り広げることが可能かもしれません。

今日話したいのはカナリアはカナリアでもカナリアリリースの話です。
SageMakerエンドポイントってカナリアリリースできるの?

SageMakerエンドポイントのデプロイ

結論からいえば、頑張ればカナリアリリースできます。
そもそもSageMakerエンドポイントってデプロイ作業とかみなさんどうしてますか?

ECSと違って絶妙にデプロイで困りませんか?

※現状CodeDeployはSageMakerエンドポイントに対応していない之図

※API的にはようやくDeployConfigに対応したっぽい之図

そうなんですよ、CodeDeployは使えないしSageMakerエンドポイント自体にデプロイ設定の項目はないし、ほんとつい先日までAPIでもDeployConfigはnot supportedだったみたいな状態で、SageMakerエンドポイントってCIからデプロイしたいときとかどうするの?みたいな疑問にぶち当たっていました。

※APIでもnot supportedだったとき之図

幸い10月にStepFunctionsがAWS SDK統合したことで、StepFunctions上からできることがものすごく増えたので、今回はそれを使ってSageMakerのデプロイを実行するなにかをつくってみたという話です。

StepFunctionsでデプロイ管理

この チャンピオンズロードのマップかな? みたいに長いのがStepFunctionsを使ったデプロイ装置です。
定義のjsonはこちらにあるので使ってみたい方はどうぞ。
ポイントをいくつか説明していきたいと思います。

①モデルの作成とエンドポイントの存在確認

なにはともかくモデルを作成しないと始まりません。
ここでいうモデルとは、機械学習モデルファイルのことではなく、SageMakerの「モデル」です。(これめっちゃわかりづらい…)
簡単に言えばエンドポイントで使用するECRのURIやモデルファイルの場所が詰まった設定みたいなものですね。
(SageMakerエンドポイントはモデル→エンドポイント設定(どのモデルを使うかなど)→エンドポイントという流れになっていて、デプロイの際はそれぞれ作成したり更新したりする必要があります)

エンドポイントの存在確認は単純に既にエンドポイントがあるかどうか確認しているだけです。
なかったら初回デプロイとみなして右側のルートに進みます。

②カナリアリリース

ここでは最初にいまエンドポイントで使われているモデルの情報を取得しておきます。
カナリアはもちろん、ロールバックする際にも必要になるので、この情報はしばらく引き回していきます。

その後、新しいモデルと2:8でトラフィックを割り振ったエンドポイント設定を作成し、エンドポイントをアップデートしていきます。
エンドポイント設定にはVariantごとの「重み」というのを設定できるのですが、この重みがトラフィックの割合になっているようです。
0.2と0.8とかで設定すれば2:8で割り振れます。
ドキュメント的にはこの辺です。

なお、UpdateEndpointを実行しても InService まで待ってくれるわけではないので、Choiceを使って InService 状態になるまでエンドポイントの更新を待つループを入れてます。

③更新 or ロールバックの判断

端折ってる部分でカナリアの方(2:8の2の方)のVariantで4xxと5xxのエラー監視のアラームを入れているのですが、これを元にエンドポイントを新しいモデルで完全に更新するのか、あるいはロールバックするのかの判断をしています。
ここは本当は手動で管理者の承認とかにするのがいいんでしょうけど、今回は ~面倒だったから~ とりあえずアラームが出ていなかったらGoサイン、アラーム出てたらロールバックという自動判断にしています。(途中の Wait は判断するまで20分ほど待ち時間をおいている感じです、その間に4xxや5xxが起きなかったらGoサイン)

無事にGoサインが出た場合は、新しいモデルを使ったエンドポイント設定を作成してエンドポイントを更新します。
ここで「え?UpdateEndpointでカナリアのVariantの重みを1にして元のVariantの重みを0にすればよくない?」と思ったりもするんですが、重みを0にしただけではインスタンスは稼働したままなので費用が発生すると思われます。(たぶん、普通にかかるはず)
かといってUpdateEndpointでVariantのDelete的なことはできそうにないので、結局新しいエンドポイント設定を作ってそれを当てる方法をとることに。

④オートスケール設定をあてる

最後に、無事に更新されたエンドポイントにオートスケールの設定をあてておきます。
もちろんオートスケール必要ない方にはいらないステップなんですが。

SageMakerエンドポイントのオートスケール設定は、Application Auto Scalingというサービスでまかなわれており、Variantごとにこれを設定してスケーリングポリシーをあてることで実現できます。
この際、あてるスケーリングポリシーの名前を SageMakerEndpointInvocationScalingPolicy にしておくことをオススメします。

これはAWSコンソールからSageMakerエンドポイントのオートスケール設定をする際の画面ですが、このように「組み込みのスケーリングポリシー」と「カスタムスケーリングポリシー」という2種類があります。
正直なにが違うのかよくわかってないんですが、Application Auto ScalingのPutScalingPolicyを実行する際、任意の名前を付けると「カスタムスケーリングポリシー」に、 SageMakerEndpointInvocationScalingPolicy という名前にすると「組み込みのスケーリングポリシー」になります。

そして当方の環境だけかわかりませんが「カスタムスケーリングポリシー」だとオートスケールが発火した際に正常に完了せずエンドポイントが「Updating」ステータスのままスタックする事象に遭遇します…
なにが原因かわかってませんが「組み込みのスケーリングポリシー」の場合はそういった事象が起きていないので、一応いまは名前を SageMakerEndpointInvocationScalingPolicy とすることで難を逃れる形にしています。

以上が「StepFunctionsでつくったSageMakerエンドポイントのデプロイをするなにか」です。
StepFunctionsに慣れた人なら、定義のjsonを見ればだいたいなにやってるかわかるかと思います。

CIから実行する

最後に、もちろんこれをいちいち手で実行していてはなんのためにStepFunctionsをつくったのかわからないので、CIから実行するように設定していきます。
たとえばCircleCIで実行する場合は

    steps:
      - checkout
      - aws-cli/install
      - aws-cli/setup
      - run:
          command: |
            aws stepfunctions start-execution --state-machine-arn "なにかしらのarn" --input "なにかしらのexecution_input的json"

こんな感じで、aws-cliのOrb使って aws stepfunctions start-execution で実行するだけです。
デプロイが完了したかわかるように、StepFunctionsの最後にSlack通知とか入れるともっとよさそうですね。

まとめ

というわけで今日はStepFunctionsを使ってSageMakerエンドポイントのデプロイを実行する話でした。
「推論コードに変更があったらGitHubでマージするだけで自動でデプロイされるぜ〜」というのがやりたかっただけなんですが、ずいぶんと大仰なものに仕上がってしまいました。

もちろんこれがベストプラクティスとも思ってなくて、APIを見る限り今後SageMakerエンドポイント自体にもECSのように気軽にBlue/Greenできるような機能が追加されていくんだろうなーと思いますし(想像でしかない)、早くそういうのが使えるようになるといいなと思ってます。個人的にはCodeDeployが使えるようになると嬉しいんですが…
現状ではデプロイ方法を探してもCloudFormation等しか出てこず、あまり自分が求めていたものがなかったので、今回はStepFunctionsでつくってみた次第です。
自動でデプロイできるようになったことはもちろん、改めてAWS SDK統合したStepFunctionsの強力さもわかってとてもいい経験になりました。

AI推進部ではこのようにCI/CD/CTといったDevOps/MLOpsが好きな方を募集しています。
ご応募お待ちしてます。


マネーフォワードでは、エンジニアを募集しています。
ご応募お待ちしています。

【サイトのご案内】
マネーフォワード採用サイト
Wantedly
京都開発拠点

【プロダクトのご紹介】
お金の見える化サービス 『マネーフォワード ME』 iPhone,iPad Android

ビジネス向けバックオフィス向け業務効率化ソリューション 『マネーフォワード クラウド』

おつり貯金アプリ 『しらたま』

お金の悩みを無料で相談 『マネーフォワード お金の相談』

だれでも貯まって増える お金の体質改善サービス 『マネーフォワード おかねせんせい』

金融商品の比較・申し込みサイト 『Money Forward Mall』

くらしの経済メディア 『MONEY PLUS』

Pocket