buildspec.yaml

version: 0.2

env:
  variables:
    CUDA_VERSION: "12.8.1"
    GDRCOPY_VERSION: "v2.5.1"
    EFA_INSTALLER_VERSION: "1.43.2"
    AWS_OFI_NCCL_VERSION: "v1.16.3"
    NCCL_VERSION: "v2.27.7-1"
    NCCL_TESTS_VERSION: "v2.16.9"
  exported-variables:
    - CUDA_VERSION
    - GDRCOPY_VERSION
    - EFA_INSTALLER_VERSION
    - AWS_OFI_NCCL_VERSION
    - NCCL_VERSION
    - NCCL_TESTS_VERSION
phases:
  pre_build:
    commands:
      - export TAG="cuda${CUDA_VERSION}-efa${EFA_INSTALLER_VERSION}-ofi${AWS_OFI_NCCL_VERSION}-nccl${NCCL_VERSION}-tests${NCCL_TESTS_VERSION}"
      - echo "TAG=$TAG"
      - export REPO_COUNT="$(aws ecr describe-repositories | grep repositoryName | grep \"${ECR_REPOSITORY_NAME}\" | wc -l)"
      - if [ "$REPO_COUNT" == 0 ]; then aws ecr create-repository --repository-name ${ECR_REPOSITORY_NAME}; else echo "Repository ${ECR_REPOSITORY_NAME} already exists"; fi
      - aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws
  build:
    commands:
      - export REPO_URI="$(aws ecr describe-repositories | grep repositoryUri | grep /${ECR_REPOSITORY_NAME}\" | cut -d '"' -f 4)"
      - echo "REPO_URI=$REPO_URI"
      - echo "Building ${REPO_URI}:${TAG} ..."
      - docker buildx create --use --name multiarch
      - cd micro-benchmarks/nccl-tests && docker buildx build --push --platform=linux/amd64,linux/arm64 --build-arg CUDA_VERSION=${CUDA_VERSION} --build-arg GDRCOPY_VERSION=$GDRCOPY_VERSION --build-arg EFA_INSTALLER_VERSION=$EFA_INSTALLER_VERSION --build-arg AWS_OFI_NCCL_VERSION=$AWS_OFI_NCCL_VERSION --build-arg NCCL_VERSION=$NCCL_VERSION --build-arg NCCL_TESTS_VERSION=$NCCL_TESTS_VERSION -t public.ecr.aws/hpc-cloud/${ECR_REPOSITORY_NAME}:${TAG} -t public.ecr.aws/hpc-cloud/${ECR_REPOSITORY_NAME}:latest -f ./nccl-tests.Dockerfile .