複数のサービスか関係している、かつジョブコンテナを自分の使いたい環境にしてAWS Batchを構築するのは結構大変でしたので、ポイントをメモしておきます。

想定仕様

• EC2タイプ（GPUを利用したいから）

• Spot Fleetでコスト削減

• ECRからコンテナイメージを取得

• GPUを利用するジョブを実行

• コンテナにEC2のディレクトリをマウント

• マウントするEC2のディレクトリにEFSを利用

全体イメージ

だいぶ大まかですが、全体イメージです。これがイメージ出来ていないと詰まった時に解決方法を探すのも大変です。

確認ポイント

• AWS Batch ベストプラクティスまとめ

• AWS Batch のトラブルシューティング

• 構築初期段階はCloudTrailの方を確認するといいかも

起動テンプレート

• EC2起動時にスクリプトを実行したい場合は必要

• AmazonLinux2のAMIが使われていること

• 更新する際は、コンピューティング環境も作り直す

• ルートデバイスの容量が十分に設定されていること

スポットインスタンスで、EC2起動時にスクリプトを実行したい場合などは起動テンプレートを使う必要があります（AMIを作らない場合）

launch_template {
  launch_template_id = aws_launch_template.xxx.id
}

AmazonLinux2以外のAMIを指定していると、ECSクラスターにインスタンスを追加するのが複雑になりそうです。
AWS Batchで起動されるECSクラスター名などを、EC2の起動時にecs.configに正しく定義するのに手間がかかりそうでした。

コンピューティング環境も作り直さないと反映されないようでした。

バッチの学習や推論用のコンテナイメージはサイズが大きくなりがちなので、余裕を見たルートデバイスの容量にしておく方が良さそうです。

block_device_mappings {
  device_name = "/dev/xvda"

  ebs {
    volume_size = 100
  }
}

ちなみに、起動テンプレートのスクリプトの内容にTerraformの変数を含める場合は、filebase64で別ファイルでスクリプトを用意するのではなく、base64encodeでインラインで記述する必要があります。

user_data = base64encode(<<EOF
MIME-Version: 1.0
Content-Type: multipart/mixed; boundary="==MYBOUNDARY=="

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

#!/bin/bash
mkdir /usr/local/share/models /usr/local/share/input /usr/local/share/output
aws s3 sync s3://${aws_s3_bucket.xxx.bucket}/models /usr/local/share/models --delete
aws s3 sync /usr/local/share/models s3://${aws_s3_bucket.xxx.bucket}/models
aws s3 sync s3://${aws_s3_bucket.xxx.bucket}/input /usr/local/share/input --delete
aws s3 sync /usr/local/share/output s3://${aws_s3_bucket.xxx.bucket}/output

--==MYBOUNDARY==
EOF
)

コンピューティング環境

• GPU対応のインスタンスタイプが選択されていること

• スポットフリートのロールが指定されていること

GPUを使うインスタンスタイプ（g4dn.xlargeなど）を指定した場合には「amzn2-ami-ecs-gpu-hvm-2.0.20230428-x86_64-ebs」のAMIで起動されるようです。

スポットインスタンスを利用する際には、スポットフリートのロールも指定しておく必要があります（権限は後述）

type = "SPOT"
spot_iam_fleet_role = aws_iam_role.xxx.arn

コンテナイメージの作成

• GPU対応のインスタンスタイプ（EC2）が選択されていること

Amazon ECS での GPU の使用に記載のインスタンスタイプが選択されていればコンテナ内でも利用できます。

nvcr.io/nvidia/pytorch:22.01-py3 など、GPUも利用できるように用意されているイメージを利用すると手間は少ないです。

コンテナ内でのGPUの確認もコードに含めておくと把握しやすいです。

import torch

print(torch.cuda.is_available())

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"デフォルトのデバイス: {device}")

if device.type == "cuda":
    print(torch.cuda.get_device_name(device))
    print("GPUメモリの総量:", torch.cuda.get_device_properties(device).total_memory)
    print("GPUメモリの使用量:", torch.cuda.memory_allocated(device))
    print(
        "GPUメモリの空き容量:",
        torch.cuda.memory_reserved(device) - torch.cuda.memory_allocated(device),
    )

AWS Batchサービスロールの権限設定

• AWSBatchServiceRoleサービスロールが付与されていること

• （ECRへのアクセス権が付与されていること）

resource "aws_iam_role_policy_attachment" "xxx" {
  role       = aws_iam_role.xxx.name
  policy_arn = "arn:aws:iam::aws:policy/service-role/AWSBatchServiceRole"
}

ECRへのアクセス権はAWS Batchのサービスロールか、EC2のインスタンスロールかどちらに必要か確認しきれていないので両方に付与してみました。

data "aws_iam_policy_document" "xxx" {
  statement {
    effect = "Allow"
    actions = [
      "ecr:GetAuthorizationToken",
      "ecr:BatchCheckLayerAvailability",
      "ecr:GetDownloadUrlForLayer",
      "ecr:GetRepositoryPolicy",
      "ecr:DescribeRepositories",
      "ecr:ListImages",
      "ecr:DescribeImages",
      "ecr:BatchGetImage"
    ]
    resources = [aws_ecr_repository.xxx.arn]
  }
}

EC2インスタンスロールの権限設定

• AmazonEC2ContainerServiceforEC2Roleが付与されていること

• （ECRへのアクセス権が付与されていること）

resource "aws_iam_role_policy_attachment" "xxx" {
  role       = aws_iam_role.xxx.name
  policy_arn = "arn:aws:iam::aws:policy/service-role/AmazonEC2ContainerServiceforEC2Role"
}

スポットフリートロールの権限設定

• AmazonEC2SpotFleetTaggingRoleが付与されていること

resource "aws_iam_role_policy_attachment" "xxx" {
  role       = aws_iam_role.xxx.name
  policy_arn = "arn:aws:iam::aws:policy/service-role/AmazonEC2SpotFleetTaggingRole"
}

ジョブ定義

• GPUが指定されていること

• EFSが指定されていること

GPUを使う場合はここで指定します（指定しない場合はコンテナ内でGPUを認識しませんでした）

"resourceRequirements" : [
  {
    "type": "GPU",
    "value": "number"
  }
]

GPUを指定して起動されたインスタンスのecs.configです。

$ cat /etc/ecs/ecs.config
ECS_CLUSTER=clusterName
ECS_DISABLE_IMAGE_CLEANUP=false
ECS_ENGINE_TASK_CLEANUP_WAIT_DURATION=2m
ECS_IMAGE_CLEANUP_INTERVAL=10m
ECS_IMAGE_MINIMUM_CLEANUP_AGE=10m
ECS_NUM_IMAGES_DELETE_PER_CYCLE=5
ECS_RESERVED_MEMORY=32
ECS_INSTANCE_ATTRIBUTES={"com.amazonaws.batch.base-instance-type":"g4dn.xlarge"}
ECS_CLUSTER=clusterName

コンテナの確認結果、正しくGPUが認識されているようです。

'"============="
'"== PyTorch =="
'"============="
NVIDIA Release 22.01 (build 31424411)
PyTorch Version 1.11.0a0+bfe5ad2
"Container image Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved."
Copyright (c) 2014-2022 Facebook Inc.
Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
Copyright (c) 2011-2013 NYU                      (Clement Farabet)
"Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)"
Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
"Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)"
Copyright (c) 2015      Google Inc.
Copyright (c) 2015      Yangqing Jia
Copyright (c) 2013-2016 The Caffe contributors
All rights reserved.
Various files include modifications (c) NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
This container image and its contents are governed by the NVIDIA Deep Learning Container License.
"By pulling and using the container, you accept the terms and conditions of this license:"
https://developer.nvidia.com/ngc/nvidia-deep-learning-container-license
NOTE: CUDA Forward Compatibility mode ENABLED.
  Using CUDA 11.6 driver version 510.39.01 with kernel driver version 470.182.03.
  See https://docs.nvidia.com/deploy/cuda-compatibility/ for details.
NOTE: MOFED driver for multi-node communication was not detected.
      Multi-node communication performance may be reduced.
NOTE: The SHMEM allocation limit is set to the default of 64MB.  This may be
   insufficient for PyTorch.  NVIDIA recommends the use of the following flags:
   docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 ...
True
デフォルトのデバイス: cuda
Tesla T4
GPUメモリの総量: 15843721216

EFSにモデルなどを格納する場合、ここで指定しておきます。

volumes = [
  {
    name = "xxx-models-efs"
    efsVolumeConfiguration = {
      fileSystemId = aws_efs_file_system.xxx.id
    }
  }
]

mountPoints = [
  {
    sourceVolume  = "xxx-models-efs"
    containerPath = "/app/models"
    readOnly      = false
  }
]

ジョブキュー

• コンピューティング環境を変更する際はジョブキューを一旦削除する

ジョブキューが存在している場合にはコンピューティング環境を変更できません。

ECS

• ECS container agentがEC2インスタンス上で動作していること

AmazonLinux2以外のAMIを使っている場合などに、設定が正しく行われていないとECS container agentが起動していない、もしくは起動しているがECSクラスター名がECS container agentに正しく設定されていないと、EC2インスタンスは起動しているけれどもAWS BatchのジョブはRUNNABLEのまま進まない、という状況になってしまいます。

ストレージ

• セキュリティグループでEC2からのアクセスを許可されていること

• マウントターゲットがサブネット分作成されていること

当初S3を活用してモデルの格納（データの永続化）を行おうかと考えていましたが、EFSを利用する方が手順を少なく出来そうだったためEFSを選択しました。
Amazon EFS ボリューム
セキュリティグループの作成

セキュリティグループの設定でソースをEC2のセキュリティグループを指定すれば大丈夫でした。バッチのサービスロールにもインスタンスロールにもEFSに対する権限の追加は必要ありませんでした。

resource "aws_security_group" "a" {
  name = "a"

  ingress {
    description      = "SSH"
    from_port        = 22
    to_port          = 22
    protocol         = "tcp"
    cidr_blocks      = local.limited_cidr_blocks
  }

  egress {
    from_port   = 0
    to_port     = 0
    protocol    = "-1"
    cidr_blocks = ["0.0.0.0/0"]
    ipv6_cidr_blocks = ["::/0"]
  }
}

resource "aws_security_group" "b" {
  name = "b"

  ingress {
    description      = "NFS"
    from_port        = 2049
    to_port          = 2049
    protocol         = "tcp"
    security_groups  = [aws_security_group.a.id]
  }

  egress {
    from_port   = 0
    to_port     = 0
    protocol    = "-1"
    cidr_blocks = ["0.0.0.0/0"]
    ipv6_cidr_blocks = ["::/0"]
  }
}

スポットフリートで複数のサブネットを指定している場合は、全てのサブネット対するマウントターゲットを指定します。

resource "aws_efs_mount_target" "xxx" {
  count           = length(local.subnet_ids)
  file_system_id  = aws_efs_file_system.xxx.id
  subnet_id       = local.subnet_ids[count.index]
  security_groups = [
      aws_security_group.b.id,
  ]
}

各定義の更新の際に気を付けること

定義の種類は以下の5種類を使います。

• 起動テンプレート

• コンピューティング環境

• ジョブ定義

• ジョブキュー

このうち、コンピューティング環境を変更する場合は、ジョブキューを一旦削除する必要があります。

起動テンプレートを変更する場合は、ジョブキューとコンピューティング環境を一旦削除する必要があります。
こちらは正しくは、一旦削除しないと起動テンプレートの変更が反映されない、という状態になります。

ジョブ定義、ジョブキューは気にせず変更可能です。

さいごに

常時最速でアクセスする必要が無い、かつGPUを利用したい場合はAWS Batchが選択肢に入ると思いますが、思いのほか構築に苦労したので、ある程度把握していないと運用も少し苦労するのではと感じました。
TerraformなどでIaCとしていても、反映ミス（定義ミス）なども怖いなと思いました。
ただ、GPUインスタンスを使う場合などのコストメリットは大きいので活用できるに越したことはないでしょう。

※ちなみにこの文章量になってくるとnoteの編集（特に文字選択）のストレスが半端ないです