利用Amazon SageMaker AI上的BoltzGen加速蛋白質設計

在Amazon SageMaker AI上部署BoltzGen可以加速蛋白質結合體的設計過程，通過端到端管理GPU計算基礎設施。BoltzGen是一種基於擴散的生成模型，能夠設計能夠與特定生物分子靶標結合的蛋白質和肽。典型的設計流程涉及多個GPU密集型步驟：主鏈生成、反向摺疊、結構驗證和候選排序。這些步驟在數百、數千甚至數百萬的設計候選項中運行時，會帶來實例預置、數據步驟間移動和成本跟蹤等操作開銷。SageMaker AI管理從實例預置到結果交付和資源清理的整個計算生命週期，使您能夠專注於設計迭代而非基礎設施操作。

本文展示瞭如何在SageMaker AI上部署BoltzGen並運行端到端的蛋白質設計實驗。完成本教程後，您將擁有一個從快速驗證運行到生產批處理均可擴展的工作設置。該設置提供兩種執行模式以適應不同的研究階段，並使用步驟級緩存來減少迭代工作流中的計算成本。

本教程適用於學術研究實驗室、生物技術初創公司、製藥研發團隊以及教育項目，無論您從事蛋白質結合體設計、治療性蛋白質工程還是從頭蛋白質架構。

SageMaker AI如何解決蛋白質設計瓶頸

BoltzGen流程中的每一步都在GPU硬件上運行，每次處理一個設計規範。根據存儲庫的基準數據，在4-GPU實例（ml.g5.12xlarge）上，1000個樣本的流程大約需要375小時完成。運營此基礎設施涉及構建CUDA環境、協調GPU實例生命週期、構建步驟間數據管道以及從長時間運行的作業故障中恢復。

SageMaker AI直接解決了這些瓶頸。提交作業後，SageMaker AI預置GPU實例並在容器內執行BoltzGen。它將結果寫入Amazon Simple Storage Service（Amazon S3），並在處理完成後釋放實例。計費按秒計算，因此不會產生閒置GPU成本。在ml.g4dn.xlarge上運行2小時的設計費用約為1.50美元（按需定價）。

該實現支持單個實例內的多GPU並行化和跨管道的多實例擴展。在管道模式下，每個步驟的輸出在Amazon S3中緩存，有效期為7天，因此當您迭代過濾參數時，約佔總計算成本90%的設計生成步驟不會重新運行。

存儲庫中的設置腳本構建容器並將其推送到Amazon Elastic Container Registry（Amazon ECR），工作示例可讓您在幾分鐘內提交第一個設計作業。驗證10個候選設計的相同配置可擴展到更大規模的流程——只需更改參數值。從ml.g4dn（最低成本T4 GPU）到ml.g6e（NVIDIA L40S GPU）的實例類型為您提供了靈活性，可根據預算匹配吞吐量。

BoltzGen的功能

BoltzGen是一種全原子生成模型，用於設計能夠結合廣泛生物分子靶標的蛋白質和肽。該模型通過擴散過程生成結合體主鏈結構。然後使用稱為BoltzIF的反向摺疊模型生成氨基酸序列，並通過Boltz2摺疊預測驗證結構。

該實現可在GitHub上的Boltzgen on SageMaker存儲庫中找到。存儲庫包含設置腳本、兩種執行模式、多GPU和多實例擴展支持以及快速入門指南。

兩種執行模式

我們在實現中提供了兩種執行模式，針對不同研究階段進行了優化，使您可以從快速實驗開始，隨着需求增長過渡到生產工作流。

SageMaker AI處理作業提供直接的批處理執行，適用於快速實驗。提交作業後，SageMaker AI預置GPU實例，運行容器，並在完成後關閉實例。單個Python腳本觸發整個工作流，無需多步編排。對於早期工作，目標是在擴展前測試設計規範，處理作業最大限度地減少了設置時間。

SageMaker AI Pipelines提供了一個5步編排工作流，具有步驟級緩存、自動擴展和SageMaker AI Studio中的可視化監控。生產工作負載受益於參數迭代，無需重新運行昂貴的計算步驟。每個步驟（設計、反向摺疊、摺疊、分析、過濾）可獨立擴展，當輸入未更改時，緩存的結果會被重用。

兩種模式遵循相同的工作流。來自Amazon ECR的容器在預置的GPU實例上運行，BoltzGen處理設計規範，SageMaker AI將結果存儲在Amazon S3中，然後關閉計算資源。

入門指南

要開始使用，您需要一個AWS賬户、一些本地安裝的工具以及一個用於數據的Amazon S3存儲桶。以下部分將引導您完成每個先決條件，然後指導您完成第一個蛋白質設計實驗。

先決條件

確保您擁有一個已啓用計費的AWS賬户，並已安裝和配置AWS CLI（aws configure）。驗證您是否具有足夠的服務配額（例如，本教程為1），用於您目標AWS區域（例如us-east-1）中的GPU實例（ml.g4dn.xlarge或更高）。

接下來，為SageMaker AI創建AWS Identity and Access Management（IAM）執行角色。附加AmazonSageMakerFullAccess託管策略，授予對您的數據存儲桶的Amazon S3讀寫訪問權限，並授予對BoltzGen容器鏡像的Amazon ECR拉取權限。

在您的本地機器上，安裝Python 3.11+和pip、最新的Boto3和SageMaker SDK（pip install boto3 sagemaker），以及用於構建容器鏡像的Docker。最後，在您的目標區域中創建一個Amazon S3存儲桶：

aws s3 mb s3://amzn-s3-demo-bucket --region

運行您的第一個蛋白質設計

為獲得最簡化的設置體驗，我們建議使用Amazon SageMaker Studio或Amazon SageMaker筆記本實例。這些環境預配置了AWS憑證、SageMaker SDK和Docker支持。

啓動SageMaker AI Studio

首先打開SageMaker AI Studio，在其中運行設置和實驗命令。

打開SageMaker AI控制台。

從左側導航中選擇Studio。

創建一個新的Studio域（如果不存在）或打開現有域，通過AWS CLI啓用Docker訪問（本地模式）。

啓動一個JupyterLab空間，實例類型例如ml.m5.2xlarge。

將ml.g4dn.xlarge和ml.g5.xlarge處理作業使用的配額限制提高到大於0。

克隆存儲庫

下載BoltzGen on SageMaker AI代碼，其中包括設置腳本、兩種執行模式和示例設計規範。

git clone https://github.com/aws-samples/sample-biofm-quickstart.git cd sample-biofm-quickstart/models/boltzgen/inference/amazon-sagemaker/boltzgen-on-sagemaker-processing-job-cli

構建並推送容器鏡像

您的容器鏡像將BoltzGen及其依賴項和GPU驅動程序打包在一起。構建一次並在實驗中重複使用。設置區域並運行構建腳本：

export AWS_REGION=./sagemaker/build_and_push.sh

配置AWS憑證

SageMaker AI需要您的AWS賬户詳細信息來預置計算資源並將結果存儲在Amazon S3中。複製環境模板並填寫您的設置：

cp sagemaker/pipeline/.env.example sagemaker/pipeline/.env vim sagemaker/pipeline/.env

示例.env：

AWS_REGION=us-east-1 AWS_S3_BUCKET=amzn-s3-demo-bucket AWS_ROLE_ARN=arn:aws:iam:::role/SageMakerExecutionRole AWS_IMAGE_URI=.dkr.ecr..amazonaws.com/boltzgen-sagemaker:latest

現在您已準備好運行第一個設計作業，有兩種選項：

選項A：在單個處理作業中快速運行所有步驟：

python sagemaker/run_processing_job.py \ --design-spec example/vanilla_protein/1g13prot.yaml \ --s3-bucket amzn-s3-demo-bucket \ --instance-type ml.g4dn.xlarge \ --num-designs 10 \ --budget 2 \ --wait

您將生成10箇中間候選設計，這些是擴散模型生成的原始主鏈結構，隨後經過反向摺疊、再摺疊、分析和過濾。實際運行中，此數字應更高，如10,000至60,000。最終多樣性優化輸出集包含兩個設計。因此，在10箇中間設計中，只有兩個針對質量和序列多樣性進行了優化的設計最終出現在輸出文件夾中。

選項B：一個更接近生產的工作流，在管道中編排不同步驟。

切換到管道目錄：

cd sagemaker/pipeline

創建管道：

python run_pipeline.py --config pipeline_config.yaml create

使用設計規範運行：

python run_pipeline.py --config pipeline_config.yaml run

檢查狀態（將ARN替換為執行ARN）：

python run_pipeline.py --region us-east-1 status --execution-arn

下載結果

對於處理作業：

aws s3 sync s3://amzn-s3-demo-bucket/boltzgen/output/boltzgen-TIMESTAMP ./results

對於管道：

aws s3 sync s3://amzn-s3-demo-bucket/boltzgen-pipeline/output/TIMESTAMP ./results

輸出目錄包含多個子目錄。final_ranked_designs/保存按預測結合親和力排名的最佳設計。intermediate_designs/包含生成的候選結構。metrics/存儲結構分析數據，包括均方根偏差（RMSD）和接觸分數。job_metadata.json文件記錄了完整的作業參數和狀態。

以下目錄樹顯示了使用--num-designs 100 --budget 10針對1G13靶標運行的輸出結構：

results/1g13prot/ ├── final_ranked_designs/ │ ├── final_10_designs/ # 針對多樣性和質量選擇的最佳設計（.cif） │ ├── intermediate_ranked_10_designs/ # 按質量分數排名的前十設計（.cif） │ ├── all_designs_metrics.csv # 所有設計的完整指標表 │ ├── final_designs_metrics_10.csv # 最終選擇的十項設計指標 │ └── results_overview.pdf # 包含過濾標準和圖表的摘要報告 ├── intermediate_designs/ # 所有生成的主鏈結構（.cif） ├── intermediate_designs_inverse_folded/ # 序列設計和再摺疊後的結構 └── config/ # 用於重現性的步驟配置

BoltzGen自動過濾和排名所有生成的設計。每個設計必須通過結構質量閾值——包括再摺疊RMSD低於2.5埃和平衡的氨基酸組成——才能進入排名階段。在此次運行中，100個設計中有3個通過了所有過濾器。然後管道按複合質量分數對設計進行排名，並應用多樣性優化以選擇最終集合。

下表顯示了前五個設計及其關鍵指標。較高的design_ptm和design_iptm分數表明更高的結構置信度和結合強度。較低的filter_rmsd表明再摺疊後更好的結構準確性，較高的delta_sasa表明更大的結合界面。

id Rank design_ptm design_iptm filter_rmsd delta_sasa 1g13prot_24 1 0.760 0.304 2.24 584.7 1g13prot_37 2 0.732 0.275 1.98 657.4 1g13prot_21 3 0.748 0.159 2.49 535.7 1g13prot_06 4 0.736 0.426 1.34 1598.0 1g13prot_59 5 0.763 0.527 1.05 2369.0

每個最終設計保存為.cif結構文件，前綴為排名（例如rank001_1g13prot_24.cif），可直接用於下游分析或實驗驗證。

現在您的第一個實驗已完成，以下部分將探討BoltzGen on SageMaker AI的內部工作原理，以便您優化架構。