Tutoriel : configurer W&B Launch sur SageMaker

Ce tutoriel montre aux ingénieurs ML et aux administrateurs de plateforme comment configurer W&B Launch pour soumettre des tâches d’entraînement à Amazon SageMaker. À la fin, vous disposerez des ressources AWS, des rôles IAM, de la configuration de la file d’attente et de l’agent Launch nécessaires pour exécuter des tâches d’entraînement SageMaker depuis W&B. Vous pouvez utiliser W&B Launch pour soumettre des jobs Launch à Amazon SageMaker afin d’entraîner des modèles de machine learning à l’aide d’algorithmes fournis ou personnalisés sur la plateforme SageMaker. SageMaker se charge de l’approvisionnement et de la libération des ressources de calcul, ce qui peut en faire un bon choix pour les équipes qui ne disposent pas de cluster EKS. Une file d’attente W&B Launch connectée à Amazon SageMaker exécute les jobs Launch en tant que SageMaker Training Jobs via l’API CreateTrainingJob. Utilisez la configuration de la file d’attente Launch pour contrôler les arguments envoyés à l’API CreateTrainingJob. Amazon SageMaker utilise des images Docker pour exécuter les tâches d’entraînement. Les images récupérées par SageMaker doivent être stockées dans Amazon Elastic Container Registry (ECR). Autrement dit, l’image que vous utilisez pour l’entraînement doit être stockée sur ECR.

Ce guide montre comment exécuter des SageMaker Training Jobs. Pour savoir comment déployer des modèles pour l’inférence sur Amazon SageMaker, consultez cet exemple de job Launch.

Prérequis

Avant de commencer, vous devez remplir les prérequis suivants :

Déterminez si vous souhaitez que l’agent Launch génère une image Docker à votre place.
Configurez les ressources AWS et rassemblez les informations relatives à S3, à ECR et aux rôles IAM SageMaker.
Créez un rôle IAM pour l’agent Launch.

Les sections suivantes expliquent comment remplir chacun de ces prérequis.

Déterminer si vous souhaitez que l’agent Launch crée une image Docker

Déterminez si vous souhaitez que l’agent W&B Launch crée une image Docker à votre place. Deux options s’offrent à vous :

Autoriser l’agent Launch à créer une image Docker, à la téléverser vers Amazon ECR et à soumettre des jobs SageMaker Training à votre place. Cette option peut simplifier la tâche des ingénieurs ML qui itèrent rapidement sur leur code d’entraînement.
Utiliser une image Docker existante contenant vos scripts d’entraînement ou d’inférence. Cette option s’intègre bien aux systèmes de CI existants. Si vous la choisissez, vous devez téléverser manuellement votre image Docker vers votre registre de conteneurs sur Amazon ECR.

Configurer les ressources AWS

Les ressources AWS suivantes doivent être configurées dans la région AWS de votre choix :

Un dépôt ECR pour stocker les images de conteneur.
Un ou plusieurs buckets S3 pour stocker les entrées et les sorties de vos tâches SageMaker Training.
Un rôle IAM pour Amazon SageMaker qui autorise SageMaker à exécuter des tâches d’entraînement et à interagir avec Amazon ECR et Amazon S3.

Notez les ARN de ces ressources. Vous en aurez besoin pour définir la configuration de la file d’attente Launch.

Créer une politique IAM pour l’agent Launch

L’agent Launch a besoin d’une politique IAM qui accorde les autorisations nécessaires pour soumettre des tâches d’entraînement SageMaker et, éventuellement, pour téléverser des images vers ECR. Pour créer la politique, procédez comme suit :

Dans l’écran IAM d’AWS, créez une nouvelle politique.
Basculez vers l’éditeur de politique JSON, puis collez la politique correspondant à votre cas d’usage. Remplacez les espaces réservés entre [BRACKETS] par vos propres valeurs :

L'agent soumet une image Docker préconstruite
L'agent construit et soumet une image 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"
          }
        }
      }
    ]
  }

Cliquez sur Next.
Attribuez un nom et une description à la politique.
Cliquez sur Create policy.

Vous disposez maintenant d’une politique IAM que vous pourrez joindre au rôle de l’agent Launch dans la section suivante.

Créer un rôle IAM pour l’agent Launch

L’agent Launch doit disposer de l’autorisation de créer des tâches d’entraînement Amazon SageMaker. En joignant la politique créée dans la section précédente à un rôle dédié, vous permettez à l’agent d’assumer ces autorisations lors de l’exécution. Pour créer un rôle IAM, procédez comme suit :

Depuis l’écran IAM d’AWS, créez un nouveau rôle.
Pour Trusted Entity, sélectionnez AWS Account (ou une autre option adaptée aux politiques de votre organisation).
Faites défiler l’écran des autorisations et sélectionnez le nom de la politique créée dans la section précédente.
Attribuez un nom et une description au rôle.
Sélectionnez Create role.
Notez l’ARN du rôle. Vous devrez le spécifier lors de la configuration de l’agent Launch.

Pour en savoir plus sur la création de rôles IAM, consultez la documentation AWS Identity and Access Management.

Si vous souhaitez que l’agent Launch génère des images, consultez la page Configuration avancée de l’agent pour connaître les autorisations supplémentaires requises.
L’autorisation kms:CreateGrant pour les files d’attente SageMaker n’est requise que si la ResourceConfig associée spécifie un VolumeKmsKeyId et que le rôle associé ne dispose d’aucune politique autorisant cette action.

Configurer la file d’attente Launch pour SageMaker

Une fois les prérequis AWS en place, vous pouvez créer la file d’attente W&B Launch qui achemine les jobs vers SageMaker. Créez dans la W&B App une file d’attente qui utilise SageMaker comme ressource de calcul :

Accédez à l’application Launch.
Cliquez sur Create Queue.
Sélectionnez l’Entity dans laquelle vous souhaitez créer la file d’attente.
Saisissez un nom pour votre file d’attente dans le champ Name.
Sélectionnez SageMaker comme Resource.

Dans le champ Configuration, fournissez les informations relatives à votre job SageMaker. Par défaut, W&B préremplit un corps de requête CreateTrainingJob au format YAML et JSON :

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

Vous devez au minimum spécifier :

RoleArn : ARN du rôle IAM d’exécution SageMaker (voir les prérequis). Ne le confondez pas avec le rôle IAM de l’agent Launch.
OutputDataConfig.S3OutputPath : URI Amazon S3 indiquant l’emplacement où SageMaker stocke les sorties.
ResourceConfig : spécification obligatoire de la configuration des ressources. Pour connaître les options disponibles, consultez la documentation AWS ResourceConfig.
StoppingCondition : spécification obligatoire des conditions d’arrêt de la tâche d’entraînement. Pour connaître les options disponibles, consultez la documentation AWS StoppingCondition.

Cliquez sur Create Queue.

Votre file d’attente est créée. Elle sera prête à recevoir des jobs une fois que vous aurez configuré un agent Launch pour l’interroger.

Configurer l’agent Launch

Les sections suivantes décrivent où vous pouvez déployer votre agent et comment le configurer en fonction de l’endroit où vous le déployez. Vous disposez de plusieurs options pour déployer l’agent Launch pour une file d’attente Amazon SageMaker : sur une machine locale, sur une instance EC2 ou dans un cluster EKS. Configurez votre agent Launch en fonction de l’endroit où vous le déployez.

Choisir où exécuter l’agent Launch

Pour les charges de travail de production et pour les clients qui disposent déjà d’un cluster EKS, W&B recommande de déployer l’agent Launch sur le cluster EKS à l’aide de ce chart Helm. Pour les charges de travail de production sans cluster EKS existant, une instance EC2 constitue une bonne option. Même si l’instance de l’agent Launch fonctionne en permanence, l’agent n’a pas besoin de plus qu’une instance EC2 de taille t2.micro, dont le coût reste abordable. Pour un usage expérimental ou individuel, vous pouvez exécuter l’agent Launch sur votre machine locale afin de démarrer rapidement. Selon votre cas d’usage, suivez les instructions des onglets ci-dessous pour configurer votre agent Launch :

EKS
EC2
Machine locale

W&B recommande d’utiliser le chart Helm géré par W&B pour installer l’agent dans un cluster EKS.

Accédez au tableau de bord Amazon EC2 et effectuez les étapes suivantes :

Cliquez sur Launch instance.
Saisissez un nom dans le champ Name. Vous pouvez également ajouter un tag.
Dans Instance type, sélectionnez un type d’instance pour votre conteneur EC2. 1 vCPU et 1 Gio de mémoire suffisent (par exemple, une t2.micro).
Créez une paire de clés pour votre organisation dans le champ Key pair (login). Vous utiliserez cette paire de clés lors d’une étape ultérieure pour vous connecter à votre instance EC2 avec un client SSH.
Dans Network settings, sélectionnez un groupe de sécurité pour votre organisation.
Développez Advanced details. Pour IAM instance profile, sélectionnez le rôle IAM de l’agent Launch que vous avez créé dans la section précédente.
Vérifiez le champ Summary. Si tout est correct, sélectionnez Launch instance.

Accédez à Instances dans le volet de gauche du tableau de bord EC2 sur AWS. Assurez-vous que l’instance EC2 que vous avez créée est en cours d’exécution (voir la colonne Instance state). Une fois ce point vérifié, ouvrez le terminal de votre machine locale et procédez comme suit :

Sélectionnez Connect.
Sélectionnez l’onglet SSH client et suivez les instructions pour vous connecter à votre instance EC2.

Dans votre instance EC2, installez les paquets suivants :

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

Ensuite, installez et démarrez Docker dans votre instance EC2 :

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

Vous pouvez maintenant configurer l’agent Launch.

Utilisez les fichiers de configuration AWS situés dans ~/.aws/config et ~/.aws/credentials pour associer un rôle à un agent qui interroge les files d’attente sur une machine locale. Fournissez l’ARN du rôle IAM que vous avez créé pour l’agent Launch à l’étape précédente.

~/.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]

Les jetons de session ont une durée de validité maximale de 1 heure ou de 3 jours, selon le principal associé.

Configurer un agent Launch

Après avoir choisi où exécuter l’agent, configurez-le afin qu’il puisse interroger votre file d’attente SageMaker et s’authentifier auprès d’AWS. Configurez l’agent Launch à l’aide d’un fichier de configuration YAML nommé launch-config.yaml. Par défaut, W&B recherche le fichier de configuration dans ~/.config/wandb/launch-config.yaml. Vous pouvez aussi indiquer un autre répertoire avec l’option -c lorsque vous démarrez l’agent Launch. L’extrait YAML suivant montre comment spécifier les principales options de configuration de l’agent :

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

Démarrez maintenant l’agent avec wandb launch-agent. Votre agent Launch est maintenant en cours d’exécution et interroge la file d’attente SageMaker pour y récupérer des jobs.

Facultatif : téléverser l’image Docker de votre job Launch vers Amazon ECR

Cette section s’applique uniquement si votre agent Launch utilise des images Docker existantes contenant votre logique d’entraînement ou d’inférence. Votre agent Launch prend en charge deux options de comportement.

Téléversez l’image Docker contenant votre job Launch vers votre dépôt Amazon ECR. Si vous utilisez des jobs basés sur des images, votre image Docker doit se trouver dans votre registre ECR avant de soumettre de nouveaux jobs Launch.

​Prérequis

​Déterminer si vous souhaitez que l’agent Launch crée une image Docker

​Configurer les ressources AWS

​Créer une politique IAM pour l’agent Launch

​Créer un rôle IAM pour l’agent Launch

​Configurer la file d’attente Launch pour SageMaker

​Configurer l’agent Launch

​Choisir où exécuter l’agent Launch

​Configurer un agent Launch

​Facultatif : téléverser l’image Docker de votre job Launch vers Amazon ECR

Prérequis

Déterminer si vous souhaitez que l’agent Launch crée une image Docker

Configurer les ressources AWS

Créer une politique IAM pour l’agent Launch

Créer un rôle IAM pour l’agent Launch

Configurer la file d’attente Launch pour SageMaker

Configurer l’agent Launch

Choisir où exécuter l’agent Launch

Configurer un agent Launch

Facultatif : téléverser l’image Docker de votre job Launch vers Amazon ECR