一鍵構(gòu)建云上高可用蛋白質(zhì)結(jié)構(gòu)預測平臺
簡介
Deepmind團隊于2021年7月公開了Alphafold2算法源代碼和相關論文，這一方案被認為可解決困擾生物學長達半個世紀的蛋白質(zhì)折疊結(jié)構(gòu)預測難題，其準確度高達92.4(百分制)。Alphafold2基于新穎的機器學習方法，以原子精度預測蛋白質(zhì)結(jié)構(gòu)，在大部分情況下表現(xiàn)出與實驗相媲美的準確性，且大大優(yōu)于其他方法。但截至目前，其官方的單機部署方式需要人工下載數(shù)據(jù)集并解壓、處理輸入輸出數(shù)據(jù)以及執(zhí)行任務腳本，無法實現(xiàn)彈性按需使用，不適用于大規(guī)模的蛋白質(zhì)結(jié)構(gòu)預測。
在本文的案例中，我們首先設計了基于API+HPC的任務調(diào)度邏輯，由Amazon Lambda作為無服務器后端，處理用戶提交的任務請求并通過Amazon SQS消息隊列服務對其進行解耦；然后將任務請求提交到Amazon Batch調(diào)用GPU計算實例，并分析Alphafold2算法的源代碼結(jié)構(gòu)進行分析，并將其優(yōu)化為Amazon Batch可使用的容器鏡像格式，存儲于Amazon ECR托管容器存儲庫中；Alphafold2算法所使用的蛋白質(zhì)數(shù)據(jù)庫文件均位于Amazon Fsx for Lustre高性能分布式文件系統(tǒng)中；用戶提交的任務過程信息使用Amazon Dynamodb作為NoSQL數(shù)據(jù)庫存儲，所有輸入的氨基酸序列文件和輸出的蛋白質(zhì)結(jié)構(gòu)預測文件均使用Amazon S3對象存儲；算法科學家可使用基于NICE DCV的高性能遠程桌面，通過goofys協(xié)議掛載Amazon S3目錄，直接在云上可視化分析Alphafold2生成的蛋白質(zhì)結(jié)構(gòu)預測文件。最后，針對以上部署過程，通過Amazon CDK打包實現(xiàn)基礎設施即代碼化，方便客戶一鍵部署。
注：本文所用代碼目前托管于：
https://github.com/wttat/alphafold
https://github.com/wttat/af2-batch-cdk
整體架構(gòu)
亞馬遜云科技服務需求
Amazon API Gateway，可以幫助開發(fā)人員輕松創(chuàng)建、發(fā)布、維護、監(jiān)控和保護任意規(guī)模的 API。
Amazon Lambda，幾乎可以為任何類型的應用程序或后端服務運行代碼，而且完全無需管理。
Amazon S3，對象存儲服務，提供行業(yè)領先的可擴展性、數(shù)據(jù)可用性、安全性和性能。
Amazon SQS， 消息隊列服務，可以幫助分離和擴展微服務、分布式系統(tǒng)和無服務器應用程序。
Amazon DynamoDB，快速、靈活的 NoSQL 數(shù)據(jù)庫服務，可在任何規(guī)模下實現(xiàn)個位數(shù)毫秒級的性能。
Amazon Batch，任意規(guī)模完全受管的批處理調(diào)度系統(tǒng)，讓開發(fā)人員、科學家和工程師能夠輕松高效地運行成千上萬個批處理計算作業(yè)。
Amazon ECR，完全托管的容器注冊表，在任何地方輕松存儲、管理、共享和部署容器映像和構(gòu)件。
Amazon SNS，完全托管的發(fā)布/訂閱消息收發(fā)、SMS、電子郵件和移動推送通知。
Amazon EventBridge，大規(guī)模構(gòu)建事件驅(qū)動型應用程序，涵蓋亞馬遜云科技、現(xiàn)有系統(tǒng)或 SaaS 應用程序
架構(gòu)說明
API 的設計
為向算法科學家屏蔽底層GPU算力調(diào)度過程，提供一個統(tǒng)一易用的Alphafold2任務調(diào)度入口，我們使用API Gateway創(chuàng)建RESTful API，對外暴露單一URL接口，實現(xiàn)對Alphafold2服務的應用級封裝。API Gateway為云原生托管API管理工具，對用戶提交的任務查詢或調(diào)度請求鑒權(quán)，通過后將其按照配置的路由發(fā)送到不同的Amazon Lambda后端。Amazon Lambda函數(shù)對請求進行解析，將任務信息寫入數(shù)據(jù)庫，并將任務參數(shù)投遞到HPC集群。共有四個Amazon Lambda函數(shù)與API Gateway直接或間接相關聯(lián)，包括身份鑒權(quán)函數(shù)、查詢請求處理函數(shù)、調(diào)度請求處理函數(shù)以及HPC集群調(diào)度函數(shù)，此處我們將調(diào)度請求處理函數(shù)與HPC集群調(diào)度函數(shù)通過SQS消息隊列服務進行解耦，從而提高了API架構(gòu)的可用性。
API參數(shù)的所有接口如下，算法科學家可直接通過API的形式對Alphafold2任務進行操作。
-
GET / ：查詢所有任務信息。
-
GET /{id} ：查詢單條任務詳細信息，包含最近的Cloudwatch log日志。
-
POST / ：投遞任務，任務信息包含在HTTP的body中。
-
DELETE /：刪除所有已完成或失敗的任務信息及其對應的Amazon S3資源（如有）。
-
DELETE /{id}：刪除單條所有已完成或失敗的任務信息及其對應的S3資源（如有）。
-
CANCEL /：取消所有正在運行的任務。
-
CANCEL /{id}：取消單條正在運行的任務。
注：ANY方法對應CANCEL方法。
云上 HPC 批處理平臺
本次方案的云上HPC批處理平臺共由四部分組成：
-
Amazon S3作為云原生的對象存儲服務，擁有11個9的高持久性以及完善的安全和合規(guī)機制，因此我們使用Amazon S3持久化存儲Alphafold2任務的輸入輸出文件，保障所有數(shù)據(jù)文件在云上的安全。在Amazon S3桶中，共有input和output兩個文件夾，分別用于存放要預測的氨基酸序列文件和2.Alphafold2預測產(chǎn)生的蛋白質(zhì)結(jié)構(gòu)文件。用戶可通過多種方式方便的檢索、下載上述文件，同時可以選擇Amazon S3中的多種存儲類，降低云上存儲成本。
-
為實現(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容器鏡像地址進行拉取。
-
Alphafold2每次預測蛋白質(zhì)結(jié)構(gòu)均需要通過Multiple Sequence Alignment (多重序列比對, MSA)，對已知蛋白質(zhì)結(jié)構(gòu)數(shù)據(jù)庫進行掃描，目前數(shù)據(jù)集總?cè)萘考s為2.2T。為了實現(xiàn)數(shù)據(jù)集的高性能共享訪問，我們將其存儲在FSx for Lustre上，使其能夠被所有任務所同時訪問并高速讀取。Fsx for Lustre可提供百兆級別吞吐、百萬級別IOPS以及穩(wěn)定的亞豪秒延遲，從而滿足Alphafold2在內(nèi)的HPC計算工作負載。
-
為實現(xiàn)計算資源的彈性擴展、降低整體計算和運維成本，我們使用Amazon Batch搭建蛋白質(zhì)結(jié)構(gòu)預測HPC批處理集群，由Batch內(nèi)置的調(diào)度器根據(jù)用戶提交的蛋白質(zhì)結(jié)構(gòu)預測任務和資源需求自動化預置所需要的GPU計算資源。我們目前針對亞馬遜云科技可提供的GPU實例類型，選擇了p3.2xlarge、p3.8xlarge、p3.16xlarge以及p4d.24xlarge四種機型作為計算環(huán)境，用戶可根據(jù)單條預測任務使用的GPU數(shù)量、任務并行度和成本綜合考慮選擇投遞任務的作業(yè)隊列，而無需關心其底層實例的部署以及環(huán)境的搭建。所有計算環(huán)境均可配置為在任務結(jié)束后即自動關閉所有計算實例節(jié)約成本。
任務信息數(shù)據(jù)存儲
Amazon DynamoDB作為云原生的無服務器NoSQL數(shù)據(jù)庫，無需預置服務器，按使用量付費，同時擁有一致性的毫秒級訪問性能以及近乎無限的吞吐量和存儲空間。本次方案使用DynamoDB存儲Alphafold2過程中產(chǎn)生的任務信息，滿足用戶并發(fā)提交或查詢?nèi)蝿諘r的性能需求，同時優(yōu)化成本，簡化數(shù)據(jù)庫配置流程。在處理任務提交的Amazon Lambda函數(shù)中，使用了UUID算法隨機生成每條任務的id，作為DynamoDB分區(qū)鍵，提高任務查詢效率。
消息與結(jié)果通知
單機部署Alphafold2時，其運行過程的所有日志均位于Docker中，如果運行環(huán)節(jié)中出現(xiàn)任何問題，均需要手動登陸實例后，通過docker log等命令進行查看；同時任務完成后缺乏相關通知推送，無法及時獲取其運行狀態(tài)，造成計算資源的浪費。
我們使用Amazon S3的事件通知機制以及Amazon EventBridge服務，自動獲取Amazon S3中壓縮文件的上傳文件事件以及Amazon Batch任務失敗事件——分別對應Alphafold2任務的成功和報錯，自動觸發(fā)不同的Amazon Lambda函數(shù)處理相關事件，更新任務信息數(shù)據(jù)庫，最后通過Amazon SNS——即亞馬遜云科技云上托管的消息收發(fā)服務，通過郵箱或者短信等方式向算法科學家推送相關信息，實現(xiàn)Alphafold2作業(yè)信息的即時通知。消息通知中包含任務信息、所需時間、下載鏈接以及Amazon CLI命令。算法科學家無需登陸亞馬遜云科技控制臺，即可直接獲取Alphafold2任務結(jié)果的下載方式。
便捷可視化終端
為滿足算法科學家遠程分析蛋白質(zhì)結(jié)構(gòu)的需求，我們使用NICE DCV遠程桌面作為蛋白質(zhì)結(jié)構(gòu)便捷可視化終端，啟動后可以直接通過瀏覽器或NICE DCV客戶端，使用圖形化界面遠程訪問。NICE DCV基于優(yōu)化的NICE DCV協(xié)議，根據(jù)網(wǎng)絡情況自適應調(diào)整傳輸帶寬，同時可利用預裝了NVIDIA T4顯卡的Amazon G4實例，提高圖形密集性軟件的渲染性能，從而使算法科學家遠程通過如Pymol或VMD等蛋白質(zhì)三維結(jié)構(gòu)可視化軟件進行分析。
部署方法
初始化 Amazon CDK 環(huán)境
cdk bootstrap aws://{ACCOUNTID}/{REGION}
通過 Amazon CDK生成
Amazon CloudFormation 模板
cdk synth
一鍵部署所有資源
cdk deploy –all
需要手動確認，輸入y回車即可，中間會輸出一些過程信息，并且在郵箱中點擊確認SNS訂閱。
部署完成
約需要四個小時將所有數(shù)據(jù)下載并解壓到FSx for Lustre，解壓完成后，會收到電子郵件通知，此時即可根據(jù)使用說明，開始投遞Alphafold2任務
算法科學家使用過程
完成上述方案構(gòu)建后，算法科學家即可開始使用Alphafold2算法預測蛋白質(zhì)結(jié)構(gòu)。
總共可分別四個步驟：
1.上傳氨基酸序列文件
通過Amazon CLI或者控制臺將氨基酸序列文件上傳到Amazon S3存儲桶的input文件夾。
2.提交任務或查詢?nèi)蝿者M度
通過終端或者HTTP調(diào)試軟件操作任務。本文以Postman調(diào)試軟件為例，演示如何基于API提交、查詢?nèi)蝿?。取消和刪除任務操作類似，可參考API Gateway中的路由設置進行提交。
首先編輯任務請求JSON文件，例：
可以參考:
https://github.com/wttat/af2-batch-cdk/blob/main/command.json
其中，
fasta 為本次預測的蛋白質(zhì)名稱，由用戶自定義，必選；
file_name為氨基酸序列名稱，必須與S3存儲桶中的文件名稱對應，必選；model_names為所使用的Alphafold2模型，共有model1-model5五個可選，缺省值為五個模型全部使用；
preset為使用的預設配置，full對應完整數(shù)據(jù)集，reduced對應壓縮數(shù)據(jù)集，CASP14對應deepmind團隊參賽時所使用的配置，缺省值為full；
max_template_date為掃描已知數(shù)據(jù)庫中蛋白質(zhì)結(jié)構(gòu)文件的截止日期，缺省值為2020-05-14；
que為batch使用的作業(yè)隊列，low對應batch的計算環(huán)境為p3.2xlarge實例，mid對應計算環(huán)境為p3.8xlarge實例，high對應計算環(huán)境為p3.16xlarge實例，p4對應計算環(huán)境為p4d.24xlarge實例，必選；
comment為用戶自定義的注釋語句，可選，缺省值為空；
gpu為使用GPU數(shù)量，并且成比例對應vCPU和內(nèi)存，可選，缺省值為1。例：p3.8xlarge實例中共有4塊NVIDIA Tesla V100顯卡，32vCPU和244G內(nèi)存，默認情況下，即gpu為1時，每個任務分配8vCPU、61G內(nèi)存和一塊V100顯卡，可同時運行四個任務，如果gpu設置為2時，該任務分配到16vCPU、122G內(nèi)存和兩塊V100顯卡，此時p3.8xlarge實例剩下兩塊顯卡，可分配給兩個gpu為1的任務或者一個gpu為2的任務。
首先配置Headers，鍵為Authorization，值為用戶自己定義的密鑰。
（1）提交任務
將HTTP方法設置為post，URL填入API Gateway的URL，Body選擇raw，格式為JSON，粘貼任務請求JSON文件，點擊Send即可提交任務，后臺即會根據(jù)選擇的作業(yè)隊列自動拉起對應的計算資源運行預測任務。API Gateway會返回每個預測任務對應的ID。
（2）查詢?nèi)蝿?/strong>
將方法修改為GET，即可查詢所有任務信息。
若想查詢單條任務詳細信息，則在URL后面加上對應任務的ID，可通過查詢所有任務信息得到，也可以在提交任務后直接復制返回的ID。
其中job info為任務相關參數(shù)，job status為任務當前狀態(tài)，若有日志產(chǎn)生，則會自動抓取最新的Amazon CloudWatch日志并返回。
3.Email自動化發(fā)送任務結(jié)果
任務完成后，會收到郵件通知提醒，告知其對應預測任務已完成。郵件正文中，包括完成預測所需時間，即可根據(jù)使用的機器類型計算成本，以及所有文件的下載鏈接HTML文件，同時提供了Amazon cli命令行，用戶可自行選擇下載方式。
打開下載鏈接HTML后，即可根據(jù)需求自行選擇下載整體壓縮包或者pdb文件進行分析。
4.可視化分析
以pymol為例，連接到NICE DCV上后，安裝并運行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進行分析。
成本估算
（僅供參考）
我們基于P3實例和P4實例，測試了不同序列長度的氨基酸序列，經(jīng)過Alphafold2算法預測其蛋白質(zhì)結(jié)構(gòu)所需時間，并且根據(jù)對應區(qū)域GPU實例單價，按比例換算為每條任務所需成本。可以看到：隨著序列長度的增加，所需時間快速增加，主要是由于GPU推理時間占比顯著增加；在序列長度超過1000時，一塊V100會提示顯存不足的問題，需要兩塊V100顯卡才能跑完完整任務；在序列長度越長的情況下，P4實例掛載的A100顯卡帶來的性能提升越明顯。
中國寧夏區(qū)域(cn-northwest-1)：
美國弗吉尼亞北部區(qū)域(us-east-1)：
注：FSx for Lustre費用未考慮在內(nèi)，需要按月收取。
優(yōu)化方向
目前Alphafold2官方源代碼還在不斷優(yōu)化迭代之中，還存在許多可以性能提升的方向。如：
-
前期掃描數(shù)據(jù)集的MSA過程只消耗CPU資源，并不占用GPU，因此可以據(jù)此將CPU和GPU運算過程接耦，拆分為兩個子任務，降低成本；
-
通常情況下，我們只需要評分最好的模型，而Alphafold2算法目前的機制默認為完整推理五個模型，可以將其修改為只預測評分最高的模型。
-
MSA中多個模塊可以CPU并行運行，提升掃描速度。
用戶可以根據(jù)自己的需要，對其進行二次開發(fā)，自行修改并構(gòu)建Alaphfold2鏡像，只需將Amazon Batch作業(yè)定義中的容器鏡像地址替換為用戶指定的鏡像地址即可。
總結(jié)
我們首先基于Amazon Batch的作業(yè)需求，容器化改造了Alphafold2鏡像，使其能夠接受來自Amazon Batch標準化的輸入?yún)?shù)，并且將輸出文件直接上傳到Amazon S3上。接下來，使用根據(jù)負載可彈性擴展的GPU計算集群和高性能共享存儲，可滿足大規(guī)模高并發(fā)的預測需求。并且，基于Amazon Batch提供完善的集群資源調(diào)度機制，降低計算資源成本以及IT部門運維復雜度，實現(xiàn)秒級的成本監(jiān)控。算法科學家基于Amazon API GW提交/查詢/刪除任務，無需接觸亞馬遜云科技資源，同于由于所有操作均基于API實現(xiàn)，因此可方便集成至現(xiàn)有工作流中。其次，基于Amazon SNS提供作業(yè)狀態(tài)提醒機制，針對任務的成功和失敗均有郵件提醒，并在任務成功的通知郵件集成了蛋白質(zhì)結(jié)構(gòu)文件的下載地址?；贏mazon CDK實現(xiàn)方案的靈活快速部署，用戶可在半小時內(nèi)將完整架構(gòu)部署于自己的亞馬遜云科技資源中，四個小時內(nèi)預測服務即可上線使用。