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

ECR、S3、IAM の設定手順に沿って、トレーニングジョブを Amazon SageMaker に送信するよう W&B Launch を設定します。

W&B Launch を使用すると、SageMaker プラットフォーム上で提供されるアルゴリズムまたはカスタムアルゴリズムを使って機械学習モデルをトレーニングする Launch ジョブ を Amazon SageMaker に送信できます。SageMaker はコンピュートリソースの起動と解放を自動で管理するため、EKS クラスターを持たないチームにとって適した選択肢となる場合があります。

Amazon SageMaker に接続された W&B Launch キューに送信された Launch ジョブ は、CreateTrainingJob API を使う SageMaker Training ジョブ として実行されます。CreateTrainingJob API に送信される引数は、Launch キュー設定で制御します。

Amazon SageMaker は、Docker イメージ を使用してトレーニングジョブを実行します。SageMaker が取得するイメージは、Amazon Elastic Container Registry (ECR) に保存されている必要があります。つまり、トレーニングに使用するイメージは ECR に保存されている必要があります。

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

​

前提条件

開始する前に、次の前提条件を満たしていることを確認してください。

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

​

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

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

• Launch エージェントが Docker イメージ をビルドし、その image を Amazon ECR にプッシュして、SageMaker Training ジョブを送信できるようにします。このオプションでは、トレーニングコードをすばやく繰り返し改善したい ML エンジニアの作業を簡略化できます。
• Launch エージェントが、トレーニングまたは推論スクリプトを含む既存の Docker イメージ を使用します。このオプションは、既存の CI システムとの相性が良好です。このオプションを選択する場合は、Docker イメージ を Amazon ECR 上のコンテナー レジストリに手動でアップロードする必要があります。

​

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

使用する AWS リージョンで、次の AWS リソースが設定されていることを確認してください。

1. コンテナーイメージを保存するための ECR リポジトリ。
2. SageMaker Training ジョブの入力と出力を保存するための、1 つ以上の S3 バケット。
3. SageMaker がトレーニングジョブを実行し、Amazon ECR および Amazon S3 と連携できるようにする、Amazon SageMaker 用の IAM ロール。

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

​

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

1. AWS の IAM 画面で、新しいポリシーを作成します。
2. JSON ポリシーエディタに切り替え、ユースケースに応じて以下のポリシーを貼り付けます。<> で囲まれた値は実際の値に置き換えてください。

エージェントが事前にビルド済みの 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/<RoleArn-from-queue-config>"
      },
    {
        "Effect": "Allow",
        "Action": "kms:CreateGrant",
        "Resource": "<ARN-OF-KMS-KEY>",
        "Condition": {
          "StringEquals": {
            "kms:ViaService": "SageMaker.<region>.amazonaws.com",
            "kms:GrantIsForAWSResource": "true"
          }
        }
      }
    ]
  }

エージェントが 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/<RoleArn-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"
          }
        }
      }
    ]
  }

3. Next をクリックします。
4. ポリシー名と説明を入力します。
5. Create policy をクリックします。

​

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

Launch エージェントが Amazon SageMaker のトレーニングジョブを作成するには、権限が必要です。以下の手順に従って IAM ロールを作成してください。

1. AWS の IAM 画面で、新しいロールを作成します。
2. Trusted Entity で、AWS Account (または組織のポリシーに適した別のオプション) を選択します。
3. 権限画面をスクロールし、先ほど作成したポリシー名を選択します。
4. ロールの名前と説明を入力します。
5. Create role を選択します。
6. ロールの ARN を控えます。Launch エージェントを設定するときに、この ARN を指定します。

IAM ロールの作成については、AWS Identity and Access Management Documentation を参照してください。

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

​

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

次に、SageMaker をコンピュートリソースとして使用するキューを W&B App で作成します。

1. Launch App にアクセスします。
2. Create Queue ボタンをクリックします。
3. キューを作成する Entity を選択します。
4. Name フィールドにキューの名を入力します。
5. Resource として SageMaker を選択します。
6. 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 (prerequisites を参照) 。Launch agent の IAM ロールと混同しないでください。
• OutputDataConfig.S3OutputPath : SageMaker の出力の保存先を指定する Amazon S3 URI。
• ResourceConfig: 必須のリソース設定です。リソース設定のオプションは こちら を参照してください。
• StoppingCondition: トレーニング ジョブの停止条件の必須設定です。オプションは こちら を参照してください。

7. Create Queue ボタンをクリックします。

​

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

このセクションでは、エージェントをどこにデプロイできるかと、デプロイ先に応じたエージェントの設定方法について説明します。

Amazon SageMaker のキューで Launch エージェントを実行する方法には、ローカルマシン、EC2 インスタンス、または EKS クラスター上で実行するなど、いくつかの選択肢があります。エージェントのデプロイ先に応じて、Launch エージェントを適切に設定してください。

​

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

本番ワークロードや、すでに EKS クラスターをお持ちのお客様には、この Helm チャートを使用して EKS クラスターに Launch エージェントをデプロイすることを W&B は推奨します。

現在 EKS クラスターがない本番ワークロードでは、EC2 インスタンスが適した選択肢です。Launch エージェント用のインスタンスは常時稼働しますが、必要なのは比較的低コストな t2.micro サイズの EC2 インスタンス程度です。

実験用途や個人利用では、ローカルマシンで Launch エージェントを実行するのが、すばやく使い始める方法として適しています。

ユースケースに応じて、以下のタブにある手順に従って Launch エージェントを適切に設定してください。

EKS

EKS クラスターにエージェントをインストールするには、W&B managed helm chart を使用することを W&B は強く推奨します。

EC2

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

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

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

1. Connect を選択します。
2. SSH client タブを選択し、表示される手順に従って EC2 インスタンスに接続します。
3. EC2 インスタンス内で、次のパッケージをインストールします。
   sudo yum install python311 -y && python3 -m ensurepip --upgrade && pip3 install wandb && pip3 install wandb[launch]
   
4. 次に、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 設定ファイルを使用します。前の step で作成した 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 エージェントを設定する

launch-config.yaml という名前の YAML 設定ファイルを使って Launch エージェントを設定します。

デフォルトでは、W&B は ~/.config/wandb/launch-config.yaml で設定ファイルを探します。必要に応じて、-c フラグを使用して Launch エージェントを起動する際に、別のディレクトリを指定できます。

以下の YAML スニペットは、主要な Launch エージェントの設定オプションを指定する方法を示しています。

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 ジョブ の Docker イメージを Amazon ECR にプッシュする

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

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