カスタマイズしたSageMaker Distributionコンテナイメージの作成パイプラインをGitHub Actionsで実現する
データ事業本部の鈴木です。
SageMaker StudioのJupyterLab Spaceはカスタムイメージで起動することが可能です。
今回はGitHubでSageMaker Distributionのイメージのカスタマイズ用Dockerfileを管理しているようなケースで、GitHub Actionsでビルドし、ECRリポジトリにPUSHし、SageMaker StudioのJupyterLab Spaceで利用するところまでを確認します。
今回の仕組み
以下の仕組みを検証しました。
検証時には以下の観点がポイントとなりました。
- GitHub ActionsでカスタマイズしたSageMaker Distributionのイメージをビルドする
- GitHub ActionsからビルドしたイメージをECRにPUSHする
- イメージをECRでスキャンする
- イメージをSageMaker StudioのJupyterLab Spaceで利用する
イメージ作成までの検証
1.ECRリポジトリ作成
ECRのコンソールより、custom-sagemaker-distribution-cicdリポジトリを作成しました。
2.AWS側の権限設定
以下を行いました。
- IAMでのOIDCプロバイダ追加
- IAMロール作成（ECRリポジトリへのPUSH権限、アシュームロールが可能な信頼ポリシーの設定済み）
OIDCプロバイダは、以下のブログを参考に追加しました。
IAMロールは以下のブログを参考に作成しましたが、後で紹介するBuildxを使った方法の場合、ECR向けの権限が追加で必要になります。
具体的には、ecr:BatchGetImageを追加して以下のようにしました。
{
"Version": "2012-10-17",
"Statement": [
{
"Sid": "GetAuthorizationToken",
"Effect": "Allow",
"Action": [
"ecr:GetAuthorizationToken"
],
"Resource": "*"
},
{
"Sid": "PushImageOnly",
"Effect": "Allow",
"Action": [
"ecr:BatchCheckLayerAvailability",
"ecr:InitiateLayerUpload",
"ecr:UploadLayerPart",
"ecr:CompleteLayerUpload",
"ecr:PutImage",
"ecr:BatchGetImage"
],
"Resource": "arn:aws:ecr:ap-northeast-1:<AWSアカウントID>:repository/<リポジトリ名>"
}
]
}
3.GitHubリポジトリの作成・設定
今回は新規にプライベートリポジトリを作成しました。
後ほどワークフロー定義から参照するシークレットおよび環境をリポジトリのSettingsよりActions secrets and variablesを開いて入力しておきました。
- Secrets
- AWS_ACCOUNT_ID：ECRリポジトリを作成するAWSアカウントのID
- AWS_IAM_ROLE_ARN：先に作成したアシュームロール用のIAMロールのARN
- Variables
- AWS_REGION：ap-northeast-1
- ECR_REGISTRY：custom-sagemaker-distribution-cicd
4.必要ファイルの作成
リポジトリは以下の構成としました。
.
├── .github
│ └── workflows
│ └── sample-workflow.yml
├── Dockerfile
├── README.md
└── requirements.txt
Dockerfile
SageMaker Distributionのイメージをベースに使用する方法を記載しました。以下のブログで紹介されています。
FROM public.ecr.aws/sagemaker/sagemaker-distribution:latest-cpu
ARG NB_USER="sagemaker-user"
ARG NB_UID=1000
ARG NB_GID=100
ENV MAMBA_USER=$NB_USER
USER root
RUN apt-get update
RUN micromamba install sagemaker-inference --freeze-installed --yes --channel conda-forge --name base
COPY ./requirements.txt requirements.txt
RUN pip install -U pip && pip install --no-cache-dir -r requirements.txt
USER $MAMBA_USER
ENTRYPOINT ["jupyter-lab"]
CMD ["--ServerApp.ip=0.0.0.0", "--ServerApp.port=8888", "--ServerApp.allow_origin=*", "--ServerApp.token=''", "--ServerApp.base_url=/jupyterlab/default"]
ワークフロー定義
以下のブログで紹介されていた、Build and push Docker imagesを使用する方法を使いました。
このActionはBuildKitによる拡張ビルド機能のための Docker CLI プラグインであるBuildxを使用してコンテナイメージをビルドします。
検証のため手動でトリガーしたかったので、workflow_dispatchイベントにしてあります。
name: Build And Push
on:
workflow_dispatch:
env:
AWS_REGION: ap-northeast-1
ECR_REGISTRY: ${{ secrets.AWS_ACCOUNT_ID }}.dkr.ecr.ap-northeast-1.amazonaws.com
jobs:
build:
name: Build And Push
runs-on: ubuntu-latest
permissions:
id-token: write
contents: read
steps:
- name: Set up QEMU
uses: docker/setup-qemu-action@v3
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v3
- uses: aws-actions/configure-aws-credentials@v4
with:
role-to-assume: ${{ secrets.AWS_IAM_ROLE_ARN }}
aws-region: ${{ env.AWS_REGION }}
- name: Login to Docker Hub
uses: docker/login-action@v3
with:
registry: ${{ env.ECR_REGISTRY }}
- name: Build and push API
uses: docker/build-push-action@v5
with:
push: true
provenance: false
tags: ${{ env.ECR_REGISTRY }}/custom-sagemaker-distribution-cicd:latest
タグはECRに作成したリポジトリ名と同じにしておくことにご注意ください。（今回だとcustom-sagemaker-distribution-cicd）
追加するPythonライブラリ
requirements.txtを用意しました。
ライブラリはなんでもよいのですが、私がよく使用するSnowpark ML向けのライブラリを指定しました。
snowflake-ml-python
なお、このときJupyterスペースでデフォルトとなっていたSageMaker Distribution 2.4.1のイメージでは、Snowpark MLはインポートできませんでした。
5.動作確認
ワークフローよりビルドとPUSHの実行
GitHubよりワークフローを実行しました。
ワークフローが成功すると、以下のようにECRリポジトリにイメージがPUSHされました。
SageMaker Studioでの利用
1.ドメインへのイメージのアタッチ
作成したイメージをSageMakerドメインにアタッチします。これにより、スペースでイメージを利用できるようにします。
まず、アタッチしたいドメインを開き、環境タブからイメージをアタッチを押します。
新しいイメージを選択し、ECRイメージURIを入力に作成したリポジトリのURLを入力しました。
イメージのプロパティを入力しました。JupyterLabスペースで使いたいのでイメージタイプはJupyterLabイメージにしました。
以下のようにカスタムイメージがアタッチされれば完了です。
なお、アタッチしたイメージにはバージョンがあります。アタッチ後再度ワークフローを実行して新しいバージョンのイメージをPUSHしましたが、自動ではバージョンは更新されず、以前作成したときのダイジェストのものを指していました。イメージを更新した場合は、特段作り込みをしていないのであれば手動で操作することになりそうです。
2.JupyterLabスペースの起動
ユーザープロファイルを指定してStudioにログインすると、スペースでドメインにアタッチしたイメージを利用できるようになっています。
このイメージを指定してJupyterLabを起動し、snowflakeをimportするとたしかに読み込めることが確認できました。
補足
今回紹介した仕組みはStudioで使用するイメージ用のパイプラインの例でした。
イメージ向けのパイプラインの重要な要素として、イメージの脆弱性検査があります。
ECRでは基本と拡張の2種類のスキャンが提供されています。
継続的・定期的な検査が望ましいと思いますが、まずはPUSH時の検査がECRでは簡単な設定でできるためご紹介します。
プライベートリポジトリのFeatures & Settingsを押すと、スキャンに関する設定ができます。
基本スキャンではPUSH時にフィルタを設定することで指定したリポジトリの脆弱性スキャンが可能です。
スキャンされると対象のリポジトリで以下のように結果が分かります。
なお、有料ですが拡張スキャンで継続的にスキャンを行うこともできます。
スキャン結果の通知は別途実装が必要です。EventBridgeで拾うことができます。EventBridgeは多くのAPIと統合されており、後続の通知処理をキックできます。
以下はChatbotの例です。
最後に
GitHubリポジトリに格納したDockerfileより、SageMaker Studioで利用するカスタマイズされたSageMaker Distributionイメージをビルドし、ECRリポジトリにPUSHする例をご紹介しました。
