SageMaker HyperPod で FSx Lustre をマウントしてみた

###################################################
# Security Group for Lustre File System
###################################################
resource "aws_security_group" "lustre" {
  name        = "${local.prefix}-lustre-sg"
  vpc_id      = module.vpc.vpc_id
  description = "${local.prefix}-hyperpod-sg"

  tags = {
    Name = "${local.prefix}-lustre-sg"
  }
}
# Ingress
resource "aws_vpc_security_group_ingress_rule" "lustre_all_traffic_self" {
  security_group_id            = aws_security_group.lustre.id
  referenced_security_group_id = aws_security_group.lustre.id
  ip_protocol                  = "-1"
}

resource "aws_vpc_security_group_ingress_rule" "lustre_all_traffic_hyperpod" {
  security_group_id            = aws_security_group.lustre.id
  referenced_security_group_id = aws_security_group.hyperpod.id
  ip_protocol                  = "-1"
}

# Egress
resource "aws_vpc_security_group_egress_rule" "lustre_all_traffic_self" {
  security_group_id            = aws_security_group.lustre.id
  referenced_security_group_id = aws_security_group.lustre.id
  ip_protocol                  = "-1"
}

resource "aws_vpc_security_group_egress_rule" "lustre_all_traffic_hyperpod" {
  security_group_id            = aws_security_group.lustre.id
  referenced_security_group_id = aws_security_group.hyperpod.id
  ip_protocol                  = "-1"
}

###################################################
# Security Group for SageMaker HyperPod Cluster
###################################################
resource "aws_security_group" "hyperpod" {
  name        = "${local.prefix}-hyperpod-sg"
  vpc_id      = module.vpc.vpc_id
  description = "${local.prefix}-hyperpod-sg"

  tags = {
    Name = "${local.prefix}-hyperpod-sg"
  }
}

# Ingress
resource "aws_vpc_security_group_ingress_rule" "hyperpod_all_traffic_self" {
  security_group_id            = aws_security_group.hyperpod.id
  referenced_security_group_id = aws_security_group.hyperpod.id
  ip_protocol                  = "-1"
}

+ resource "aws_vpc_security_group_ingress_rule" "hyperpod_all_traffic_lustre" {
+   security_group_id            = aws_security_group.hyperpod.id
+   referenced_security_group_id = aws_security_group.lustre.id
+   ip_protocol                  = "-1"
+ }

# Egress
resource "aws_vpc_security_group_egress_rule" "hyperpod_all_traffic_self" {
  security_group_id            = aws_security_group.hyperpod.id
  referenced_security_group_id = aws_security_group.hyperpod.id
  ip_protocol                  = "-1"
}

+ resource "aws_vpc_security_group_egress_rule" "hyperpod_all_traffic_lustre" {
+   security_group_id            = aws_security_group.hyperpod.id
+   referenced_security_group_id = aws_security_group.lustre.id
+   ip_protocol                  = "-1"
+ }

resource "aws_vpc_security_group_egress_rule" "hyperpod_all_traffic_ipv4" {
  security_group_id = aws_security_group.hyperpod.id
  cidr_ipv4         = "0.0.0.0/0"
  ip_protocol       = "-1"
}

###################################################
# Data Repository for Lustre File System
###################################################
resource "aws_s3_bucket" "data_repository" {
  bucket        = "${local.prefix}-hyperpod-data-${local.account_id}"
  force_destroy = true
}

resource "aws_s3_bucket_public_access_block" "data_repository" {
  bucket                  = aws_s3_bucket.data_repository.bucket
  block_public_acls       = true
  block_public_policy     = true
  ignore_public_acls      = true
  restrict_public_buckets = true
}

resource "aws_s3_bucket_ownership_controls" "data_repository" {
  bucket = aws_s3_bucket.data_repository.bucket
  rule {
    object_ownership = "BucketOwnerPreferred"
  }
}

###################################################
# Lustre File System
###################################################
resource "aws_fsx_lustre_file_system" "this" {
  storage_type                = "SSD"
  file_system_type_version    = "2.15"
  storage_capacity            = 1200
  security_group_ids          = [aws_security_group.lustre.id]
  subnet_ids                  = [module.vpc.private_subnets[0]]
  data_compression_type       = "LZ4"
  deployment_type             = "PERSISTENT_2"
  per_unit_storage_throughput = 250

  metadata_configuration {
    mode = "AUTOMATIC"
  }
}

resource "aws_fsx_data_repository_association" "this" {
  file_system_id       = aws_fsx_lustre_file_system.this.id
  data_repository_path = "s3://${aws_s3_bucket.data_repository.bucket}"
  file_system_path     = "/"

  s3 {
    auto_export_policy {
      events = ["NEW", "CHANGED", "DELETED"]
    }

    auto_import_policy {
      events = ["NEW", "CHANGED", "DELETED"]
    }
  }
}

###################################################
# Lifecycle Script Objects for SageMaker HyperPod Cluster
###################################################
resource "aws_s3_object" "life_cycle_scripts" {
  for_each = fileset("./config/", "**")
  bucket   = aws_s3_bucket.life_cycle_scripts.bucket
  key      = "config/${each.value}"
  source   = "./config/${each.value}"
}

resource "aws_s3_object" "life_cycle_scripts_provisioning_parameters" {
  bucket = aws_s3_bucket.life_cycle_scripts.bucket
  key    = "config/provisioning_parameters.json"

  content = jsonencode(
    {
      "version"          = "1.0.0"
      "workload_manager" = "slurm"
      "controller_group" = "controller-group"
      "worker_groups"    = []
+      "fsx_dns_name"     = aws_fsx_lustre_file_system.this.dns_name
+      "fsx_mountname"    = aws_fsx_lustre_file_system.this.mount_name
    }
  )
}

###################################################
# SageMaker HyperPod Cluster
###################################################
resource "awscc_sagemaker_cluster" "this" {
  cluster_name = "${local.prefix}-cluster"

  instance_groups = [
    {
      execution_role      = aws_iam_role.hyperpod.arn
      instance_count      = 1
      instance_group_name = "controller-group"
      instance_type       = "ml.t3.medium"
      life_cycle_config = {
        source_s3_uri = "s3://${aws_s3_bucket.life_cycle_scripts.id}/config/"
        on_create     = "on_create.sh"
      }
    }
  ]

  vpc_config = {
    security_group_ids = [aws_security_group.hyperpod.id]
    subnets            = [module.vpc.private_subnets[0]]
  }

  depends_on = [
+    aws_iam_role_policy_attachment.hyperpod_vpc,
+    aws_iam_role_policy_attachment.hyperpod_s3,
+    aws_vpc_endpoint.s3,
+    aws_s3_object.life_cycle_scripts_provisioning_parameters
  ]
}

aws ssm start-session \
  --target sagemaker-cluster:je7plekuheau_controller-group-i-02cb085c72ba19ca6

[cloudshell-user@ip-10-132-77-98 ~]$ aws ssm start-session \
> --target sagemaker-cluster:je7plekuheau_controller-group-i-02cb085c72ba19ca6

Starting session with SessionId: cm-takakuni.shinnosuke-fsfoepusyfrug2ch99lopbchtu
# whoami
root

# df -h
Filesystem               Size  Used Avail Use% Mounted on
/dev/root                 97G   51G   46G  53% /
devtmpfs                 1.9G     0  1.9G   0% /dev
tmpfs                    1.9G     0  1.9G   0% /dev/shm
tmpfs                    386M  1.2M  385M   1% /run
tmpfs                    5.0M     0  5.0M   0% /run/lock
tmpfs                    1.9G     0  1.9G   0% /sys/fs/cgroup
/dev/loop0                64M   64M     0 100% /snap/core20/2379
/dev/loop2                75M   75M     0 100% /snap/core22/1621
/dev/loop1                64M   64M     0 100% /snap/core20/2434
/dev/loop5                92M   92M     0 100% /snap/lxd/29619
/dev/loop3                74M   74M     0 100% /snap/core22/1663
/dev/loop4                39M   39M     0 100% /snap/snapd/21759
/dev/nvme0n1p15          105M  6.1M   99M   6% /boot/efi
tmpfs                    386M   24K  386M   1% /run/user/128
10.0.1.35@tcp:/vtmdtbev  1.2T   38M  1.2T   1% /fsx
#

# cd /fsx
# ls
enroot  ubuntu
# echo hello! from hyperpod node! > hello.txt
# cat hello.txt
hello! from hyperpod node!
#

hello! from S3!

# ls
enroot  hello.txt  hello_from_s3.txt  ubuntu
# cat hello_from_s3.txt
hello! from S3!#