自定义 AMI

AMI 是如何构建的

我们将了解预部署的 x86 AMI 是如何使用 EC2 Image Builder 构建的。理解这一过程非常重要,因为我们需要为生产环境的 PCS 集群构建和维护自己的 AMI。

EC2 Image Builder 概述

EC2 Image Builder 通过流水线自动化 AMI 的创建:

  1. Component(组件) — 一个 YAML 文档,定义安装步骤(安装什么软件以及如何安装)
  2. Recipe(配方) — 将基础 AMI 与一个或多个组件结合,并指定根卷大小
  3. Infrastructure Configuration(基础设施配置) — 定义构建实例类型、子网、安全组和 IAM 配置文件
  4. Distribution Configuration(分发配置) — 控制 AMI 分发到哪里(哪些区域)
  5. Image(镜像) — 触发构建并生成 AMI ID

Image Builder 启动一个临时的 EC2 实例,运行所有组件步骤,根据结果创建 AMI,然后终止该实例。构建日志存储在 S3 中以供调试。

x86 AMI 的组件步骤

预部署的 AMI 是基于 Amazon Linux 2023 x86_64 构建的,包含以下步骤:

步骤 安装内容 原因
UpdateOS 最新的操作系统补丁 安全基线
TurnOffSRSO 禁用推测性返回栈溢出缓解措施 为 HPC 优化性能
PCS Agent AWS PCS 节点注册代理 PCS 必需 — 处理节点生命周期
Slurm 作业调度器 PCS 必需 — 必须与集群版本匹配
EFA Drivers Elastic Fabric Adapter 低延迟的节点间网络;EFA 安装程序还捆绑了 OpenMPI 5 + environment-modules(module load openmpi5
Lustre Client FSx for Lustre 挂载支持 高性能共享存储
EFS Utils Amazon EFS 挂载支持 共享主目录
CloudWatch Agent 指标和日志收集 监控
Development Tools gcc, g++, make, htop, vim, jq 编译 MPI 程序,交互式使用
Enroot + Pyxis 容器运行时 + Slurm 集成 运行容器化工作负载
SSSD/LDAP Client 用户认证 通过 LDAP 进行多用户访问
Enroot 和 Pyxis 是 NVIDIA 开源的一对工具,用来在 Slurm 集群上跑容器

Enroot:HPC 版的容器运行时

普通 Docker 在 HPC 集群上不好用:需要 root 守护进程(安全隐患)、隔离太重(网络/文件系统隔离反而妨碍 MPI 和共享存储)。Enroot 是 NVIDIA 做的轻量替代

以下是 Image Builder 组件 YAML 中的关键步骤。完整的组件位于 pcs-workshop-infra.yaml 的 CloudFormation 模板中。

PCS Agent 安装:

- name: GetPCSAgent
  action: WebDownload
  inputs:
    - source: https://aws-pcs-repo-<REGION>.s3.<REGION>.amazonaws.com/aws-pcs-agent/aws-pcs-agent-v1-latest.tar.gz
      destination: /tmp/aws-pcs-agent-v1-latest.tar.gz

- name: InstallPCSAgent
  action: ExecuteBash
  inputs:
    commands:
      - |
        cd /tmp && tar xf aws-pcs-agent-v1-latest.tar.gz
        cd aws-pcs-agent && ./installer.sh

Slurm 安装(独立于 PCS 代理):

- name: GetSlurm
  action: WebDownload
  inputs:
    - source: https://aws-pcs-repo-<REGION>.s3.<REGION>.amazonaws.com/aws-pcs-slurm/aws-pcs-slurm-${SlurmVersion}-installer-latest.tar.gz
      destination: /tmp/aws-pcs-slurm-${SlurmVersion}-installer-latest.tar.gz

- name: InstallSlurm
  action: ExecuteBash
  inputs:
    commands:
      - |
        cd /tmp && tar xf aws-pcs-slurm-${SlurmVersion}-installer-latest.tar.gz
        cd aws-pcs-slurm-${SlurmVersion}-installer && ./installer.sh -y

Enroot + Pyxis 安装:

- name: InstallPyxisEnroot
  action: ExecuteBash
  inputs:
    commands:
      - |
        arch=$(uname -m)
        dnf install -y https://github.com/NVIDIA/enroot/releases/download/v4.0.1/enroot-4.0.1-1.el8.${arch}.rpm
        dnf install -y https://github.com/NVIDIA/enroot/releases/download/v4.0.1/enroot+caps-4.0.1-1.el8.${arch}.rpm
        # Clone and compile Pyxis against PCS Slurm headers
        git clone https://github.com/NVIDIA/pyxis
        cd pyxis
        make CFLAGS=\"-fPIC -I/opt/aws/pcs/scheduler/slurm-${SlurmVersion}/include\"
        cp spank_pyxis.so /opt/aws/pcs/scheduler/slurm-${SlurmVersion}/lib/slurm/
        # Register as Slurm SPANK plugin
        echo \"required /opt/aws/pcs/scheduler/slurm-${SlurmVersion}/lib/slurm/spank_pyxis.so\" > \\
          /etc/aws/pcs/scheduler/slurm-${SlurmVersion}/plugstack.conf.d/pyxis.conf

用于容器支持的 Enroot 和 Pyxis

该 AMI 包含 Enroot 和 Pyxis,用于通过 Slurm 运行容器化工作负载:

  • Enroot — 一个为 HPC 设计的轻量级容器运行时。它导入容器镜像(来自 Docker Hub、ECR 或本地文件),并将其作为非特权用户命名空间运行。与 Docker 不同,Enroot 不需要守护进程,并且能自然地与 Slurm 集成。
  • Pyxis — 一个 Slurm SPANK 插件,让我们可以通过 srun --container-image 直接运行容器。它在底层使用 Enroot,但提供了无缝的 Slurm 体验 — 无需手动导入或管理容器镜像。

Pyxis 是针对 PCS Slurm 头文件从源代码编译的,并注册为 SPANK 插件。这是在 AMI 构建期间完成的,而不是在运行时。

GPU AMI

对于 GPU 工作负载,我们应该从 AWS Deep Learning AMI(DLAMI)开始,而不是 Amazon Linux 2023。DLAMI 预装了 GPU 驱动程序、CUDA 和 ML 框架。然后我们在此基础上添加相同的 PCS 组件(代理、Slurm、EFA、SSSD、Enroot+Pyxis)。

构建 Arm64 AMI 并添加 Graviton 节点

我们将使用 EC2 Image Builder 构建一个 Arm64 AMI,然后将 Graviton 计算节点添加到我们的集群中。这与构建预部署的 x86 AMI 所使用的流程相同 —— 现在我们要为不同的架构亲自完成这一过程。

为什么选择 Graviton?

AWS Graviton 处理器基于 Arm 架构,可为许多工作负载提供更优的性价比。并非所有 HPC 代码都针对 Arm 进行了优化,但许多科学应用程序(尤其是那些用 C/Fortran 编写并使用标准 MPI 的应用程序)只需极少的改动即可在 Graviton 上编译并良好运行。

同时提供两种架构可让用户为任务选择合适的工具 —— hpc7a 适用于紧耦合的 EFA 工作负载,Graviton 适用于面向吞吐量或成本敏感的工作负载。

构建 Arm64 AMI

static/cfn/pcs-arm64-ami.yaml 提供了一个 CloudFormation 模板,用于为 Arm64 AMI 创建所有 Image Builder 资源(组件、配方、基础设施配置、分发和镜像)。该组件安装与 x86 AMI 相同的软件\:PCS agent、Slurm :param{key="slurmVersion”}、EFA 驱动程序、Lustre 客户端、SSSD/LDAP、Enroot+Pyxis 以及开发工具。

我们仍然通过前面练习连接到登录节点。在运行以下命令之前,请返回到我们的 CloudShell 标签页 —— 这些命令会部署 CloudFormation 堆栈,应该从 CloudShell 运行,而不是从集群内部运行。\n\n反复输入 exit,直到我们回到 CloudShell 提示符(我们可能需要退出 user1、退出 root 并退出 SSM 会话)。我们的 CloudShell 提示符看起来像 [cloudshell-user@ip-xx-xx-xx-xx ~]$

首先,从基础设施堆栈中收集参数值\:

export INFRA_STACK="pcs-workshop-infra"
export AWS_DEFAULT_REGION="us-east-2"

export LDAP_URI=$(aws cloudformation describe-stacks \
  --stack-name $INFRA_STACK \
  --query "Stacks[0].Outputs[?OutputKey=='LdapServerUri'].OutputValue" \
  --output text --region $AWS_DEFAULT_REGION)
export LDAP_IP=$(echo $LDAP_URI | sed 's|ldap://||')

export LDAP_PWD=$(aws secretsmanager get-secret-value \
  --secret-id pcs-workshop-ldap-password-$INFRA_STACK \
  --query SecretString --output text --region $AWS_DEFAULT_REGION)

export PUBLIC_SUBNET=$(aws cloudformation describe-stacks \
  --stack-name $INFRA_STACK \
  --query "Stacks[0].Outputs[?OutputKey=='PublicSubnetId1'].OutputValue" \
  --output text --region $AWS_DEFAULT_REGION)

export PUBLIC_SG=$(aws cloudformation describe-stacks \
  --stack-name $INFRA_STACK \
  --query "Stacks[0].Outputs[?OutputKey=='PublicSecurityGroupId'].OutputValue" \
  --output text --region $AWS_DEFAULT_REGION)

export IB_PROFILE_NAME=$(aws cloudformation describe-stacks \
  --stack-name $INFRA_STACK \
  --query "Stacks[0].Outputs[?OutputKey=='ImageBuilderInstanceProfileName'].OutputValue" \
  --output text --region $AWS_DEFAULT_REGION)

export S3_BUCKET=$(aws cloudformation describe-stacks \
  --stack-name $INFRA_STACK \
  --query "Stacks[0].Outputs[?OutputKey=='ImageBuilderLogBucketName'].OutputValue" \
  --output text --region $AWS_DEFAULT_REGION)

echo "LDAP IP: $LDAP_IP"
echo "Subnet: $PUBLIC_SUBNET"
echo "SG: $PUBLIC_SG"
echo "IB Profile: $IB_PROFILE_NAME"
echo "S3 Bucket: $S3_BUCKET"

我们可以在此处下载并查看该模板\::link[下载 pcs-arm64-ami.yaml]{href=”:assetUrl{path=’/cfn/pcs-arm64-ami.yaml’}” download="pcs-arm64-ami.yaml”}

创建参数文件并部署堆栈\:

cat > arm64-params.json << PARAMSEOF
[
  {"ParameterKey": "LdapServerIp", "ParameterValue": "$LDAP_IP"},
  {"ParameterKey": "LdapPassword", "ParameterValue": "$LDAP_PWD"},
  {"ParameterKey": "BuildSubnetId", "ParameterValue": "$PUBLIC_SUBNET"},
  {"ParameterKey": "BuildSecurityGroupId", "ParameterValue": "$PUBLIC_SG"},
  {"ParameterKey": "ImageBuilderInstanceProfileName", "ParameterValue": "$IB_PROFILE_NAME"},
  {"ParameterKey": "LogBucketName", "ParameterValue": "$S3_BUCKET"}
]
PARAMSEOF

curl ':assetUrl{path="/cfn/pcs-arm64-ami.yaml"}' --output pcs-arm64-ami.yaml

aws cloudformation create-stack \
  --stack-name pcs-arm64-ami \
  --template-body file://pcs-arm64-ami.yaml \
  --parameters file://arm64-params.json \
  --region $AWS_DEFAULT_REGION

AMI 构建需要 30-45 分钟。Image Builder 会启动一个 Graviton 实例(c7g.xlarge),运行所有组件步骤,创建 AMI,然后终止该实例。

监控构建:

watch -n 30 "aws cloudformation describe-stacks \
  --stack-name pcs-arm64-ami \
  --query 'Stacks[0].StackStatus' \
  --output text \
  --region us-east-2"

当它显示 CREATE_COMPLETE 时,按 Ctrl+C

获取 Arm64 AMI ID

export ARM_AMI=$(aws cloudformation describe-stacks \
  --stack-name pcs-arm64-ami \
  --query "Stacks[0].Outputs[?OutputKey=='Arm64AmiId'].OutputValue" \
  --output text --region $AWS_DEFAULT_REGION)
echo "Arm64 AMI: $ARM_AMI"

创建 Arm64 启动模板

Arm64 启动模板比 hpc7a 模板更简单 —— 没有 EFA 网络,只有文件系统挂载\:

export PRIVATE_SG=$(aws cloudformation describe-stacks \
  --stack-name $INFRA_STACK \
  --query "Stacks[0].Outputs[?OutputKey=='PrivateSecurityGroupId'].OutputValue" \
  --output text --region $AWS_DEFAULT_REGION)

export EFS_ID=$(aws cloudformation describe-stacks \
  --stack-name $INFRA_STACK \
  --query "Stacks[0].Outputs[?OutputKey=='EfsFileSystemId'].OutputValue" \
  --output text --region $AWS_DEFAULT_REGION)

export FSX_ID=$(aws cloudformation describe-stacks \
  --stack-name $INFRA_STACK \
  --query "Stacks[0].Outputs[?OutputKey=='FsxLustreFileSystemId'].OutputValue" \
  --output text --region $AWS_DEFAULT_REGION)

export FSX_MOUNT=$(aws cloudformation describe-stacks \
  --stack-name $INFRA_STACK \
  --query "Stacks[0].Outputs[?OutputKey=='FsxLustreMountName'].OutputValue" \
  --output text --region $AWS_DEFAULT_REGION)

export FSX_DNS="${FSX_ID}.fsx.${AWS_DEFAULT_REGION}.amazonaws.com"
cat > arm64-userdata.txt << USERDATA
MIME-Version: 1.0
Content-Type: multipart/mixed; boundary="==//=="

--==//==
Content-Type: text/x-shellscript; charset="us-ascii"
MIME-Version: 1.0

#!/bin/bash

echo "${EFS_ID}.efs.${AWS_DEFAULT_REGION}.amazonaws.com:/ /home nfs4 nfsvers=4.1,rsize=1048576,wsize=1048576,hard,timeo=600,retrans=2,_netdev 0 0" >> /etc/fstab
mount /home

echo "${FSX_DNS}@tcp:/${FSX_MOUNT} /fsx lustre defaults,_netdev,flock,user_xattr,noatime 0 0" >> /etc/fstab
mkdir -p /fsx
chmod a+rwx /fsx
mount /fsx
chmod 777 /fsx

# Configure Enroot for container support
# Cache on Lustre (shared across nodes), everything else on EBS (needs overlayfs xattr support)
cat > /etc/enroot/enroot.conf << 'ENROOTCFG'
ENROOT_RUNTIME_PATH /var/tmp/enroot/runtime
ENROOT_CONFIG_PATH  /var/tmp/enroot/config
ENROOT_CACHE_PATH   /fsx/enroot/cache
ENROOT_DATA_PATH    /var/tmp/enroot/data
ENROOT_TEMP_PATH    /var/tmp/enroot/tmp
ENROOTCFG
mkdir -p /var/tmp/enroot/runtime /var/tmp/enroot/config /var/tmp/enroot/data /var/tmp/enroot/tmp /fsx/enroot/cache
chmod 1777 /var/tmp/enroot /var/tmp/enroot/runtime /var/tmp/enroot/config /var/tmp/enroot/data /var/tmp/enroot/tmp
chmod 1777 /fsx/enroot/cache

# Remove NVIDIA container hook (CPU-only workshop)
rm -f /etc/enroot/hooks.d/98-nvidia.sh

--==//==
USERDATA
为我们即将创建的启动模板设置 EC2 密钥对。Workshop Studio 会在我们的账户中自动创建 ws-default-keypair,因此下面的默认值对于 Workshop Studio 运行是正确的。如果我们自行部署了本 workshop 并使用了不同的密钥对名称,请在运行此代码块之前覆盖 KEY_NAME。该变量不能为空 —— 空的 KeyName 会生成一个在 PCS 启动实例时无法通过验证的启动模板。
export KEY_NAME="${KEY_NAME:-ws-default-keypair}"

创建启动模板\:

aws ec2 create-launch-template \
  --launch-template-name arm64-compute-lt \
  --launch-template-data "{
    \"KeyName\": \"${KEY_NAME}\",
    \"MetadataOptions\": {
      \"HttpEndpoint\": \"enabled\",
      \"HttpPutResponseHopLimit\": 2,
      \"HttpTokens\": \"required\"
    },
    \"Monitoring\": {\"Enabled\": true},
    \"NetworkInterfaces\": [
      {
        \"DeviceIndex\": 0,
        \"NetworkCardIndex\": 0,
        \"Groups\": [\"${PRIVATE_SG}\"]
      }
    ],
    \"UserData\": \"$(base64 -w 0 arm64-userdata.txt 2>/dev/null || base64 -i arm64-userdata.txt | tr -d '\n')\"
  }" \
  --region $AWS_DEFAULT_REGION

验证启动模板已创建\:

aws ec2 describe-launch-templates \
  --launch-template-names arm64-compute-lt \
  --query "LaunchTemplates[0].{Name:LaunchTemplateName,Id:LaunchTemplateId,Version:LatestVersionNumber}" \
  --output table \
  --region $AWS_DEFAULT_REGION

我们应该看到一行内容,包含 Name: arm64-compute-ltVersion: 1

创建 Graviton 计算节点组和队列

export INSTANCE_PROFILE_ARN=$(aws cloudformation describe-stacks \
  --stack-name $INFRA_STACK \
  --query "Stacks[0].Outputs[?OutputKey=='PcsInstanceProfileArn'].OutputValue" \
  --output text --region $AWS_DEFAULT_REGION)

export ARM_LT_ID=$(aws ec2 describe-launch-templates \
  --launch-template-names arm64-compute-lt \
  --query "LaunchTemplates[0].LaunchTemplateId" \
  --output text --region $AWS_DEFAULT_REGION)

export HPC_SUBNET=$(aws cloudformation describe-stacks \
  --stack-name $INFRA_STACK \
  --query "Stacks[0].Outputs[?OutputKey=='HpcSubnetId'].OutputValue" \
  --output text --region $AWS_DEFAULT_REGION)

aws pcs create-compute-node-group \
  --cluster-identifier hpc-cluster \
  --compute-node-group-name graviton \
  --ami-id $ARM_AMI \
  --subnet-ids $HPC_SUBNET \
  --purchase-option ONDEMAND \
  --custom-launch-template id=$ARM_LT_ID,version=1 \
  --iam-instance-profile-arn $INSTANCE_PROFILE_ARN \
  --scaling-configuration minInstanceCount=0,maxInstanceCount=2 \
  --instance-configs '[{"instanceType":"c8g.16xlarge"}]' \
  --region $AWS_DEFAULT_REGION

等待计算节点组变为活动状态\:

watch -n 10 "aws pcs get-compute-node-group \
  --cluster-identifier hpc-cluster \
  --compute-node-group-identifier graviton \
  --query 'computeNodeGroup.status' \
  --output text \
  --region us-east-2"

创建 arm 队列\:

export ARM_CNG_ID=$(aws pcs get-compute-node-group \
  --cluster-identifier hpc-cluster \
  --compute-node-group-identifier graviton \
  --query "computeNodeGroup.id" \
  --output text --region $AWS_DEFAULT_REGION)

aws pcs create-queue \
  --cluster-identifier hpc-cluster \
  --queue-name arm \
  --compute-node-group-configurations "computeNodeGroupId=$ARM_CNG_ID" \
  --region $AWS_DEFAULT_REGION

从登录节点验证并提交测试作业

通过 SSM 进入登录节点,然后切换到 user1,以便 Slurm 命令在一个可通过 LDAP 解析、且在 EFS 上拥有主目录的账户下运行\:

sudo su - user1

检查两个队列都可见\:

sinfo -a

我们应该看到两个分区\:

PARTITION AVAIL  TIMELIMIT  NODES  STATE NODELIST
x86          up   infinite      2  idle~ hpc7a-[1-2]
arm          up   infinite      2  idle~ graviton-[1-2]

向 Graviton 提交测试作业\:

sbatch -p arm --wrap="hostname && uname -m && echo 'Hello from Graviton'" \
  -o /fsx/arm-test.%j.out

uname -m 命令将显示 aarch64,确认作业在 Arm64 实例上运行。等待作业完成并检查\:

watch -n 5 squeue
cat /fsx/arm-test.*.out

我们应该看到类似如下的输出\:

graviton-1
aarch64
Hello from Graviton

我们的集群现在拥有两种架构。对 hpc7a EFA 工作负载使用 -p x86,对 Graviton 工作负载使用 -p arm

AMI 最佳实践

构建和维护 AMI 是运营 PCS 集群的关键部分。以下是最重要的实践。

从最小化开始,逐步添加

构建我们的第一个 AMI 时,只包含必需的组件(PCS agent、Slurm、EFA、SSSD)。用一个简单的作业进行测试。然后一次添加一层软件——container runtime、应用程序库、监控代理。如果出现问题,我们可以确切地知道是哪次添加导致的。

为我们的 AMI 设置版本

使用一致的命名方案,包含日期和更改内容:

pcs-workshop-x86-2026-03-03-base
pcs-workshop-x86-2026-03-10-enroot
pcs-workshop-arm64-2026-03-03-base

Image Builder 可以在分发配置中使用 {{imagebuilder:buildDate}} 自动附加构建日期。

使用 Image Builder 管道自动化

不要手动构建 AMI。使用带有定时构建(每周或在依赖项更新时)的 Image Builder 管道。管道会重新构建 AMI,然后我们更新计算节点组以使用新版本。

https://github.com/aws-samples/sample-parallel-computing-service/tree/main/modules/ami Terraform 仓库(引用 main 分支以获取最新版本,或固定到发布标签以实现可重现性)包含按每周计划构建的管道定义。

保持 Slurm 版本同步

我们的 AMI 中的 Slurm 版本必须与集群的 Slurm 版本完全匹配。如果我们将集群升级到新的 Slurm 版本,则必须先使用匹配的版本重新构建所有 AMI,然后再更新计算节点组。

在生产环境前进行测试

在将新 AMI 推广到生产队列之前,请始终使用小型计算节点组进行测试。创建一个 maxInstanceCount=1 的测试节点组,提交作业,验证一切正常,然后更新生产节点组。

安全补丁

AMI 是快照——它们在创建后不会接收更新。安排定期重建以获取操作系统安全补丁。Image Builder 组件中的 UpdateOS 步骤会在每次构建时自动处理此事。

在生产环境中,像对待软件发布一样对待 AMI 管理——为它们设置版本、测试它们,并自动化构建管道。