跳到正文
AWS Machine Learning Blog· Malav Shastri·· 3 小时前AI 评分36

OpenFold3 上线 Amazon SageMaker JumpStart,可预测生物分子结构

OpenFold3 on Amazon SageMaker JumpStart: Predict biomolecular structures

AI 导读

OpenFold3-preview 现可从 Amazon SageMaker JumpStart 部署,在 AWS 托管容器与异步推理端点上预测蛋白质、DNA、RNA、小分子配体及其复合物的三维结构。

正文 · AI 翻译

你可以从 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 异步推理上:

  1. 您将请求上传到 Amazon S3 并调用 InvokeEndpointAsync,该调用会立即返回。
  2. 端点将请求排队并处理,调用超时最长可达 1 小时。
  3. 端点将结果或错误详情写入您选择的 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

  1. 打开 SageMaker Studio,在导航窗格中选择 JumpStart。
  2. 搜索 OpenFold3。
  3. 打开模型卡片以查看模型描述、许可证和支持的实例类型。
  4. 从模型卡中打开 示例 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()

OpenFold3 predicted structure of ubiquitin colored by per-atom pLDDT confidence, from red for low to blue for high

图 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 没有明显改变预测时间。

Carbonic anhydrase II predictions without an MSA on the left and with an MSA on the right, each overlaid on the gray PDB 3HS4 crystal structure

图 2:相同的碳酸酐酶 II、锌和乙酰唑胺查询在不使用 MSA(左)和使用 MSA(右)的情况下折叠。每个预测按 pLDDT 着色,并叠加在实验晶体结构上(灰色,PDB 3HS4;Sippel 等,2009)。

将预测结果叠加到晶体结构上,基于 MSA 的预测与晶体结构高度吻合:蛋白质骨架偏差为 0.2 Å(Cα RMSD),锌离子偏差小于 0.1 Å,乙酰唑胺重原子偏差为 0.4 Å RMSD。单序列预测则发生错误折叠,Cα RMSD 为 22 Å。

Close-up of the predicted binding site showing the zinc ion coordinated by His94, His96, His119, and the acetazolamide sulfonamide nitrogen

图 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

Malav Shastri 是 AWS 的软件开发工程师,他负责构建基础设施,通过 Amazon SageMaker JumpStart 将开源和专有基础模型(文本、视觉、音频和视频)提供给客户。在过去的 5 年多里,他一直在 ML 技术栈的各个层面工作,从 SageMaker 工作流和容器到模型部署、微调流水线,以及实现模型接入自动化的多智能体系统。他最感兴趣的是模型服务以及让 ML 在生产中可靠且快速的基础设施。

Manu Pillai

Manu Pillai

Manu 是 AWS 的资深 专业解决方案架构师,专注于为生命科学组织提供高性能计算和云解决方案。他在生命科学行业拥有超过 12 年的经验,带来深厚的领域专业知识,帮助组织利用云技术推进科学发现。

Kevin Wang

Kevin Wang

Kevin 是 AWS 深度学习容器/AMI 团队的软件开发工程师,他负责构建和维护客户在 Amazon Elastic Compute Cloud (Amazon EC2)、Amazon SageMaker AI 和 Amazon Elastic Kubernetes Service (Amazon EKS) 上运行的深度学习镜像。他的团队在 GitHub 上以开源方式管理这些镜像及其构建系统。

Anamaria Todor

Anamaria Todor

Anamaria 是常驻丹麦哥本哈根的首席解决方案架构师。她 4 岁时第一次见到计算机,此后便一直投身于计算机科学、电子游戏和工程领域。她曾在丹麦多家公司担任过各种技术职务,从自由职业者、全栈开发人员到数据工程师、技术负责人和 CTO,主要专注于游戏和广告行业。她在 AWS 工作超过 5 年,担任首席解决方案架构师,主要专注于生命科学和 AI/ML。Anamaria 拥有应用工程与计算机科学学士学位、计算机科学硕士学位,以及超过 10 年的 AWS 经验。在不工作或不玩电子游戏时,她会指导女孩和女性专业人士理解和找到自己在技术领域的道路。

Zhuofu Bai

Zhuofu Bai

Zhuofu 是 AWS 的软件开发经理。

来源:AWS Machine Learning Blog · aws.amazon.com