亚马逊AWS官方博客
利用 AWS Batch 来为容器化负载调用海量云端算力
概述
Batch是AWS托管的一个批量计算服务,用户可以通过它运行任意规模的容器化工作负载,目前已经广泛应用于药物研发、基因分析、建模仿真、金融模拟等高性能计算(HPC)的场景。在HPC的场景,一个计算作业可能会用到几千核甚至上万核CPU的算力,那我们如何通过Batch来调用海量的云端算力呢?主要分两种情况,松耦合应用和紧耦合应用。
对于松耦合应用,我们可以通过Batch的数组作业来完成,数组作业是共享通用参数 (如任务定义、vCPU 和内存) 的作业,它以一系列相关但独立的基本作业的形式运行,这些作业可以跨多个主机分布并且同时运行。数组作业通过每台主机的环境变量AWS_BATCH_JOB_ARRAY_INDEX为每个作业从通用参数获取不同的参数值来运行独立的作业,比如在蛋白质药物研发的结构预测场景中,输入文件是一个文本文件,每一行包含一个蛋白质序列, 数组作业的每个子作业处理同一个输入文件的不同行,通过这种方式在云端启动大量的计算实例来达到同时调用成千上万千核来进行蛋白质序列分析的目的。数组作业的具体使用方式,在官方Batch文档中有一个示例可供参考,本篇就不再赘述。
对于紧耦合应用,我们可以通过Batch的多节点并行作业来完成。多节点并行作业能够跨多个EC2实例运行单个作业,Batch 的多节点并行作业与支持基于 IP 的节点间通信的任何框架均可兼容,例如 Apache MXNet、TensorFlow、Caffe2 或消息传递接口 (MPI)。本篇接下来通过一个简单的基于openmpi的程序来演示如何使用Batch的多节点并行作业来调用海量计算资源。
先决条件
您需要一个AWS帐户来完成本演练,其它条件包括(本篇不再描述其具体使用方式):
- 开发环境中完成AWS CLI v2安装和配置,本篇使用us-east-1区域
- 完成VPC的设置,包括私有子网和NAT的配置
- 创建一个EFS用作共享文件系统,用于存放程序脚本和计算的中间数据
- 熟悉Batch计算环境、作业队列、任务定义的配置
- 开发环境中需要安装openmpi用于后续测试代码的编译
- 熟悉进程管理程序supervisor的使用
应用部署
容器化
在 AWS Batch 上运行多节点并行作业,应用程序代码必须包含进行分布式通信所需的框架和库,并且完成容器化镜像的制作。为多节点并行作业构建容器镜像的几个关键点是:
- 设置容器之间可以通过无密码ssh连接
- 下载并编译安装openmpi
- 设置supervisor在Docker容器启动时运行ssh服务,并运行run-mpi.sh
在开发环境中构建容器镜像,Dockerfile如下:
在ECR中创建镜像仓库mpi-demo,构建镜像后推送到该仓库(123456789012替换成您自己的AWS账号ID),确保开发测试机具有推送ECR镜像的权限。
任务脚本
AWS Batch多节点并行作业跨节点的容器之间使用 Amazon ECS awsvpc 网络模式进行通信,在作业运行时,ECS容器从后端接收环境变量,这些变量可用于确定哪个运行容器是主容器,哪个是子容器,以下是使用的4个环境变量:
- AWS_BATCH_JOB_MAIN_NODE_INDEX -此变量为作业的主节点的索引号,应用程序代码可以将 AWS_BATCH_JOB_MAIN_NODE_INDEX 与单个节点上的 AWS_BATCH_JOB_NODE_INDEX 进行比较,以确定它是否为主节点
- AWS_BATCH_JOB_MAIN_NODE_PRIVATE_IPV4_ADDRESS -此变量为作业的主节点的私有 IPv4 地址,仅在子节点中存在,子节点的应用程序代码可以使用此地址与主节点进行通信
- AWS_BATCH_JOB_NODE_INDEX -此变量为节点的唯一索引号,从 0 开始
- AWS_BATCH_JOB_NUM_NODES -此变量为多节点并行作业请求的节点数
以下代码是一个MPI同步脚本run-mpi.sh,容器启动后执行该脚本,子节点上的容器会调用report_to_master()函数把自身的ip和CPU核数报告给主节点,主节点在所有子节点初始化完成之后通过mpirun来调用计算程序mpi_hello_world
supervisor是基于Python开发的一套通用的进程管理程序,能将一个普通的命令行进程以Daemon方式后台运行,并监控进程状态,异常退出时能自动重启,创建以下配置文件supervisord.conf:
容器启动时运行脚本entry-point.sh,该脚本会通过supervisor在容器启动时运行ssh服务,并运行run-mpi.sh
openmpi 测试程序
在开发环境创建如下程序mpi_hello_world.c(详细可参考此教程)
通过以下命令编译,生成可执行文件mpi_hello_world
mpicc -o mpi_hello_world mpi_hello_world.c
EFS
开发环境挂载EFS文件系统,在文件系统中创建mpi-demo目录,保存entry-point.sh,run-mpi.sh,mpi_hello_world和supervisord.conf 到该目录下,并且把entry-point.sh和run-mpi.sh修改为可执行权限(这些文件也可以通过Dockerfile打包进容器镜像,放在容器外面主要是为了调试及后期修改方便,不用每次都重新构建镜像)
置放群组(可选)
对于多节点并行作业,为了提高节点间通信的网络性能,可考虑在单个可用区中创建集群 置放群组,并与计算资源进行关联。在EC2控制台创建一个策略为集群的置放群组:
Batch设置
可通过Batch控制台或者命令行工具来进行Batch环境的配置,以下均通过命令行工具来演示如何进行Batch计算环境、作业队列、任务定义的设置(123456789012替换成您自己的AWS账号id)
计算环境
创建配置文件env-mpi.json:
subnets:配置了NAT网关路由的私有子网
securityGroupIds:默认安全组id
placementGroup:置放群组名称
instanceTypes:使用的实例类型
tags:EC2标签
创建计算环境
作业队列
配置文件q-mpi.json:
computeEnvironment:上一步创建的计算环境的arn
创建作业队列
任务定义
配置文件jd-mpi.json
numNodes:节点数量,本次演示使用10个节点
mainNode:主节点索引
targetNodes:节点索引范围
volumes和mountPoints:将efs的mpi-demo目录挂载到容器内的/data目录下
environment:环境变量
注册任务定义
测试
提交任务
可以看到任务提交后启动了一个具有10个节点的并行任务
在EC2控制台也能看到启动了10台EC实例
作业执行完成之后,最终的作业状态(SUCCEEDED 或 FAILED)由主节点的最终作业状态决定
在CloudWatch查看主节点的日志
结论
在这篇文章中,我们演示了在仅有一次任务提交的情况下,如何通过AWS Batch的多节点并行作业为紧耦合的mpi工作负载调用大量的CPU计算资源。在实际的场景中,您可以将mpi_hello_world改为自己的程序或者与其它分析软件结合使用。除了CPU之外,您也可以使用各种类型的GPU资源,无需担心算力问题,在AWS上尽情运行您的HPC工作负载吧!
参考文档
AWS Batch用户指南:https://docs.aws.amazon.com/zh_cn/batch/latest/userguide/what-is-batch.html
aws-mnpbatch-template:https://github.com/aws-samples/aws-mnpbatch-template
MPI教程:https://github.com/aws-samples/aws-mnpbatch-template
在AWS Batch中构建具有多节点并行作业的紧耦合分子动力学工作流:https://aws.amazon.com/cn/blogs/compute/building-a-tightly-coupled-molecular-dynamics-workflow-with-multi-node-parallel-jobs-in-aws-batch/