CircleCIからAWS SageMakerを利用する(1)

昨日に引き続き、CircleCI Advent Calendar 2023の12/2を担当します CircleCI Developer Advocate の @mfunaki です。

AWS re:Invent 2023 Las Vegas に関連して

After Corona では久々の「超」大規模のイベント、AWS re:Invent 2023が、11/27~12/1まで開催されていました。日本からの参加者もたくさんいらっしゃった様で、弊社 CircleCI ブースにもお立ち寄りいただいた方がいらっしゃるという話は聞いております。ありがとうございます。

AWS re:Invent 2023 Las Vegas では、SageMakerに関しても数多くの発表がありました。また、AI学習用のアクセラレータ Trainium2 を搭載するインスタンスも2024年には出てきそうです。CircleCI でもこれまで以上にパワフルなGPUインスタンスを提供できるのではないかと思います(GPUインスタンスを自分でEC2上に立ち上げなくても、CircleCIでジョブを実行して、そのインスタンスにSSHでログインして使う、ポートフォワードして使うのも、私個人のお勧めです)。

• Your guide to generative AI and ML at AWS re:Invent 2023

CircleCI SageMaker Orb リリース

CircleCI も今回、SageMaker Orb をリリースいたしました。

• setup
• create_model
• create_endpoint_configuration
• deploy_endpoint

といった Orb が提供するコマンドから見ても、一通りのことができそう感はあります。インバウンド Webhook を利用することで、トレーニング用のデータセットが更新されたり、テストデータが追加されたのを自動検出して、(リポジトリに空コミット＆プッシュしなくても)ワークフローが回せそうです。

SageMaker Orbを使って SageMaker カンタン自動処理！

さて、SageMaker Orbの使い方に関しては、CircleCI 公式ブログがすでに公開されています。

• Using Amazon SageMaker orb to orchestrate model deployment across environments

このブログの内容を日本語で追っていきたいと思います。

ワークフローの流れ

今回実行するのは、画像によるアワビの年齢推定です。具体的な処理内容は、AmazonからXGBoost for regressionサンプルが公開されているので、そちらを参考にしていただければと思います。また、classmethodの大澤さんによる「AmazonSageMakerのXGBoostでアワビの年齢を予測してみた」が参考になります。
以下の説明では自動化に絞った形で説明させていただきます。

• 更新されたのが model-train ブランチなら

  • kitten-model-train ジョブを実行する
    • CircleCIのPythonコンテナ上でトレーニングデータを取得(gather_data.py)する
    • SageMaker上でモデルのトレーニングと登録(train_register.py)を行う

• 更新されたのが main ブランチなら

  • create_modelコマンドを実行して、SageMaker上でモデルを生成する
  • create_endpoint_configurationコマンドを実行して、devデプロイ環境を用意する
  • deploy_endpointコマンドを実行して、モデルをdev環境にデプロイする
  • prod環境にデプロイしていいか、CircleCIの approval 機能を使って確認プロセスを挟む(手動承認後に処理を続行するワークフロー)
  • create_endpoint_configurationコマンドを実行して、prodデプロイ環境を用意する
  • deploy_endpointコマンドを実行して、モデルをprod環境にデプロイする

ワークフロー実行に必要な設定

OIDCトークンの設定

AWS SageMakerを利用するにあたり、(AWSシークレットを環境変数に保持しておくのではなく)OIDCトークンを使用します。しくみ等は私が以前書いた下のブログをご覧いただくとして、ここではどんな作業が必要かを説明していきます。

• ブログ「AWS へのアクセスに Open ID Connect(OIDC) を使用する(https://circleci.com/ja/blog/circleci-aws-oidc/)」

まずは、CircleCIのOrganization SettingsからOrganization IDをコピーしておきます。

次に、AWSマネージメントコンソールから、IAM > (アクセス管理の下の)IDプロバイダ > プロバイダを追加 の順に選択し、IDプロバイダを追加します。プロバイダのタイプ(Provider type)は OpenID Connect を選び、プロバイダ用のURL(Provider URL)には、https://oidc.circleci.com/org/<先にコピーしておいた Organization ID を指定します。対象者(Audience)には、先にコピーしておいた Organization ID を指定します。ここまでが済んだら、プロバイダを追加(Add provider)ボタンをクリックします。

ロールの設定

次に、IAM > ロール > ロールを作成 の順に選択し、ロールを追加します。信頼されたエンティティタイプ にウェブアイデンティティを選択すると、アイデンティティプロバイダーを選択できるようになるので、先ほど追加したIDプロバイダを選択します。また Audienceに、先ほど指定した Organization ID (=対象者) を選択しておきます。

次へを押すと許可ポリシーの選択画面になりますが、ここでは特に何も選択しないで、次へを押します。すると、ロール名を設定する画面になるので、(私の場合は) masahiko-sagemaker-deploy-example-policy と指定しました。信頼ポリシーには次のようなポリシーを設定します(英語ブログのままでは動作しなかったので一部追加してあります)。リソース名は適宜、ご自身の環境に置き換えてください。

{
	"Version": "2012-10-17",
	"Statement": [
		{
			"Sid": "OrbPermissions",
			"Effect": "Allow",
			"Action": [
				"sagemaker:AddTags",
				"sagemaker:CreateEndpoint",
				"sagemaker:CreateEndpointConfig",
				"sagemaker:CreateModel",
				"sagemaker:DescribeEndpoint",
				"sagemaker:DescribeEndpointConfig",
				"sagemaker:ListEndpoints",
				"sagemaker:ListModelPackages",
				"sagemaker:ListTags",
				"sagemaker:UpdateEndpoint",
				"iam:PassRole"
			],
			"Resource": "*"
		},
		{
			"Sid": "S3Access",
			"Effect": "Allow",
			"Action": [
				"s3:GetObject",
				"s3:ListBucket"
			],
			"Resource": [
				"arn:aws:s3:::masahiko-sagemaker-deploy-examples/*"
			]
		},
		{
			"Sid": "S3AccessTrainModel",
			"Effect": "Allow",
			"Action": [
				"s3:GetObject",
				"s3:ListBucket",
				"s3:PutObject"
			],
			"Resource": [
				"arn:aws:s3:::sagemaker-sample-files/*",
				"arn:aws:s3:::masahiko-cci-sagemaker-pipeline",
				"arn:aws:s3:::masahiko-cci-sagemaker-pipeline/*"
			]
		},
		{
			"Sid": "SageMakerTrainModel",
			"Effect": "Allow",
			"Action": [
				"sagemaker:CreateTrainingJob",
				"sagemaker:DescribeTrainingJob",
				"logs:DescribeLogStreams",
				"sagemaker:ListModelPackageGroups",
				"sagemaker:CreateModelPackage",
				"sagemaker:CreateModelPackageGroup",
				"sagemaker:UpdateModelPackage"
			],
			"Resource": "*"
		}
	]
}

信頼されたエンティティ(Trust relationship)には、次のようなポリシーを設定します(こちらは英語ブログのとおり)。

{
	"Version": "2012-10-17",
	"Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "Federated": "arn:aws:iam::<AWS-ACCOUNT-ID>:oidc-provider/oidc.circleci.com/org/<CIRCLECI-ORG-ID>"
            },
            "Action": "sts:AssumeRoleWithWebIdentity",
            "Condition": {
                "StringLike": {
                    "oidc.circleci.com/org/<CIRCLECI-ORG-ID>:sub": "org/<CIRCLECI-ORG-ID>/project/<CIRCLECI-PROJECT-ID>/user/*"
                }
            }
        },
        {
			"Effect": "Allow",
			"Principal": {
				"Service": "sagemaker.amazonaws.com"
			},
			"Action": "sts:AssumeRole"
		}

    ]
}

設定が完了すると、ロールのARNが表示されるので、コピーしておきます。

おわりに

ここまでで AWS 側の仕込みは完了です。明日はCircleCI側の設定を行い、実際にワークフローを実行してみましょう！