AWS Batchでハマったときに確認する6つのこと（Inferentia向け）

先端技術開発グループ（WAND）の小島です。前回の記事で解説したAWS Inferentiaモデルのコンパイルに続き、この記事ではAWS BatchでInferentiaインスタンスを動かす際の注意点と、汎用的なデバッグテクニックを解説します。

AWS Batchとは

AWS Batchは、フルマネージド型のバッチ処理サービスです。ECR（Elastic Container Registry）に登録したDockerジョブを、指定したコンピューティング環境で実行します。ジョブキューを用いて、「SQS + ECS on EC2/Fargate」のような構成をシンプルに管理できます。

AWS バッチ処理ベースのアーキテクチャ より

本記事では、AWS BatchとInferentiaインスタンスを使い、Hugging FaceからダウンロードしたLLMをNeuron用にコンパイルする場合の、AWS Batchのハマりどころについて解説します。

1. そもそもCPUクォーターに抵触してEC2が起動できなくて「Runnable」で止まる

対策：EC2単独で起動してみる

AWS Batchのジョブを起動すると、このように「Runnable」になりますが、数十分待っても「Starting」にならずにRunnableで止まることがあります。これには複数の理由がありますが、Neuron特有の理由としてinf2.48xlargeのような大きなインスタンスを確保しようとしたときに、クォーターに抵触してしまいEC2が起動できないというのがハマりがちな点です。

このようなときは、まず「EC2単独で起動できるかを確かめてみる」というのが一つの解決法です。もしクォーターに抵触している場合、以下のようなメッセージが表示されます。

エラー文のURLでは、サポートチケットに誘導されますが、実は「Service Quotas」から申請できます。Service Quotasから「ダッシュボード」→「EC2」と開き、利用したいインスタンスタイプについて（Inf2なら「Inf」で検索し）、「Running On-Demand ◯◯ instances」の上限解放申請をします。

この値はインスタンスタイプのvCPU数以上にする必要があります。例えば、inf2.48xlargeならvCPUが192であるため、この値は「192以上」で申請します。申請が承認されても、反映まで30分程度遅延する場合があります。

2. コンピューティング環境の最大vCPUがインスタンスタイプより低く、EC2が起動できなくて「Runnable」で止まる

対策：インスタンスタイプとコンピューティング環境の最大vCPUの数を揃える

「EC2単独では起動できるが、Batch経由では起動できずにジョブが開始されない」というパターンです。これはコンピューティング環境の設定が適切でない可能性があります。

例えば、上記の設定はジョブが「Runnable」のまま進行しません。この例では、コンピューティング環境の最大vCPU値「95」が、inf2.24xlargeインスタンスのvCPU数「96」を下回っており、1台もインスタンスを起動できないためです。「96」にすると起動できます。

最小vCPUと最大vCPUは、オートスケーリンググループにおけるvCPU総数の最小値・最大値と同様の意味です。ジョブの需要に応じて最大vCPU数までインスタンスを起動します。最小vCPUを0より大きくすると、ジョブが来ない状態でもインスタンスを起動し続けます。

3. AMIがECSに対応しておらずEC2内でタスクを開始できない

対策：コンピューティング環境においてECS対応のAMIに変える

NeuronやGPUを使う場合は、AWS Batch内でECS on EC2を動かす場合が多いです。Batchの裏側で動いているのがECSの場合は、ECSに最適化されたAMIを明示的に指定する必要があります。NeuronであってもECS最適化されたAMIはあり、以下は「Amazon Linux 2023でNeuronに対応した推奨AMIのID」を示す、SSMパラメーターストアのパスです（このような便利なパラメーターが用意されています）。

/aws/service/ecs/optimized-ami/amazon-linux-2023/neuron/recommended

「パラメーターストア」→「パブリックパラメーター」→サービス名：「ecs」で一覧を取得できますが、数千～1万のパラメーターがあるので、一覧をプログラムで取得してテキストに出力する方法が考えられます。AWS CLIなどで検索するとレート制限に抵触する可能性があります。

4. ボリュームの拡張をどこに書くか迷う

対策：ボリュームの拡張だけ起動テンプレートに書く

イメージサイズが大きい場合、ルートボリュームを拡張しないとインスタンスの空き容量不足問題が発生します。ECS on EC2の場合は、起動テンプレートとコンピューティング環境の使い分けが悩ましいですが、ルートボリュームのサイズ変更以外は基本的にコンピューティング環境で指定します。Terraformでは以下のようになります。

resource "aws_launch_template" "this" {
  name                   = "${var.prefix}-lt"
  update_default_version = true

  block_device_mappings {
    device_name = "/dev/xvda" # AL2023のルートデバイス
    ebs {
      volume_size           = 200
      volume_type           = "gp3"
      delete_on_termination = true
      encrypted            = true
    }
  }

  tag_specifications {
    resource_type = "instance"
    tags          = var.tags
  }
}

resource "aws_batch_compute_environment" "this" {
  name = "${var.prefix}-ce"
  type = "MANAGED"
  compute_resources {
    type                = local.ce_type
    allocation_strategy = var.allocation_strategy
    min_vcpus           = var.min_vcpus
    max_vcpus           = var.max_vcpus
    desired_vcpus       = var.desired_vcpus
    
    instance_type       = var.instance_types
    subnets             = var.subnet_ids
    security_group_ids  = [aws_security_group.this.id]
    instance_role       = aws_iam_instance_profile.batch_instance_profile.arn
    image_id            = jsondecode(data.aws_ssm_parameter.neuron_optimized.value).image_id
    tags                = var.tags
    ec2_key_pair        = var.key_pair_name # デバッグ用で本番は非推奨
    
    launch_template {
      launch_template_id = aws_launch_template.this.id
      version            = "$Latest"
    }
  }
  service_role = aws_iam_role.batch_service.arn
  tags         = var.tags

  depends_on = [
    aws_iam_role_policy_attachment.batch_service,
    aws_iam_role_policy.batch_service_ecs
  ]
}

5. ジョブが開始されたがログが出ず、永遠に終わらない

対策：EC2にSSHログインしてDockerのログを見る

これはできれば回避したいやり方ですが、場合によっては「ECSが実行されているEC2にログインしてログを見たい」というケースはあると思います。「空き容量不足でECRからのイメージの取得に失敗した」などの問題は、この方法で原因を特定できます。テスト時は何らかの形で、ECSが稼働するEC2インスタンスにSSHログインできるようにしておくと、デバッグが容易になります。

ただし、この方法で確認できるのはコンテナの起動プロセスに関するログ（Dockerデーモンのログなど）であり、ジョブ内のアプリケーションログはCloudWatch Logsで確認します。ログの書き込み権限はBatchの実行ロールに持たせておきましょう。

ECSが実行されているEC2にSSHログインできたら、以下のコマンドを実行します。docker psでジョブを実行しているコンテナのIDを確認し、そのIDをdocker logsコマンドの<CONTAINER_ID_OR_NAME>に指定します。

# コンテナを特定
docker ps --format 'table {{.ID}}\t{{.Names}}\t{{.Label "com.amazonaws.ecs.container-name"}}'
# ログを表示
docker logs -f <CONTAINER_ID_OR_NAME>

1つ目のコマンドを実行すると以下のようにコンテナIDが表示されます。

$ docker ps --format 'table {{.ID}}\t{{.Names}}\t{{.Label "com.amazonaws.ecs.container-name"}}'
CONTAINER ID   NAMES       container name
8af190cceaa2   ecs-agent

6. コンピューティング環境とジョブキューに依存関係があり、更新時に削除する必要がある

対策：Terraformの場合はジョブキューのreplace_triggered_byを指定する

コンピューティング環境（vCPUやインスタンスタイプなど）を更新する際は、リソースの再作成が必要です。AWS Batchの仕様では、コンピューティング環境とジョブキューに依存関係があるため、環境の再作成と同時にジョブキューも再作成（削除→作成）しなければなりません。

この挙動をIaCで明示的に定義する必要があり、例えばTerraformでは、aws_batch_job_queueリソースのreplace_triggered_byにコンピューティング環境を指定します。

resource "aws_batch_job_queue" "this" {
  name     = "${var.prefix}-jq"
  state    = "ENABLED"
  priority = 1

  compute_environment_order {
    compute_environment = aws_batch_compute_environment.this.arn
    order               = 1
  }
  
  tags = var.tags

  lifecycle {
    # CE が置換されたら queue も置換（delete->create）させる
    replace_triggered_by = [aws_batch_compute_environment.this]
  }  
}

おわりに

AWS BatchはNeuronの例に限らず、機械学習のジョブで登場頻度が多いように思えます。Batchの魅力は、Fargateと比べてやや敷居が高いECS on EC2を抽象化できる点で、スポットインスタンスも柔軟に利用できます。SageMakerの追加コストを抑えたい場合、GPUのジョブ実行では有力な選択肢になるので、ぜひ使いこなしていきたいところです。

本記事がAWS Batchでハマっている方の参考になれば幸いです。