Skip to main content

Tutorial: Set up W&B Launch on SageMaker

You can use W&B Launch to submit launch jobs to Amazon SageMaker to train machine learning models using provided or custom algorithms on the SageMaker platform. SageMaker takes care of spinning up and releasing compute resources, so it can be a good choice for teams without an EKS cluster.

Launch jobs sent to a W&B Launch queue connected to Amazon SageMaker are executed as SageMaker Training Jobs with the CreateTrainingJob API. Use the launch queue configuration to control arguments sent to the CreateTrainingJob API.

Amazon SageMaker uses Docker images to execute training jobs. Images pulled by SageMaker must be stored in the Amazon Elastic Container Registry (ECR). This means that the image you use for training must be stored on ECR.

note

This guide shows how to execute SageMaker Training Jobs. For information on how to deploy to models for inference on Amazon SageMaker, see this example Launch job.

Prerequisites​

Before you get started, ensure you satisfy the following prerequisites:

Decide if you want the Launch agent to build a Docker images​

Decide if you want the W&B Launch agent to build a Docker image for you. There are two options you can choose from:

  • Permit the launch agent build a Docker image, push the image to Amazon ECR, and submit SageMaker Training jobs for you. This option can offer some simplicity to ML Engineers rapidly iterating over training code.
  • The launch agent uses an existing Docker image that contains your training or inference scripts. This option works well with existing CI systems. If you choose this option, you will need to manually upload your Docker image to your container registry on Amazon ECR.

Set up AWS resources​

Ensure you have the following AWS resources configured in your preferred AWS region:

  1. An ECR repository to store container images.
  2. One or more S3 buckets to store inputs and outputs for your SageMaker Training jobs.
  3. An IAM role for Amazon SageMaker that permits SageMaker to run training jobs and interact with Amazon ECR and Amazon S3.

Make a note of the ARNs for these resources. You will need the ARNs when you define the Launch queue configuration.

Create a IAM Policy for Launch agent​

  1. From the IAM screen in AWS, create a new policy.
  2. Toggle to the JSON policy editor, then paste the following policy based on your use case. Substitute values enclosed with <> with your own values:
{
  "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"
        }
      }
    }
  ]
}
  1. Click Next.
  2. Give the policy a name and description.
  3. Click Create policy.

Create an IAM role for Launch agent​

The Launch agent needs permission to create Amazon SageMaker training jobs. Follow the procedure below to create an IAM role:

  1. From the IAM screen in AWS, create a new role.
  2. For Trusted Entity, select AWS Account (or another option that suits your organization's policies).
  3. Scroll through the permissions screen and select the policy name you just created above.
  4. Give the role a name and description.
  5. Select Create role.
  6. Note the ARN for the role. You will specify the ARN when you set up the launch agent.

For more information on how to create IAM role, see the AWS Identity and Access Management Documentation.

info
  • If you want the launch agent to build images, see the Advanced agent set up for additional permissions required.
  • The kms:CreateGrant permission for SageMaker queues is required only if the associated ResourceConfig has a specified VolumeKmsKeyId and the associated role does not have a policy that permits this action.

Configure launch queue for SageMaker​

Next, create a queue in the W&B App that uses SageMaker as its compute resource:

  1. Navigate to the Launch App.
  2. Click on the Create Queue button.
  3. Select the Entity you would like to create the queue in.
  4. Provide a name for your queue in the Name field.
  5. Select SageMaker as the Resource.
  6. Within the Configuration field, provide information about your SageMaker job. By default, W&B will populate a YAML and JSON CreateTrainingJob request body:
{
  "RoleArn": "<REQUIRED>", 
  "ResourceConfig": {
      "InstanceType": "ml.m4.xlarge",
      "InstanceCount": 1,
      "VolumeSizeInGB": 2
  },
  "OutputDataConfig": {
      "S3OutputPath": "<REQUIRED>"
  },
  "StoppingCondition": {
      "MaxRuntimeInSeconds": 3600
  }
}

You must at minimum specify:

  • RoleArn : ARN of the SageMaker execution IAM role (see prerequisites). Not to be confused with the launch agent IAM role.
  • OutputDataConfig.S3OutputPath : An Amazon S3 URI specifying where SageMaker outputs will be stored.
  • ResourceConfig: Required specification of a resource config. Options for resource config are outlined here.
  • StoppingCondition: Required specification of the stopping conditions for the training job. Options outlined here.
  1. Click on the Create Queue button.

Set up the launch agent​

The following section describes where you can deploy your agent and how to configure your agent based on where it is deployed.

There are several options for how the Launch agent is deployed for a Amazon SageMaker queue: on a local machine, on an EC2 instance, or in an EKS cluster. Configure your launch agent appropriately based on the where you deploy your agent.

Decide where to run the Launch agent​

For production workloads and for customers who already have an EKS cluster, W&B recommends deploying the Launch agent to the EKS cluster using this Helm chart.

For production workloads without an current EKS cluster, an EC2 instance is a good option. Though the launch agent instance will keep running all the time, the agent doesn't need more than a t2.micro sized EC2 instance which is relatively affordable.

For experimental or solo use cases, running the Launch agent on your local machine can be a fast way to get started.

Based on your use case, follow the instructions provided in the following tabs to properly configure up your launch agent:

W&B strongly encourages that you use the W&B managed helm chart to install the agent in an EKS cluster.

Configure a launch agent​

Configure the launch agent with a YAML config file named launch-config.yaml.

By default, W&B will check for the config file in ~/.config/wandb/launch-config.yaml. You can optionally specify a different directory when you activate the launch agent with the -c flag.

The following YAML snippet demonstrates how to specify the core config agent options:

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

Now start the agent with wandb launch-agent

(Optional) Push your launch job Docker image to Amazon ECR​

info

This section applies only if your launch agent uses existing Docker images that contain your training or inference logic. There are two options on how your launch agent behaves.

Upload your Docker image that contains your launch job to your Amazon ECR repo. Your Docker image needs to be in your ECR registry before you submit new launch jobs if you are using image-based jobs.

On this page
Was this page helpful?👍👎