OpenFold3 on Amazon SageMaker JumpStart: Predict biomolecular structures

Amazon Web Services added the OpenFold3-preview biomolecular structure prediction model to Amazon SageMaker JumpStart, allowing computational biology teams to deploy the system into their own AWS accounts using managed infrastructure. Developed by the OpenFold Consortium and the AlQuraishi Lab at Columbia University, the open source model generates joint 3D structure predictions for proteins, DNA, RNA, and small-molecule ligands. The integration runs inside an AWS-managed container on an asynchronous SageMaker endpoint backed by network isolation, as detailed in AWS Machine Learning Blog's report.

This walkthrough details how life sciences teams configure the JumpStart model, submit asynchronous inference jobs through Amazon S3, run multi-GPU parallel predictions, supply custom multiple sequence alignments, and configure scale-to-zero autoscaling. The setup serves computational drug discovery teams, protein engineers, and screening groups that need private infrastructure to handle long-running folding tasks without managing custom PyTorch and CUDA driver stacks.

01All-atom biomolecular modeling and licensing

OpenFold3 operates as an all-atom structure prediction engine that maps molecular sequences directly to 3D atomic coordinates. Rather than restricting predictions to isolated single-chain proteins, the model predicts multi-chain protein complexes, protein–nucleic acid assemblies, and protein–ligand binding configurations within a single execution pass. Small molecules enter the pipeline formatted as SMILES strings alongside standard FASTA sequences.

The OpenFold Consortium and Columbia University release the system as a research preview under the Apache 2.0 license. This license covers code, training methods, data definitions, and model weights, granting pharmaceutical researchers full rights to inspect architecture layers or fine-tune parameters on proprietary data. The JumpStart model ID huggingface-proteinfolding-openfold3 packages the OpenFold3 0.4.1 release using the of3-p2-155k checkpoint.

Results write out in standard mmCIF format alongside three confidence scores: predicted local distance difference test (pLDDT), predicted aligned error (PAE), and predicted distance error (PDE). Computational teams evaluate target structures lacking experimental data, score candidate stability for antibody design, or generate binding poses to guide structure-based virtual screening campaigns.

02Asynchronous inference pipeline and timeouts

Standard real-time SageMaker endpoints enforce a 60-second execution cutoff, which fails for complex biomolecular folding jobs that routinely run for several minutes. OpenFold3 addresses this workload profile by deploying through SageMaker Asynchronous Inference, raising the individual invocation timeout ceiling up to one full hour.

The inference cycle begins by uploading an input JSON payload to an Amazon S3 request prefix and invoking InvokeEndpointAsync. The API call returns an immediate confirmation token, allowing client processes to proceed without maintaining open socket connections. The managed endpoint queues requests internally, buffers high-volume bursts, executes the folding pipeline, and deposits structured mmCIF outputs or error logs into designated S3 destination buckets.

Teams track job completions by polling the target S3 output prefix or subscribing downstream automation directly to Amazon Simple Notification Service notifications. The managed container handles initial GPU warmup cycles automatically, ensuring pipeline stability during sudden traffic surges without requiring custom server-side task managers or third-party queue systems.

03Hardware options and environment configuration

The JumpStart implementation defaults to the ml.g6.4xlarge compute profile, while supporting ml.g6.2xlarge and multi-GPU ml.g6.12xlarge instances. Deploying these instances requires sufficient service quotas within the target AWS Region under Amazon SageMaker AI. Because endpoints bill continuously per instance-hour while active, infrastructure teams manage costs through explicit teardown or autoscaling configurations.

Prerequisites include an AWS account running the SageMaker Python SDK v3 alongside py3Dmol for visualization, installable via pip. The host execution identity requires AWS Identity and Access Management permissions to create, invoke, and delete SageMaker endpoints, accompanied by iam:PassRole rights. Security policies restrict the endpoint's execution role to reading and writing solely within the designated S3 workflow bucket prefixes.

To deploy via Python, engineers build a ModelBuilder instance referencing JumpStartConfig with accept_eula=True to acknowledge the Apache 2.0 license. Passing an AsyncInferenceConfig specifying output_path and failure_path parameters initializes the deployment call. S3 data security rules recommend using customer-managed AWS Key Management Service keys, bucket-level permission controls, and automated S3 Lifecycle policies.

04Custom sequence alignments and privacy

OpenFold3 improves prediction precision when supplied with rich evolutionary data through multiple sequence alignments. The model accepts user-generated alignments as part of the invocation payload, avoiding reliance on external alignment services and preventing proprietary amino acid sequences from leaking outside the customer's Virtual Private Cloud boundary.

Because the SageMaker endpoint operates with complete network isolation enabled, external web access is blocked during model inference. Proprietary protein engineering targets and competitive therapeutic sequences stay confined to the customer's virtual storage boundary. Input payloads pair alignment data with chemical structures defined as SMILES strings, allowing co-folding analysis without exposing early-stage intellectual property.

Output payloads deliver structural metrics alongside atomic coordinates. Structural engineers evaluate pLDDT arrays to assess regional backbone confidence, while PAE matrices indicate relative orientation certainty across interacting protein domains. PDE metrics provide complementary distance estimates, helping researchers distinguish high-confidence binding interactions from unverified structural artifacts before initiating laboratory synthesis.

05Multi-GPU parallelism and zero scaling

For large screening batches, the ml.g6.12xlarge instance type provides multi-GPU compute capacity that processes multiple structural candidates simultaneously. Multi-GPU parallel runs split inference queues across separate hardware engines, lowering turnaround latency during screening passes that evaluate hundreds of candidate variants across candidate antibody libraries.

To eliminate compute costs between experimental campaigns, teams bind AWS Application Auto Scaling policies to the SageMaker asynchronous endpoint. The scaling configuration attaches to the endpoint's ApproximateBacklogSize CloudWatch metric, instructing the cluster to scale down to zero active instances whenever the incoming queue remains empty.

When a researcher uploads new sequence batches to S3 and calls InvokeEndpointAsync, the arriving requests repopulate the CloudWatch backlog metric. SageMaker responds by spinning up instance hardware, loading the of3-p2-155k checkpoint into memory, and executing the processing queue before scaling back down to zero instances upon queue exhaustion.