OpenFold3 on Amazon SageMaker JumpStart: Predict biomolecular structures

You can deploy OpenFold3-preview from Amazon SageMaker JumpStart . The OpenFold Consortium and the AlQuraishi Lab at Columbia University developed OpenFold3 as an open source biomolecular structure prediction model. With a few lines of code, you can typically deploy OpenFold3 into your own AWS account. You can then predict the 3D structures of proteins, DNA, RNA, small-molecule ligands, and the complexes they form. The endpoint runs with network isolation, which you can use to help keep your sequences within your AWS account. Computational drug discovery teams now use structure prediction as a core part of their work. Teams use it in target identification, antibody and protein engineering, and structure-based virtual screening. Teams have found running these models challenging. You need GPU infrastructure, a working CUDA and PyTorch stack, large model checkpoints, and a serving layer that handles requests that take minutes rather than milliseconds. SageMaker JumpStart can handle much of that work for you. OpenFold3 runs in an AWS-maintained container on a managed Amazon SageMaker AI endpoint. The endpoint is designed to handle request queuing and GPU warmup, and you can add autoscaling. In this post, we show you how to: Deploy OpenFold3-preview on a SageMaker asynchronous inference endpoint. Submit structure prediction jobs and interpret the results, including predicted structures and confidence scores. Run several predictions in parallel on multi-GPU instances . Supply your own multiple sequence alignments (MSAs) to improve prediction accuracy while helping keep your sequences private. Help reduce costs by scaling your endpoint to zero when it’s idle. Overview of OpenFold3 OpenFold3 is an all-atom structure prediction model. It predicts the joint 3D structure of a set of molecules from their sequences. It isn’t limited to single protein chains: It can fold multi-chain protein complexes, protein–nucleic acid assemblies, and protein–ligand complexes in one prediction. Attribute Details Developer OpenFold Consortium / AlQuraishi Lab, Columbia University Model type All-atom biomolecular structure prediction Inputs Protein, DNA, and RNA sequences. Small-molecule ligands as SMILES Outputs Predicted structures in mmCIF format, with confidence metrics (pLDDT, PAE, PDE) Version OpenFold3 0.4.1, of3-p2-155k checkpoint License Apache 2.0 JumpStart model ID huggingface-proteinfolding-openfold3 Supported instances ml.g6.2xlarge , ml.g6.4xlarge (default), ml.g6.12xlarge The OpenFold3 developers release its code and model weights under the Apache 2.0 license, and they publish the training data and methods. Research organizations and pharmaceutical companies can inspect the model, fine-tune it, and build on it under the terms of that license. The developers describe OpenFold3-preview as a research preview. Use cases Target and complex structure prediction – Predict structures for targets that have no experimental structure, including multi-chain assemblies and protein–nucleic acid complexes. Protein–ligand co-folding – Fold a protein together with a small molecule supplied as SMILES to generate binding-pose hypotheses for structure-based design. Antibody and protein engineering – Evaluate designed variants by predicting their structures and comparing confidence across candidates. High-throughput screening pipelines – Queue large batches of prediction jobs against an autoscaling endpoint and process the results directly from Amazon Simple Storage Service (Amazon S3). Why OpenFold3 uses asynchronous inference Most JumpStart models, such as large language models and embedding models, deploy to real-time endpoints. A real-time endpoint returns its response within a single request, and that request is capped at 60 seconds. Structure prediction works differently. Folding time grows with the number and length of the chains in a query, and real-world proteins and complexes often take several minutes. For this reason, OpenFold3 on JumpStart runs on SageMaker Asynchronous Inference : You upload the request to Amazon S3 and call InvokeEndpointAsync , which returns immediately. The endpoint queues the request and processes it, with an invocation timeout of up to 1 hour. The endpoint writes the result, or the error details, to an S3 location you choose. You can poll for it or get a notification through Amazon Simple Notification Service (Amazon SNS). Two properties of asynchronous endpoints suit scientific workloads especially well. First, they queue bursts of requests internally, so you can submit a large batch of sequences without building your own job queue. Second, they can scale to zero instances when the queue is empty, which helps you avoid paying for idle instances between experiments. Prerequisites To follow this walkthrough, you need: An AWS account, and a SageMaker AI execution role for the endpoint. Following least privilege, this AWS Identity and Access Management (IAM) role needs read and write access only to the S3 prefix that you use for requests and results. For more information, see How to use SageMaker AI execution roles . IAM permissions for the identity that runs the code: to create, describe, invoke, and delete the model, endpoint configuration, and endpoint, plus iam:PassRole for the execution role. For the optional scale-to-zero section, it also needs AWS Application Auto Scaling and Amazon CloudWatch alarm permissions. In production, scope these permissions to your resources rather than using broad managed policies such as AmazonSageMakerFullAccess . For more information, see Complete the prerequisites for asynchronous inference. Access to Amazon SageMaker Studio , a SageMaker notebook instance, or a local IDE. The SageMaker Python SDK v3 and py3Dmol, which you can install with pip install "sagemaker>=3, . Service quota for at least one ml.g6.4xlarge endpoint instance in your AWS Region. You also need quota for ml.g6.12xlarge if you follow the multi-GPU section. You can check your quotas and request increases in the Service Quotas console under Amazon SageMaker . Important: SageMaker endpoints are billed per instance-hour while they’re running, including while they’re idle. To avoid unexpected charges, follow the Clean up section, or configure scale-to-zero as described later in this post. For pricing details, see Amazon SageMaker AI pricing . Discover OpenFold3 in SageMaker JumpStart Open SageMaker Studio and choose JumpStart in the navigation pane. Search for OpenFold3 . Open the model card to review the model description, license, and supported instance types. Open the example notebook from the model card. It contains the complete asynchronous deployment walkthrough shown in this post. Deploy OpenFold3 with the SageMaker Python SDK First, build the model from its JumpStart configuration and define an AsyncInferenceConfig . The output and failure paths are the S3 locations where the endpoint writes results and error details. 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}") Setting accept_eula=True confirms that you’ve reviewed and accepted the m