蛋白设计正在从“单次手工跑模型”走向“可复现、可扩展、可缓存”的工程化工作流。BoltzGen 部署到 Amazon SageMaker AI 后，可以把快速验证和生产批处理放在同一套基础设施里运行：研究早期用小规模任务验证输入、参数和模型行为，进入筛选阶段后再扩展到批量生成与评估。
这类工作流的关键不只是把模型跑起来，而是让每一步都能被追踪、复用和按需扩容。尤其在蛋白设计中，输入预处理、候选生成、结构或属性评估往往会反复迭代；如果没有缓存机制，同样的中间结果会反复消耗 GPU 时间。
从一次实验到可扩展流水线
BoltzGen 在 SageMaker AI 上的价值，可以理解为把蛋白设计实验拆成几个稳定阶段：
- 环境构建：准备容器镜像、依赖、模型权重和运行脚本。
- 快速验证：用少量序列、短任务或小批次确认流程可运行。
- 批量处理：把候选输入扩展到更大的数据集，交给 SageMaker 作业调度。
- 步骤级缓存：对预处理、生成、评估等阶段建立缓存，避免重复计算。
这带来的直接收益是：研究人员不用在本地 GPU、临时脚本和手动文件夹之间来回切换；工程团队也能把实验参数、输入数据、输出结果和成本边界纳入统一管理。
两种执行模式：验证和生产不要混在一起
来源摘要提到该方案提供两种执行模式，适合研究工作的不同阶段。可以这样理解：
- 快速验证模式：目标是尽快发现输入格式、依赖、权限、模型参数是否正确。它通常使用小数据、小实例、短运行时间。
- 生产批处理模式：目标是稳定吞吐和可重复执行。它通常使用批量输入、明确的输出路径、可监控的作业状态，并且更依赖缓存降低重复成本。
一个常见误区是直接把探索性 notebook 改成生产脚本。更稳妥的做法是从一开始就把参数、输入路径、输出路径和运行模式显式化，让同一份代码能在不同规模下运行。
可以这样实践：用 SageMaker ProcessingJob 封装 BoltzGen 任务
下面示例展示一个可改造的最小脚本：用 SageMaker Python SDK 启动一个处理作业，把输入数据从 S3 挂载进容器，并把输出写回 S3。实际使用时，需要把镜像地址、角色 ARN、S3 路径和 BoltzGen 命令替换成你的环境配置。
# run_boltzgen_processing.py
import argparse
import sagemaker
from sagemaker.processing import ProcessingInput, ProcessingOutput, ScriptProcessor
parser = argparse.ArgumentParser()
parser.add_argument("--mode", choices=["validation", "batch"], default="validation")
parser.add_argument("--image-uri", required=True)
parser.add_argument("--role", required=True)
parser.add_argument("--input-s3", required=True)
parser.add_argument("--output-s3", required=True)
args = parser.parse_args()
session = sagemaker.Session()
instance_type = "ml.g5.xlarge" if args.mode == "validation" else "ml.g5.12xlarge"
instance_count = 1 if args.mode == "validation" else 2
processor = ScriptProcessor(
image_uri=args.image_uri,
command=["python3"],
role=args.role,
instance_type=instance_type,
instance_count=instance_count,
sagemaker_session=session,
)
processor.run(
code="boltzgen_entrypoint.py",
arguments=["--mode", args.mode, "--cache-dir", "/opt/ml/processing/cache"],
inputs=[
ProcessingInput(
source=args.input_s3,
destination="/opt/ml/processing/input",
)
],
outputs=[
ProcessingOutput(
source="/opt/ml/processing/output",
destination=args.output_s3,
)
],
)
配套的入口脚本可以把“验证”和“批处理”的差异限制在参数层，而不是写两套逻辑：
# boltzgen_entrypoint.py
import argparse
import hashlib
import json
import subprocess
from pathlib import Path
parser = argparse.ArgumentParser()
parser.add_argument("--mode", choices=["validation", "batch"], required=True)
parser.add_argument("--cache-dir", required=True)
args = parser.parse_args()
input_dir = Path("/opt/ml/processing/input")
output_dir = Path("/opt/ml/processing/output")
cache_dir = Path(args.cache_dir)
output_dir.mkdir(parents=True, exist_ok=True)
cache_dir.mkdir(parents=True, exist_ok=True)
input_files = sorted(input_dir.glob("*.fasta"))
if args.mode == "validation":
input_files = input_files[:2]
for fasta in input_files:
digest = hashlib.sha256(fasta.read_bytes()).hexdigest()[:16]
cached_result = cache_dir / f"{digest}.json"
final_result = output_dir / f"{fasta.stem}.json"
if cached_result.exists():
final_result.write_text(cached_result.read_text())
continue
# 假设容器内已安装 boltzgen CLI；请替换为你的实际命令和参数。
subprocess.run(
[
"boltzgen",
"design",
"--input",
str(fasta),
"--output",
str(final_result),
],
check=True,
)
cached_result.write_text(final_result.read_text())
metadata = {
"mode": args.mode,
"processed_files": [path.name for path in input_files],
}
(output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2))
运行方式示例：
python run_boltzgen_processing.py \
--mode validation \
--image-uri 123456789012.dkr.ecr.us-east-1.amazonaws.com/boltzgen:latest \
--role arn:aws:iam::123456789012:role/SageMakerExecutionRole \
--input-s3 s3://my-protein-design/input/ \
--output-s3 s3://my-protein-design/output/validation/
进入批处理阶段时，只需要切换模式和输出路径：
python run_boltzgen_processing.py \
--mode batch \
--image-uri 123456789012.dkr.ecr.us-east-1.amazonaws.com/boltzgen:latest \
--role arn:aws:iam::123456789012:role/SageMakerExecutionRole \
--input-s3 s3://my-protein-design/input/ \
--output-s3 s3://my-protein-design/output/batch-001/
步骤级缓存为什么重要
蛋白设计的迭代通常不是线性的。你可能会调整候选序列集合、修改生成参数、替换评估阈值，或者只重跑某个失败样本。如果每次都从头开始，成本会被重复计算吞掉。
步骤级缓存的设计原则很简单：
- 缓存键必须稳定：由输入文件内容、模型版本、关键参数共同决定。
- 缓存粒度不要过粗：按样本或步骤缓存，比整批缓存更容易复用。
- 缓存结果要可解释：输出里保留模式、参数、输入摘要和模型版本。
- 缓存失效要明确：模型权重、依赖版本或核心参数变化时，应该主动生成新缓存。
上面的示例只用 FASTA 内容生成缓存键，适合说明思路。生产环境里应把模型版本、BoltzGen 参数、容器镜像 digest 等也纳入哈希，否则可能误用旧结果。
采用建议：先让实验可复现，再追求吞吐
把 BoltzGen 放到 SageMaker AI 上，不应该只看“能不能跑更大的 GPU”。更重要的是把实验流程变成可以复现、可以扩容、可以控制成本的系统。
落地时可以按这个清单推进：
- 先用验证模式跑通最小 FASTA 输入和输出结构。
- 把容器镜像、模型版本、参数文件和 S3 路径全部显式记录。
- 为高成本步骤设计缓存键，并测试缓存命中和失效逻辑。
- 批处理前设置预算、日志、失败重试和输出校验。
- 不要把探索性参数直接带入生产批量任务，先做小样本回归验证。
这套方式的边界也很清楚：SageMaker AI 解决的是运行、扩展、管理和成本控制问题；蛋白设计结果是否具有生物学意义，仍然需要领域评估、后续筛选和实验验证。工程平台能加快迭代，但不能替代科学判断。