Tutorial: Set up W&B Launch on SageMaker

W&B Launch 를 사용하여 제공된 또는 사용자 지정 알고리즘을 사용하여 Amazon SageMaker 에 launch 작업을 제출하여 SageMaker 플랫폼에서 기계 학습 모델을 트레이닝할 수 있습니다. SageMaker 는 컴퓨팅 리소스를 가동 및 해제하는 작업을 처리하므로 EKS 클러스터가 없는 팀에게 적합한 선택이 될 수 있습니다.

Amazon SageMaker 에 연결된 W&B Launch 대기열로 전송된 Launch 작업은 CreateTrainingJob API를 통해 SageMaker 트레이닝 작업으로 실행됩니다. launch 대기열 설정을 사용하여 CreateTrainingJob API 로 전송되는 인수를 제어합니다.

Amazon SageMaker 는 Docker 이미지를 사용하여 트레이닝 작업을 실행합니다. SageMaker 가 가져오는 이미지는 Amazon Elastic Container Registry (ECR)에 저장해야 합니다. 즉, 트레이닝에 사용하는 이미지는 ECR 에 저장해야 합니다.

전제 조건

시작하기 전에 다음 전제 조건을 충족하는지 확인하십시오.

• Launch 에이전트가 Docker 이미지를 빌드하도록 할지 결정합니다.
• AWS 리소스를 설정하고 S3, ECR 및 Sagemaker IAM 역할에 대한 정보를 수집합니다.
• Launch 에이전트에 대한 IAM 역할을 만듭니다.

Launch 에이전트가 Docker 이미지를 빌드하도록 할지 결정

W&B Launch 에이전트가 Docker 이미지를 빌드하도록 할지 결정합니다. 다음 두 가지 옵션 중에서 선택할 수 있습니다.

• Launch 에이전트가 Docker 이미지를 빌드하고, 이미지를 Amazon ECR 에 푸시하고, 사용자를 위해 SageMaker 트레이닝 작업을 제출하도록 허용합니다. 이 옵션은 ML 엔지니어가 트레이닝 코드를 빠르게 반복하는 데 약간의 단순성을 제공할 수 있습니다.
• Launch 에이전트는 트레이닝 또는 추론 스크립트가 포함된 기존 Docker 이미지를 사용합니다. 이 옵션은 기존 CI 시스템과 잘 작동합니다. 이 옵션을 선택하는 경우 Docker 이미지를 Amazon ECR 의 컨테이너 레지스트리에 수동으로 업로드해야 합니다.

AWS 리소스 설정

선호하는 AWS 리전에서 다음 AWS 리소스가 구성되어 있는지 확인합니다.

1. 컨테이너 이미지를 저장할 ECR 레포지토리.
2. SageMaker 트레이닝 작업에 대한 입력 및 출력을 저장할 하나 이상의 S3 버킷.
3. SageMaker 가 트레이닝 작업을 실행하고 Amazon ECR 및 Amazon S3 와 상호 작용할 수 있도록 허용하는 Amazon SageMaker 에 대한 IAM 역할.

이러한 리소스에 대한 ARN 을 기록해 두십시오. Launch 대기열 설정을 정의할 때 ARN 이 필요합니다.

Launch 에이전트에 대한 IAM 정책 만들기

1. AWS 의 IAM 화면에서 새 정책을 만듭니다.
2. JSON 정책 편집기로 전환한 다음 사용 사례에 따라 다음 정책을 붙여넣습니다. <> 로 묶인 값을 자신의 값으로 대체합니다.

• 에이전트가 미리 빌드된 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/<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"
        }
      }
    }
  ]
}

3. 다음을 클릭합니다.
4. 정책에 이름과 설명을 지정합니다.
5. 정책 생성을 클릭합니다.

Launch 에이전트에 대한 IAM 역할 만들기

Launch 에이전트는 Amazon SageMaker 트레이닝 작업을 생성할 수 있는 권한이 필요합니다. 아래 절차에 따라 IAM 역할을 만듭니다.

1. AWS 의 IAM 화면에서 새 역할을 만듭니다.
2. 신뢰할 수 있는 엔터티의 경우 AWS 계정(또는 조직의 정책에 적합한 다른 옵션)을 선택합니다.
3. 권한 화면을 스크롤하여 위에서 방금 만든 정책 이름을 선택합니다.
4. 역할에 이름과 설명을 지정합니다.
5. 역할 생성을 선택합니다.
6. 역할에 대한 ARN 을 기록해 둡니다. launch 에이전트를 설정할 때 ARN 을 지정합니다.

IAM 역할을 만드는 방법에 대한 자세한 내용은 AWS Identity and Access Management 설명서를 참조하십시오.

SageMaker 에 대한 Launch 대기열 구성

다음으로 SageMaker 를 컴퓨팅 리소스로 사용하는 W&B App 에서 대기열을 만듭니다.

1. Launch App으로 이동합니다.
2. 대기열 생성 버튼을 클릭합니다.
3. 대기열을 만들려는 Entities를 선택합니다.
4. 이름 필드에 대기열 이름을 제공합니다.
5. 리소스로 SageMaker 를 선택합니다.
6. 설정 필드 내에서 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: 리소스 구성에 대한 필수 사양입니다. 리소스 구성에 대한 옵션은 여기에 설명되어 있습니다.
• StoppingCondition: 트레이닝 작업에 대한 중지 조건에 대한 필수 사양입니다. 옵션은 여기에 설명되어 있습니다.

7. 대기열 생성 버튼을 클릭합니다.

Launch 에이전트 설정

다음 섹션에서는 에이전트를 배포할 수 있는 위치와 배포 위치에 따라 에이전트를 구성하는 방법을 설명합니다.

Amazon SageMaker 에 대해 Launch 에이전트를 배포하는 방법에 대한 몇 가지 옵션이 있습니다. 대기열: 로컬 시스템, EC2 인스턴스 또는 EKS 클러스터에서. 에이전트를 배포하는 위치에 따라 Launch 에이전트를 적절하게 구성합니다.

Launch 에이전트를 실행할 위치 결정

프로덕션 워크로드 및 이미 EKS 클러스터가 있는 고객의 경우 W&B 는 이 Helm 차트를 사용하여 Launch 에이전트를 EKS 클러스터에 배포하는 것이 좋습니다.

현재 EKS 클러스터가 없는 프로덕션 워크로드의 경우 EC2 인스턴스가 좋은 옵션입니다. launch 에이전트 인스턴스가 항상 실행되지만 에이전트는 상대적으로 저렴한 t2.micro 크기의 EC2 인스턴스 이상이 필요하지 않습니다.

실험적 또는 단독 사용 사례의 경우 로컬 시스템에서 Launch 에이전트를 실행하는 것이 시작하는 빠른 방법이 될 수 있습니다.

사용 사례에 따라 다음 탭에 제공된 지침에 따라 Launch 에이전트를 올바르게 구성합니다.

• EKS
• EC2
• 로컬 시스템

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

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

[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>

Launch 에이전트 구성

launch-config.yaml 이라는 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 이미지를 Amazon ECR 레포지토리로 업로드합니다. 이미지 기반 작업을 사용하는 경우 새 launch 작업을 제출하기 전에 Docker 이미지가 ECR 레지스트리에 있어야 합니다.