0.aws-batch-distributed-training.yaml

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

AWSTemplateFormatVersion: '2010-09-09'
Description: >
 Setup for large scale AWS Batch with CEs, JQ, ECR and Job Definition.
 Author: Pierre-Yves Aquilanti - pierreya@

Parameters:
  VPCStackParameter:
    Type: String
    Description: Private subnets will be retrieved for the compute environment
    Default: 'LargeScaleVPC'
  AMI:
    Description: Default ECS Optimized AMI or your own AMI ID for the Compute Environment
    Type: AWS::EC2::Image::Id
    Default: ''
  CapacityReservationId:
    Description: Leave empty if unused or if using Capacity Reservation Resource Group
    Type: String
    Default: ''
  CapacityReservationResourceGroup:
    Description: Leave empty if unused or if using Capacity Reservation ID
    Type: String
    Default: ''
  EC2KeyPair:
    Description: Select your keypair or leave as default, useful for debug
    Type: AWS::EC2::KeyPair::KeyName
  CreatePlacementGroup:
    AllowedValues:
      - 'true'
      - 'false'
    Default: 'false'
    Description:
      Create a new placement group with cluster placement
    Type: String

Metadata:
  AWS::CloudFormation::Interface:
    ParameterGroups:
      - Label:
          default: General configuration
        Parameters:
          - VPCStackParameter
      - Label:
          default: AWS Batch Configuration
        Parameters:
          - AMI
          - CapacityReservationId
          - CapacityReservationResourceGroup
          - CreatePlacementGroup
          - EC2KeyPair
    ParameterLabels:
      VPCStackParameter:
        default: Name of the VPC Stack
      AMI:
        default: ECS Compatible AMI
      CapacityReservationId:
        default: Capacity Reservation ID
      CapacityReservationResourceGroup:
        default: Capacity Reservation Resource Group ARN
      CreatePlacementGroup:
        default: Create Placement Group
      EC2KeyPair:
        default: EC2 Keypair


Conditions:
  ImageID: !Equals [ !Ref AMI, '' ]
  CrID: !Equals [ !Ref CapacityReservationId, '' ]
  CrRG: !Equals [ !Ref CapacityReservationResourceGroup, '' ]
  PG: !Equals [ !Ref CreatePlacementGroup, 'true' ]
  KeyPair: !Equals [ !Ref EC2KeyPair, '' ]

Resources:
  ###################
  ## EC2 Resources ##
  ###################
  DistributedDeepLearningLT:
    Type: AWS::EC2::LaunchTemplate
    Properties:
      LaunchTemplateData:
        CapacityReservationSpecification:
          CapacityReservationTarget:
              CapacityReservationId: !If [ CrID, !Ref 'AWS::NoValue', !Ref CapacityReservationId]
              CapacityReservationResourceGroupArn:  !If [ CrRG, !Ref 'AWS::NoValue', !Ref CapacityReservationResourceGroup]
        NetworkInterfaces:
          - Description: EFA Interface 0
            Groups: !Split
              - ','
              - Fn::ImportValue: !Sub ${VPCStackParameter}-SecurityGroup
            NetworkCardIndex: 0
            DeviceIndex: 0
            DeleteOnTermination: true
            InterfaceType: efa
          - Description: EFA Interface 1
            Groups: !Split
              - ','
              - Fn::ImportValue: !Sub ${VPCStackParameter}-SecurityGroup
            NetworkCardIndex: 1
            DeviceIndex: 1
            DeleteOnTermination: true
            InterfaceType: efa
          - Description: EFA Interface 2
            Groups: !Split
              - ','
              - Fn::ImportValue: !Sub ${VPCStackParameter}-SecurityGroup
            NetworkCardIndex: 2
            DeviceIndex: 2
            DeleteOnTermination: true
            InterfaceType: efa
          - Description: EFA Interface 3
            Groups: !Split
              - ','
              - Fn::ImportValue: !Sub ${VPCStackParameter}-SecurityGroup
            NetworkCardIndex: 3
            DeviceIndex: 3
            DeleteOnTermination: true
            InterfaceType: efa
        UserData:
         Fn::Base64:
          !Sub |
            MIME-Version: 1.0
            Content-Type: multipart/mixed; boundary="==MYBOUNDARY=="

            --==MYBOUNDARY==

            Content-Type: text/x-shellscript; charset="us-ascii"

            #!/bin/bash

            # # log userdata to console
            # exec >>(tee /var/log/user-data.log|logger -t user-data -s 2>/dev/console) 2>&1

            # # Format NVMe drives as RAID0 volume and mount on /scratch
            # N_DRIVES=8 # max number of drives

            # for DISK in $(seq 1 $N_DRIVES)
            # do
            # sed -e 's/\s*\([\+0-9a-zA-Z]*\).*/\1/' << EOF | sudo fdisk /dev/nvme${!DISK}n1
            #   o # clear the in memory partition table
            #   n # new partition
            #   p # primary partition
            #   1 # partition number 1
            #     # default - start at beginning of disk
            #     # default - stop at end of disk
            #   p # print the in-memory partition table
            #   t # change partition type
            #   fd #
            #   p # print the in-memory partition table
            #   w # save
            #   q # and we're done
            # EOF
            # done

            # # do some checking
            # mdadm --examine /dev/nvme[1-$N_DRIVES]n1

            # # create the raid devise, echo y for the create array validation
            # echo y | mdadm --create /dev/md0 -l raid0 -n $N_DRIVES /dev/nvme[1-$N_DRIVES]n1

            # # display some information
            # cat /proc/mdstat
            # mdadm --examine /dev/nvme[1-$N_DRIVES]n1
            # mdadm --detail /dev/md0

            # # format the file system
            # yum install xfsprogs -y

            # # use 256k block size:
            # mkfs.xfs -d su=256k -d sw=$N_DRIVES /dev/md0

            # # and mount
            # mkdir /scratch
            # mount /dev/md0 /scratch
            # echo '/dev/md0    /scratch  xfs    defaults        0   0' | sudo tee -a /etc/fstab

            # # saves the raid configurations
            # mdadm --detail --scan --verbose | sudo tee -a /etc/mdadm.conf
            # cat /etc/mdadm.conf
            # # dirty
            # chmod 777 /scratch
            # echo ECS_ENABLED_GPU_SUPPORT=true>>/etc/ecs/ecs.config

            --==MYBOUNDARY==--

  ClusterPlacementGroup:
    Type: AWS::EC2::PlacementGroup
    Properties:
      Strategy: cluster

  ########################
  ## Batch Architecture ##
  ########################

  ##
  ## IAM Roles for AWS Batch
  ##
  # Configure IAM roles for Batch
  BatchInstanceRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          -
            Effect: Allow
            Principal:
              Service:
                - batch.amazonaws.com
            Action:
              - sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSBatchServiceRole

  ECSTaskServiceRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          -
            Effect: Allow
            Principal:
              Service:
                - ec2.amazonaws.com
            Action:
              - sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AmazonEC2ContainerServiceforEC2Role

  ECSTaskInstanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties:
      Path: /
      Roles:
        - !Ref ECSTaskServiceRole
      InstanceProfileName: !Join [ "", [ "ECSTaskInstanceProfileIAM-", !Ref AWS::StackName ] ]

  ##
  ## Compute Environment and Job Definition
  ##
  # Build the AWS Batch CEs
  DistributedDeepLearningCE:
    Type: AWS::Batch::ComputeEnvironment
    Properties:
      Type: MANAGED
      ServiceRole: !Ref BatchInstanceRole
      ComputeResources:
        AllocationStrategy: BEST_FIT
        MaxvCpus: 100000
        DesiredvCpus: 0
        MinvCpus: 0
        PlacementGroup: !If [ PG, !Ref AWS::NoValue, !Ref ClusterPlacementGroup]
        Subnets: !Split
          - ','
          - Fn::ImportValue: !Sub ${VPCStackParameter}-PrivateSubnet
        Type: EC2
        Ec2KeyPair: !If [ KeyPair, !Ref AWS::NoValue, !Ref EC2KeyPair ]
        ImageId: !If [ ImageID, '/aws/service/ecs/optimized-ami/amazon-linux-2/gpu/recommended/image_id', !Ref AMI]
        InstanceRole: !Ref ECSTaskInstanceProfile
        LaunchTemplate:
          LaunchTemplateId: !Ref DistributedDeepLearningLT
          Version: $Latest
        InstanceTypes:
          - p4d.24xlarge
      State: ENABLED
      Tags:
        Name: Batch Deep Learning

  # Job Queue
  DistributedDeepLearningJQ:
    Type: AWS::Batch::JobQueue
    Properties:
      ComputeEnvironmentOrder:
        - ComputeEnvironment: !Ref DistributedDeepLearningCE
          Order: 1
      Priority: 1
      State: "ENABLED"

  ##
  ## ECR and AWS Batch Job definition
  ##
  DistributedDeepLearningRepository:
    Type: AWS::ECR::Repository

  StressJDSingleInstance:
    Type: AWS::Batch::JobDefinition
    Properties:
      Type: container
      ContainerProperties:
        Command:
          - /workspace/gpu-nccl-sanity.sh
        Image: !Join ['', [!Ref 'AWS::AccountId','.dkr.ecr.', !Ref 'AWS::Region', '.amazonaws.com/', !Ref DistributedDeepLearningRepository ] ]
        ResourceRequirements:
          - Type: VCPU
            Value: 1
          - Type: GPU
            Value: 8
          - Type: MEMORY
            Value: 16384
        Ulimits:
          - Name: memlock
            HardLimit: -1
            SoftLimit: -1
          - Name: stack
            HardLimit: 67108864
            SoftLimit: 67108864
          - Name: nofile
            HardLimit: 1024000
            SoftLimit: 1024000
        LinuxParameters:
          SharedMemorySize: 8192
          Devices:
            - HostPath: /dev/infiniband/uverbs0
              ContainerPath: /dev/infiniband/uverbs0
              Permissions:
                - READ
                - WRITE
                - MKNOD
            - HostPath: /dev/infiniband/uverbs1
              ContainerPath: /dev/infiniband/uverbs1
              Permissions:
                - READ
                - WRITE
                - MKNOD
            - HostPath: /dev/infiniband/uverbs2
              ContainerPath: /dev/infiniband/uverbs2
              Permissions:
                - READ
                - WRITE
                - MKNOD
            - HostPath: /dev/infiniband/uverbs3
              ContainerPath: /dev/infiniband/uverbs3
              Permissions:
                - READ
                - WRITE
                - MKNOD
      PropagateTags: true
      RetryStrategy:
        Attempts: 10


  StressJDMultiInstances:
    Type: AWS::Batch::JobDefinition
    Properties:
      Type: multinode
      NodeProperties:
        MainNode: 0
        NumNodes: 2
        NodeRangeProperties:
          - TargetNodes: '0:'
            Container:
              InstanceType: p4d.24xlarge
              Image: !Join ['', [!Ref 'AWS::AccountId','.dkr.ecr.', !Ref 'AWS::Region', '.amazonaws.com/', !Ref DistributedDeepLearningRepository ] ]
              Command:
                - /workspace/gpu-nccl-sanity.sh
              ResourceRequirements:
                - Type: VCPU
                  Value: 1
                - Type: GPU
                  Value: 8
                - Type: MEMORY
                  Value: 16384
              Ulimits:
                - Name: memlock
                  HardLimit: -1
                  SoftLimit: -1
                - Name: stack
                  HardLimit: 67108864
                  SoftLimit: 67108864
                - Name: nofile
                  HardLimit: 1024000
                  SoftLimit: 1024000
              LinuxParameters:
                SharedMemorySize: 8192
                Devices:
                  - HostPath: /dev/infiniband/uverbs0
                    ContainerPath: /dev/infiniband/uverbs0
                    Permissions:
                      - READ
                      - WRITE
                      - MKNOD
                  - HostPath: /dev/infiniband/uverbs1
                    ContainerPath: /dev/infiniband/uverbs1
                    Permissions:
                      - READ
                      - WRITE
                      - MKNOD
                  - HostPath: /dev/infiniband/uverbs2
                    ContainerPath: /dev/infiniband/uverbs2
                    Permissions:
                      - READ
                      - WRITE
                      - MKNOD
                  - HostPath: /dev/infiniband/uverbs3
                    ContainerPath: /dev/infiniband/uverbs3
                    Permissions:
                      - READ
                      - WRITE
                      - MKNOD
      PropagateTags: true
      RetryStrategy:
        Attempts: 10

Outputs:

  ECRRepository:
    Description: ECR Repository for the containers
    Value: !Ref DistributedDeepLearningRepository

  ECRRepositoryUrl:
    Description: ECR Repository for the containers
    Value: !Join ['', [!Ref 'AWS::AccountId','.dkr.ecr.', !Ref 'AWS::Region', '.amazonaws.com/', !Ref DistributedDeepLearningRepository ] ]

  JobDefinitionSingleInstance:
    Description: Job definition for single instance Jobs
    Value: !Ref StressJDSingleInstance

  JobDefinitionMultiInstance:
    Description: Job definition for Multi-node Parallel Jobs
    Value: !Ref StressJDMultiInstances

  DistributedDeepLearningJQ:
    Description: Job Queue
    Value: !Ref DistributedDeepLearningJQ