POC for Batch

batch.tf

data "template_file" "container_properties" {
  template = file("templates/container_properties.yaml")
  vars = {
    bucket_name =  var.bucket_name
  }
}

data "aws_ssm_parameter" "image_id" {
  name = "/aws/service/ecs/optimized-ami/amazon-linux-2/recommended/image_id"
}

resource "aws_launch_template" "batch_launch_template" {

  block_device_mappings {
    device_name = "/dev/xvda"

    ebs {
      volume_size = 100
      encrypted   = true
    }
  }

  image_id = data.aws_ssm_parameter.image_id.value
}

resource "aws_batch_compute_environment" "spot" {
  compute_environment_name = "spot-fleet"

  compute_resources {
    allocation_strategy = "SPOT_CAPACITY_OPTIMIZED"

    instance_role = aws_iam_instance_profile.ecs_instance_role.arn
    instance_type = [
      "optimal",
    ]


    max_vcpus     = 256
    min_vcpus     = 0
    desired_vcpus = 4

    security_group_ids = [
      aws_security_group.this.id,
    ]

    subnets = data.aws_subnet_ids.private.ids

    type = "SPOT"

    launch_template {
      launch_template_id = aws_launch_template.batch_launch_template.id
      version            = "$Latest"
    }
  }

  service_role = aws_iam_role.aws_batch_service_role.arn
  type         = "MANAGED"
}

resource "aws_batch_job_queue" "this" {
  name     = "queue"
  state    = "ENABLED"
  priority = "1"
  compute_environments = [
    aws_batch_compute_environment.spot.arn,
  ]
}

resource "aws_batch_job_definition" "example" {
  name = "batch-job-definition"
  type = "container"
  container_properties = jsonencode(yamldecode(
    data.template_file.container_properties.rendered
  ))
}

container_properties.yaml

command:
  - aws
  - s3
  - cp
  - /etc/motd
  - Ref::BUCKET_NAME
image: "<AWS ACCOUNT ID>.dkr.ecr.<REGION>.amazonaws.com/<AWSCLI REPO>"
memory: 128
vcpus: 1

security_groups.tf

resource "aws_security_group" "this" {

  name   = "batch_compute_env"
  vpc_id = data.aws_vpc.selected.id
  
  # egress only to VPC + S3 buckets for ECR pulling/test bucket
  egress {
    from_port = 443
    to_port   = 443
    protocol  = "tcp"
    cidr_blocks = [
      data.aws_vpc.selected.cidr_block,
      "54.231.0.0/17",
      "52.216.0.0/15",
      "3.5.16.0/21",
      "52.92.16.0/20",
      "3.5.0.0/20",
    ]
  }
}

service_execution_policy.tf

data "aws_iam_policy_document" "uploader_policy" {
  statement {
    actions = [
      "s3:PutObject",
    ]

    resources = [
      "arn:aws:s3:::${var.bucket_name}/",
    ]
  }
}

resource "aws_iam_policy" "ecs_batch_uploader" {
  name   = "ecs_batch_uploader"
  path   = "/"
  policy = data.aws_iam_policy_document.uploader_policy.json
}

resource "aws_iam_role_policy_attachment" "ecs_uploader_policy" {
  role       = aws_iam_role.ecs_instance_role.name
  policy_arn = aws_iam_policy.ecs_batch_uploader.arn
}

service_iam.tf

# Batch Service Role
resource "aws_iam_role" "aws_batch_service_role" {
  name = "aws_batch_service_role"

  assume_role_policy = <<EOF
{
        "Version": "2012-10-17",
        "Statement": [{
            "Action": "sts:AssumeRole",
            "Effect": "Allow",
            "Principal": {
                "Service": "batch.amazonaws.com"
            }
        }]
}
EOF
}

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

# Batch Instance Role
resource "aws_iam_role" "ecs_instance_role" {
  name = "ecs_instance_role"

  assume_role_policy = <<EOF
{
    "Version": "2012-10-17",
    "Statement": [{
        "Action": "sts:AssumeRole",
        "Effect": "Allow",
        "Principal": {
            "Service": "ec2.amazonaws.com"
        }
    }]
}
EOF
}

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

resource "aws_iam_role_policy_attachment" "ecs_ssm_policy" {
  role       = aws_iam_role.ecs_instance_role.name
  policy_arn = "arn:aws:iam::aws:policy/service-role/AmazonEC2RoleforSSM"
}

resource "aws_iam_instance_profile" "ecs_instance_role" {
  name = "ecs_instance_role"
  role = aws_iam_role.ecs_instance_role.name
}

vpc_data.tf

data "aws_vpc" "selected" {
  id = var.vpc_id
}

data "aws_subnet_ids" "private" {
  vpc_id = data.aws_vpc.selected.id

  tags = {
    subnet = "private"
  }
}