Terraformで構築する機械学習ワークロード（Batch on EC2編）

nokomoro3

2023.09.19

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

こんちには。

データアナリティクス事業本部インテグレーション部機械学習チームの中村です。

今回も「Terraformで構築する機械学習ワークロード」ということで、今回はその処理をBatch on EC2に載せてみたいと思います。

いままでの記事は以下です。

構成イメージ

構成としては、前回とほぼ変わらないので割愛します（FargateがEC2となったのみの差分）。

動作環境

Docker、Terraformはインストール済みとします。

Terraformを実行する際のAWSリソースへの権限は、aws-vaultで環境構築をしておきます。

aws-vaultについては以下も参考にされてください。

ホストPCとしてのバージョン情報配下です。

OS: Windows 10 バージョン22H2
docker: Docker version 24.0.2-rd, build e63f5fa
terraform: Terraform v1.4.6 (on windows_amd64)

またコンテナ内のバージョン情報配下となります。

Python: 3.10.12 (main, Aug 15 2023, 15:43:05) [GCC 7.3.1 20180712 (Red Hat 7.3.1-15)]
PyTorch: 2.0.1+cu117
OpenCV: 4.8.0
MMEngine: 0.8.4
MMDetection: 3.1.0+

成果物

成果物はGitHub上にあげておきましたので、詳細は必要に応じてこちらを参照ください。

フォルダ構成は以下のようになっています。

├─asset
│      yolox_l_8x8_300e_coco_20211126_140236-d3bd2b23.pth
│
├─docker
│  └─batch_fargate
│          push_ecr.sh
│
├─python
│      Dockerfile
│      mmdetect_handler.py
│      run.py
│      s3_handler.py
│
└─terraform
    ├─environments
    │  └─dev
    │          locals.tf
    │          main.tf
    │          variables.tf
    │
    └─modules
        ├─batch
        │      main.tf
        │      outputs.tf
        │      variables.tf
        │
        ├─ecr
        │      main.tf
        │      variables.tf
        │
        ├─event_bridge
        │      main.tf
        │      variables.tf
        │
        ├─iam
        │      main.tf
        │      outputs.tf
        │      variables.tf
        │
        ├─s3
        │      main.tf
        │      variables.tf
        │
        └─vpc
                main.tf
                outputs.tf
                variables.tf

前回記事と特に変化なしです。

コードの説明

ここからはコードを説明します。Fargateの時と違う点をメインに説明していきます。

コンテナイメージのビルド

DockerfileはFargateの時と全く同じです。

Pythonファイルについて

PythonファイルもFargateの時と全く同じとなります。

tfファイル

tfファイルは以下のようなフォルダ構成となっています。

    ├─environments
    │  └─dev
    │          locals.tf
    │          main.tf
    │          variables.tf
    │
    └─modules
        ├─batch
        │      main.tf
        │      outputs.tf
        │      variables.tf
        │
        ├─ecr
        │      main.tf
        │      variables.tf
        │
        ├─event_bridge
        │      main.tf
        │      variables.tf
        │
        ├─iam
        │      main.tf
        │      outputs.tf
        │      variables.tf
        │
        ├─s3
        │      main.tf
        │      variables.tf
        │
        └─vpc
                main.tf
                outputs.tf
                variables.tf

構成は全く同じですが、前回と異なる点は以下です。

batchの修正
- ここをFargateからEC2用に修正しています。本記事のメイン部分です
iamの修正
- EC2用のインスタンスロール、インスタンスプロファイルを追加しています
- ジョブ実行ロールは削除しています（インスタンスロールで十分なため）

以降、変更のある部分のみ説明します。

modules/batch/

batchについては以下のようになっています。

data "aws_region" "current" {}

// ジョブ定義
resource "aws_batch_job_definition" "main" {
  name                 = "${var.project_prefix}-job-definition"
  type                 = "container"
  container_properties = jsonencode({
    command = [
      "python", "run.py",
      "--input-bucket-name", "Ref::input_bucket_name",
      "--input-object-key", "Ref::input_object_key"
    ]
    image = "${var.image_uri}:latest"
    jobRoleArn = "${var.job_role_arn}"
    resourceRequirements = [
      {
        type = "VCPU"
        value = "1"
      },
      {
        type = "MEMORY"
        value = "2048"
      }
    ]
    environment = "${var.environments}"

    logConfiguration = {
      logDriver     = "awslogs",
      options = {
        awslogs-group = "/aws/batch/job/${var.project_prefix}"
      }
    }
  })
}

resource "aws_cloudwatch_log_group" "log" {
  name = "/aws/batch/job/${var.project_prefix}"
}

// コンピューティング環境
resource "aws_batch_compute_environment" "ec2" {
  compute_environment_name = "${var.project_prefix}-compute-environment"

  compute_resources {
    instance_role = var.instance_profile_arn
    instance_type = ["g4dn.xlarge"]
    max_vcpus = 16

    security_group_ids = [
      var.security_group_id
    ]

    subnets = [
      var.subnet_id
    ]

    type = "EC2"
  }

  type         = "MANAGED"
  service_role = var.service_role_arn
}

# ジョブキュー
resource "aws_batch_job_queue" "job_queue" {
  name                 = "${var.project_prefix}-job-queue"
  state                = "ENABLED"
  priority             = 0
  compute_environments = [aws_batch_compute_environment.ec2.arn]
}

変更としては以下です。

ジョブ定義
- executionRoleArnを削除
- logConfigurationとして、ロググループを設定
コンピューティング環境
- typeをEC2に変更
- compute_resourcesにinstance_roleとinstance_typeを追加
- instance_typeはGPUの使用を今後想定してg4dnを指定してみました

executionRoleArnは本来コンテナエージェントが必要とする権限のあるロールですが、コンテナエージェントは今回の場合EC2上で動作するため、EC2のインスタンロール（プロファイル）に権限があるため、あらためて指定する必要がなくなっています。

実際公式ドキュメントにも以下のような記載があります。

https://docs.aws.amazon.com/batch/latest/APIReference/API_ContainerProperties.html

executionRoleArn The Amazon Resource Name (ARN) of the execution role that AWS Batch can assume. For jobs that run on Fargate resources, you must provide an execution role. For more information, see AWS Batch execution IAM role in the AWS Batch User Guide.

modules/iam/

iamはjob_execution_roleを削除する代わりに、instance_roleとそれを含むinstance_profileを追加します。

# 信頼ポリシー
data "aws_iam_policy_document" "trust_ecs_tasks" {
  statement {
    effect = "Allow"
    principals {
      type        = "Service"
      identifiers = ["ecs-tasks.amazonaws.com"]
    }
    actions = ["sts:AssumeRole"]
  }
}

# 信頼ポリシー
data "aws_iam_policy_document" "trust_ec2" {
  statement {
    effect = "Allow"
    principals {
      type        = "Service"
      identifiers = ["ec2.amazonaws.com"]
    }
    actions = ["sts:AssumeRole"]
  }
}

# ジョブロール
resource "aws_iam_role" "job_role" {
  name = "${var.project_prefix}-job-role"
  assume_role_policy = data.aws_iam_policy_document.trust_ecs_tasks.json
}

# IAMポリシーデータ
data "aws_iam_policy_document" "job_role" {
  statement {
    effect = "Allow"
    actions = [
      "logs:CreateLogGroup",
      "logs:CreateLogStream",
      "logs:PutLogEvents"
    ]
    resources = ["*"]
  }

  statement {
    effect = "Allow"
    actions = [
      "s3:*",
      "s3-object-lambda:*"
    ]
    resources = [
      "arn:aws:s3:::${var.project_prefix}-${var.account_id}",
      "arn:aws:s3:::${var.project_prefix}-${var.account_id}/*"
    ]
  }
}

# IAMポリシー
resource "aws_iam_policy" "job_role" {
  name = "${var.project_prefix}-job-role-policy"
  policy = data.aws_iam_policy_document.job_role.json
}

# IAMポリシーのアタッチ
resource "aws_iam_role_policy_attachment" "job_role" {
  role       = aws_iam_role.job_role.name
  policy_arn = aws_iam_policy.job_role.arn
}


# インスタンスロール
resource "aws_iam_role" "instance_role" {
  name = "${var.project_prefix}-instance-role"
  assume_role_policy = data.aws_iam_policy_document.trust_ec2.json
  managed_policy_arns = [
    "arn:aws:iam::aws:policy/service-role/AmazonEC2ContainerServiceforEC2Role"
  ]
}

# インスタンスプロファイル
resource "aws_iam_instance_profile" "instance_profile" {
  name = "${var.project_prefix}-instance-role-profile"
  role = aws_iam_role.instance_role.name
}

# 信頼ポリシー
data "aws_iam_policy_document" "trust_batch" {
  statement {
    effect = "Allow"

    principals {
      type        = "Service"
      identifiers = ["batch.amazonaws.com"]
    }

    actions = ["sts:AssumeRole"]
  }
}

# サービスロール
resource "aws_iam_role" "batch_service_role" {
  name = "${var.project_prefix}-batch-service-role"
  assume_role_policy = data.aws_iam_policy_document.trust_batch.json
  managed_policy_arns = [
    "arn:aws:iam::aws:policy/service-role/AWSBatchServiceRole"
  ]
}

# 信頼ポリシー
data "aws_iam_policy_document" "trust_events" {
  statement {
    effect = "Allow"

    principals {
      type        = "Service"
      identifiers = ["events.amazonaws.com"]
    }
    actions = ["sts:AssumeRole"]
  }
}

# 実行ロール
resource "aws_iam_role" "events_execution_role" {
  name = "${var.project_prefix}-events-execution-role"
  assume_role_policy = data.aws_iam_policy_document.trust_events.json
}


# IAMポリシーデータ
data "aws_iam_policy_document" "events_execution_role" {
  statement {
    effect = "Allow"
    actions = [
      "batch:SubmitJob"
    ]
    resources = [
      "arn:aws:batch:ap-northeast-1:${var.account_id}:job/${var.project_prefix}-job",
      "arn:aws:batch:ap-northeast-1:${var.account_id}:job-definition/${var.project_prefix}-job-definition:*",
      "arn:aws:batch:ap-northeast-1:${var.account_id}:job-queue/${var.project_prefix}-job-queue"
    ]
  }
}

# IAMポリシー
resource "aws_iam_policy" "events_execution_role" {
  name = "${var.project_prefix}-events-execution-role-policy"
  policy = data.aws_iam_policy_document.events_execution_role.json
}

# IAMポリシーのアタッチ
resource "aws_iam_role_policy_attachment" "events_execution_role" {
  role       = aws_iam_role.events_execution_role.name
  policy_arn = aws_iam_policy.events_execution_role.arn
}

インスタンスロールについては、managed_policy_arnsで直接以下のポリシーを指定しています。

arn:aws:iam::aws:policy/service-role/AmazonEC2ContainerServiceforEC2Role

またインスタンスロールは信頼ポリシーもec2となりますので、以下も追加しています。

# 信頼ポリシー
data "aws_iam_policy_document" "trust_ec2" {
  statement {
    effect = "Allow"
    principals {
      type        = "Service"
      identifiers = ["ec2.amazonaws.com"]
    }
    actions = ["sts:AssumeRole"]
  }
}

手順

手順も前回と同様ですが再掲しておきます。

Terraformでリソースを全て構築

最初にリソースをすべて作成します。

# 作業フォルダ: terraform/environments/dev/
aws-vault exec {プロファイル名} -- terraform apply -var 'project_prefix={任意のプレフィックス}' # aws-vault経由で実行

Lambdaの時と異なり、ECRレポジトリにimageがpushされていなくてもBatchのジョブ定義は作成できますので、最初にすべて作ることが可能です。

イメージのビルド

まずdocker/lambda/.envというファイルを作成して、環境変数を入力しておきます。

PROJECT_PREFIX="{任意のプレフィックス}"

{任意のプレフィックス}は、後述のterraform applyの際に指定したものと合致するようにしておいてください。

その後は以下でビルドができます。

# 作業フォルダ: docker/lambda/
docker compose build

ECRへコンテナイメージをpush

push_ecr.shというスクリプトを準備していますのでそちらを実行してください。

.\push_ecr.sh {プロファイル名} {任意のプレフィックス}

{任意のプレフィックス}は、後述のterraform applyの際に指定したものと合致するようにしておいてください。

push_ecr.shの内容は以下です。

ProfileName=$1
ProjectPrefix=$2

REGION=$(aws --profile $ProfileName configure get region)
ACCOUNT_ID=$(aws --profile $ProfileName sts get-caller-identity --query 'Account' --output text )

REPOSITORY_NAME=$ProjectPrefix
ECR_BASE_URL="${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com"
ECR_IMAGE_URI="${ECR_BASE_URL}/${REPOSITORY_NAME}"

echo "ECR_BASE_URL: ${ECR_BASE_URL}"
echo "ECR_IMAGE_URI: ${ECR_IMAGE_URI}"

# ECRへのログイン
aws --profile $ProfileName ecr get-login-password --region ${REGION} | docker login --username AWS --password-stdin ${ECR_BASE_URL}

# tagの付け替え
docker tag "${REPOSITORY_NAME}:latest" "${ECR_IMAGE_URI}:latest"

# ECRへのpush
docker push "${ECR_IMAGE_URI}:latest"

今回からPowershellではなくGit Bashなどからshellを使う形にしています。

以上で構築の準備が完了しました。

動作確認

AWS CLIでjpgファイルをアップロードしてみます。

前回同様にasset/demo.jpgにサンプル画像を配置してありますので良ければお試しください。

（今回はオブジェクトキーは時刻情報が付与するようにしています）

aws s3 cp asset/demo.jpg s3://{バケット名}/input/$(date "+%Y%m%d-%H%M%S").jpg --profile {プロファイル名}

処理が終わると、output/に結果が配置されます。

aws s3 ls s3://{バケット名}/output/ --profile {プロファイル名}

# 2023-09-18 02:04:30          0
# 2023-09-18 12:32:38      78343 20230918-122754.jpg

マネジメントコンソールで確認するとジョブのログが以下のように確認できます。

インスタンスタイプにもよりそうですが、作成から開始まで4分30秒とFargateより時間が掛かっています。

処理した結果については前回と同様ですので割愛いたします。

まとめ

いかがでしたでしょうか。

今回はBach on EC2で機械学習のワークロードを構築する方法を見ていきました。

本記事が皆様のお役に立てば幸いです。

Terraformで構築する機械学習ワークロード（Batch on EC2編）

構成イメージ

動作環境

成果物

コードの説明

コンテナイメージのビルド

Pythonファイルについて

tfファイル

modules/batch/

modules/iam/

手順

Terraformでリソースを全て構築

イメージのビルド

ECRへコンテナイメージをpush

動作確認

まとめ

関連記事

主なカテゴリ

AWS

おすすめ

プロダクト

コンテンツ

お問い合わせ

運営会社