チュートリアル: SageMaker で W&B Launch を設定する

このチュートリアルでは、ML エンジニアとプラットフォーム管理者を対象に、Amazon SageMaker にトレーニングジョブを送信するように W&B Launch を設定する方法を説明します。チュートリアルを完了すると、W&B から SageMaker Training ジョブを実行するのに必要な AWS リソース、IAM ロール、キュー設定、Launch エージェントがすべて揃います。 W&B Launch を使用すると、Amazon SageMaker に Launch job を送信し、SageMaker プラットフォーム上で組み込みのアルゴリズムまたはカスタムアルゴリズムを使って機械学習モデルをトレーニングできます。コンピュートリソースのプロビジョニングと解放は SageMaker が処理するため、EKS クラスターを持たないチームに適した選択肢です。 Amazon SageMaker に接続された W&B Launch キューは、CreateTrainingJob API を使用して Launch job を SageMaker Training ジョブとして実行します。CreateTrainingJob API に送信する引数は、Launch キュー設定で制御します。 Amazon SageMaker は Docker イメージを使用してトレーニングジョブを実行します。SageMaker がプルするイメージは Amazon Elastic Container Registry (ECR) に保存する必要があるため、トレーニングに使用するイメージも ECR に保存してください。

このガイドでは、SageMaker Training ジョブの実行方法を説明します。Amazon SageMaker で推論用にモデルをデプロイする方法については、こちらの Launch job の例を参照してください。

前提条件

作業を始める前に、以下の前提条件を満たしておく必要があります。

Launch エージェントに Docker イメージをビルドさせるかどうかを決めます。
AWS リソースをセットアップし、S3、ECR、SageMaker の IAM ロールに関する情報を収集します。
Launch エージェント用の IAM ロールを作成します。

以下のセクションでは、各前提条件を満たす方法を説明します。

Launch エージェントに Docker イメージをビルドさせるかどうかを決める

W&B Launch エージェントに Docker イメージをビルドさせるかどうかを決めます。次の 2 つのオプションから選択できます。

Launch エージェントに、Docker イメージのビルド、Amazon ECR へのイメージのプッシュ、SageMaker Training ジョブの送信を任せます。トレーニングコードを短いサイクルで反復改善する ML エンジニアにとっては、このオプションで作業を簡素化できます。
トレーニングスクリプトまたは推論スクリプトを含む既存の Docker イメージを使用します。このオプションは既存の CI システムと組み合わせやすい方法です。このオプションを選択する場合は、Amazon ECR 上のコンテナーレジストリに Docker イメージを手動でアップロードする必要があります。

AWS リソースをセットアップする

使用する AWS リージョンで、次の AWS リソースを設定しておく必要があります。

コンテナーイメージを保存するための ECR リポジトリ
SageMaker Training ジョブの入力と出力を保存するための S3 バケット (1 つ以上)
SageMaker によるトレーニングジョブの実行と、Amazon ECR および Amazon S3 へのアクセスを許可する Amazon SageMaker 用の IAM ロール

これらのリソースの ARN を控えておいてください。ARN は Launch キュー設定を定義する際に必要です。

Launch エージェント用の IAM ポリシーを作成する

Launch エージェントには、SageMaker のトレーニングジョブを送信するための権限と、必要に応じて ECR にイメージをプッシュするための権限を付与する IAM ポリシーが必要です。次の手順に従ってポリシーを作成します。

AWS の IAM 画面で、新しいポリシーを作成します。
JSON ポリシーエディターに切り替え、ユースケースに応じて次のいずれかのポリシーを貼り付けます。[BRACKETS] 内のプレースホルダーは、実際の値に置き換えてください。

エージェントがビルド済みの Docker イメージを送信する場合
エージェントが Docker イメージをビルドして送信する場合

  {
    "Version": "2012-10-17",
    "Statement": [
      {
        "Effect": "Allow",
        "Action": [
          "logs:DescribeLogStreams",
          "SageMaker:AddTags",
          "SageMaker:CreateTrainingJob",
          "SageMaker:DescribeTrainingJob"
        ],
        "Resource": "arn:aws:sagemaker:[REGION]:[ACCOUNT-ID]:*"
      },
      {
        "Effect": "Allow",
        "Action": "iam:PassRole",
        "Resource": "arn:aws:iam::[ACCOUNT-ID]:role/[ROLE-ARN-FROM-QUEUE-CONFIG]"
      },
    {
        "Effect": "Allow",
        "Action": "kms:CreateGrant",
        "Resource": "[ARN-OF-KMS-KEY]",
        "Condition": {
          "StringEquals": {
            "kms:ViaService": "SageMaker.[REGION].amazonaws.com",
            "kms:GrantIsForAWSResource": "true"
          }
        }
      }
    ]
  }

  {
    "Version": "2012-10-17",
    "Statement": [
      {
        "Effect": "Allow",
        "Action": [
          "logs:DescribeLogStreams",
          "SageMaker:AddTags",
          "SageMaker:CreateTrainingJob",
          "SageMaker:DescribeTrainingJob"
        ],
        "Resource": "arn:aws:sagemaker:[REGION]:[ACCOUNT-ID]:*"
      },
      {
        "Effect": "Allow",
        "Action": "iam:PassRole",
        "Resource": "arn:aws:iam::[ACCOUNT-ID]:role/[ROLE-ARN-FROM-QUEUE-CONFIG]"
      },
       {
      "Effect": "Allow",
      "Action": [
        "ecr:CreateRepository",
        "ecr:UploadLayerPart",
        "ecr:PutImage",
        "ecr:CompleteLayerUpload",
        "ecr:InitiateLayerUpload",
        "ecr:DescribeRepositories",
        "ecr:DescribeImages",
        "ecr:BatchCheckLayerAvailability",
        "ecr:BatchDeleteImage"
      ],
      "Resource": "arn:aws:ecr:[REGION]:[ACCOUNT-ID]:repository/[REPOSITORY]"
    },
    {
      "Effect": "Allow",
      "Action": "ecr:GetAuthorizationToken",
      "Resource": "*"
    },
    {
        "Effect": "Allow",
        "Action": "kms:CreateGrant",
        "Resource": "[ARN-OF-KMS-KEY]",
        "Condition": {
          "StringEquals": {
            "kms:ViaService": "SageMaker.[REGION].amazonaws.com",
            "kms:GrantIsForAWSResource": "true"
          }
        }
      }
    ]
  }

Next をクリックします。
ポリシーの名前と説明を入力します。
Create policy をクリックします。

これで IAM ポリシーが作成されました。このポリシーは、次のセクションで Launch エージェントのロールにアタッチします。

Launch エージェント用の IAM ロールを作成する

Launch エージェントには、Amazon SageMaker のトレーニングジョブを作成する権限が必要です。前のセクションで作成したポリシーを専用のロールにアタッチしておくと、エージェントは実行時にそのロールを引き受けて、これらの権限を利用できます。次の手順に従って IAM ロールを作成します。

AWS の IAM 画面で、新しいロールを作成します。
Trusted Entity で AWS Account (または組織のポリシーに合った別のオプション) を選択します。
権限の画面をスクロールし、前のセクションで作成したポリシー名を選択します。
ロールの名前と説明を入力します。
Create role を選択します。
ロールの ARN を控えておきます。この ARN は Launch エージェントを設定する際に指定します。

IAM ロールの作成方法については、AWS Identity and Access Management のドキュメントを参照してください。

Launch エージェントでイメージをビルドする場合に必要な追加の権限については、エージェントの高度な設定を参照してください。
SageMaker キューに対する kms:CreateGrant 権限が必要になるのは、関連付けられた ResourceConfig で VolumeKmsKeyId が指定されており、かつ関連付けられたロールにこのアクションを許可するポリシーがない場合のみです。

SageMaker 用の Launch キューを設定する

AWS の前提条件が整ったら、ジョブを SageMaker にルーティングする W&B Launch キューを作成できます。W&B アプリで、SageMaker をコンピュートリソースとして使用するキューを作成します。

Launch App にアクセスします。
Create Queue をクリックします。
キューを作成する Entity を選択します。
Name フィールドにキューの名前を入力します。
Resource で SageMaker を選択します。

Configuration フィールドに、SageMaker ジョブに関する情報を入力します。デフォルトでは、W&B によって YAML および JSON 形式の CreateTrainingJob リクエストボディが自動入力されます。

{
  "RoleArn": "[REQUIRED]", 
  "ResourceConfig": {
      "InstanceType": "ml.m4.xlarge",
      "InstanceCount": 1,
      "VolumeSizeInGB": 2
  },
  "OutputDataConfig": {
      "S3OutputPath": "[REQUIRED]"
  },
  "StoppingCondition": {
      "MaxRuntimeInSeconds": 3600
  }
}

少なくとも以下を指定する必要があります。

RoleArn: SageMaker 実行 IAM ロールの ARN (前提条件を参照) 。Launch エージェントの IAM ロールと混同しないでください。
OutputDataConfig.S3OutputPath: SageMaker が出力を保存する場所を指定する Amazon S3 URI。
ResourceConfig: リソース設定 (必須) 。リソース設定のオプションについては、AWS ResourceConfig ドキュメントを参照してください。
StoppingCondition: トレーニングジョブの停止条件 (必須) 。オプションについては、AWS StoppingCondition ドキュメントを参照してください。

Create Queue をクリックします。

これでキューが作成されます。このキューをポーリングするように Launch エージェントを設定すると、ジョブを受け付けられるようになります。

Launch エージェントを設定する

以下のセクションでは、エージェントのデプロイ先と、デプロイ先に応じたエージェントの設定方法について説明します。 Amazon SageMaker キュー用の Launch エージェントは、ローカルマシン、EC2 インスタンス、EKS クラスターなど、複数の場所にデプロイできます。エージェントのデプロイ先に応じて、Launch エージェントを設定してください。

Launch エージェントの実行場所を決める

本番ワークロードの場合や、すでに EKS クラスターをお持ちのお客様の場合、W&B では、この Helm チャートを使用して Launch エージェントを EKS クラスターにデプロイすることをお勧めします。本番ワークロードで EKS クラスターをお持ちでない場合は、EC2 インスタンスが適しています。Launch エージェントのインスタンスは常時稼働しますが、エージェントには t2.micro サイズ程度の EC2 インスタンスがあれば十分なため、低コストで運用できます。実験的な用途や個人での利用の場合は、ローカルマシンで Launch エージェントを実行すると手早く始められます。ユースケースに応じて、以下のタブの手順に従って Launch エージェントを設定してください。

EKS
EC2
ローカルマシン

W&B では、W&B が管理する Helm チャートを使用して EKS クラスターにエージェントをインストールすることをお勧めします。

Amazon EC2 ダッシュボードにアクセスし、次の手順を実行します。

Launch instance をクリックします。
Name フィールドに名前を入力します。必要に応じてタグを追加します。
Instance type で、EC2 コンテナーのインスタンスタイプを選択します。1 vCPU と 1 GiB のメモリ (例: t2.micro) があれば十分です。
Key pair (login) フィールドで、組織用のキーペアを作成します。このキーペアは、後のステップで SSH クライアントを使用して EC2 インスタンスに接続する際に使用します。
Network settings で、組織のセキュリティグループを選択します。
Advanced details を展開します。IAM instance profile で、前のセクションで作成した Launch エージェントの IAM ロールを選択します。
Summary フィールドを確認します。問題がなければ、Launch instance を選択します。

AWS の EC2 ダッシュボードの左パネルにある Instances にアクセスします。作成した EC2 インスタンスが実行中であることを確認します (Instance state 列を参照) 。EC2 インスタンスが実行中であることを確認したら、ローカルマシンのターミナルで次の手順を実行します。

Connect を選択します。
SSH client タブを選択し、手順に従って EC2 インスタンスに接続します。

EC2 インスタンス内で、次のパッケージをインストールします。

sudo yum install python311 -y && python3 -m ensurepip --upgrade && pip3 install wandb && pip3 install wandb[launch]

次に、EC2 インスタンス内で Docker をインストールして起動します。

sudo yum update -y && \
sudo yum install -y docker python3 && \
sudo systemctl start docker && \
sudo systemctl enable docker && \
sudo usermod -a -G docker ec2-user

newgrp docker

これで、Launch エージェント設定を行えるようになりました。

~/.aws/config および ~/.aws/credentials にある AWS 設定ファイルを使用して、ローカルマシン上でポーリングするエージェントにロールを関連付けます。前のステップで Launch エージェント用に作成した IAM ロールの ARN を指定します。

~/.aws/config

[profile SageMaker-agent]
role_arn = arn:aws:iam::[ACCOUNT-ID]:role/[AGENT-ROLE-NAME]
source_profile = default                                                                   

~/.aws/credentials

[default]
aws_access_key_id=[ACCESS-KEY-ID]
aws_secret_access_key=[SECRET-ACCESS-KEY]
aws_session_token=[SESSION-TOKEN]

セッショントークンの最大有効期間は、関連付けられたプリンシパルに応じて 1 時間または 3 日間です。

Launch エージェントを設定する

エージェントを実行する場所を決めたら、SageMaker キューをポーリングし、AWS で認証できるようにエージェントを設定します。Launch エージェントの設定には、launch-config.yaml という名前の YAML 設定ファイルを使用します。デフォルトでは、W&B は ~/.config/wandb/launch-config.yaml にある設定ファイルを参照します。Launch エージェントを有効化する際に -c フラグを使用すると、別のディレクトリを指定することもできます。次の YAML スニペットは、エージェントの主要な設定オプションを指定する方法を示しています。

launch-config.yaml

max_jobs: -1
queues:
  - [QUEUE-NAME]
environment:
  type: aws
  region: [YOUR-REGION]
registry:
  type: ecr
  uri: [ECR-REPO-ARN]
builder: 
  type: docker

次に、wandb launch-agent でエージェントを起動します。これで Launch エージェントが実行され、SageMaker キューにジョブがないかポーリングするようになります。

オプション: Launch job の Docker イメージを Amazon ECR にプッシュする

このセクションは、トレーニングまたは推論のロジックを含む既存の Docker イメージを Launch エージェントが使用する場合にのみ該当します。Launch エージェントでは、2 つの動作オプションがサポートされています。

Launch job を含む Docker イメージを Amazon ECR リポジトリにアップロードします。イメージベースのジョブを使用する場合は、新しい Launch job を送信する前に、Docker イメージを ECR レジストリに格納しておく必要があります。

​前提条件

​Launch エージェントに Docker イメージをビルドさせるかどうかを決める

​AWS リソースをセットアップする

​Launch エージェント用の IAM ポリシーを作成する

​Launch エージェント用の IAM ロールを作成する

​SageMaker 用の Launch キューを設定する

​Launch エージェントを設定する

​Launch エージェントの実行場所を決める

​Launch エージェントを設定する

​オプション: Launch job の Docker イメージを Amazon ECR にプッシュする

前提条件

Launch エージェントに Docker イメージをビルドさせるかどうかを決める

AWS リソースをセットアップする

Launch エージェント用の IAM ポリシーを作成する

Launch エージェント用の IAM ロールを作成する

SageMaker 用の Launch キューを設定する

Launch エージェントを設定する

Launch エージェントの実行場所を決める

Launch エージェントを設定する

オプション: Launch job の Docker イメージを Amazon ECR にプッシュする