利用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的內部工作原理，以便您最佳化架構。