Amazon SageMaker StudioでMLOpsを始めました

こんにちは！スマートバンクでSREをしている @maaaato です。

今回はスマートバンクが提供するB/43の裏側にAmazon SageMaker Studio（以下SageMaker Studio）の利用を開始したのでSageMaker Studioにまつわる話をしたいと思います。

筆者はこれまでにMLOpsの運用は行った経験がなく初めての試みでした。同じくSageMakerの構築・運用も初めてだったためいくつか悩んだポイントがありました。

特にSageMakerにはいろんな機能がありどれを使うのがベストなのか迷ったり、MLOpsを設計するにあたり、AWS Step Functions（以下Step Functions）を採用している事例などもあることからAmazon SageMaker Studio Pipelinesとどちらがベストなのか？Step Functionsとの違いはなにか考えることが多々ありました。

そんな経験を公開することで少しでもみなさんのお役に立てられれば幸いです。

この記事の想定読者は以下の通りです。

これからSageMakerの運用を始める人
AWSでMLOpsを始めるにあたってSageMaker StudioがいいのかAWS Step Functionsがいいのか判断材料がほしい人

B/43の裏側を支えるSageMakerについて

改めてSageMakerについて簡単に説明すると、機械学習（Machine Learning）のためのマネージドサービスです。機械学習モデルの構築、トレーニング、デプロイメントなどの一連の作業を簡単に実行できます。

初めてSageMakerの説明を読んだときは全く理解ができず、そもそもの前提知識が無さすぎでした。そこでAWSのBlack BeltのYoutubeの動画やクラスメソッドさんの記事を参考にしつつ、解像度を上げていきました。

アーキテクチャー概要

ざっくりとしたフローは以下の通りです。

学習データ取得バッチによりDBから必要なデータを取得しS3にアップロード
学習データ取得バッチが学習用パイプラインを実行
新しいモデルが作成される
作成されたモデルの評価を行い基準を超えていたらApproveする
ApproveされたことをEventBridgeで検知しデプロイ用パイプラインを実行
推論エンドポイントが作成される

当初、上記フローの4,5はLambdaを利用することを考えていました。公式ドキュメントに目を通すとEventBridgeからパイプラインを実行できる説明があったため採用しました。

以下のようなルールを設定し、ターゲットに対象のパイプラインを指定しています。ModelPackageGroupNameで絞ることで他のモデルがApproveされた時に実行されないようにしています。

{
  "source": ["aws.sagemaker"],
  "detail-type": ["SageMaker Model Package State Change"],
  "detail": {
    "ModelApprovalStatus": ["Approved"],
    "ModelPackageGroupName": ["xxxxx"]
  }
}

SageMakerのイベントについては公式のドキュメントを参照しルールのテストを行いました

SageMakerとの通信の部分にVPCエンドポイントを利用することも可能でしたが利用していません。理由は推論エンドポイントを本番環境以外の開発環境用にも用意しているのですが、SageMakerの設定をVPC Onlyで開発環境にもセットアップする必要があるため運用コストが高くなるという結論に至りました。

仮にSageMakerのセットアップを開発環境に行った場合、推論エンドポイントとの通信をVPCエンドポイント経由にすることで通信速度の向上が見込めると思います。

Amazon SageMaker StudioとAWS Step FunctionsのどちらがMLOpsを始めやすいのか

MLOpsの第一歩としてSageMaker Studioを採用していますが、他にStep Functionsを使ったMLOpsもあります。

SageMaker Studioは以下のような特徴があります。

IDEを提供してくれるのでモデルの確認やパイプラインの編集・実行をStudio上で完結できる
DomainにUserProfileを用意することで複数人それぞれの環境の用意が可能

アーキテクチャについては公式のブログがわかりやすいかと思います。

Step FunctionsはAWSのサービスなどをワークフロー（Step Functionsではステートマシンと呼ぶ）として設定できるサービスです。これを利用してSageMakerのProcessing Job とTraining Job を呼び出すことでMLOpsを組むことが出来ます。

Step Functionsを利用したMLOpsを運用するにあたって以下の点が気になりました。

Processing Job, Training Jobも利用する際に実行スクリプトを埋め込んだコンテナイメージが必要になるため、Dockerfile, スクリプト, Build環境, ECRが必要になる
Step Functionsのステートマシン（どういうフローにするかまとめたもの）の管理が必要

当プロジェクトは期間の制約は強くはありませんでした。しかし、今回は筆者含め一緒にプロジェクトを進めていたサーバーサイドエンジニアもMLOpsが初めてだったため運用イメージを掴む方を優先し構築の手間がかからない方を選択した方がよいと思いました。

SageMaker Studioを触れる環境を筆者が用意し、サーバーサイドエンジニアに触ってもらいました。並行してStep Functionsの調査を行い、上記の1に上げたような作業が必要になると判断しランニングコストがかかる印象を受けました。

SageMaker Studioの「IDEを提供してくれるのでモデルの確認やパイプラインの編集・実行をStudio上で完結できる」がMLOps初心者には理解しやすく構築のランニングも低かったため採用することにしました。

ちなみにAWSのSAの方にどちらがオススメか確認したところ、SageMakerだけで完結させたいか、他サービスと連結させたいのかが判断基準になるという回答でした。

他社の事例を見るとStep Functionsを利用されているケースが多い印象でした。今後はStep Functionsを利用して細かな処理を行う必要も出てくるかもしれません。

IaCとパイプラインのコードの棲み分け

B/43の9割以上のインフラの構成がTerraformで定義されています。

SageMakerに関するコードはSREが管理し、パイプラインのコードは別のリポジトリでサーバーサイドエンジニアが管理しています。

以下のリソースをTerraformで管理しています。

SageMaker Domain
UserProfile
IAMロール（SageMakerが利用）

DomainとUserProfileは以下のようなコードで定義しています。今後UserProfileが増える可能性があることを考慮してVariableにプロジェクトをKeyとしてユーザー名を定義しています。

resource "aws_sagemaker_domain" "project-sb" {
  domain_name = "project-sb"
  auth_mode   = "IAM"
  vpc_id      = aws_vpc.vpc.id
  subnet_ids = [
    aws_subnet.subnet-a.id,
    aws_subnet.subnet-d.id
  ]

  default_user_settings {
    execution_role = aws_iam_role.project-sb-role.arn
  }
}

resource "aws_sagemaker_user_profile" "project-sb" {
  for_each          = { for v in var.user_profile.project_sb : v => v }
  domain_id         = aws_sagemaker_domain.project-sb.id
  user_profile_name = each.value
  user_settings {
    execution_role = aws_iam_role.aws_iam_role.project-sb-role.arn
  }
}