簡介

Deepmind團(tuán)­隊(duì)于2021年7月公開了Alphafold2算法源代碼和相關(guān)論文，這一方案被認(rèn)為可解決困擾生物學(xué)長達(dá)半個世紀(jì)的蛋白質(zhì)折疊結(jié)構(gòu)預(yù)測難題，其準(zhǔn)確度高達(dá)92.4(百分制)。Alphafold2基于新穎的機(jī)器學(xué)習(xí)方法，以原子精度預(yù)測蛋白質(zhì)結(jié)構(gòu)，在大部分情況下表現(xiàn)出與實(shí)驗(yàn)相媲美的準(zhǔn)確性，且大大優(yōu)于其他方法。但截至目前，其官方的單機(jī)部署方式需要人工下載數(shù)據(jù)集并解壓、處理輸入輸出數(shù)據(jù)以及執(zhí)行任務(wù)腳本，無法實(shí)現(xiàn)彈性按需使用，不適用于大規(guī)模的蛋白質(zhì)結(jié)構(gòu)預(yù)測。

在本文的案例中，我們首先設(shè)計(jì)了基于API+HPC的任務(wù)調(diào)度邏輯，由Amazon Lambda作為無服務(wù)器后端，處理用戶提交的任務(wù)請求并通過Amazon SQS消息隊(duì)列服務(wù)對其進(jìn)行解耦；然后將任務(wù)請求提交到Amazon Batch調(diào)用GPU計(jì)算實(shí)例，并分析Alphafold2算法的源代碼結(jié)構(gòu)進(jìn)行分析，并將其優(yōu)化為Amazon Batch可使用的容器鏡像格式，存儲于Amazon ECR托管容器存儲庫中；Alphafold2算法所使用的蛋白質(zhì)數(shù)據(jù)庫文件均位于Amazon Fsx for Lustre高性能分布式文件系統(tǒng)中；用戶提交的任務(wù)過程信息使用Amazon Dynamodb作為NoSQL數(shù)據(jù)庫存儲，所有輸入的氨基酸序列文件和輸出的蛋白質(zhì)結(jié)構(gòu)預(yù)測文件均使用Amazon S3對象存儲；算法科學(xué)家可使用基于NICE DCV的高性能遠(yuǎn)程桌面，通過goofys協(xié)議掛載Amazon S3目錄，­直接在云上可視化分析Alphafold2生成的蛋白質(zhì)結(jié)構(gòu)預(yù)測文件。最后，針對以上部署過程，通過Amazon CDK打包實(shí)現(xiàn)基礎(chǔ)設(shè)施即代碼化，方便客戶一鍵部署。

注：本文所用代碼目前托管于：

https://github.com/wttat/alphafold

https://github.com/wttat/af2-batch-cdk

整體架構(gòu)

亞馬遜云科技服務(wù)需求

Amazon API Gateway，可以幫助開發(fā)人員輕松創(chuàng)建、發(fā)布、維護(hù)、監(jiān)控和保護(hù)任意規(guī)模的 API。

Amazon Lambda，幾乎可以為任何類型的應(yīng)用程序或后端服務(wù)運(yùn)行代碼，而且完全無需管理。

Amazon S3，對象存儲服務(wù)，提供行業(yè)領(lǐng)先的可擴(kuò)展性、數(shù)據(jù)可用性、安全性和性能。

Amazon SQS， 消息隊(duì)列服務(wù)，可以幫助分離和擴(kuò)展微服務(wù)、分布式系統(tǒng)和無服務(wù)器應(yīng)用程序。

Amazon DynamoDB，快速、靈活的 NoSQL 數(shù)據(jù)庫服務(wù)，可在任何規(guī)模下實(shí)現(xiàn)個位數(shù)毫秒級的性能。

Amazon Batch，任意規(guī)模完全受管的批處理調(diào)度系統(tǒng)，讓開發(fā)人員、科學(xué)家和工程師能夠輕松高效地運(yùn)行成千上萬個批處理計(jì)算作業(yè)。

Amazon ECR，完全托管的容器注冊表，在任何地方輕松存儲、管理、共享和部署容器映像和構(gòu)件。

Amazon SNS，完全托管的發(fā)布/訂閱消息收發(fā)、SMS、電子郵件和移動推送通知。

Amazon EventBridge，大規(guī)模構(gòu)建事件驅(qū)動型應(yīng)用程序，涵蓋亞馬遜云科技、現(xiàn)有系統(tǒng)或 SaaS 應(yīng)用程序

架構(gòu)說明

API 的設(shè)計(jì)

為向算法科學(xué)家屏蔽底層GPU算力調(diào)度過程，提供一個統(tǒng)一易用的Alphafold2任務(wù)調(diào)度入口，我們使用API Gateway創(chuàng)建RESTful API，對外暴露單一URL接口，實(shí)現(xiàn)對Alphafold2服務(wù)的應(yīng)用級封裝。API Gateway為云原生托管API管理工具，對用戶提交的任務(wù)查詢或調(diào)度請求鑒權(quán)，通過后將其按照配置的路由發(fā)送到不同的Amazon Lambda后端。Amazon Lambda函數(shù)對請求進(jìn)行解析，將任務(wù)信息寫入數(shù)據(jù)庫，并將任務(wù)參數(shù)投遞到HPC集群。共有四個Amazon Lambda函數(shù)與API Gateway直接或間接相關(guān)聯(lián)，包括身份鑒權(quán)函數(shù)、查詢請求處理函數(shù)、調(diào)度請求處理函數(shù)以及HPC集群調(diào)度函數(shù)，此處我們將調(diào)度請求處理函數(shù)與HPC集群調(diào)度函數(shù)通過SQS消息隊(duì)列服務(wù)進(jìn)行解耦，從而提高了API架構(gòu)的可用性。

API參數(shù)的所有接口如下，算法科學(xué)家可直接通過API的形式對Alphafold2任務(wù)進(jìn)行操作。

• GET / ：查詢所有任務(wù)信息。

• GET /{id} ：查詢單條任務(wù)詳細(xì)信息，包含最近的Cloudwatch log日志。

• POST / ：投遞任務(wù)，任務(wù)信息包含在HTTP的body中。

• DELETE /：刪除所有已完成或失敗的任務(wù)信息及其對應(yīng)的Amazon S3資源（如有）。

• DELETE /{id}：刪除單條所有已完成或失敗的任務(wù)信息及其對應(yīng)的S3資源（如有）。

• CANCEL /：取消所有正在運(yùn)行的任務(wù)。

• CANCEL /{id}：取消單條正在運(yùn)行的任務(wù)。

注：ANY方法對應(yīng)CANCEL方法。

云上 HPC 批處理平臺

本次方案的云上HPC批處理平臺共由四部分組成：

• Amazon S3作為云原生的對象存儲服務(wù)，擁有11個9的高持久性以及完善的安全和合規(guī)機(jī)制，因此我們使用Amazon S3持久化存儲Alphafold2任務(wù)的輸入輸出文件，保障所有數(shù)據(jù)文件在云上的安全。在Amazon S3桶中，共有input和output兩個文件夾，分別用于存放要預(yù)測的氨基酸序列文件和2.Alphafold2預(yù)測產(chǎn)生的蛋白質(zhì)結(jié)構(gòu)文件。用戶可通過多種方式方便的檢索、下載上述文件，同時(shí)可以選擇Amazon S3中的多種存儲類，降低云上存儲成本。

• 為實(shí)現(xiàn)對Alphafold2優(yōu)化鏡像的統(tǒng)一托管，我們使用Amazon ECR作為容器鏡像存儲倉庫，免去了用戶自建容器鏡像倉庫的繁瑣流程。我們對Alphafold2代碼做了一定的優(yōu)化，使其能夠接受Amazon Batch作業(yè)的參數(shù)和環(huán)境變量輸入，并且使其能夠使用Amazon S3作為文件存儲，并重新打包為Docker鏡像，Amazon ECR支持與Amazon Batch無縫集成，將打包后Docker鏡像上傳至Amazon ECR后，即可以直接在Batch的作業(yè)定義中指定存儲在Amazon ECR中的Alphafold2容器鏡像地址進(jìn)行拉取。

• Alphafold2每次預(yù)測蛋白質(zhì)結(jié)構(gòu)均需要通過Multiple Sequence Alignment (多重序列比對, MSA)，對已知蛋白質(zhì)結(jié)構(gòu)數(shù)據(jù)庫進(jìn)行掃描，目前數(shù)據(jù)集總?cè)萘考s為2.2T。為了實(shí)現(xiàn)數(shù)據(jù)集的高性能共享訪問，我們將其存儲在FSx for Lustre上，使其能夠被所有任務(wù)所同時(shí)訪問并高速讀取。Fsx for Lustre可提供百兆級別吞吐、百萬級別IOPS以及穩(wěn)定的亞豪秒延遲，從而滿足Alphafold2在內(nèi)的HPC計(jì)算工作負(fù)載。

• 為實(shí)現(xiàn)計(jì)算資源的彈性擴(kuò)展、降低整體計(jì)算和運(yùn)維成本，我們使用Amazon Batch搭建蛋白質(zhì)結(jié)構(gòu)預(yù)測HPC批處理集群，由Batch內(nèi)置的調(diào)度器根據(jù)用戶提交的蛋白質(zhì)結(jié)構(gòu)預(yù)測任務(wù)和資源需求自動化預(yù)置所需要的GPU計(jì)算資源。我們目前針對亞馬遜云科技可提供的GPU實(shí)例類型，選擇了p3.2xlarge、p3.8xlarge、p3.16xlarge以及p4d.24xlarge四種機(jī)型作為計(jì)算環(huán)境，用戶可根據(jù)單條預(yù)測任務(wù)使用的GPU數(shù)量、任務(wù)并行度和成本綜合考慮選擇投遞任務(wù)的作業(yè)隊(duì)列，而無需關(guān)心其底層實(shí)例的部署以及環(huán)境的搭建。所有計(jì)算環(huán)境均可配置為在任務(wù)結(jié)束后即自動關(guān)閉所有計(jì)算實(shí)例節(jié)約成本。

任務(wù)信息數(shù)據(jù)存儲

Amazon DynamoDB作為云原生的無服務(wù)器NoSQL數(shù)據(jù)庫，無需預(yù)置服務(wù)器，按使用量付費(fèi)，同時(shí)擁有一致性的毫秒級訪問性能以及近乎無限的吞吐量和存儲空間。本次方案使用DynamoDB存儲Alphafold2過程中產(chǎn)生的任務(wù)信息，滿足用戶并發(fā)提交或查詢?nèi)蝿?wù)時(shí)的性能需求，同時(shí)優(yōu)化成本，簡化數(shù)據(jù)庫配置流程。在處理任務(wù)提交的Amazon Lambda函數(shù)中，使用了UUID算法隨機(jī)生成每條任務(wù)的id，作為DynamoDB分區(qū)鍵，提高任務(wù)查詢效率。

消息與結(jié)果通知

單機(jī)部署Alphafold2時(shí)，其運(yùn)行過程的所有日志均位于Docker中，如果運(yùn)行環(huán)節(jié)中出現(xiàn)任何問題，均需要手動登陸實(shí)例后，通過docker log等命令進(jìn)行查看；同時(shí)任務(wù)完成后缺乏相關(guān)通知推送，無法及時(shí)獲取其運(yùn)行狀態(tài)，造成計(jì)算資源的浪費(fèi)。

我們使用Amazon S3的事件通知機(jī)制以及Amazon EventBridge服務(wù)，自動獲取Amazon S3中壓縮文件的上傳文件事件以及Amazon Batch任務(wù)失敗事件——分別對應(yīng)Alphafold2任務(wù)的成功和報(bào)錯，自動觸發(fā)不同的Amazon Lambda函數(shù)處理相關(guān)事件，更新任務(wù)信息數(shù)據(jù)庫，最后通過Amazon SNS——即亞馬遜云科技云上托管的消息收發(fā)服務(wù)，通過郵箱或者短信等方式向算法科學(xué)家推送相關(guān)信息，實(shí)現(xiàn)Alphafold2作業(yè)信息的即時(shí)通知。消息通知中包含任務(wù)信息、所需時(shí)間、下載鏈接以及Amazon CLI命令。算法科學(xué)家無需登陸亞馬遜云科技控制臺，即可直接獲取Alphafold2任務(wù)結(jié)果的下載方式。

便捷可視化終端

為滿足算法科學(xué)家遠(yuǎn)程分析蛋白質(zhì)結(jié)構(gòu)的需求，我們使用NICE DCV遠(yuǎn)程桌面作為蛋白質(zhì)結(jié)構(gòu)便捷可視化終端，啟動后可以直接通過瀏覽器或NICE DCV客戶端，使用圖形化界面遠(yuǎn)程訪問。NICE DCV基于優(yōu)化的NICE DCV協(xié)議，根據(jù)網(wǎng)絡(luò)情況自適應(yīng)調(diào)整傳輸帶寬，同時(shí)可利用預(yù)裝了NVIDIA T4顯卡的Amazon G4實(shí)例，提高圖形密集性軟件的渲染性能，從而使算法科學(xué)家遠(yuǎn)程通過如Pymol或VMD等蛋白質(zhì)三維結(jié)構(gòu)可視化軟件進(jìn)行分析。

部署方法

初始化 Amazon CDK 環(huán)境

cdk bootstrap aws://{ACCOUNTID}/{REGION}

通過 Amazon CDK生成

Amazon CloudFormation 模板

cdk synth

一鍵部署所有資源

cdk deploy –all

需要手動確認(rèn)，輸入y回車即可，中間會輸出一些過程信息，并且在郵箱中點(diǎn)擊確認(rèn)SNS訂閱。

部署完成

約需要四個小時(shí)將所有數(shù)據(jù)下載并解壓到FSx for Lustre，解壓完成后，會收到電子郵件通知，此時(shí)即可根據(jù)使用說明，開始投遞Alphafold2任務(wù)

算法科學(xué)家使用過程

完成上述方案構(gòu)建后，算法科學(xué)家即可開始使用Alphafold2算法預(yù)測蛋白質(zhì)結(jié)構(gòu)。

總共可分別四個步驟：

1.上傳氨基酸序列文件

通過Amazon CLI或者控制臺將氨基酸序列文件上傳到Amazon S3存儲桶的input文件夾。

2.提交任務(wù)或查詢?nèi)蝿?wù)進(jìn)度

通過終端或者HTTP調(diào)試軟件操作任務(wù)。本文以Postman調(diào)試軟件為例，演示如何基于API提交、查詢?nèi)蝿?wù)。取消和刪除任務(wù)操作類似，可參考API Gateway中的路由設(shè)置進(jìn)行提交。

首先編輯任務(wù)請求JSON文件，例：

可以參考:

https://github.com/wttat/af2-batch-cdk/blob/main/command.json

其中，

fasta 為本次預(yù)測的蛋白質(zhì)名稱，由用戶自定義，必選；

file_name為氨基酸序列名稱，必須與S3存儲桶中的文件名稱對應(yīng)，必選；model_names為所使用的Alphafold2模型，共有model1-model5五個可選，缺省值為五個模型全部使用；

preset為使用的預(yù)設(shè)配置，full對應(yīng)完整數(shù)據(jù)集，reduced對應(yīng)壓縮數(shù)據(jù)集，CASP14對應(yīng)deepmind團(tuán)隊(duì)參賽時(shí)所使用的配置，缺省值為full；

max_template_date為掃描已知數(shù)據(jù)庫中蛋白質(zhì)結(jié)構(gòu)文件的截止日期，缺省值為2020-05-14；

que為batch使用的作業(yè)隊(duì)列，low對應(yīng)batch的計(jì)算環(huán)境為p3.2xlarge實(shí)例，mid對應(yīng)計(jì)算環(huán)境為p3.8xlarge實(shí)例，high對應(yīng)計(jì)算環(huán)境為p3.16xlarge實(shí)例，p4對應(yīng)計(jì)算環(huán)境為p4d.24xlarge實(shí)例，必選；

comment為用戶自定義的注釋語句，可選，缺省值為空；

gpu為使用GPU數(shù)量，并且成比例對應(yīng)vCPU和內(nèi)存，可選，缺省值為1。例：p3.8xlarge實(shí)例中共有4塊NVIDIA Tesla V100顯卡，32vCPU和244G內(nèi)存，默認(rèn)情況下，即gpu為1時(shí)，每個任務(wù)分配8vCPU、61G內(nèi)存和一塊V100顯卡，可同時(shí)運(yùn)行四個任務(wù)，如果gpu設(shè)置為2時(shí)，該任務(wù)分配到16vCPU、122G內(nèi)存和兩塊V100顯卡，此時(shí)p3.8xlarge實(shí)例剩下兩塊顯卡，可分配給兩個gpu為1的任務(wù)或者一個gpu為2的任務(wù)。

首先配置Headers，鍵為Authorization，值為用戶自己定義的密鑰。

（1）提交任務(wù)

將HTTP方法設(shè)置為post，URL填入API Gateway的URL，Body選擇raw，格式為JSON，粘貼任務(wù)請求JSON文件，點(diǎn)擊Send即可提交任務(wù)，后臺即會根據(jù)選擇的作業(yè)隊(duì)列自動拉起對應(yīng)的計(jì)算資源運(yùn)行預(yù)測任務(wù)。API Gateway會返回每個預(yù)測任務(wù)對應(yīng)的ID。

（2）查詢?nèi)蝿?wù)

將方法修改為GET，即可查詢所有任務(wù)信息。

若想查詢單條任務(wù)詳細(xì)信息，則在URL后面加上對應(yīng)任務(wù)的ID，可通過查詢所有任務(wù)信息得到，也可以在提交任務(wù)后直接復(fù)制返回的ID。

其中job info為任務(wù)相關(guān)參數(shù)，job status為任務(wù)當(dāng)前狀態(tài)，若有日志產(chǎn)生，則會自動抓取最新的Amazon CloudWatch日志并返回。

3.Email自動化發(fā)送任務(wù)結(jié)果

任務(wù)完成后，會收到郵件通知提醒，告知其對應(yīng)預(yù)測任務(wù)已完成。郵件正文中，包括完成預(yù)測所需時(shí)間，即可根據(jù)使用的機(jī)器類型計(jì)算成本，以及所有文件的下載鏈接HTML文件，同時(shí)提供了Amazon cli命令行，用戶可自行選擇下載方式。

打開下載鏈接HTML后，即可根據(jù)需求自行選擇下載整體壓縮包或者pdb文件進(jìn)行分析。

4.可視化分析

以pymol為例，連接到NICE DCV上后，安裝并運(yùn)行pymol：

wget https://pymol.org/installers/PyMOL-2.5.2_293-Linux-x86_64-py37.tar.bz2
tar -jxf PyMOL-2.5.2_293-Linux-x86_64-py37.tar.bz2
./pymol/pymol

*左滑查看更多

打開桌面上Amazon S3文件夾，找到需要分析的蛋白質(zhì)文件夾，將pdb文件拖拽到pymol中打開即可直接在云上分析蛋白質(zhì)結(jié)構(gòu)。

也可以自行安裝VMD進(jìn)行分析。

成本估算

（僅供參考）

我們基于P3實(shí)例和P4實(shí)例，測試了不同序列長度的氨基酸序列，經(jīng)過Alphafold2算法預(yù)測其蛋白質(zhì)結(jié)構(gòu)所需時(shí)間，并且根據(jù)對應(yīng)區(qū)域GPU實(shí)例單價(jià)，按比例換算為每條任務(wù)所需成本?？梢钥吹剑弘S著序列長度的增加，所需時(shí)間快速增加，主要是由于GPU推理時(shí)間占比顯著增加；在序列長度超過1000時(shí)，一塊V100會提示顯存不足的問題，需要兩塊V100顯卡才能跑完完整任務(wù)；在序列長度越長的情況下，P4實(shí)例掛載的A100顯卡帶來的性能提升越明顯。

中國寧夏區(qū)域(cn-northwest-1)：

美國弗吉尼亞北部區(qū)域(us-east-1)：

注：FSx for Lustre費(fèi)用未考慮在內(nèi)，需要按月收取。

優(yōu)化方向

目前Alphafold2官方源代碼還在不斷優(yōu)化迭代之中，還存在許多可以性能提升的方向。如：

• 前期掃描數(shù)據(jù)集的MSA過程只消耗CPU資源，并不占用GPU，因此可以據(jù)此將CPU和GPU運(yùn)算過程接耦，拆分為兩個子任務(wù)，降低成本；

• 通常情況下，我們只需要評分最好的模型，而Alphafold2算法目前的機(jī)制默認(rèn)為完整推理五個模型，可以將其修改為只預(yù)測評分最高的模型。

• MSA中多個模塊可以CPU并行運(yùn)行，提升掃描速度。

用戶可以根據(jù)自己的需要，對其進(jìn)行二次開發(fā)，自行修改并構(gòu)建Alaphfold2鏡像，只需將Amazon Batch作業(yè)定義中的容器鏡像地址替換為用戶指定的鏡像地址即可。

總結(jié)

我們首先基于Amazon Batch的作業(yè)需求，容器化改造了Alphafold2鏡像，使其能夠接受來自Amazon Batch標(biāo)準(zhǔn)化的輸入?yún)?shù)，并且將輸出文件直接上傳到Amazon S3上。接下來，使用根據(jù)負(fù)載可彈性擴(kuò)展的GPU計(jì)算集群和高性能共享存儲，可滿足大規(guī)模高并發(fā)的預(yù)測需求。并且，基于Amazon Batch提供完善的集群資源調(diào)度機(jī)制，降低計(jì)算資源成本以及IT部門運(yùn)維復(fù)雜度，實(shí)現(xiàn)秒級的成本監(jiān)控。算法科學(xué)家基于Amazon API GW提交/查詢/刪除任務(wù)，無需接觸亞馬遜云科技資源，同于由于所有操作均基于API實(shí)現(xiàn)，因此可方便集成至現(xiàn)有工作流中。其次，基于Amazon SNS提供作業(yè)狀態(tài)提醒機(jī)制，針對任務(wù)的成功和失敗均有郵件提醒，并在任務(wù)成功的通知郵件集成了蛋白質(zhì)結(jié)構(gòu)文件的下載地址?；贏mazon CDK實(shí)現(xiàn)方案的靈活快速部署，用戶可在半小時(shí)內(nèi)將完整架構(gòu)部署于自己的亞馬遜云科技資源中，四個小時(shí)內(nèi)預(yù)測服務(wù)即可上線使用。