Tutorial: Set up W&B Launch on SageMaker

4 minute read

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

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

Amazon SageMaker は、Docker イメージを使用して Training ジョブを実行します。SageMaker によってプルされるイメージは、Amazon Elastic Container Registry (ECR) に保存する必要があります。これは、トレーニングに使用するイメージが ECR に保存されている必要があることを意味します。

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

前提条件

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

Launch エージェントに Docker イメージを構築させるかどうかを決定します。
AWS リソースをセットアップし、S3、ECR、および Sagemaker IAM ロールに関する情報を収集します。
Launch エージェントの IAM ロールを作成します。

Launch エージェントに Docker イメージを構築させるかどうかを決定します。

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

Launch エージェントが Docker イメージを構築し、イメージを Amazon ECR にプッシュして、SageMaker Training ジョブを送信できるようにします。このオプションは、機械学習エンジニアがトレーニングコードを迅速に反復処理する際に、ある程度の簡素化をもたらすことができます。
Launch エージェントは、トレーニングスクリプトまたは推論スクリプトを含む既存の Docker イメージを使用します。このオプションは、既存の CI システムとうまく連携します。このオプションを選択した場合は、Docker イメージを Amazon ECR のコンテナーレジストリに手動でアップロードする必要があります。

AWS リソースのセットアップ

優先する AWS リージョンで次の AWS リソースが構成されていることを確認してください。

コンテナーイメージを保存する ECR リポジトリ。
SageMaker Training ジョブの入出力を保存する 1 つ以上の S3 バケット。
SageMaker が Training ジョブを実行し、Amazon ECR および Amazon S3 とやり取りすることを許可する Amazon SageMaker の IAM ロール。

これらのリソースの ARN をメモしておきます。Launch キューの設定を定義する際に、ARN が必要になります。

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

AWS の IAM 画面から、新しいポリシーを作成します。
JSON ポリシーエディターに切り替え、ユースケースに基づいて次のポリシーを貼り付けます。<> で囲まれた値を独自の値に置き換えます。

{
  "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"
        }
      }
    }
  ]
}

{
  "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"
        }
      }
    }
  ]
}

[Next] をクリックします。
ポリシーに名前と説明を付けます。
[Create policy] をクリックします。

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

Launch エージェントに Amazon SageMaker Training ジョブを作成する権限が必要です。次の手順に従って、IAM ロールを作成します。

AWS の IAM 画面から、新しいロールを作成します。
[Trusted Entity] で、[AWS Account] (または組織のポリシーに適した別のオプション) を選択します。
権限画面をスクロールして、上記で作成したポリシー名を選択します。
ロールに名前と説明を付けます。
[Create role] を選択します。
ロールの ARN をメモします。Launch エージェントを設定する際に、ARN を指定します。

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

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

SageMaker の Launch キューを設定する

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

Launch アプリに移動します。
[Create Queue] ボタンをクリックします。
キューを作成する [Entity] を選択します。
[Name] フィールドにキューの名前を入力します。
[Resource] として [SageMaker] を選択します。
[Configuration] フィールド内で、SageMaker ジョブに関する情報を提供します。デフォルトでは、W&B は YAML および JSON の CreateTrainingJob リクエスト本文を生成します。

{
  "RoleArn": "<必須>", 
  "ResourceConfig": {
      "InstanceType": "ml.m4.xlarge",
      "InstanceCount": 1,
      "VolumeSizeInGB": 2
  },
  "OutputDataConfig": {
      "S3OutputPath": "<必須>"
  },
  "StoppingCondition": {
      "MaxRuntimeInSeconds": 3600
  }
}

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

RoleArn: SageMaker 実行 IAM ロールの ARN (前提条件を参照)。Launch エージェント IAM ロールと混同しないようにしてください。
OutputDataConfig.S3OutputPath: SageMaker の出力が保存される Amazon S3 URI。
ResourceConfig: リソース設定に必要な仕様。リソース設定のオプションはこちらに概説されています。
StoppingCondition: Training ジョブの停止条件に必要な仕様。オプションはこちらに概説されています。

[Create Queue] ボタンをクリックします。

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

次のセクションでは、エージェントをデプロイできる場所と、デプロイ場所に基づいてエージェントを構成する方法について説明します。

Amazon SageMaker の Launch エージェントをデプロイする方法には、いくつかのオプションがあります。ローカルマシン、EC2 インスタンス、または EKS クラスターです。エージェントをデプロイする場所に基づいて、Launch エージェントを適切に構成します。

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

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

現在の EKS クラスターを使用しない本番環境のワークロードの場合、EC2 インスタンスは優れたオプションです。Launch エージェントインスタンスは常に実行され続けますが、エージェントには t2.micro サイズの EC2 インスタンス以上のものは必要ありません。これは比較的安価です。

実験的なユースケースや個人のユースケースの場合、ローカルマシンで Launch エージェントを実行すると、すばやく開始できます。

ユースケースに基づいて、次のタブに記載されている手順に従って、Launch エージェントを適切に構成してください。

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 を指定します。

[profile SageMaker-agent]
role_arn = arn:aws:iam::<account-id>:role/<agent-role-name>
source_profile = default

[default]
aws_access_key_id=<access-key-id>
aws_secret_access_key=<secret-access-key>
aws_session_token=<session-token>

セッショントークンには、関連付けられているプリンシパルに応じて、最大長が 1 時間または 3 日であることに注意してください。

Launch エージェントを構成する

YAML 構成ファイル launch-config.yaml を使用して Launch エージェントを構成します。

デフォルトでは、W&B は ~/.config/wandb/launch-config.yaml で構成ファイルを確認します。必要に応じて、-c フラグを使用して Launch エージェントをアクティブ化するときに、別のディレクトリーを指定できます。

次の 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 エージェントがトレーニングロジックまたは推論ロジックを含む既存の Docker イメージを使用する場合にのみ適用されます。Launch エージェントの動作方法には 2 つのオプションがあります。

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

[i18n] feedback_title

[i18n] feedback_question

Glad to hear it! Please tell us how we can improve.

Sorry to hear that. Please tell us how we can improve.

最終更新 April 3, 2025

ページの編集ドキュメントのissueを作成セクション全体を印刷

Tutorial: Set up W&B Launch on SageMaker

前提条件

Launch エージェント に Docker イメージを構築させるかどうかを決定します。

AWS リソースのセットアップ

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

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