基于 AWS Step Functions 和 AWS Batch 服务构建宏基因分析工作负载

概述

本文旨在介绍一个在 AWS 上构建宏基因分析工作负载的方案，该方案对常用的一些宏基因分析软件完成了容器化，并以批量计算服务 AWS Batch 为核心，构建容器化的宏基因分析集群和自动化工作流。本方案已经在多个客户的实际宏基因分析业务中得到应用，如土壤微生物/病毒宏基因分析，海洋微生物宏基因分析等。根据客户的实际业务场景，本方案也可以扩展到通用的基因分析场景。

宏基因分析业务

宏基因分析是一种利用生物信息学技术来研究微生物群落物种和功能的方法。它通过对微生物群落的全部 DNA 进行高通量测序，然后将测序序列与公共数据库进行比对或从头组装出微生物基因组，从而识别微生物群落的物种和功能基因，以了解微生物的多样性、种群结构、功能活性、相互协作关系以及与环境之间的关系。宏基因分析的一般流程包括以下步骤：

数据预处理：对原始序列进行去接头、质量剪切以及去除污染等优化处理。
组装和基因预测：使用优质序列进行拼接组装和基因预测，得到各样本预测得到的基因集合。
基因集合的处理：将各样本预测得到的基因集合进行去冗余，得到非冗余基因集。
物种注释：将非冗余基因集序列与参考数据库比对，得到物种信息，分析得到整个基因集的物种信息。
功能注释：将非冗余基因集序列与功能数据库比对，获得基因集的功能注释信息，如 COG 功能注释、KEGG 功能注释等。
特殊数据库注释：将非冗余基因集序列与特殊数据库比对，如 CAZy 数据库（碳水化合物活性酶）和 CARD 数据库（抗生素抗性基因），获得基因集特殊的功能信息。

宏基因组业务在环境科学、生态学、微生物学等领域具有广泛的应用价值，可以帮助我们更好地理解微生物的生态功能和环境适应性，为环境保护、生态修复和农业生产等提供科学依据。

基于 AWS 服务的宏基因分析方案

本文的方案基于一个常见的宏基因分析流程所设计，流程如下图所示：

不同的功能用到了以下软件：

质控：fastqc，multiqc，trimmomatic
组装/分箱/提炼：megahit，metawrap
预测：prodigal
去冗余：dRep，cd-hit， mummer
注释：gtdbtk，metaphlan， diamond
丰度计算：coverm， checkm，minimap2
分析及处理：seqkit，salmon

本方案对宏基因分析流程中所用到的软件做了容器化适配，基于 AWS 的计算和存储等服务构建了高性能计算集群，支持大规模任务的批量分析；此外还构建了自动化工作流，将分析流程解耦，用不同计算实例适配不同任务类型，旨在优化云上的分析成本和减少运维操作。AWS 服务如何支撑具体业务如下图所示：

业务场景

本方案主要适合两个业务场景：

大规模批量任务分析：有大量的样本需要进行批量处理，比如有几千组样本要进行质控处理
自动化流程：自动化运行某个（或多个）分析流程中的每一个任务，比如质控→组装→分箱提炼→注释

大规模批量任务分析

本业务场景以批量计算服务 AWS Batch 为核心，构建了一个高性能计算集群

支持同时进行数千、数万任务的分析
方案内置了质控、分箱提炼、注释等分析脚本
提供内置任务提交程序的 Amazon SageMaker Notebook 作为用户工作台，用户只需要准备输入文件和文件列表，在 Amazon SageMaker Notebook 中通过 Python 提交任务即可
计算资源根据实际的任务状态动态生成和销毁

该业务场景的 AWS 架构图如下所示：

自动化流程

在高性能集群之上，基于可视化的工作流服务 AWS Step Functions + Batch，构建了一个分析流程（质控→组装→分箱提炼→注释），自动化运行分析流程中的每一个任务：

根据流程中不同任务对计算资源的不同需求做了解耦，用不同计算实例适配不同任务类型，以达到优化成本的目的，如分别使用不同的实例类型 c6g.4xlarge(Arm-based)、c6i.24xlarge、c6i.8xlarge、c6i.12xlarge 来运行质控、组装、分箱、注释任务
用户只需一次性地启动流程，无需手动执行每个步骤
通过订阅邮件来接收流程执行状态：成功或失败

该业务场景的 AWS 架构图如下所示：

主要用到了以下 AWS 服务：

AWS Batch：批量计算，自动计算资源调度
Amazon EC2：适合各种计算需求的实例类型，Arm 架构的 Graviton 实例显著提升性价比
Amazon S3：数据存储，通过智能分层优化成本
Amazon ECR：容器镜像仓库
Amazon DynamoDB：存储任务状态信息
Amazon EFS：共享文件系统，存储任务运行脚本
Amazon SageMaker Notebook：用户工作台，任务投递
AWS Step Functions：自动化工作流
Amazon SNS：流程执行状态邮件通知

方案部署

本方案通过 CDK 进行安装部署，篇幅所限，详细操作步骤请参考此链接。

配置工作

完成部署之后，还需要一些其它配置工作，包括容器镜像，任务脚本和参考数据库的配置。

Amazon SageMaker Notebook 工作台

本方案部署了一个 Amazon SageMaker Notebook 作为工作台方便用户操作，登录 Amazon SageMaker Notebook 控制台，打开笔记本实例的 JupyterLab 进行操作：

metagenomic-analysis-on-aws/deployment 目录下为项目部署的 cdk 项目
metagenomic-analysis-on-aws/notebook 目录下为内置的三个笔记本，分别为 container-images.ipynb（内置镜像拉取），job-submit-batch.ipynb（大规模批量任务），job-submit-stepfunctions.ipynb（自动化流程），以及输入文件列表示例 sample-info.xlsx
metagenomic-analysis-on-aws/scripts 目录下为内置的 4 个任务脚本

容器镜像

根据业务流程，将本方案用到的软件分为 4 个容器镜像进行构建，每个镜像包含的软件如下：

qc: fastqc, multiqc, trimmomatic, parallel, awscli
metawrap: metawrap, mount-s3, awscli
annotation: seqkit, seqtk, drep, coverm, cd-hit, minimap2, checkm, mummer, metaphlan, prodigal, gtdbtk, parallel, mamba, mount-s3, awscli
diamond: diamond

可以使用本方案预置的镜像，也可以参考此链接的 Dockerfile 分别构建容器镜像，完成之后推送到 ECR 的镜像仓库保存，简便起见，4 个镜像放在同一个镜像仓库中，通过标签来区分：

任务脚本

本方案有 4 个脚本，对应自动化流程中的 4 个步骤，参考此链接复制任务脚本到 EFS 中：

qc.sh：通过 fastqc/multiqc/trimmomatic 实现质量控制
metawrap-assembly.sh：通过 metawrap assembly 实现组装功能，目前支持单样本和 3 样本组装
metawrap-binning.sh：通过 metawrap 的 binning 和 bin_finement 实现分箱提炼，目前支持单样本和 3 样本的分箱提炼
annotation.sh：通过 dRep/coverm/gtdbtk/metaphlan/prodigal/cd-hit/seqkit/salmon 等实现去冗余/丰度计算/物种注释/基因预测等功能

参考数据库

所用到的软件 checkm、gtdbtk、metaphlan 依赖于参考数据库，参考软件官方文档及以下链接下载对应的参考数据库并解压保存到 Amazon S3 上：

也可联系本文作者，通过 request-payer 的方式快速拉取这些参考数据库到您自己的 AWS 账号。

任务测试

本节通过 2 个示例演示如何运行大规模批量任务和自动化分析流程。

大规模批量任务

参考此链接进行大规模批量任务投递，首先上传源数据到 Amazon S3 存储桶中的 sources 目录下：

编辑一个 excel 文件 sample-info.xlsx，每一行都包含需要处理的一个样本信息，上传该文件到 notebook/job-submit-batch.ipynb 对应的目录下，打开 job-submit-batch.ipynb，运行 python 提交任务，针对每一个样本会提交一个任务，如下图：

查看 Amazon Batch 中的任务：

查看 Amazon DynamoDB 表中的任务执行状态信息：

查看 Amazon S3 桶中的结果文件：

自动化流程

参考此链接提交自动化流程任务，同样上传数据到 Amazon S3，上传 sample-info.xlsx 到 notebook 对应目录，打开 notebook/job-submit-stepfunctions.ipynb，运行 python 提交任务，查看 Amazon Step Functions 中状态机执行状态，可以看到流程已经启动：

任务完成之后，能够收到邮件通知：

查看 Amazon Step Functions 中状态机执行状态，所有流程运行完成：

选择一个流程查看每个执行步骤详细情况：

查看 Amazon DynamoDB 表中对应样本每个步骤的任务执行时间和状态，如果某个步骤的执行时间很少，一般就是有问题，可以去查看 Amazon Batch 任务中的日志进行分析：

查看 Amazon S3 桶中的每个步骤的结果文件：

总结

本方案对宏基因分析流程中所用到的常用软件做了容器化适配，基于 AWS 的计算和存储等服务构建了高性能计算集群，旨在帮助用户最小化运维管理，能够以一种简单的、成本优化的方式运行大规模的宏基因分析任务和优化云上任务流程。主要针对两个分析场景：大规模批量单任务分析和自动化工作流（也可二者结合），将分析流程解耦，用不同计算实例适配不同任务类型，优化云上的分析成本和减少运维操作。方案提供 CDK 一键部署脚本，方便客户构建测试环境；方案提供了完整的 Workshop（参考文档 1），可帮助客户熟悉方案并构建自己的自定义分析环境。

参考文档

在 AWS 上构建大规模宏基因分析流程：

https://catalog.us-east-1.prod.workshops.aws/workshops/59cc73fe-fd92-4445-b048-deb54d5c38f4/zh-CN

基于 Graviton2 处理器构建容器化基因分析工作负载：

https://thinkwithwp.com/cn/blogs/china/construction-of-containerized-gene-analysis-workload-based-on-graviton2-processor/

亚马逊AWS官方博客