AWS BatchのAI環境構築について
複数のサービスか関係している、かつジョブコンテナを自分の使いたい環境にしてAWS Batchを構築するのは結構大変でしたので、ポイントをメモしておきます。
想定仕様
EC2タイプ（GPUを利用したいから）
Spot Fleetでコスト削減
ECRからコンテナイメージを取得
GPUを利用するジョブを実行
コンテナにEC2のディレクトリをマウント
マウントするEC2のディレクトリにEFSを利用
全体イメージ
だいぶ大まかですが、全体イメージです。これがイメージ出来ていないと詰まった時に解決方法を探すのも大変です。
確認ポイント
構築初期段階は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で起動されるようです。
GPU (Kernel 4.14)AMI name: amzn2-ami-ecs-gpu-hvm-2.0.20230428-x86_64-ebs
ECS Agent version: 1.71.0
Docker version: 20.10.22
Containerd version: 1.6.19
NVIDIA driver version: 470.182.03
CUDA version: 11.4
Source AMI name: amzn2-ami-minimal-hvm-2.0.20230418.0-x86_64-ebs
スポットインスタンスを利用する際には、スポットフリートのロールも指定しておく必要があります（権限は後述）
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」
コンテナ用に予約する物理 GPU の数。ジョブ内のすべてのコンテナ用に予約されている GPU の数は、ジョブが起動されたコンピューティングリソースで使用できる GPU の数以下である必要があります。
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メモリの総量: 15843721216EFSにモデルなどを格納する場合、ここで指定しておきます。
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の編集（特に文字選択）のストレスが半端ないです