利用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的内部工作原理，以便您优化架构。