OpenFold3 上线 Amazon SageMaker JumpStart,可预测生物分子结构
OpenFold3 on Amazon SageMaker JumpStart: Predict biomolecular structures
OpenFold3-preview 现可从 Amazon SageMaker JumpStart 部署,在 AWS 托管容器与异步推理端点上预测蛋白质、DNA、RNA、小分子配体及其复合物的三维结构。
你可以从 Amazon SageMaker JumpStart 部署 OpenFold3-preview。OpenFold Consortium 和哥伦比亚大学 AlQuraishi 实验室开发了 OpenFold3,这是一个开源生物分子结构预测模型。只需几行代码,你通常就可以将 OpenFold3 部署到你自己的 AWS 账户中。然后,你可以预测蛋白质、DNA、RNA、小分子配体及其形成的复合物的 3D 结构。该端点以网络隔离方式运行,你可以利用它帮助将序列保留在你的 AWS 账户内。
计算药物发现团队现在将结构预测作为其工作的核心部分。团队将其用于靶点识别、抗体和蛋白质工程,以及基于结构的虚拟筛选。团队发现运行这些模型颇具挑战性。你需要 GPU 基础设施、可用的 CUDA 和 PyTorch 技术栈、大型模型检查点,以及一个能够处理耗时数分钟而非数毫秒请求的服务层。SageMaker JumpStart 可以为你处理其中大部分工作。OpenFold3 在由 AWS 维护的容器中运行,部署在托管的 Amazon SageMaker AI 端点上。该端点旨在处理请求排队和 GPU 预热,你还可以添加自动扩缩容。
在本文中,我们将向你展示如何:
- 在 SageMaker 异步推理端点上部署 OpenFold3-preview。
- 提交结构预测任务并解读结果,包括预测结构和置信度分数。
- 在多 GPU 实例上并行运行多个预测。
- 提供你自己的多序列比对(MSA)以提高预测准确性,同时帮助保持序列私密。
- 通过在空闲时将端点缩容至零来帮助降低成本。
OpenFold3 概述
OpenFold3 是一个全原子结构预测模型。它根据序列预测一组分子的联合 3D 结构。它不限于单条蛋白质链:它可以在一次预测中折叠多链蛋白质复合物、蛋白质–核酸组装体以及蛋白质–配体复合物。
| 属性 | 详情 |
| 开发者 | OpenFold Consortium / 哥伦比亚大学 AlQuraishi 实验室 |
| 模型类型 | 全原子生物分子结构预测 |
| 输入 | 蛋白质、DNA 和 RNA 序列。小分子配体以 SMILES 表示 |
| 输出 | mmCIF 格式的预测结构,附带置信度指标(pLDDT、PAE、PDE) |
| 版本 | OpenFold3 0.4.1,of3-p2-155k 检查点 |
| 许可证 | Apache 2.0 |
| JumpStart 模型 ID | huggingface-proteinfolding-openfold3 |
| 支持的实例 | ml.g6.2xlarge、ml.g6.4xlarge(默认)、ml.g6.12xlarge |
OpenFold3 开发者根据 Apache 2.0 许可证发布其代码和模型权重,并公开训练数据和方法。研究机构和制药公司可以根据该许可证条款检查模型、微调模型并在此基础上构建。开发者将 OpenFold3-preview 描述为研究预览版。
使用场景
- 靶点和复合物结构预测 – 为没有实验结构的靶点预测结构,包括多链组装体和蛋白质–核酸复合物。
- 蛋白质–配体共折叠 – 将蛋白质与以 SMILES 提供的小分子一起折叠,为基于结构的设计生成结合构象假设。
- 抗体和蛋白质工程 – 通过预测设计变体的结构并比较候选之间的置信度来评估这些变体。
- 高通量筛选流水线 – 向自动扩缩容端点排队提交大批量预测任务,并直接从 Amazon Simple Storage Service(Amazon S3)处理结果。
OpenFold3 为何使用异步推理
大多数 JumpStart 模型(例如大语言模型和嵌入模型)部署到实时端点。实时端点会在单个请求内返回响应,且该请求上限为 60 秒。结构预测的工作方式不同。折叠时间随查询中链的数量和长度增长,真实世界的蛋白质和复合物往往需要数分钟。
因此,JumpStart 上的 OpenFold3 运行在 SageMaker 异步推理上:
- 您将请求上传到 Amazon S3 并调用
InvokeEndpointAsync,该调用会立即返回。 - 端点将请求排队并处理,调用超时最长可达 1 小时。
- 端点将结果或错误详情写入您选择的 S3 位置。您可以轮询获取,或通过 Amazon Simple Notification Service(Amazon SNS)获取通知。
异步端点的两个特性尤其适合科学工作负载。第一,它们在内部对突发请求进行排队,因此您可以提交大批量序列,而无需自行构建任务队列。第二,当队列为空时,它们可以缩减到零实例,这有助于避免在实验之间为空闲实例付费。
先决条件
要完成本演练,您需要:
- 一个 AWS 账户,以及用于该端点的 SageMaker AI 执行角色。遵循最小权限原则,此 AWS Identity and Access Management(IAM)角色只需对用于请求和结果的 S3 前缀具有读写访问权限。有关更多信息,请参阅如何使用 SageMaker AI 执行角色。
- 运行代码的身份所需的 IAM 权限:创建、描述、调用和删除模型、端点配置和端点,以及为执行角色使用
iam:PassRole。对于可选的缩减到零部分,还需要 AWS Application Auto Scaling 和 Amazon CloudWatch 告警权限。在生产环境中,请将这些权限限定到您的资源,而不是使用AmazonSageMakerFullAccess等宽泛的托管策略。有关更多信息,请参阅异步推理的完成先决条件。 - 访问 Amazon SageMaker Studio、SageMaker 笔记本实例或本地 IDE。
- SageMaker Python SDK v3 和 py3Dmol,可使用
pip install "sagemaker>=3,<4" py3Dmol安装。 - 您的 AWS 区域中至少一个
ml.g6.4xlarge端点实例的服务配额。如果遵循多 GPU 部分,还需要ml.g6.12xlarge的配额。您可以在 Service Quotas 控制台的 Amazon SageMaker 下检查配额并请求增加。
重要:SageMaker 端点运行期间按实例小时计费,包括空闲期间。为避免意外费用,请遵循清理部分,或按本文后面所述配置缩减到零。有关定价详情,请参阅 Amazon SageMaker AI 定价。
在 SageMaker JumpStart 中发现 OpenFold3
- 打开 SageMaker Studio,在导航窗格中选择 JumpStart。
- 搜索 OpenFold3。
- 打开模型卡片以查看模型描述、许可证和支持的实例类型。
- 从模型卡中打开 示例 notebook。其中包含本文所示的完整异步部署演练。
使用 SageMaker Python SDK 部署 OpenFold3
首先,根据其 JumpStart 配置构建模型,并定义一个 AsyncInferenceConfig。输出路径和失败路径是端点写入结果和错误详情的 S3 位置。
import json, time, uuid
import boto3
from sagemaker.core.helper.session_helper import Session
from sagemaker.core.inference_config import AsyncInferenceConfig
from sagemaker.core.jumpstart.configs import JumpStartConfig
from sagemaker.core.training.configs import Compute
from sagemaker.serve.model_builder import ModelBuilder
session = Session()
bucket = session.default_bucket() # sagemaker-{region}-{account-id}
prefix = "openfold3-async"
model_builder = ModelBuilder.from_jumpstart_config(
jumpstart_config=JumpStartConfig(
model_id="huggingface-proteinfolding-openfold3",
accept_eula=True,
),
compute=Compute(instance_type="ml.g6.4xlarge"),
)
model = model_builder.build()
async_config = AsyncInferenceConfig(
output_path=f"s3://{bucket}/{prefix}/output",
failure_path=f"s3://{bucket}/{prefix}/failure",
)
endpoint = model_builder.deploy(inference_config=async_config)
endpoint_name = endpoint.endpoint_name
print(f"Deployed asynchronous endpoint: {endpoint_name}")设置 accept_eula=True 即确认你已查看并接受该模型的许可条款,这是部署此模型所必需的。如果你在 SageMaker Studio 之外运行此代码,还需通过 role_arn="<role-arn>" 将你的 SageMaker 执行角色传给 ModelBuilder.from_jumpstart_config。
端点会将你的请求和结果存储在 Amazon S3 中,其中包含你的序列和预测结构。对于专有序列,建议使用专用的 S3 存储桶或前缀,并使用由你管理的 AWS Key Management Service(AWS KMS)密钥进行服务端加密。要使用你的密钥加密结果,请将其 ARN 作为 kms_key_id 传给 AsyncInferenceConfig。通过存储桶策略限制访问,仅允许运行此工作流的角色访问,并添加 S3 生命周期规则,使不再需要保留的请求和结果对象过期。有关更多信息,请参阅 Amazon S3 安全最佳实践。
部署通常在几分钟内完成。这包括一个预热步骤,容器会在每个 GPU 上运行一次简短预测,以便端点达到 InService 时 GPU 已预热。
提示: 在某些区域,G6 容量可能受限。如果部署因
InsufficientInstanceCapacity失败,请稍后重试、选择其他受支持的实例大小,或在模型可用的其他 AWS 区域中部署。
运行结构预测
请求是一个 JSON 文档,包含一个或多个具名 查询,每个查询列出其 链。一个查询中的所有链会作为一个复合物一起折叠。不同查询则独立折叠。
每条链包含以下字段:
molecule_type–protein、dna、rna或ligand。chain_ids– 链标识符列表,例如["A"]。要指定相同副本(例如同源二聚体),请使用["A", "B"]。sequence– 残基序列,用于蛋白质、DNA 和 RNA 链。smiles– SMILES 字符串,用于配体链,替代sequence。
以下示例折叠人类泛素,这是一种特征明确的 76 残基蛋白质:
s3 = boto3.client("s3")
smr = boto3.client("sagemaker-runtime")
payload = {
"queries": {
"ubiquitin": {
"chains": [
{
"molecule_type": "protein",
"chain_ids": ["A"],
"sequence": "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG",
}
]
}
}
}
# 1. Stage the request in S3
input_key = f"{prefix}/input/{uuid.uuid4().hex}.json"
s3.put_object(Bucket=bucket, Key=input_key, Body=json.dumps(payload).encode())
# 2. Submit asynchronously -- returns immediately
resp = smr.invoke_endpoint_async(
EndpointName=endpoint_name,
ContentType="application/json",
InputLocation=f"s3://{bucket}/{input_key}",
InvocationTimeoutSeconds=3600,
)
output_location, failure_location = resp["OutputLocation"], resp["FailureLocation"]然后轮询 S3 以获取结果或失败记录:
def split(uri):
b, _, k = uri.removeprefix("s3://").partition("/")
return b, k
result = None
for _ in range(120): # up to ~30 minutes
for loc, is_failure in ((output_location, False), (failure_location, True)):
try:
body = s3.get_object(Bucket=split(loc)[0], Key=split(loc)[1])["Body"].read()
except s3.exceptions.NoSuchKey:
continue
if is_failure:
raise RuntimeError(body.decode())
result = json.loads(body)
if result:
break
time.sleep(15)
if result is None:
raise TimeoutError(f"No result yet; check {output_location}")
if result.get("status") == "error":
raise RuntimeError(result["error"])输入无效的请求(例如链格式错误或 MSA 名称不受支持)仍会生成结果对象,但会包含 "status": "error" 和一条 error 消息,而不是结构。最后一项检查会显示该消息。
预测时间随输入规模增长。我们在 ml.g6.4xlarge 上的测试中,较小的蛋白质比较大型蛋白质完成得更快。
解读结果
响应包含一个 status、一个以 mmCIF 格式表示的预测 structures 列表,以及与之匹配的 confidences。默认情况下,OpenFold3 会为每个查询生成五个扩散样本,因此每个查询返回五个候选结构。每个样本附带两个置信度文件:一个详细文件(*_confidences.json)和一个摘要文件(*_confidences_aggregated.json)。摘要中包含一个 sample_ranking_score,可用于挑选最佳样本。
import os
os.makedirs("structures", exist_ok=True)
for s in result["structures"]:
path = os.path.join("structures", os.path.basename(s["filename"]))
with open(path, "w") as f:
f.write(s["content"])
print("saved", path, f'({s["size_bytes"]/1024:.0f} KB)')
summaries = {c["filename"]: c["data"] for c in result["confidences"]
if c["filename"].endswith("_aggregated.json")}
for name, m in sorted(summaries.items(), key=lambda kv: -kv[1]["sample_ranking_score"]):
print(f'{name}: ranking={m["sample_ranking_score"]:.2f} pLDDT={m["avg_plddt"]:.1f} '
f'pTM={m["ptm"]:.2f} ipTM={m["iptm"]:.2f}')你可以在任何结构查看器中打开 .cif 文件,例如 PyMOL、ChimeraX 或 Mol*。要在 notebook 中内联查看结构,请使用 py3Dmol:
import py3Dmol # pip install py3Dmol
view = py3Dmol.view(width=600, height=450)
view.addModel(result["structures"][0]["content"], "cif")
# pLDDT is stored in the B-factor field: red = low confidence, blue = high
view.setStyle({"cartoon": {"colorscheme": {"prop": "b", "gradient": "roygb", "min": 50, "max": 90}}})
view.zoomTo(); view.show()图 1:OpenFold3 预测的泛素结构,使用上述代码渲染,并按每个原子的置信度(pLDDT)着色。本文所有图均使用 3Dmol.js 渲染(Rego and Koes, Bioinformatics, 2015)。
预测结果显示泛素具有紧凑的折叠结构:一个混合β折叠片层紧贴着一个α螺旋。结构的大部分呈蓝色(高置信度)。只有图像顶部的柔性C端尾部偏向黄色和红色,这与实验结构中已知的尾部柔性相符。
置信度指标告诉你每个预测结果的可信程度:
- pLDDT(预测局部距离差异测试)——局部结构中每个原子的置信度,范围为0–100。数值越高越好,通常高于约70的值表示局部结构可靠。摘要文件报告的平均值为
avg_plddt。 - pTM和ipTM(预测TM分数和界面pTM)——对折叠整体以及各链和配体相对位置的置信度,范围为0–1。对于复合物,ipTM是关键数值。摘要文件还会报告每对链的
chain_pair_iptm。 - PAE(预测对齐误差)和PDE(预测距离误差)——详细的置信度文件中包含的逐残基对误差估计。用它们来判断结构域和链的相对位置。
折叠复合物
要预测复合物,请在同一查询中列出多条链。以下示例将人碳酸酐酶II与其催化锌离子和药物乙酰唑胺共同折叠。这是一个研究充分的复合物,蛋白质数据库中已有其实验结构(例如PDB 3HS4):
ca2_sequence = (
"MSHHWGYGKHNGPEHWHKDFPIAKGERQSPVDIDTHTAKYDPSLKPLSVSYDQATSLRILNNGHAFNVEFDDSQDKAVLKGG"
"PLDGTYRLIQFHFHWGSLDGQGSEHTVDKKKYAAELHLVHWNTKYGDFGKAVQQPDGLAVLGIFLKVGSAKPGLQKVVDVLDS"
"IKTKGKSADFTNFDPRGLLPESLDYWTYPGSLTTPPLLECVTWIVLKEPISVSSEQVLKFRKLNFNGEGEPEELMVDNWRPAQ"
"PLKNRQIKASFK"
)
payload = {
"queries": {
"ca2_acetazolamide": {
"chains": [
{"molecule_type": "protein", "chain_ids": ["A"], "sequence": ca2_sequence},
{"molecule_type": "ligand", "chain_ids": ["B"], "smiles": "[Zn+2]"},
{"molecule_type": "ligand", "chain_ids": ["C"], "smiles": "CC(=O)Nc1nnc(s1)S(=O)(=O)N"},
]
}
}
}这些示例中的蛋白质序列来自UniProt(P0CG48、P00918)。提交和读取此请求的方式与泛素示例相同。对于此类复合物,预测质量在很大程度上取决于是否提供MSA,如下文各节所示。
在多GPU实例上并行运行预测
单个OpenFold3预测在一个GPU上运行,因此多GPU实例并不会让单个预测更快。相反,你可以同时运行多个预测。在诸如ml.g6.12xlarge这样具有四个GPU的多GPU实例上,容器会为每个并发请求分配各自的GPU。要使用所有GPU,请将max_concurrent_invocations_per_instance设置为实例上的GPU数量(对于ml.g6.12xlarge为四个):
multi_gpu_builder = ModelBuilder.from_jumpstart_config(
jumpstart_config=JumpStartConfig(
model_id="huggingface-proteinfolding-openfold3",
accept_eula=True,
),
compute=Compute(instance_type="ml.g6.12xlarge"),
)
multi_gpu_model = multi_gpu_builder.build()
async_config = AsyncInferenceConfig(
output_path=f"s3://{bucket}/{prefix}/output",
failure_path=f"s3://{bucket}/{prefix}/failure",
max_concurrent_invocations_per_instance=4, # = number of GPUs on the instance
)
multi_gpu_endpoint = multi_gpu_builder.deploy(inference_config=async_config)这会创建第二个端点,因此完成后记得也将其删除(参见清理)。
在我们的测试中,四个并发预测各自运行在独立的GPU上,吞吐量增益随蛋白质规模增大而增长。对于较大的蛋白质,我们基于按需实例定价在ml.g6.12xlarge上测得的每次预测成本接近ml.g6.4xlarge。你的结果取决于输入和区域。有关定价,请参阅Amazon SageMaker AI定价。
注意:
ml.g6.12xlarge容量可能有限。如果你的区域不可用,可以通过使用本文后面描述的自扩缩配置横向扩展多个ml.g6.4xlarge实例来提高吞吐量。
使用下表选择实例:
| 工作负载 | 推荐实例 | 原因 |
| 单个用户的交互式探索 | ml.g6.2xlarge或ml.g6.4xlarge |
一个GPU一次处理一个预测,在受支持的实例中每小时成本最低 |
| 团队使用或并发请求突发 | ml.g6.12xlarge搭配max_concurrent_invocations_per_instance=4 |
四个预测并行运行而非排队 |
| 高通量筛选 | ml.g6.12xlarge搭配自扩缩 |
在受支持的实例中每实例吞吐量最高。随队列深度扩展并缩减至零 |
使用你自己的MSA提高预测准确性
OpenFold3 可以为每条蛋白质链额外接收一个多序列比对(MSA)作为输入。MSA 是一组与你的查询序列比对上的进化相关序列。这些序列之间的共变信号指示哪些残基可能相互接触,因此 MSA 可以显著提高许多蛋白质和复合物的预测准确度。
为了帮助将你的序列保留在 AWS 账户内,JumpStart 端点以网络隔离方式运行,不会调用外部 MSA 服务器。相反,它在请求中接受预计算的 MSA。你使用已有的流程和序列数据库,在你控制的基础设施上生成比对,并将其传递给端点。这种设计有助于需要保护专有序列隐私的药物发现项目。
如果你不提供 MSA,端点仅根据序列进行预测。这对于小型、特征明确的蛋白质(如泛素)效果良好。对于大多数真实世界的靶标和复合物,请包含 MSA。本节其余部分将展示如何操作,以及它能带来多大改善。
自带 MSA
每条蛋白质链接受一个 precomputed_msa 字段:一个比对列表,每个比对包含一个 name 以及 A3M 或 Stockholm 格式的比对 content。对于单个合并的比对,例如 ColabFold MMseqs2 搜索的输出,请使用名称 colabfold_main:
with open("ca2_colabfold.a3m") as f:
ca2_msa = f.read()
payload["queries"]["ca2_acetazolamide"]["chains"][0]["precomputed_msa"] = [
{"name": "colabfold_main", "content": ca2_msa}
]此示例只有一条蛋白质链。如果查询包含多条不同的蛋白质链,请为每条链计算 MSA,并将其附加到该链的 precomputed_msa 字段。相同的副本列在一个 chain_ids 条目下,例如同源二聚体,共享单个 MSA。
如果你使用不同的流程生成 MSA,端点也接受其他比对名称,例如 uniref90_hits、mgnify_hits 和 bfd_uniclust_hits。由于请求暂存在 S3 中,你可以发送大型比对。在我们的测试中,几兆字节的比对可以正常工作。
你可以使用 ColabFold(MMseqs2)或 HHblits 或 JackHMMER 等工具生成 MSA。为了帮助将专有序列保留在你的环境内,请在你自己的 AWS 基础设施上针对本地托管的序列数据库运行搜索,而不是将序列发送到公共服务器。例如,ColabFold 的 colabfold_batch 命令默认使用公共 ColabFold MSA 服务器,因此对于专有序列,请改用 colabfold_search 配合本地托管的数据库。
MSA 能带来多大差异?
我们在同一端点上对上一节中的碳酸酐酶 II、锌和乙酰唑胺复合物进行了两次折叠:一次在单序列模式下,一次使用来自 ColabFold MMseqs2 搜索的 MSA。
| 指标(5 个样本中的第一个) | 无 MSA | 有 MSA |
| 平均 pLDDT | 30 | 98 |
| pTM | 0.22 | 0.98 |
| ipTM | 0.49 | 0.98 |
| 锌–His94/His96/His119 距离 | 2.3–9.4 Å | 各 2.0 Å |
| 乙酰唑胺氮–锌距离 | 5.6 Å | 2.0 Å |
| 预测时间 | 约 1.5 分钟 | 约 1.5 分钟 |
使用 MSA 后,OpenFold3 恢复了已知的结合模式。锌离子由三个活性位点组氨酸配位,乙酰唑胺磺酰胺氮直接与锌结合,与实验结构一致。不使用 MSA 时,蛋白质折叠置信度低,结合位点未形成。添加 MSA 没有明显改变预测时间。
图 2:相同的碳酸酐酶 II、锌和乙酰唑胺查询在不使用 MSA(左)和使用 MSA(右)的情况下折叠。每个预测按 pLDDT 着色,并叠加在实验晶体结构上(灰色,PDB 3HS4;Sippel 等,2009)。
将预测结果叠加到晶体结构上,基于 MSA 的预测与晶体结构高度吻合:蛋白质骨架偏差为 0.2 Å(Cα RMSD),锌离子偏差小于 0.1 Å,乙酰唑胺重原子偏差为 0.4 Å RMSD。单序列预测则发生错误折叠,Cα RMSD 为 22 Å。
图 3:基于 MSA 的预测中结合位点的特写。预测的锌离子(紫色)由 His94、His96 和 His119 以及乙酰唑胺的磺酰胺氮(绿色)配位。晶体构象(灰色)几乎完全被预测结果遮挡。
碳酸酐酶 II 与乙酰唑胺的复合物是蛋白质数据库中研究最广泛的复合物之一,密切相关的结构很可能包含在 OpenFold3 的训练数据中。请将此示例视为对工作流程以及 MSA 重要程度的演示,而非对新颖靶点准确性的基准测试。
当你将 OpenFold3 的结果与其他工具或已发表的基准进行比较时,请确保同类比较:单序列结果与单序列结果比较,基于 MSA 的结果与基于 MSA 的结果比较。
帮助降低成本:空闲时缩减至零
异步端点在没有请求排队时可以缩减至零实例,并在新请求到达时重新扩容。这种模式适合结构预测工作,因为这类工作通常以突发方式运行。
aas = boto3.client("application-autoscaling")
resource_id = f"endpoint/{endpoint_name}/variant/AllTraffic"
aas.register_scalable_target(
ServiceNamespace="sagemaker",
ResourceId=resource_id,
ScalableDimension="sagemaker:variant:DesiredInstanceCount",
MinCapacity=0,
MaxCapacity=4,
)
# Scale out and in to keep about 4 queued requests per instance
aas.put_scaling_policy(
PolicyName="openfold3-backlog",
ServiceNamespace="sagemaker",
ResourceId=resource_id,
ScalableDimension="sagemaker:variant:DesiredInstanceCount",
PolicyType="TargetTrackingScaling",
TargetTrackingScalingPolicyConfiguration={
"TargetValue": 4.0,
"CustomizedMetricSpecification": {
"MetricName": "ApproximateBacklogSizePerInstance",
"Namespace": "AWS/SageMaker",
"Dimensions": [{"Name": "EndpointName", "Value": endpoint_name}],
"Statistic": "Average",
},
"ScaleInCooldown": 600,
"ScaleOutCooldown": 300,
},
)目标跟踪会在队列空闲一段时间后将端点缩减至零。在我们的测试中,这大约需要 15–20 分钟。然而,目标跟踪不会将端点从零重新扩容。要在请求到达且没有实例运行时启动实例,请添加一个由 HasBacklogWithoutCapacity 指标触发的步进扩缩策略:
cw = boto3.client("cloudwatch")
step = aas.put_scaling_policy(
PolicyName="openfold3-scale-from-zero",
ServiceNamespace="sagemaker",
ResourceId=resource_id,
ScalableDimension="sagemaker:variant:DesiredInstanceCount",
PolicyType="StepScaling",
StepScalingPolicyConfiguration={
"AdjustmentType": "ChangeInCapacity",
"MetricAggregationType": "Average",
"Cooldown": 300,
"StepAdjustments": [{"MetricIntervalLowerBound": 0, "ScalingAdjustment": 1}],
},
)
cw.put_metric_alarm(
AlarmName=f"{endpoint_name}-has-backlog-without-capacity",
Namespace="AWS/SageMaker",
MetricName="HasBacklogWithoutCapacity",
Dimensions=[{"Name": "EndpointName", "Value": endpoint_name}],
Statistic="Average",
Period=60,
EvaluationPeriods=2,
DatapointsToAlarm=2,
Threshold=1,
ComparisonOperator="GreaterThanOrEqualToThreshold",
TreatMissingData="missing",
AlarmActions=[step["PolicyARN"]],
)当请求到达零实例的端点时,它会在队列中等待,同时启动一个实例。从零扩容后的第一个请求将需要更长时间完成,因为这包括告警触发、实例启动以及预测本身的时间。同一突发中的后续请求将以正常速度运行。这种权衡对于批处理工作负载效果很好。对于交互式使用,请设置 MinCapacity=1 以保持一个实例处于热备状态。更多详情,请参阅自动扩缩异步端点。
清理
为避免产生费用,请在完成后删除您创建的资源。如果您配置了自动扩缩,请先移除它:删除 Amazon CloudWatch 告警,然后注销可扩缩目标。接着删除每个端点、其端点配置及其模型,包括您部署的多 GPU 端点(如果已部署):
# 1. If you configured autoscaling, remove it before deleting the endpoint
cw.delete_alarms(AlarmNames=[f"{endpoint_name}-has-backlog-without-capacity"])
aas.deregister_scalable_target(
ServiceNamespace="sagemaker",
ResourceId=resource_id,
ScalableDimension="sagemaker:variant:DesiredInstanceCount",
)
# 2. Delete the endpoints, endpoint configurations, and models
# (include only the ones you deployed)
from sagemaker.core.resources import EndpointConfig
for m, ep in ((model, endpoint), (multi_gpu_model, multi_gpu_endpoint)):
endpoint_config = EndpointConfig.get(endpoint_config_name=ep.endpoint_name)
m.delete()
ep.delete()
endpoint_config.delete()如果您不再需要请求和结果对象,也可以从 s3://<bucket>/openfold3-async/ 中删除它们。
结论
借助 SageMaker JumpStart 上的 OpenFold3,研究和药物发现团队只需几行代码即可将开放、可检查的结构预测模型部署到自己的 AWS 账户中。异步推理可处理长时间运行的折叠和突发性批处理工作负载,多 GPU 实例可并行运行预测。网络隔离有助于保护专有序列的私密性,并且由于您可以自带 MSA,因此无需将序列发送到外部服务器即可受益于 MSA。通过缩减至零,当队列空闲足够长时间使端点缩容后,您将停止为端点实例付费。其他费用(例如 Amazon S3 存储费用)仍然适用。
本文中的代码为用于实验的示例代码。在将其用于专有或受监管数据之前,请与您的安全和法务团队合作,以满足您组织的安全、监管和合规要求。
要开始使用,请在 SageMaker JumpStart 中搜索 OpenFold3,或从模型卡片中打开示例笔记本。要了解本文中使用的功能的更多信息,请参阅异步推理和为异步端点自动扩缩。如需讨论您的结构生物学工作负载,请联系您的 AWS 客户团队。
关于作者
Malav Shastri
Malav Shastri 是 AWS 的软件开发工程师,他负责构建基础设施,通过 Amazon SageMaker JumpStart 将开源和专有基础模型(文本、视觉、音频和视频)提供给客户。在过去的 5 年多里,他一直在 ML 技术栈的各个层面工作,从 SageMaker 工作流和容器到模型部署、微调流水线,以及实现模型接入自动化的多智能体系统。他最感兴趣的是模型服务以及让 ML 在生产中可靠且快速的基础设施。
Manu Pillai
Manu 是 AWS 的资深 专业解决方案架构师,专注于为生命科学组织提供高性能计算和云解决方案。他在生命科学行业拥有超过 12 年的经验,带来深厚的领域专业知识,帮助组织利用云技术推进科学发现。
Kevin Wang
Kevin 是 AWS 深度学习容器/AMI 团队的软件开发工程师,他负责构建和维护客户在 Amazon Elastic Compute Cloud (Amazon EC2)、Amazon SageMaker AI 和 Amazon Elastic Kubernetes Service (Amazon EKS) 上运行的深度学习镜像。他的团队在 GitHub 上以开源方式管理这些镜像及其构建系统。
Anamaria Todor
Anamaria 是常驻丹麦哥本哈根的首席解决方案架构师。她 4 岁时第一次见到计算机,此后便一直投身于计算机科学、电子游戏和工程领域。她曾在丹麦多家公司担任过各种技术职务,从自由职业者、全栈开发人员到数据工程师、技术负责人和 CTO,主要专注于游戏和广告行业。她在 AWS 工作超过 5 年,担任首席解决方案架构师,主要专注于生命科学和 AI/ML。Anamaria 拥有应用工程与计算机科学学士学位、计算机科学硕士学位,以及超过 10 年的 AWS 经验。在不工作或不玩电子游戏时,她会指导女孩和女性专业人士理解和找到自己在技术领域的道路。
Zhuofu Bai
Zhuofu 是 AWS 的软件开发经理。
来源:AWS Machine Learning Blog · aws.amazon.com


