sagemaker-studio-image-buildでSageMaker用のカスタムコンテナイメージをビルド&プッシュして利用する

データアナリティクス事業本部の貞松です。

Amazon SageMakerではデータの前処理やモデル学習・推論のジョブ実行時に任意のライブラリや独自アルゴリズム(カスタムコード)を含むカスタムDockerコンテナイメージを使用することができます。

以前(概ねSageMaker Studio登場以前)は、カスタムコンテナイメージをSageMaker用にビルド→ECRにプッシュする為に、自前でビルド&プッシュ処理用のシェルスクリプトを実装する必要がありました。

現在はsagemaker-studio-image-buildというCLIツールを使用することで、Dockerfileを含むコンテナイメージに含めるリソースを準備してツールを実行するだけでSageMaker用にコンテナイメージをビルド&プッシュすることができます。

本記事では、sagameker-studio-image-buildの概要と使用方法を解説し、実際にビルド&プッシュしたカスタムコンテナイメージをSageMakerで利用できることを確認します。

sagemaker-studio-image-buildの概要

sagemaker-studio-image-buildは、PyPIに登録されたPythonパッケージであり、pipコマンドでインストールが可能なCLIツールです。 その名の通りSageMaker Studio上で実行することを想定されており、ターミナルやNotebook上でコマンドライン実行することによりSageMaker用のカスタムコンテナイメージをビルドしてリポジトリにプッシュします。

コンテナイメージのビルドはAWS CodeBuild上で実行され、ビルド済みコンテナイメージをAmazon ECR上のリポジトリにプッシュします。

sagemaker-studio-image-buildの実行に必要なIAMロールの設定

前述の通り、sagemaker-studio-image-buildはCodebuildでコンテナイメージのビルドを実行し、ビルド済みコンテナイメージをECRのリポジトリにプッシュするので、それぞれのサービスにアクセスする為のIAMロール設定が必要です。

CodeBuildの信頼ポリシーを追加

実行ロールに対して、以下の信頼ポリシー設定を追加します。

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Principal": {
        "Service": [
          "codebuild.amazonaws.com"
        ]
      },
      "Action": "sts:AssumeRole"
    }
  ]
}

コンテナイメージをビルド、プッシュするための権限設定を追加

実行ロールに対して、以下の権限設定を追加します。

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": [
        "codebuild:DeleteProject",
        "codebuild:CreateProject",
        "codebuild:BatchGetBuilds",
        "codebuild:StartBuild"
      ],
      "Resource": "arn:aws:codebuild:*:*:project/sagemaker-studio*"
    },
    {
      "Effect": "Allow",
      "Action": "logs:CreateLogStream",
      "Resource": "arn:aws:logs:*:*:log-group:/aws/codebuild/sagemaker-studio*"
    },
    {
      "Effect": "Allow",
      "Action": [
        "logs:GetLogEvents",
        "logs:PutLogEvents"
      ],
      "Resource": "arn:aws:logs:*:*:log-group:/aws/codebuild/sagemaker-studio*:log-stream:*"
    },
    {
      "Effect": "Allow",
      "Action": "logs:CreateLogGroup",
      "Resource": "*"
    },
    {
      "Effect": "Allow",
      "Action": [
        "ecr:CreateRepository",
        "ecr:BatchGetImage",
        "ecr:CompleteLayerUpload",
        "ecr:DescribeImages",
        "ecr:DescribeRepositories",
        "ecr:UploadLayerPart",
        "ecr:ListImages",
        "ecr:InitiateLayerUpload",
        "ecr:BatchCheckLayerAvailability",
        "ecr:PutImage"
      ],
      "Resource": "arn:aws:ecr:*:*:repository/sagemaker-studio*"
    },
    {
      "Effect": "Allow",
      "Action": "ecr:GetAuthorizationToken",
      "Resource": "*"
    },
    {
      "Effect": "Allow",
      "Action": [
        "s3:GetObject",
        "s3:DeleteObject",
        "s3:PutObject"
      ],
      "Resource": "arn:aws:s3:::sagemaker-*/*"
    },
    {
      "Effect": "Allow",
      "Action": [
        "s3:CreateBucket"
      ],
      "Resource": "arn:aws:s3:::sagemaker*"
    },
    {
      "Effect": "Allow",
      "Action": [
        "iam:GetRole",
        "iam:ListRoles"
      ],
      "Resource": "*"
    },
    {
      "Effect": "Allow",
      "Action": "iam:PassRole",
      "Resource": "arn:aws:iam::*:role/*",
      "Condition": {
        "StringLikeIfExists": {
          "iam:PassedToService": "codebuild.amazonaws.com"
        }
      }
    }
  ]
}

sagemaker-studio-image-buildのインストール

SageMaker Studio上のターミナルあるいはNotebookで以下のコマンドを実行してインストールします。

!pip install sagemaker-studio-image-build

カスタムコンテナイメージ用のリソース準備

Dockerfileの作成

なにはなくともDockerfileを作成します。
SageMaker用のコンテナでは、以下のDockerコマンドが実行されます。

• 学習ジョブ実行時 : docker run [image_name] train
• 推論ジョブ実行時 : docker run [image_name] serve

その為、Dockerfile内で、trainとserveというファイルに対するパスと実行権限を設定する必要があります。

以下、Dockerfileの設定例です。

FROM ubuntu:16.04

RUN apt-get update
RUN apt-get install -y software-properties-common
RUN add-apt-repository ppa:deadsnakes/ppa
RUN apt-get update

RUN apt-get install -y build-essential python3.6 python3.6-dev python3-pip python3.6-venv
RUN python3.6 -m pip install pip --upgrade

RUN apt-get -y update && apt-get install -y --no-install-recommends \
         wget \
         nginx \
         libgcc-5-dev \ 
         ca-certificates \
    && rm -rf /var/lib/apt/lists/*

RUN pip install --upgrade pip
RUN pip install numpy scikit-learn pandas flask gevent gunicorn requests

ENV PYTHONUNBUFFERED=TRUE
ENV PYTHONDONTWRITEBYTECODE=TRUE

ENV PATH="/opt/ml/code:${PATH}"

COPY custom_code /opt/ml/code
WORKDIR /opt/ml/code
RUN chmod +x /opt/ml/code/train
RUN chmod +x /opt/ml/code/serve

モデル学習・推論用の最小ファイル構成

SageMaker用のコンテナ仕様を踏まえた上で、学習・推論ジョブを実行するための最小ファイル構成は以下のようになります。

container_root
├── Dockerfile
└── custom_code
    ├── nginx.conf
    ├── predictor.py
    ├── serve
    ├── train
    └── wsgi.py

参考記事

カスタムコンテナイメージの作成についての詳細は以下の記事を参考にしてください。

sagemaker-studio-image-buildの実行

SageMaker Studio上のターミナルあるいはNotebookで以下のコマンドを実行します。

sm-docker build [container-image-resource-path] [option]

[container-image-resource-path]にDockerfileが配置されたディレクトリのパスを指定します。

sagemaker-studio-image-buildのコマンドラインオプションは以下の通りです。

• --repository [repository]:[tag]
  • デフォルトでは、[repository]はsagemaker-studio-[studio-domainid]という名前を使用します。
  • このパラメータを設定して、既存のリポジトリにプッシュしたり、任意の名前で新しいリポジトリを作成したりできます。
  • デフォルトの[tag]は[user-profile-name]という名前を使用します。
  • このパラメータは、タグ付け戦略をカスタマイズするためにも使用できます。
• --role [iam-role-name]
  • デフォルトでは、CLIはSageMaker実行ロールを使用してCLIが使用するAWSサービス(CodeBuild、ECR)とやり取りします。
  • オプションで、必要な権限を持つ任意のIAMロールを指定することもできます。
• --bucket [bucket-name]
  • デフォルトでは、CLIはCodeBuildに送信されたパッケージ化された入力を保存するためにSageMakerのデフォルトのセッションバケットを使用します。
  • オプションで、使用する任意のS3バケットを指定できます。
• --no-logs
  • デフォルトでは、CLIは実行中のCodeBuildの出力ログを表示します。
  • このオプションを指定することで、ログ出力をしないように設定することができます。

SageMakerのジョブでカスタムコンテナイメージを利用

学習・推論での利用

以下、SageMakerでカスタムコンテナイメージを使用して学習・推論処理を実行するコードの例です。 Estimatorのimage_uriにカスタムコンテナイメージのURIを指定します。

import boto3
import sagemaker

sess = sagemaker.Session()
client = boto3.client('sagemaker-runtime')

instance_count = 1
instance_type = 'ml.m4.xlarge'
train_input_s3_path = [入力元S3バケットパス]
train_output_s3_path = [出力先S3バケットパス]
role = [実行IAMロール]

# Estimatorのimage_uriにカスタムコンテナイメージのURIを指定
sm = sagemaker.estimator.Estimator(
    image_uri = [ECRにプッシュしたカスタムコンテナイメージのURIを指定],
    role = role,
    instance_count = instance_count,
    instance_type =　instance_type,
    max_run = 86400,
    output_path = train_output_s3_path,
    sagemaker_session = sess
)

# 学習実行
sm.fit(
    inputs = {'training': train_input_s3_path},
    job_name = 'training_job_name'
)

# エンドポイントデプロイ
predictor = sm.deploy(
    endpoint_name = 'endpoint_name',
    initial_instance_count = instance_count,
    instance_type = instance_type
)

# デプロイしたエンドポイントで推論実行
response = client.invoke_endpoint(
    EndpointName = endpoint_name, 
    ContentType = 'application/json',
    Accept = 'application/json',
    Body = json.dumps(request_parameter)
)

まとめ

SageMaker用のカスタムコンテナイメージをCodeBuildでビルド→ECRにプッシュする為のCLIツールであるsagemaker-studio-image-buildについて、概要と使用方法について解説しました。
自前でコンテナイメージのビルド&プッシュ用シェルスクリプトを作成する手間を省くことができ、かつSageMaker Studio上で操作を完結できる為、さらに開発環境をSageMaker Studioに集約させる動きが加速していると言えそうです。