[アップデート] AWS PCS が HashiCorp Terraform でデプロイできるようになりました

terraform {
  required_providers {
    aws = {
      source  = "hashicorp/aws"
      version = "5.92.0"
    }
    awscc = {
      source  = "hashicorp/awscc"
      version = "1.34.0"
    }
  }
}

provider "aws" {
  # Configuration options
  region = "ap-northeast-1"
}

provider "awscc" {
  # Configuration options
  region = "ap-northeast-1"
}

data "aws_region" "this" {}
data "aws_caller_identity" "this" {}

locals {
  account_id = data.aws_caller_identity.this.account_id
  region     = data.aws_region.this.name
}

module "vpc" {
  source = "terraform-aws-modules/vpc/aws"

  name = "pcs"
  cidr = "10.0.0.0/16"

  azs             = ["${local.region}a", "${local.region}c"]
  private_subnets = ["10.0.1.0/24", "10.0.2.0/24"]
  public_subnets  = ["10.0.101.0/24", "10.0.102.0/24"]

  enable_nat_gateway = true
  enable_vpn_gateway = true
  single_nat_gateway = true
}

###########################################
# Node Group
###########################################
variable "my_ip" {
  description = "My IP address"
  default     = "0.0.0.0/0"
}

# クラスターのセキュリティグループ
resource "aws_security_group" "cluster" {
  name        = "pcs-cluster-sg"
  description = "pcs-cluster-sg"
  vpc_id      = module.vpc.vpc_id

  # 自己参照のインバウンドルール
  ingress {
    from_port   = 0
    to_port     = 0
    protocol    = "-1"
    self        = true
    description = "Allow all inbound from self"
  }

  ingress {
    from_port = 22
    to_port   = 22
    protocol  = "tcp"
    cidr_blocks = [
      var.my_ip
    ]
    description = "Allow all outbound to world"
  }

  # 自己参照のアウトバウンドルール
  egress {
    from_port   = 0
    to_port     = 0
    protocol    = "-1"
    self        = true
    description = "Allow all outbound to self"
  }

  # インターネットへのアウトバウンドルール
  egress {
    from_port   = 0
    to_port     = 0
    protocol    = "-1"
    cidr_blocks = ["0.0.0.0/0"]
    description = "Allow all outbound to world"
  }

  tags = {
    Name = "pcs-cluster-sg"
  }
}

###########################################
# Cluster
###########################################
resource "awscc_pcs_cluster" "this" {
  name = "pcs-cluster"

  networking = {
    subnet_ids         = [module.vpc.private_subnets[0]]
    security_group_ids = [aws_security_group.cluster.id]
  }

  scheduler = {
    type    = "SLURM"
    version = "24.05"
  }

  size = "SMALL"
}

###########################################
# Role for the cluster
###########################################
resource "aws_iam_role" "pcs_cluster" {
  name = "AWSPCS-cluster-role"
  path = "/"

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

resource "aws_iam_role_policy" "pcs_cluster" {
  name = "AWSPCS-cluster-policy"
  role = aws_iam_role.pcs_cluster.id

  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Effect   = "Allow"
        Action   = ["pcs:RegisterComputeNodeGroupInstance"]
        Resource = "*"
      }
    ]
  })
}

resource "aws_iam_role_policy_attachment" "pcs_cluster" {
  role       = aws_iam_role.pcs_cluster.name
  policy_arn = "arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore"
}

resource "aws_iam_instance_profile" "pcs_cluster" {
  name = aws_iam_role.pcs_cluster.name
  path = "/"
  role = aws_iam_role.pcs_cluster.name
}

resource "aws_launch_template" "pcs" {
  name        = "pcs-launch-template"
  description = "PCS cluster node launch template"

  # AMIとインスタンスタイプ
  # aws-pcs-sample_ami-amzn2-x86_64-slurm-24.05-2024-12-14T05-28-32.441Z at ap-northeast-1
  image_id      = "ami-0e18e980afc64cc20"
  instance_type = "c6i.xlarge"
  key_name      = "pcs-blog"

  # ネットワーク設定
  vpc_security_group_ids = [aws_security_group.cluster.id]

  # IAMインスタンスプロファイル（必要に応じて）
  iam_instance_profile {
    name = aws_iam_instance_profile.pcs_cluster.name
  }
}

resource "awscc_pcs_compute_node_group" "this" {
  cluster_id               = awscc_pcs_cluster.this.cluster_id
  iam_instance_profile_arn = aws_iam_instance_profile.pcs_cluster.arn
  subnet_ids               = [module.vpc.private_subnets[0]]
  custom_launch_template = {
    template_id = aws_launch_template.pcs.id
    version     = aws_launch_template.pcs.latest_version
  }
  scaling_configuration = {
    max_instance_count = 2
    min_instance_count = 0
  }
  instance_configs = [
    {
      instance_type = "c6i.xlarge"
    }
  ]
}

###########################################
# Queue
###########################################
resource "awscc_pcs_queue" "this" {
  cluster_id = awscc_pcs_cluster.this.cluster_id
  compute_node_group_configurations = [
    {
      compute_node_group_id = awscc_pcs_compute_node_group.this.compute_node_group_id
    }
  ]
}