diff --git a/README.md b/README.md
index ec5f8658..d5986ee7 100644
--- a/README.md
+++ b/README.md
@@ -42,6 +42,7 @@ For information on using ESM3, see the [ESM3 README](https://github.com/Biohub/e
- [ESMC](#esmc)
- [ESMC Sparse Autoencoders](#esmc-sparse-autoencoders)
- [ESMFold2](#esmfold2)
+- [Batch API](#batch-api)
- [Parallel Job Execution](#parallel-job-execution)
- [Frontier-Safety](#frontier-safety)
- [Licenses](#licenses)
@@ -50,7 +51,44 @@ For information on using ESM3, see the [ESM3 README](https://github.com/Biohub/e
## ESM on PyPI
-[`esm v3.4.post1`](https://github.com/Biohub/esm/releases/tag/v3.4.1.post1) is now available on PyPI with full support for ESMC and ESMFold2.
+[`esm v3.4.1.post1`](https://github.com/Biohub/esm/releases/tag/v3.4.1.post1) is now available on PyPI with full support for ESMC and ESMFold2.
+
+### Installation of ESM Package
+
+| Command | Contents |
+| :---- | :---- |
+| `pip install esm` | The models and the reference PyTorch path. No CUDA kernels. |
+| `pip install "esm[gpu]"` | Adds the xformers and cuequivariance kernels. Linux x86-64 only; a no-op elsewhere. |
+
+> **Changed in this release.** `pip install esm` no longer installs the
+> cuequivariance backend. Install `"esm[gpu]"` to keep it.
+
+Notes:
+
+- Kernels are optional at runtime. The default install runs everywhere, it is only slower on a GPU.
+- Triton is not one of the extras. Torch requires it on Linux, so the `"fused"` backend already works in the default install.
+
+#### Install size
+
+| Install | Size |
+| :---- | :---- |
+| `pip install esm` | about 5.6 GB |
+| `pip install "esm[gpu]"` | about 5.8 GB |
+
+The two are nearly the same size because torch on PyPI is always the CUDA
+build, about 4.6 GB of either total. The extra adds roughly 0.2 GB on top.
+
+Installing torch from its CPU-only index replaces that 4.6 GB CUDA build with
+an 0.8 GB CPU build, which is where the real saving is:
+
+```
+pip install esm --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple
+```
+
+Worth doing if you have no NVIDIA GPU, are on Apple Silicon, or want smaller CI
+images and containers.
+
+`flash-attn` and `transformer-engine` are not installable from PyPI against our torch pin. Prebuilt wheels for both are published in the [EvolutionaryScale wheels release](https://github.com/evolutionaryscale/wheels/releases).
## ESM Package vs. Hugging Face Transformers Usage
@@ -150,7 +188,7 @@ For tutorials on how to use ESMC, see our [tutorials](https://github.com/Biohub/
Sparse autoencoders (SAE) are an unsupervised method for decomposing representations of large transformer language models into interpretable units. We released SAEs trained on ESMC to reveal the interpretable units of functional organization that ESMC's world model has learned.
-The sparse autoencoder used in the Atlas and analyzed in the paper, `ESMC-6B-sae-layer60-k64-codebook16384`, is built on the ESMC 6B model. We also provide human-interpretable, agent-generated feature descriptions for this SAE's codebook.
+The sparse autoencoder used in the example below, `ESMC-6B-sae-k64-codebook16384`, is built on the ESMC 6B model and covers all layers. The SAE used in the Atlas and analyzed in the paper is `ESMC-6B-sae-layer60-k64-codebook16384`; the human-interpretable, agent-generated feature descriptions apply to its codebook.
Codebase, model weights, and model variants for ESMC SAEs are available through [Hugging Face](https://huggingface.co/collections/biohub/esmc-saes-for-hidden-states-all-layers).
@@ -266,7 +304,7 @@ with open("1mht_pred.cif", "w") as f:
f.write(result.complex.to_mmcif())
```
-> **AMD ROCm users:** use ROCm 6.4 with PyTorch 2.9 or newer.
+> **AMD ROCm users:** use ROCm 6.4 or newer. `esm` requires PyTorch 2.11.
### Running ESMFold2 Through the Biohub Platform
@@ -312,9 +350,15 @@ with open("result.cif", "w") as f:
For tutorials on how to use ESMFold2, see our [tutorials](https://github.com/Biohub/esm/tree/main/cookbook/tutorials).
+For high quality structure prediction, see the [maximum-accuracy folding SDK guide](cookbook/snippets/fold_max_accuracy.md).
+
+## Batch API
+
+The Biohub Batch API processes larger or longer-running workloads in the background and offers higher limits. See `ForgeBatchClient` and the [Batch API guide](cookbook/snippets/batch.md) for examples of submission, progress checks, result retrieval, cache control, and cancellation. Queued tasks expire after 24 hours; results are retained for 31 days after submission.
+
## Parallel Job Execution
-For jobs that require processing multiple inputs, the Parallel Executor sends parallel requests, executing them concurrently and efficiently while respecting rate limits and adapting to request latency. The following example provides an example of using the parallel_executor context manager when embedding sequences.
+Use `parallel_executor` to send multiple Sync API requests concurrently while respecting rate limits and adapting to request latency. For queued background work, use `ForgeBatchClient`. The following example uses `parallel_executor` to embed sequences through the Sync API.
```py
from esm.sdk.forge import ESMCForgeInferenceClient
@@ -334,7 +378,7 @@ sequences = ["A", "AA", "AAA"]
client = ESMCForgeInferenceClient(model="esmc-6b-2024-12", url="https://biohub.ai", token="")
# Usage Example:
-# To execute a batch job, wrap your function inside the batch executor context manager.
+# To process multiple inputs through the Sync API, wrap your function in parallel_executor.
# Syntax:
# with parallel_executor() as executor:
# outputs = executor.execute_batch(user_func=, **kwargs)
diff --git a/cookbook/foldcp/README.md b/cookbook/foldcp/README.md
index b08ec4dc..3359f9db 100644
--- a/cookbook/foldcp/README.md
+++ b/cookbook/foldcp/README.md
@@ -11,10 +11,9 @@ Tested on 1x and 4x H100 80GB GPUs.
| Script | Input | GPUs | Result |
| --- | --- | --- | --- |
| `esmfold2-none.py` | 7ysz (1 protein chain ×2) | 1 | Reference single-GPU path |
-| `esmfold2-cueq.py` | 7ysz (1 protein chain ×2) | 1 | cuEquivariance backend |
| `esmfold2-fused.py` | 7ysz (1 protein chain ×2) | 1 | Triton fused backend |
| `esmfold2-cp.py` | 7ysz (1 protein chain ×2) | 4 | Same fold API with CP setup |
-| `cuequivariance_cp.py` | 5xgo (1 protein chain ×12, CL ligand) | 4 | Larger input via CP |
+| `large_input_cp.py` | 5xgo (1 protein chain ×12, CL ligand) | 4 | Larger input via CP |
| `will_fail.py` | 5xgo (same as CP) | 1 | **OOMs** — too large for 1xH100 |
## Environment setup
@@ -24,27 +23,23 @@ Start from the NVIDIA PyTorch container, then install the dependencies:
```bash
docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:26.03-py3 bash -l
-# Main ESM package
-pip install "git+https://github.com/Biohub/esm.git@main"
-# fused dependency
-pip install "esm[fused] @ git+https://github.com/Biohub/esm.git@main"
-# cuEquivariance dependency (cueq12 also exists)
-pip install "esm[cueq13] @ git+https://github.com/Biohub/esm.git@main"
-# fold-cp dependency for larger inputs
-pip install "esm[fold-cp] @ git+https://github.com/Biohub/esm.git@main"
+# ESM plus the Triton kernels
+pip install "esm[gpu] @ git+https://github.com/Biohub/esm.git@main"
```
+Context parallelism additionally needs `transformer-engine`, which is not
+installable from PyPI against our torch pin. Take the prebuilt wheel from the
+[EvolutionaryScale wheels release](https://github.com/evolutionaryscale/wheels/releases).
+
## Single-GPU accelerated backends
-[single-gpu.md](single-gpu.md) compares the three single-GPU backend choices:
-`None` for the pure PyTorch reference path, `"cuequivariance"` for NVIDIA
-cuEquivariance triangle-multiplication kernels, and `"fused"` for Triton kernels
-that fuse several ESMFold2 hot-path operations. These backends keep the same model
+[single-gpu.md](single-gpu.md) compares the two single-GPU backend choices:
+`None` for the pure PyTorch reference path and `"fused"` for Triton kernels
+that fuse several ESMFold2 hot-path operations. Both backends keep the same model
weights and outputs while trading dependencies, warm-up behavior, and speed.
```bash
python esmfold2-none.py
-python esmfold2-cueq.py
python esmfold2-fused.py
```
@@ -62,5 +57,5 @@ torchrun --nproc-per-node=4 esmfold2-cp.py
# Larger 5xgo example.
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
- torchrun --nproc-per-node=4 cuequivariance_cp.py
+ torchrun --nproc-per-node=4 large_input_cp.py
```
diff --git a/cookbook/foldcp/esmfold2-cueq.py b/cookbook/foldcp/esmfold2-cueq.py
deleted file mode 100644
index 20a89dc0..00000000
--- a/cookbook/foldcp/esmfold2-cueq.py
+++ /dev/null
@@ -1,54 +0,0 @@
-from time import time
-
-import torch
-
-from esm.models.esmfold2 import (
- ESMFold2InputBuilder,
- EsmFold2Model,
- MolecularComplexResult,
- ProteinInput,
- StructurePredictionInput,
-)
-
-spi = StructurePredictionInput( # 7ysz
- sequences=[
- ProteinInput(
- id=["A1", "B1"],
- sequence=(
- "MDNFDNYEQVASIKVIGIGGAGNNAVNRMIEAGVQGVEFIVANTDAQIISVSKSKNKIVLGKETSKGLGA"
- "GANPDVGRQAAIESAEEIKDALKGADMVFVAAGMGGGTGTGAAPIIAKLAREQGALTVGIITTPFSFEGR"
- "ARNSYAIQGTEELRKHVDSLIIISNDRLLEVIGGVPLKDSFKEADNILRQGVQTITDLIAVPSLINLDFA"
- "DIKTVMKNKGNALFGIGIGSGKDKAIEAANKAIISPLLEASIRGARDAIINVTGGNTLTLNDANDAVDIV"
- "KQAIGGEVNIIFGTAVNEHLDDEMIVTVIATGFDGSHHHHHH"
- ),
- )
- ]
-)
-
-torch.cuda.reset_peak_memory_stats()
-
-model = (
- EsmFold2Model.from_pretrained("biohub/ESMFold2", esmc_precision="bf16")
- .cuda()
- .eval()
-)
-model.set_kernel_backend("cuequivariance")
-
-start = time()
-with torch.inference_mode():
- result = ESMFold2InputBuilder().fold(
- model, spi, num_loops=10, num_sampling_steps=50, num_diffusion_samples=1, seed=0
- )
-end = time()
-peak_mib = torch.cuda.max_memory_allocated() / (1024**2)
-
-# num_diffusion_samples=1 returns a single result with confidence heads populated.
-assert isinstance(result, MolecularComplexResult)
-assert result.plddt is not None and result.ptm is not None and result.iptm is not None
-
-print(
- f"pLDDT mean: {float(result.plddt.mean()):.3f}, "
- f"pTM: {float(result.ptm):.3f}, ipTM: {float(result.iptm):.3f}"
-)
-print(f"Elapsed: {end - start:.2f} sec")
-print(f"Max VRAM: {peak_mib:.1f} MB")
diff --git a/cookbook/foldcp/fold-cp.md b/cookbook/foldcp/fold-cp.md
index fb3bcf5a..a792be28 100644
--- a/cookbook/foldcp/fold-cp.md
+++ b/cookbook/foldcp/fold-cp.md
@@ -14,10 +14,13 @@ proteins than fit on a single GPU. Wrapping results in the same object with a fe
## Quickstart
-Install the fold-cp dependency
+Install ESM, then `transformer-engine`. PyPI only ships it as an sdist that needs
+a CUDA toolchain, and its core wants a cuBLAS symbol newer than the one torch
+pins, so take the prebuilt wheel from the
+[EvolutionaryScale wheels release](https://github.com/evolutionaryscale/wheels/releases).
```bash
-pip install "esm[fold-cp] @ git+https://github.com/Biohub/esm.git@main"
+pip install "esm[gpu] @ git+https://github.com/Biohub/esm.git@main"
```
Use [esmfold2-cp.py](esmfold2-cp.py) as the CP version of the single-GPU
@@ -43,7 +46,7 @@ unchanged.
- Python environment:
- This ESM package
- - transformer-engine 2 (installed by `esm[fold-cp]` dependency)
+ - transformer-engine 2, from the EvolutionaryScale wheels release
- Use a perfect-square number of GPUs: 1, 4, 9, 16, ...
- Keep `num_diffusion_samples=1` in `ESMFold2InputBuilder().fold(...)`. The CP
diffusion path currently expects one diffusion sample.
@@ -111,7 +114,7 @@ Two differences are expected, both small:
- **Rounding** differs at about 1e-3, because numbers are summed across GPUs in a
different order.
-Single-GPU kernel backends such as `"fused"` and `"cuequivariance"` are ignored by the CP-wrapped stages; those stages use the distributed CP implementations instead.
+The single-GPU `"fused"` kernel backend is ignored by the CP-wrapped stages; those stages use the distributed CP implementations instead.
## Larger example
@@ -125,7 +128,7 @@ Context parallelism sharding will split the memory and allow this protein to be
Launch with torchrun on a perfect-square number of GPUs:
```bash
-PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True torchrun --nproc-per-node=4 cuequivariance_cp.py
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True torchrun --nproc-per-node=4 large_input_cp.py
```
## How context parallelism works
diff --git a/cookbook/foldcp/cuequivariance_cp.py b/cookbook/foldcp/large_input_cp.py
similarity index 96%
rename from cookbook/foldcp/cuequivariance_cp.py
rename to cookbook/foldcp/large_input_cp.py
index 3bada7d7..539bf8d7 100644
--- a/cookbook/foldcp/cuequivariance_cp.py
+++ b/cookbook/foldcp/large_input_cp.py
@@ -30,7 +30,7 @@
]
)
-# torchrun entrypoint: torchrun --nproc_per_node=4 cuequivariance_cp.py
+# torchrun entrypoint: torchrun --nproc_per_node=4 large_input_cp.py
# torchrun sets LOCAL_RANK / WORLD_SIZE / RANK / MASTER_ADDR / MASTER_PORT.
# world_size must be a perfect square (the CP grid is n×n).
local_rank = int(os.environ["LOCAL_RANK"])
@@ -65,7 +65,7 @@
# Fused kernels accelerate the modules *outside* the CP region (notably the
# diffusion sampler / structure head). The CP-wrapped trunk and MSA encoder
# no-op this call, so it never touches the distributed ring math.
-model.set_kernel_backend("cuequivariance")
+model.set_kernel_backend("fused")
start = time()
with torch.inference_mode():
diff --git a/cookbook/foldcp/single-gpu.md b/cookbook/foldcp/single-gpu.md
index 49d74105..dd1aa6dd 100644
--- a/cookbook/foldcp/single-gpu.md
+++ b/cookbook/foldcp/single-gpu.md
@@ -9,12 +9,12 @@ ESMFold2 runs the expensive Pairformer/attention math through a selectable
```python
model = ESMFold2Model.from_pretrained("biohub/ESMFold2").cuda().eval()
-model.set_kernel_backend(None | "fused" | "cuequivariance")
+model.set_kernel_backend(None | "fused")
```
This propagates out to every module that runs Pairformer-style blocks (`folding_trunk`,
`lm_encoder`, `parcae_coda`, `confidence_head`, `structure_head`). It only swaps the
-*implementation* of a few hot ops. Weights are untouched and the three
+*implementation* of a few hot ops. Weights are untouched and both
backends are numerically equivalent to bf16 rounding (same pLDDT/pTM).
> This page covers the single-GPU backends. For multi-GPU context parallelism see
@@ -22,22 +22,21 @@ backends are numerically equivalent to bf16 rounding (same pLDDT/pTM).
## Table of contents
-- [The three backends](#the-three-backends)
+- [The two backends](#the-two-backends)
- [None (default experience)](#none-default-experience)
-- [cuEquivariance](#cuequivariance)
- [Fused](#fused)
- [Performance summary](#performance-summary)
- [Decision guide](#decision-guide)
-## The three backends
+## The two backends
-| | None | fused | cuequivariance |
-|---|---|---|---|
-| Implementation | PyTorch | Triton kernels | [cuEquivariance](https://github.com/nvidia/cuequivariance) |
-| Accelerated Operations | none (reference) | tri-mul + LN+SwiGLU + dropout-residual + pair-bias | tri-mul |
-| Extra dependency | none | `triton>=3` | `cuequivariance-torch` (CUDA-matched build) |
-| Training / autograd | yes | inference-only | yes |
-| First-call compilation | none | yes — Triton JIT + autotune | no — precompiled kernels |
+| | None | fused |
+|---|---|---|
+| Implementation | PyTorch | Triton kernels |
+| Accelerated Operations | none (reference) | tri-mul + LN+SwiGLU + dropout-residual + pair-bias |
+| Extra dependency | none | `triton>=3` |
+| Training / autograd | yes | inference-only |
+| First-call compilation | none | yes — Triton JIT + autotune |
## None (default experience)
@@ -70,59 +69,6 @@ Elapsed: 58.02 sec
Max VRAM: 19682.0 MB
```
-## cuEquivariance
-
-[cuEquivariance](https://github.com/nvidia/cuequivariance) is NVIDIA’s precompiled CUDA kernel backend for ESMFold2’s triangle-multiplication hot path, giving a large single-GPU speedup without Triton JIT or accuracy changes.
-
-### Dependencies
-
-* `cuequivariance-torch`
-* Depending on your CUDA version
- * `cuequivariance-ops-torch-cu13`
- * `cuequivariance-ops-torch-cu12`
-
-Install the CUDA-matched extra:
-
-```bash
-pip install "esm[cueq12]" # CUDA 12 build
-```
-or
-```bash
-pip install "esm[cueq13]" # CUDA 13 build
-```
-
-You can figure out which one you need with `nvidia-smi`
-
-```shell
-$ nvidia-smi
-
-Thu Jul 16 13:03:23 2026
-+---------------------------------------------------------------------------------------+
-| NVIDIA-SMI 535.216.03 Driver Version: 535.216.03 CUDA Version: 13.2 |
-|-----------------------------------------+----------------------+----------------------+
-```
-
-In this case, the CUDA Version is 13.2, so you would need to install `esm[cueq13]`.
-
-### Usage
-
-```python
-model = ESMFold2Model.from_pretrained("biohub/ESMFold2").cuda().eval()
-model.set_kernel_backend("cuequivariance")
-```
-
-### Performance
-
-Fold a 644-residue protein using the cuEquivariance backend with [esmfold2-cueq.py](esmfold2-cueq.py) on a single H100 SXM
-
-```shell
-python esmfold2-cueq.py
-
-pLDDT mean: 0.751, pTM: 0.459, ipTM: 0.096
-Elapsed: 19.90 sec
-Max VRAM: 19680.1 MB
-```
-
## Fused
The "fused" backend uses Triton, a Python-based language for writing custom GPU kernels, to fuse several ESMFold2 hot-path operations into fewer CUDA launches for the fastest single-GPU inference while preserving the same model weights and outputs.
@@ -132,7 +78,7 @@ The "fused" backend uses Triton, a Python-based language for writing custom GPU
The "fused" backend needs Triton 3. It's GPU-only (Triton JIT-compiles to PTX) and **inference-only** (falls back to reference under autograd).
```bash
-pip install "esm[fused]" # triton>=3,<4
+pip install "esm[gpu]" # brings triton>=3,<4
```
> If Triton is not importable, `set_kernel_backend("fused")` silently runs the reference path. Check that Triton can be imported if `"fused"` is unexpectedly slow.
@@ -164,7 +110,6 @@ examples above.
| backend | elapsed | max VRAM | pLDDT | pTM | ipTM | speedup vs `None` |
|---|---:|---:|---:|---:|---:|---:|
| `None` | 58.02 s | 19682.0 MB | 0.751 | 0.458 | 0.096 | 1.0× |
-| "cuequivariance" | 19.90 s | 19680.1 MB | 0.751 | 0.459 | 0.096 | 2.9× |
| "fused" | 16.56 s | 19795.1 MB | 0.750 | 0.458 | 0.096 | 3.5× |
## Decision guide
@@ -173,8 +118,6 @@ If you want:
- **Fastest throughput →** `"fused"`.
- Fastest steady state, and the incremental per-new-length compile is only ~0.5 s
-- **Huge length variety →** `"cuequivariance"`
- — precompiled, steady state a bit slower than fused.
- **Bit-exact reference for debugging →** `None`.
> Tip: For any benchmarking, include a warm-up fold at startup to pay the ~10s global cost
diff --git a/cookbook/foldcp/will_fail.py b/cookbook/foldcp/will_fail.py
index dd8b2d56..c10806af 100644
--- a/cookbook/foldcp/will_fail.py
+++ b/cookbook/foldcp/will_fail.py
@@ -32,7 +32,7 @@
.cuda()
.eval()
)
-model.set_kernel_backend("cuequivariance")
+model.set_kernel_backend("fused")
start = time()
with torch.inference_mode():
diff --git a/cookbook/snippets/batch.md b/cookbook/snippets/batch.md
new file mode 100644
index 00000000..8f26a211
--- /dev/null
+++ b/cookbook/snippets/batch.md
@@ -0,0 +1,217 @@
+# Batch API
+
+Use the Biohub Batch API to process larger or longer-running workloads in the background. Some models/services may only be available over the Batch API, not Sync API. The Batch API offers inference at discounted rates, and should be used when immediate responses are not required.
+
+`ForgeBatchClient` lets you submit many tasks, wait for completion, and retrieve decoded results. This guide covers submission, progress checks, result retrieval, expiry, cache control, cancellation, and Python asyncio usage.
+
+## Supported services
+
+| Service | Description | Guide |
+| --- | --- | --- |
+| `fold` | Protein structure prediction with ESMFold2 | Examples below |
+| `fold_all_atom` | Molecular-complex prediction with ESMFold2 | [Input guide](../../README.md#running-esmfold2-through-the-biohub-platform) |
+| `fold_max_accuracy` | Maximum-accuracy molecular-complex prediction | [SDK guide](fold_max_accuracy.md) |
+
+We will be launching more models/services soon!
+
+## Submit, wait, and retrieve results
+
+```python
+import os
+
+from esm.sdk.api import FoldRequest
+from esm.sdk.forge import ForgeBatchClient
+
+sequences = ["MKTAYIAKQRQISFVKSHFSRQ", "GSHMASMTGGQQMGRGS"]
+
+with ForgeBatchClient(
+ token=os.environ["ESM_API_KEY"], model="esmfold2-fast-2026-05"
+) as client:
+ submission = client.fold.submit(
+ requests=[FoldRequest(sequence=sequence) for sequence in sequences]
+ )
+ print("Batch ID:", submission.batch_id)
+ summary = client.fold.wait(submission)
+ outputs = client.fold.results(submission)
+```
+
+
+`submit()` returns a `BatchSubmission` containing the batch ID. Pass that submission to `wait()` to poll until the tasks finish, then to `results()` to download and decode their outputs.
+
+`fold.results()` returns one `ESMProtein` or `BatchError` per listed task, in server listing order. A failed item does not discard
+other results. Inspect `error.task_status` to distinguish the outcome:
+
+| `task_status` | Meaning |
+| --- | --- |
+| `"failed"` | Task execution failed. |
+| `"cancelled"` | The task was cancelled, including expiry before it started. `error_msg` includes the server's reason when provided. |
+| `"queued"` or `"in_progress"` | The result is not ready; wait and retrieve it later. |
+| `None` | The error came from a request, decoding, or waiting rather than a task state. |
+
+
+`wait()` periodically polls `/batch/list` until batch processing finishes, including failed or cancelled tasks. It raises an error if waiting times out. It returns the batch summary while `results()` retrieves the outputs.
+`results()` fetches every page automatically and collects the outputs in a list.
+Use `page_size` to control how many tasks are listed per request (default 100,
+maximum 1000).
+
+A failed task appears as a `BatchError` in the list alongside successful results.
+If the batch cannot be fully retrieved, `results()` raises a `BatchError`; outputs
+already retrieved are available in `error.partial_results` as `(task_id, output)`
+pairs.
+
+### Submission errors and retries
+
+The SDK sends each submission once. A timeout or lost response may occur after
+the server accepts the work, so submitting again can create a duplicate batch.
+Automatic retries apply to status reads, batch listing, and result downloads.
+
+`submit()` raises `BatchError` on failure; `run()` returns it. If the server
+response included usable job IDs, `error.submission` retains them. Otherwise the
+error has no recovery handle. Save the handle returned by a successful submit
+before waiting for results.
+
+### Optional: save a submission for later
+
+To continue later, save the submission and restore it:
+
+```python
+import json
+from pathlib import Path
+
+from esm.sdk.batch import BatchSubmission
+
+Path("fold-job.json").write_text(json.dumps(submission.to_dict()))
+restored = BatchSubmission.from_dict(json.loads(Path("fold-job.json").read_text()))
+```
+
+Pass `restored` to `wait()` and `results()` on a new client authenticated as the
+same user and pointed at the same service.
+
+If you only kept the batch ID, build the submission directly. `endpoint` must
+match the service you call.
+
+```python
+from esm.sdk.batch import BatchSubmission
+
+submission = BatchSubmission(endpoint="fold", batch_id="")
+summary = client.fold.wait(submission)
+outputs = client.fold.results(submission)
+```
+
+## Queue expiry and result retention
+
+Both deadlines are measured from when the batch is submitted:
+
+| Deadline | Behavior |
+| --- | --- |
+| 24 hours (`expires_at`) | Tasks that have not started are cancelled without charge. Tasks already running finish normally. |
+| 31 days (`results_expire_at`) | Results are no longer available. Download and save outputs before this deadline. |
+
+Fetching an expired result produces a `BatchError` with `error_code=410` and an expiry explanation in `error_msg`.
+
+## Bypass cached results
+
+By default, submissions allow cached inference results (`disable_cache=False`).
+To bypass the cache, pass `disable_cache=True` to `submit()` or `run()`:
+
+```python
+submission = client.fold.submit(
+ requests=[FoldRequest(sequence=sequence) for sequence in sequences],
+ disable_cache=True,
+)
+```
+
+## Process results incrementally
+
+For large batches that consume too much memory in aggregate, use `iter_results()` in place of `results()` to process outputs
+as they are retrieved. It follows pagination, downloads and decodes stored outputs,
+and yields `(task_id, result)` pairs in server listing order. `results()` uses the
+same order.
+
+```python
+from esm.sdk.batch import BatchError
+
+for task_id, result in client.fold.iter_results(
+ submission, page_size=100, max_concurrency=4
+):
+ if isinstance(result, BatchError):
+ print(task_id, result.error_msg)
+ continue
+ result.to_pdb(f"{task_id}.pdb")
+```
+
+`max_concurrency` bounds both simultaneous downloads and prefetched decoded outputs. The next page is fetched after the current page
+is consumed.
+
+Iteration does not wait for unfinished tasks; call `wait()` first when you want all tasks to finish.
+
+## List a batch
+
+Use `list_batch()` for progress inspection, status filters, manual paging, or
+selecting outputs to download. Use `iter_results()` to download and decode outputs
+incrementally, or `results()` to collect every output in the submission.
+
+```python
+summary = client.list_batch(submission.batch_id)
+page = client.list_batch(
+ submission.batch_id,
+ status=["done", "failed"],
+ include_responses=True,
+ limit=100,
+)
+while True:
+ for item in page["items"]:
+ print(item["task_id"], item["status"])
+ if page["cursor"] is None:
+ break
+ page = client.list_batch(
+ submission.batch_id,
+ status=["done", "failed"],
+ include_responses=True,
+ limit=100,
+ cursor=page["cursor"],
+ )
+```
+
+## Run the whole lifecycle
+
+`run()` owns the whole pipeline: it submits the work, waits for completion, and
+downloads and decodes the results in one call.
+
+```python
+from esm.sdk.batch import BatchError
+
+outputs = client.fold.run(
+ requests=[FoldRequest(sequence=sequence) for sequence in sequences],
+)
+if isinstance(outputs, BatchError):
+ print(outputs.error_msg)
+```
+
+On timeout or interruption, `wait()` stops waiting and leaves work running.
+`run()` attempts to cancel queued work in both cases by default
+(`cancel_on_timeout=True`). Set `cancel_on_timeout=False` to leave work running.
+
+Manual cancellation uses exactly one kind of ID:
+
+```python
+client.cancel(task_id=task_id) # An actual ID returned by list_batch() or iter_results().
+client.cancel(batch_id=submission.batch_id)
+```
+
+## Using Python asyncio
+
+Both client interfaces submit jobs for background processing. The examples above
+use regular Python methods. In an asyncio application, use the `async_*` methods
+to avoid blocking the event loop while making requests or polling for completion.
+
+```python
+async with ForgeBatchClient(
+ token=os.environ["ESM_API_KEY"], model="esmfold2-fast-2026-05"
+) as client:
+ submission = await client.fold.async_submit(
+ requests=[FoldRequest(sequence=sequence) for sequence in sequences],
+ )
+ await client.fold.async_wait(submission)
+ outputs = await client.fold.async_results(submission)
+```
diff --git a/cookbook/snippets/fold_max_accuracy.md b/cookbook/snippets/fold_max_accuracy.md
new file mode 100644
index 00000000..2e0ce0ab
--- /dev/null
+++ b/cookbook/snippets/fold_max_accuracy.md
@@ -0,0 +1,82 @@
+# Maximum-accuracy folding with the SDK
+
+Use `ForgeBatchClient.fold_max_accuracy` to predict a molecular complex with a 4096-token input limit through
+Biohub's Batch API. The service generates an MSA, runs five folds with different
+seeds, and selects the best predicted structure. The service uses ESMFold2 Pro trained on data with a 2025 cutoff, `esmfold2-2026-05-cutoff-2025`.
+
+For general batch operations, see the [Batch API guide](batch.md).
+
+## Inputs and MSA settings
+
+Provide a `StructurePredictionInput` containing protein, RNA, DNA, or ligand
+entries. The service searches UniRef30 and ColabFold to generate its own multiple sequence
+alignment (MSA), using these search settings:
+
+| Setting | Value |
+| --- | --- |
+| `max_prefilter` | 100,000 |
+| `max_accept` | 10,000 |
+| `diff` | 3,000 |
+| `expand_diff` | 3,000 |
+
+These settings are managed by the service. Leave `ProteinInput.msa` unset. Use `FoldMaxAccuracyConfig` for output options:
+
+| Option | Requested output |
+| --- | --- |
+| `include_pae` | Predicted Aligned Error (PAE) matrix. |
+| `include_pair_chains_iptm` | Pair-chain interface predicted TM scores. |
+| `include_embeddings` | Pooled pair embeddings. |
+
+## Submit and save the handle
+
+This example submits two proteins, waits for both,
+and saves each selected structure to an mmCIF file named by task ID. Each request
+describes one complex; to fold multiple chains together, put them in the same
+`StructurePredictionInput.sequences` list.
+
+```python
+import json
+import os
+from pathlib import Path
+
+from esm.sdk.api import FoldMaxAccuracyConfig, FoldMaxAccuracyRequest
+from esm.sdk.batch import BatchError
+from esm.sdk.forge import ForgeBatchClient
+from esm.utils.structure.input_builder import ProteinInput, StructurePredictionInput
+
+sequences = [
+ "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG",
+ "MTYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDDATKTFTVTE",
+]
+fold_inputs = [
+ StructurePredictionInput(sequences=[ProteinInput(id="A", sequence=sequence)])
+ for sequence in sequences
+]
+config = FoldMaxAccuracyConfig(include_pae=True)
+
+with ForgeBatchClient(token=os.environ["ESM_API_KEY"]) as client:
+ submission = client.fold_max_accuracy.submit(
+ requests=[FoldMaxAccuracyRequest(all_atom_input=item) for item in fold_inputs],
+ config=config,
+ )
+ # Optional: save the submission handle for later retrieval.
+ Path("fold-submission.json").write_text(json.dumps(submission.to_dict()))
+ print("Batch ID:", submission.batch_id)
+
+ client.fold_max_accuracy.wait(submission)
+
+ for task_id, result in client.fold_max_accuracy.iter_results(submission):
+ if isinstance(result, BatchError):
+ print(task_id, result.error_msg)
+ continue
+ Path(f"{task_id}.cif").write_text(result.complex.to_mmcif())
+ print(task_id, "pTM:", result.ptm)
+ # The requested PAE matrix is available as result.pae.
+```
+
+The shared `config` applies to every request. To customize one
+item, pass `config=FoldMaxAccuracyConfig(...)` to its `FoldMaxAccuracyRequest`.
+An item's config replaces the shared config entirely; fields are not merged.
+
+In an asyncio application, use the `async` counterpart methods as shown in the
+[asyncio example](batch.md#using-python-asyncio).
diff --git a/cookbook/tutorials/binder_design.py b/cookbook/tutorials/binder_design.py
index 7909cd6d..beddbbd3 100644
--- a/cookbook/tutorials/binder_design.py
+++ b/cookbook/tutorials/binder_design.py
@@ -48,9 +48,7 @@
from esm.models.esmfold2.experimental import EsmFold2ExperimentalModel
from esm.models.esmfold2.experimental import MSAEncoder as EsmFold2MSAEncoder
from esm.models.esmfold2.layers import (
- BACKEND_CUEQ,
BACKEND_FUSED,
- CUE_AVAILABLE,
TRITON_KERNELS_AVAILABLE,
PairUpdateBlock,
)
@@ -1239,11 +1237,7 @@ def _load_hf_model(
_ESMC_CACHE[esmc_id] = model.esmc
model.esmc = _ESMC_CACHE[esmc_id]
model.configure_lm_dropout(lm_dropout, force_lm_dropout_during_inference=True)
- kernel_backend = None
- if TRITON_KERNELS_AVAILABLE:
- kernel_backend = BACKEND_FUSED
- elif CUE_AVAILABLE:
- kernel_backend = BACKEND_CUEQ
+ kernel_backend = BACKEND_FUSED if TRITON_KERNELS_AVAILABLE else None
model.set_kernel_backend(kernel_backend)
return model.to(device=device).eval().requires_grad_(False)
@@ -1371,7 +1365,6 @@ def get_base_image():
.pip_install(
"flash-attn==2.8.3",
"transformer-engine[core-cu12,pytorch]==2.13.0",
- "xformers==0.0.32.post1",
extra_options="--no-build-isolation",
env={
"CPATH": (
@@ -1389,13 +1382,7 @@ def get_base_image():
},
)
.pip_install("abnumber", "esm", "modal")
- .env(
- {
- "HF_HOME": "/models",
- "HF_XET_HIGH_PERFORMANCE": "1",
- "XFORMERS_IGNORE_FLASH_VERSION_CHECK": "1",
- }
- )
+ .env({"HF_HOME": "/models", "HF_XET_HIGH_PERFORMANCE": "1"})
)
diff --git a/cookbook/tutorials/embed.ipynb b/cookbook/tutorials/embed.ipynb
index 6aae11e6..ad561730 100644
--- a/cookbook/tutorials/embed.ipynb
+++ b/cookbook/tutorials/embed.ipynb
@@ -40,7 +40,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "First, we need to set up the Biohub client. Generate an API key on the [Biohub platform](https://biohub.ai) and add it to your account. This API key manages your access to credits and tokens, and the term API key/token is often used interchangeably within documentation.\n",
+ "First, we need to set up the Biohub client. Generate an API key on the [Biohub platform](https://biohub.ai/?utm_source=colab&utm_medium=tutorial_notebook&utm_content=embed) and add it to your account. This API key manages your access to credits and tokens, and the term API key/token is often used interchangeably within documentation.\n",
"\n",
"Please note that your API key is like a password for your account and you should take care to protect it. For this reason it is recommended to frequently create a new API key and delete old, unused ones. It is also recommended to paste the API key directly into an environment variable or use a utility like `getpass` as shown in the example below so keys are not accidentally shared or checked into code repositories."
]
diff --git a/cookbook/tutorials/esm3_generate.ipynb b/cookbook/tutorials/esm3_generate.ipynb
index 820bdaef..da6a040d 100644
--- a/cookbook/tutorials/esm3_generate.ipynb
+++ b/cookbook/tutorials/esm3_generate.ipynb
@@ -54,7 +54,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "First, we need to set up the [Biohub client](https://biohub.ai). Generate an API key on Biohub and add it to your account. This API key manages your access to credits and tokens, and the term API key/token is often used interchangeably within documentation.\n",
+ "First, we need to set up the [Biohub client](https://biohub.ai/?utm_source=colab&utm_medium=tutorial_notebook&utm_content=esm3_generate). Generate an API key on Biohub and add it to your account. This API key manages your access to credits and tokens, and the term API key/token is often used interchangeably within documentation.\n",
"\n",
"Please note that your API key is like a password for your account and you should take care to protect it. For this reason it is recommended to frequently create a new API key and delete old, unused ones. It is also recommended to paste the API key directly into an environment variable or use a utility like `getpass` as shown in the example below so keys are not accidentally shared or checked into code repositories."
]
diff --git a/cookbook/tutorials/esmc_layer_sweep.ipynb b/cookbook/tutorials/esmc_layer_sweep.ipynb
index 3d37a1c1..8368bd02 100644
--- a/cookbook/tutorials/esmc_layer_sweep.ipynb
+++ b/cookbook/tutorials/esmc_layer_sweep.ipynb
@@ -107,7 +107,7 @@
"metadata": {},
"source": [
"## 2. Set up Biohub client. \n",
- "First, we need to set up the Biohub client. Generate an API key on the [Biohub platform](https://biohub.ai) and add it to your account. This API key manages your access to credits and tokens, and the term API key/token is often used interchangeably within documentation.\n",
+ "First, we need to set up the Biohub client. Generate an API key on the [Biohub platform](https://biohub.ai/?utm_source=colab&utm_medium=tutorial_notebook&utm_content=esmc_layer_sweep) and add it to your account. This API key manages your access to credits and tokens, and the term API key/token is often used interchangeably within documentation.\n",
"\n",
"Please note that your API key is like a password for your account and you should take care to protect it. For this reason it is recommended to frequently create a new API key and delete old, unused ones. It is also recommended to paste the API key directly into an environment variable or use a utility like `getpass` as shown in the example below so keys are not accidentally shared or checked into code repositories.\n",
"\n",
diff --git a/cookbook/tutorials/esmc_mutation_scoring.ipynb b/cookbook/tutorials/esmc_mutation_scoring.ipynb
index 99181e64..56dd8767 100644
--- a/cookbook/tutorials/esmc_mutation_scoring.ipynb
+++ b/cookbook/tutorials/esmc_mutation_scoring.ipynb
@@ -177,7 +177,7 @@
},
"source": [
"# **2. Set up ESMC client**\n",
- "First, we need to set up the Biohub client. Generate an API key on the [Biohub platform](https://biohub.ai) and add it to your account. This API key manages your access to credits and tokens, and the term API key/token is often used interchangeably within documentation.\n",
+ "First, we need to set up the Biohub client. Generate an API key on the [Biohub platform](https://biohub.ai/?utm_source=colab&utm_medium=tutorial_notebook&utm_content=esmc_mutation_scoring) and add it to your account. This API key manages your access to credits and tokens, and the term API key/token is often used interchangeably within documentation.\n",
"\n",
"Please note that your API key is like a password for your account and you should take care to protect it. For this reason it is recommended to frequently create a new API key and delete old, unused ones. It is also recommended to paste the API key directly into an environment variable or use a utility like `getpass` as shown in the example below so keys are not accidentally shared or checked into code repositories."
]
diff --git a/cookbook/tutorials/esmc_sae_feature_interpretation.ipynb b/cookbook/tutorials/esmc_sae_feature_interpretation.ipynb
index 292915c6..b18c549c 100644
--- a/cookbook/tutorials/esmc_sae_feature_interpretation.ipynb
+++ b/cookbook/tutorials/esmc_sae_feature_interpretation.ipynb
@@ -85,7 +85,7 @@
"id": "cell-5"
},
"source": [
- "First, we need to set up the Biohub client. Generate an API key on the [Biohub platform](https://biohub.ai) and add it to your account. This API key manages your access to credits and tokens, and the term API key/token is often used interchangeably within documentation.\n",
+ "First, we need to set up the Biohub client. Generate an API key on the [Biohub platform](https://biohub.ai/?utm_source=colab&utm_medium=tutorial_notebook&utm_content=esmc_sae_feature_interpretation) and add it to your account. This API key manages your access to credits and tokens, and the term API key/token is often used interchangeably within documentation.\n",
"\n",
"Please note that your API key is like a password for your account and you should take care to protect it. For this reason it is recommended to frequently create a new API key and delete old, unused ones. It is also recommended to paste the API key directly into an environment variable or use a utility like `getpass` as shown in the example below so keys are not accidentally shared or checked into code repositories."
]
diff --git a/cookbook/tutorials/esmfold2.ipynb b/cookbook/tutorials/esmfold2.ipynb
index 463606df..aa4eb749 100644
--- a/cookbook/tutorials/esmfold2.ipynb
+++ b/cookbook/tutorials/esmfold2.ipynb
@@ -235,7 +235,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "First, we need to set up the Biohub client. Generate an API key on the [Biohub platform](https://biohub.ai) and add it to your account. This API key manages your access to credits and tokens, and the term API key/token is often used interchangeably within documentation.\n",
+ "First, we need to set up the Biohub client. Generate an API key on the [Biohub platform](https://biohub.ai/?utm_source=colab&utm_medium=tutorial_notebook&utm_content=esmfold2) and add it to your account. This API key manages your access to credits and tokens, and the term API key/token is often used interchangeably within documentation.\n",
"\n",
"Please note that your API key is like a password for your account and you should take care to protect it. For this reason it is recommended to frequently create a new API key and delete old, unused ones. It is also recommended to paste the API key directly into an environment variable or use a utility like `getpass` as shown in the example below so keys are not accidentally shared or checked into code repositories."
]
diff --git a/cookbook/tutorials/esmfold2_local_applesilicon.ipynb b/cookbook/tutorials/esmfold2_local_applesilicon.ipynb
index 8a1a9307..3f610de9 100644
--- a/cookbook/tutorials/esmfold2_local_applesilicon.ipynb
+++ b/cookbook/tutorials/esmfold2_local_applesilicon.ipynb
@@ -168,8 +168,11 @@
"source": [
"## Load the model\n",
"\n",
- "Two imports: the **model** from `mlx_lm`, the **builder** from `esm`. That is the\n",
- "only line in this notebook that differs from the CUDA version.\n",
+ "Two imports: the **loader** from `mlx_lm`, the **builder** from `esm`. That is the\n",
+ "only part of this notebook that differs from the CUDA version.\n",
+ "\n",
+ "`lazy=True` reads nothing yet, so casting to bf16 first keeps the peak near the\n",
+ "13 GB the weights actually occupy instead of the 27 GB they take as fp32.\n",
"\n",
"The first load pulls ~25 GB and takes a few minutes; after that it is resident.\n"
]
@@ -181,11 +184,22 @@
"metadata": {},
"outputs": [],
"source": [
- "from mlx_lm.models.esmfold2 import ESMFold2Model # <- MLX instead of esm\n",
+ "import mlx.core as mx\n",
+ "from mlx.utils import tree_flatten\n",
+ "from mlx_lm import load # <- MLX instead of esm\n",
"\n",
"from esm.models.esmfold2 import ESMFold2InputBuilder\n",
"\n",
- "model = ESMFold2Model.from_pretrained(\"biohub/ESMFold2-Fast\")\n",
+ "\n",
+ "def load_model(repo, dtype=mx.bfloat16):\n",
+ " model, _tokenizer = load(repo, lazy=True)\n",
+ " model.set_dtype(dtype)\n",
+ " for _, param in tree_flatten(model.parameters()):\n",
+ " mx.eval(param) # one at a time, so no fp32 copy of the whole model\n",
+ " return model\n",
+ "\n",
+ "\n",
+ "model = load_model(\"biohub/ESMFold2-Fast\")\n",
"builder = ESMFold2InputBuilder()"
]
},
@@ -412,7 +426,7 @@
"metadata": {},
"outputs": [],
"source": [
- "model = ESMFold2Model.from_pretrained(\"biohub/ESMFold2\")"
+ "model = load_model(\"biohub/ESMFold2\")"
]
},
{
diff --git a/cookbook/tutorials/gfp_design.ipynb b/cookbook/tutorials/gfp_design.ipynb
index bcd6a565..3e1f1fdc 100644
--- a/cookbook/tutorials/gfp_design.ipynb
+++ b/cookbook/tutorials/gfp_design.ipynb
@@ -79,7 +79,7 @@
"source": [
"ESM3 is a frontier generative model for biology. It is scalable due to its ability to tokenize sequence, structure, and function and use a (nearly standard) transformer architecture while still being able to reason across all modalities simulateously.\n",
"\n",
- "The largest ESM3 (98 billion parameters) was trained with 1.07e24 FLOPs on 2.78 billion proteins and 771 billion unique tokens. To create esmGFP we used the 7 billion parameter variant of ESM3. We'll use this model via the [Biohub](https://biohub.ai) API.\n",
+ "The largest ESM3 (98 billion parameters) was trained with 1.07e24 FLOPs on 2.78 billion proteins and 771 billion unique tokens. To create esmGFP we used the 7 billion parameter variant of ESM3. We'll use this model via the [Biohub](https://biohub.ai/?utm_source=colab&utm_medium=tutorial_notebook&utm_content=gfp_design) API.\n",
"\n",
"Grab a token from [Biohub](https://biohub.ai/) and add it below. Note that your token is like a password for your account and you should take care to protect it. For this reason it is recommended to frequently create a new token and delete old, unused ones. It is also recommended to paste the token directly into an environment variable or use a utility like `getpass` as shown below so tokens are not accidentally shared or checked into code repositories.\n"
]
diff --git a/esm/__init__.py b/esm/__init__.py
index 405f1bf0..46aa803f 100644
--- a/esm/__init__.py
+++ b/esm/__init__.py
@@ -1 +1 @@
-__version__ = "3.4.1.post1"
+__version__ = "3.4.2"
diff --git a/esm/layers/attention.py b/esm/layers/attention.py
index 718dea5f..9669bce9 100644
--- a/esm/layers/attention.py
+++ b/esm/layers/attention.py
@@ -129,11 +129,14 @@ def forward(
self.k_ln(key_ND).to(query_ND.dtype),
)
- qkv_N3D = torch.stack([query_ND, key_ND, value_ND], dim=1)
- qkv_N3HD = einops.rearrange(
- qkv_N3D, pattern="n a (h d) -> n a h d", h=self.n_heads
+ query_NHD, key_NHD, value_NHD = (
+ einops.rearrange(t, pattern="n (h d) -> n h d", h=self.n_heads)
+ for t in (query_ND, key_ND, value_ND)
)
- qkv_N3HD = self.rotary(qkv_N3HD, cu_seqlens, max_seqlen)
+ # Rotate before packing: stacking the rotated q/k is one copy, whereas rotating
+ # a packed tensor out-of-place and re-stacking is two.
+ query_NHD, key_NHD = self.rotary(query_NHD, key_NHD, cu_seqlens, max_seqlen)
+ qkv_N3HD = torch.stack([query_NHD, key_NHD, value_NHD], dim=1)
assert flash_attn_varlen_qkvpacked_func is not None
context_NHD = flash_attn_varlen_qkvpacked_func(
diff --git a/esm/layers/rotary.py b/esm/layers/rotary.py
index 302a81aa..0f3dc4a3 100644
--- a/esm/layers/rotary.py
+++ b/esm/layers/rotary.py
@@ -26,7 +26,9 @@
from einops import rearrange, repeat
try:
- from flash_attn.ops.triton.rotary import apply_rotary as apply_triton_rotary
+ # Autograd-aware wrapper around the same Triton kernel; the bare
+ # ``flash_attn.ops.triton.rotary.apply_rotary`` driver has no ``grad_fn``.
+ from flash_attn.layers.rotary import apply_rotary_emb as apply_triton_rotary
except ImportError:
apply_triton_rotary = None # ty:ignore[invalid-assignment]
@@ -227,33 +229,40 @@ def forward(
class TritonRotaryEmbedding(RotaryEmbedding):
- def forward(self, qkv: torch.Tensor, cu_seqlens, max_seqlen) -> torch.Tensor:
+ def forward(
+ self,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ max_seqlen: int,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
"""
- qkv: (n, 3, nheads, headdim)
+ q, k: (n, nheads, headdim) varlen-flat
cu_seqlens: cumulative sequence lengths
max_seqlen: max sequence length
"""
- self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype)
+ self._update_cos_sin_cache(max_seqlen, device=q.device, dtype=q.dtype)
assert self._cos_cached is not None
assert self._sin_cached is not None
assert apply_triton_rotary is not None
- # In-place modification
- apply_triton_rotary(
- qkv[:, 0],
- self._cos_cached,
- self._sin_cached,
- cu_seqlens=cu_seqlens,
- max_seqlen=max_seqlen,
- inplace=True,
+ # The raw ``flash_attn.ops.triton.rotary.apply_rotary`` driver is not an autograd
+ # op, so its output carried no ``grad_fn`` and backward skipped RoPE entirely.
+ return (
+ apply_triton_rotary(
+ q,
+ self._cos_cached,
+ self._sin_cached,
+ interleaved=self.interleaved,
+ cu_seqlens=cu_seqlens,
+ max_seqlen=max_seqlen,
+ ),
+ apply_triton_rotary(
+ k,
+ self._cos_cached,
+ self._sin_cached,
+ interleaved=self.interleaved,
+ cu_seqlens=cu_seqlens,
+ max_seqlen=max_seqlen,
+ ),
)
- apply_triton_rotary(
- qkv[:, 1],
- self._cos_cached,
- self._sin_cached,
- cu_seqlens=cu_seqlens,
- max_seqlen=max_seqlen,
- inplace=True,
- )
-
- return qkv
diff --git a/esm/models/esmc/checkpoint_layout.py b/esm/models/esmc/checkpoint_layout.py
index ced4bcaf..ca3b5bc2 100644
--- a/esm/models/esmc/checkpoint_layout.py
+++ b/esm/models/esmc/checkpoint_layout.py
@@ -5,11 +5,8 @@
``v_proj``, ``gate_proj``, ``up_proj`` — which is the layout ``transformers``
reads and the only one its tensor-parallel plan can shard.
-This implementation instead keeps the projections fused, because
-``transformer_engine``'s ``LayerNormLinear`` / ``LayerNormMLP`` fuse each
-LayerNorm into the GEMM that follows it and take a single packed weight. That
-fusion is what the fp8 path runs on, so the module tree keeps Transformer
-Engine's parameter names and the two layouts are reconciled here instead.
+This implementation instead keeps the projections fused: one packed QKV weight
+and one packed ``[gate; up]`` weight. The two layouts are reconciled here.
The conversion is exact in both directions: a fused weight is the row-wise
concatenation of its parts, so ``cat``/``chunk`` along dim 0 round-trips
@@ -97,7 +94,7 @@ def published_to_native(raw: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]
"""Rewrite a published checkpoint onto this implementation's layout.
Concatenates the separate q/k/v and gate/up projections into the fused
- weights Transformer Engine expects. Keys already in the in-memory layout,
+ weights this implementation holds. Keys already in the in-memory layout,
and keys belonging to neither vocabulary, are returned unchanged.
"""
out: dict[str, torch.Tensor] = {}
diff --git a/esm/models/esmc/compatibility.py b/esm/models/esmc/compatibility.py
index fea1e8b5..319985d0 100644
--- a/esm/models/esmc/compatibility.py
+++ b/esm/models/esmc/compatibility.py
@@ -19,7 +19,7 @@
import torch
import torch.nn as nn
from attr import dataclass
-from transformers import PreTrainedTokenizerFast
+from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
from esm.models.esmc.config import (
ESMC_6B_HF_REPO,
diff --git a/esm/models/esmc/kernels.py b/esm/models/esmc/kernels.py
index 8aa52c90..a8c950b9 100644
--- a/esm/models/esmc/kernels.py
+++ b/esm/models/esmc/kernels.py
@@ -6,40 +6,15 @@
The kernels, in the order ESMC prefers them:
-* ``transformer_engine`` — fused LayerNorm+Linear / LayerNorm+MLP with an
- fp32 reduction inside the LayerNorm. Recommended on GPU for accurate bf16
- inference; without it the pure-PyTorch fallback drifts ~O(10) in fp32 and
- ~O(100) in bf16 on the unnormalized residual stream (perplexity stays
- within rounding noise).
-* ``xformers`` — preferred fused attention kernel with a deterministic
- reduction order.
-* ``flash-attn`` — secondary fused attention kernel (fp16 / bf16 only), plus
- the Triton RoPE kernel and the varlen packing helpers.
+* ``flash-attn`` — fused attention kernel (fp16 / bf16 only), plus the
+ Triton RoPE kernel and the varlen packing helpers.
"""
-import logging
+import warnings
from typing import Any
import torch
-logger = logging.getLogger(__name__)
-
-try:
- import transformer_engine.pytorch as te
-
- TE_INSTALLED = True
-except ImportError:
- te: Any = None
- TE_INSTALLED = False
-
-try:
- import xformers.ops as xops
-
- XFORMERS_INSTALLED = True
-except ImportError:
- xops: Any = None
- XFORMERS_INSTALLED = False
-
try:
from flash_attn import flash_attn_func, flash_attn_varlen_qkvpacked_func
from flash_attn.bert_padding import pad_input, unpad_input
@@ -53,7 +28,10 @@
FLASH_ATTN_INSTALLED = False
try:
- from flash_attn.ops.triton.rotary import apply_rotary as apply_triton_rotary
+ # ``flash_attn.layers.rotary.apply_rotary_emb`` wraps the same Triton kernel in an
+ # ``autograd.Function``; the bare ``flash_attn.ops.triton.rotary.apply_rotary`` driver
+ # returns a tensor with no ``grad_fn``, so a backward through it silently skips RoPE.
+ from flash_attn.layers.rotary import apply_rotary_emb as apply_triton_rotary
FLASH_ATTN_ROTARY_INSTALLED = True
except ImportError:
@@ -61,31 +39,23 @@
FLASH_ATTN_ROTARY_INSTALLED = False
-if torch.cuda.is_available() and not TE_INSTALLED:
- logger.warning(
- "ESMC: Transformer Engine is not installed; falling back to "
- "pure-PyTorch LayerNorm+Linear / LayerNorm+MLP. Outputs will differ "
- "numerically - measured on the unnormalized residual stream (before "
- "the final LayerNorm), ~O(10) max-diff in fp32 and ~O(100) in bf16; "
- "after the final LayerNorm these shrink to a few ULP and perplexity "
- "stays within rounding noise. Install with "
- "`pip install transformer-engine[pytorch]` to enable fused fp32-"
- "reduction LayerNorm."
- )
-
-if torch.cuda.is_available() and not (XFORMERS_INSTALLED or FLASH_ATTN_INSTALLED):
- logger.warning(
- "ESMC: neither xformers nor flash-attn is installed; falling back "
- "to PyTorch ``F.scaled_dot_product_attention``. The attention "
- "reduction order in bf16 differs from a fused kernel by ~1 bf16 "
- "ULP per attention block; compounded across the 80-block stack "
- "this reaches ~O(100) max-diff on the unnormalized residual stream. "
- "Install an xformers build matching your PyTorch and CUDA versions "
- "for a fused attention kernel."
- )
-
-if torch.cuda.is_available() and not FLASH_ATTN_ROTARY_INSTALLED:
- logger.warning(
- "ESMC: flash-attn rotary kernel not installed; falling back to "
- "pure-PyTorch RoPE. For faster GPU inference run `pip install flash-attn`."
- )
+FLASH_ATTN_WHEEL_URL = (
+ "https://github.com/evolutionaryscale/wheels/releases/download/"
+ "py312-pt211-cu13-sm80-90/flash_attn-2.7.4.post1-cp312-cp312-linux_x86_64.whl"
+)
+
+if torch.cuda.is_available() and not (
+ FLASH_ATTN_INSTALLED and FLASH_ATTN_ROTARY_INSTALLED
+):
+ # warnings, not logging, and forced: the esmfold2 Triton modules ignore
+ # every UserWarning process-wide.
+ with warnings.catch_warnings():
+ warnings.simplefilter("always")
+ warnings.warn(
+ "ESMC: flash-attn is not installed or failed to import; attention and "
+ "RoPE fall back to pure PyTorch. PyPI cannot carry our wheel, so install "
+ f'it by hand:\n\n pip install "flash_attn @ {FLASH_ATTN_WHEEL_URL}"\n\n'
+ "It needs Linux x86-64, Python 3.12, torch 2.11, CUDA 13, compute "
+ "capability 8.0-9.0 and libstdc++ from GCC 13+ (Ubuntu 24.04, or "
+ "LD_PRELOAD a newer libstdc++ on 22.04)."
+ )
diff --git a/esm/models/esmc/layers.py b/esm/models/esmc/layers.py
index 8a55589e..f5b7c3d9 100644
--- a/esm/models/esmc/layers.py
+++ b/esm/models/esmc/layers.py
@@ -1,9 +1,8 @@
"""ESMC-specific transformer layers.
RoPE, QK-LayerNorm attention, SwiGLU feed-forward and the transformer stack used
-by the ESMC encoder. Each fused module (Transformer Engine LayerNorm+Linear /
-LayerNorm+MLP, xformers / flash-attn attention, Triton RoPE) has a pure-PyTorch
-fallback with matching parameter names so state dicts are interchangeable.
+by the ESMC encoder. Each fused kernel (flash-attn attention, Triton RoPE) has a
+pure-PyTorch fallback.
"""
import math
@@ -17,12 +16,9 @@
from esm.models.esmc.kernels import (
FLASH_ATTN_INSTALLED,
FLASH_ATTN_ROTARY_INSTALLED,
- XFORMERS_INSTALLED,
apply_triton_rotary,
flash_attn_func,
flash_attn_varlen_qkvpacked_func,
- te,
- xops,
)
@@ -198,34 +194,65 @@ class EsmcTritonRotaryEmbedding(EsmcRotaryEmbedding):
Used inside :class:`EsmcFlashMultiHeadAttention` when Flash Attention 2 is
available. The ``forward`` signature differs from
- :class:`EsmcRotaryEmbedding` because Flash Attention packs Q, K, V together.
+ :class:`EsmcRotaryEmbedding` because Flash Attention uses a varlen-flat
+ layout: positions restart at zero at every ``cu_seqlens`` boundary.
"""
def forward(
- self, qkv: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int
- ) -> torch.Tensor:
- """Apply RoPE in-place to a packed ``(N, 3, n_heads, head_dim)`` tensor."""
- self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype)
+ self,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ cu_seqlens: torch.Tensor | None = None,
+ max_seqlen: int | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Apply RoPE to varlen-flat Q and K.
+
+ Parameters
+ ----------
+ q, k : torch.Tensor
+ Varlen-flat tensors of shape ``(total_tokens, n_heads, head_dim)``,
+ or batched ``(batch, seqlen, n_heads, head_dim)`` when ``cu_seqlens``
+ is ``None``.
+ cu_seqlens : torch.Tensor | None
+ Cumulative sequence lengths, shape ``(batch + 1,)``, ``int32``.
+ ``None`` selects the batched parent path, which the flash block takes
+ once the model has been moved off CUDA.
+ max_seqlen : int | None
+ Length of the longest sequence in the batch; required with
+ ``cu_seqlens``.
+
+ Returns
+ -------
+ tuple[torch.Tensor, torch.Tensor]
+ Rotated ``q`` and ``k``, same shapes as the inputs.
+ """
+ if cu_seqlens is None:
+ return super().forward(q, k)
+ assert max_seqlen is not None
+ self._update_cos_sin_cache(max_seqlen, device=q.device, dtype=q.dtype)
assert self._cos_cached is not None and self._sin_cached is not None
assert apply_triton_rotary is not None
- apply_triton_rotary(
- qkv[:, 0],
- self._cos_cached,
- self._sin_cached,
- cu_seqlens=cu_seqlens,
- max_seqlen=max_seqlen,
- inplace=True,
+ # The raw ``flash_attn.ops.triton.rotary.apply_rotary`` driver is not an autograd
+ # op, so its output carried no ``grad_fn`` and backward skipped RoPE entirely.
+ return (
+ apply_triton_rotary(
+ q,
+ self._cos_cached,
+ self._sin_cached,
+ interleaved=self.interleaved,
+ cu_seqlens=cu_seqlens,
+ max_seqlen=max_seqlen,
+ ),
+ apply_triton_rotary(
+ k,
+ self._cos_cached,
+ self._sin_cached,
+ interleaved=self.interleaved,
+ cu_seqlens=cu_seqlens,
+ max_seqlen=max_seqlen,
+ ),
)
- apply_triton_rotary(
- qkv[:, 1],
- self._cos_cached,
- self._sin_cached,
- cu_seqlens=cu_seqlens,
- max_seqlen=max_seqlen,
- inplace=True,
- )
- return qkv
def esmc_swiglu_hidden_dim(expansion_ratio: float, d_model: int) -> int:
@@ -236,10 +263,7 @@ def esmc_swiglu_hidden_dim(expansion_ratio: float, d_model: int) -> int:
class EsmcLayerNormMLP(nn.Module):
"""LayerNorm + SwiGLU MLP (no bias), as used by ESMC.
- Shares the parameter names ``layer_norm_weight``, ``layer_norm_bias``,
- ``fc1_weight`` and ``fc2_weight`` with Transformer Engine's fused
- ``LayerNormMLP`` so the two are state-dict compatible: the fused kernel is
- used when available and this pure-PyTorch module is the fallback.
+ ``fc1_weight`` packs the gate and up projections as ``[gate; up]``.
"""
def __init__(
@@ -270,21 +294,10 @@ def forward(self, x: Tensor) -> Tensor:
return F.linear(x, self.fc2_weight)
-def esmc_swiglu_ln_ffn(
- d_model: int, expansion_ratio: float, bias: bool, use_te: bool
-) -> nn.Module:
- """LayerNorm + SwiGLU MLP, fused via Transformer Engine when ``use_te``."""
+def esmc_swiglu_ln_ffn(d_model: int, expansion_ratio: float, bias: bool) -> nn.Module:
+ """LayerNorm + SwiGLU MLP."""
assert not bias, "ESMC was trained with bias=False; bias=True not supported"
hidden = esmc_swiglu_hidden_dim(expansion_ratio, d_model)
- if use_te:
- return te.LayerNormMLP(
- hidden_size=d_model,
- ffn_hidden_size=hidden,
- bias=bias,
- activation="swiglu",
- init_method=None,
- output_layer_init_method=None,
- )
return EsmcLayerNormMLP(hidden_size=d_model, ffn_hidden_size=hidden)
@@ -300,13 +313,7 @@ def esmc_gelu_ln_ffn(d_model: int, expansion_ratio: float, bias: bool) -> nn.Seq
class EsmcLayerNormLinear(nn.Module):
- """LayerNorm followed by a Linear projection (no bias).
-
- Shares the parameter names ``layer_norm_weight``, ``layer_norm_bias`` and
- ``weight`` with Transformer Engine's fused ``LayerNormLinear`` so the two
- are state-dict compatible: the fused kernel is used when available and this
- pure-PyTorch module is the fallback.
- """
+ """LayerNorm followed by a Linear projection (no bias)."""
def __init__(self, d_in: int, d_out: int, eps: float = 1e-5) -> None:
super().__init__()
@@ -324,19 +331,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
return F.linear(x, self.weight)
-def make_esmc_attn_layernorm_qkv(d_model: int, bias: bool, use_te: bool) -> nn.Module:
- """LayerNorm + fused QKV projection, fused via Transformer Engine when
- ``use_te``."""
+def make_esmc_attn_layernorm_qkv(d_model: int, bias: bool) -> nn.Module:
+ """LayerNorm + packed QKV projection."""
assert not bias, "ESMC was trained with bias=False; bias=True not supported"
- if use_te:
- return te.LayerNormLinear(d_model, d_model * 3, bias=bias, init_method=None)
return EsmcLayerNormLinear(d_model, d_model * 3)
-def make_esmc_attn_out_proj(d_model: int, bias: bool, use_te: bool) -> nn.Module:
- """Attention output projection, fused via Transformer Engine when ``use_te``."""
- if use_te:
- return te.Linear(d_model, d_model, bias=bias, init_method=None)
+def make_esmc_attn_out_proj(d_model: int, bias: bool) -> nn.Module:
+ """Attention output projection."""
return nn.Linear(d_model, d_model, bias=bias)
@@ -358,22 +360,11 @@ def esmc_scaled_dot_product_attention(
Dispatches in order of preference:
- 1. xformers ``memory_efficient_attention`` - preferred fused kernel,
- requires ``xformers``, unmasked only.
- 2. Flash Attention 2 (``flash_attn_func``) - secondary fused kernel,
- requires ``flash-attn``, unmasked only, fp16 / bf16 only.
- 3. PyTorch's ``F.scaled_dot_product_attention`` - last-resort path; the only
+ 1. Flash Attention 2 (``flash_attn_func``) - fused kernel, requires
+ ``flash-attn``, unmasked only, fp16 / bf16 only.
+ 2. PyTorch's ``F.scaled_dot_product_attention`` - last-resort path; the only
one that takes a mask, and the only fp32 one.
"""
- if seq_id is None and XFORMERS_INSTALLED and q.is_cuda:
- b, s, _ = q.shape
- q4 = q.view(b, s, n_heads, d_head)
- k4 = k.view(b, s, n_heads, d_head)
- v4 = v.view(b, s, n_heads, d_head)
- context = xops.memory_efficient_attention(
- q4, k4, v4, attn_bias=None, scale=d_head**-0.5
- )
- return context.reshape(b, s, n_heads * d_head)
if (
seq_id is None
and FLASH_ATTN_INSTALLED
@@ -411,12 +402,7 @@ class EsmcMultiHeadAttention(nn.Module):
"""
def __init__(
- self,
- d_model: int,
- n_heads: int,
- bias: bool = False,
- qk_layernorm: bool = True,
- use_te: bool = False,
+ self, d_model: int, n_heads: int, bias: bool = False, qk_layernorm: bool = True
):
super().__init__()
self.d_model = d_model
@@ -424,8 +410,8 @@ def __init__(
self.d_head = d_model // n_heads
assert not bias, "ESMC was trained with bias=False; bias=True not supported"
- self.layernorm_qkv = make_esmc_attn_layernorm_qkv(d_model, bias, use_te)
- self.out_proj = make_esmc_attn_out_proj(d_model, bias, use_te)
+ self.layernorm_qkv = make_esmc_attn_layernorm_qkv(d_model, bias)
+ self.out_proj = make_esmc_attn_out_proj(d_model, bias)
if qk_layernorm:
self.q_ln = nn.LayerNorm(d_model, bias=bias)
@@ -490,19 +476,10 @@ class EsmcFlashMultiHeadAttention(EsmcMultiHeadAttention):
"""Flash-Attention 2 variant of :class:`EsmcMultiHeadAttention`."""
def __init__(
- self,
- d_model: int,
- n_heads: int,
- bias: bool = False,
- qk_layernorm: bool = True,
- use_te: bool = False,
+ self, d_model: int, n_heads: int, bias: bool = False, qk_layernorm: bool = True
):
super().__init__(
- d_model=d_model,
- n_heads=n_heads,
- bias=bias,
- qk_layernorm=qk_layernorm,
- use_te=use_te,
+ d_model=d_model, n_heads=n_heads, bias=bias, qk_layernorm=qk_layernorm
)
self.rotary = EsmcTritonRotaryEmbedding(d_model // n_heads)
@@ -512,6 +489,10 @@ def forward(
seq_id: torch.Tensor | None,
output_attentions: bool = False,
) -> tuple[torch.Tensor, torch.Tensor | None]:
+ if not x.is_cuda:
+ # Built on CUDA then moved off it: the parent model hands over the
+ # batched layout and a padding mask, so the SDPA path applies as is.
+ return super().forward(x, seq_id, output_attentions=output_attentions)
if output_attentions:
raise ValueError(
"output_attentions=True is not supported with "
@@ -532,9 +513,11 @@ def forward(
# q/k/v are 2D (T, D) here: the parent model unpads the batch before the
# transformer stack to produce the varlen-flat layout that
# ``flash_attn_varlen_qkvpacked_func`` requires.
- T = q.shape[0]
- qkv_packed = torch.stack([q, k, v], dim=1).view(T, 3, self.n_heads, self.d_head)
- qkv_packed = self.rotary(qkv_packed, cu_seqlens, max_seqlen)
+ q, k, v = (t.unflatten(-1, (self.n_heads, self.d_head)) for t in (q, k, v))
+ # Rotate before packing: stacking the rotated q/k is one copy, whereas rotating
+ # a packed tensor out-of-place and re-stacking is two.
+ q, k = self.rotary(q, k, cu_seqlens, max_seqlen)
+ qkv_packed = torch.stack([q, k, v], dim=1)
assert flash_attn_varlen_qkvpacked_func is not None
context = flash_attn_varlen_qkvpacked_func(
@@ -567,8 +550,6 @@ class EsmcUnifiedTransformerBlock(nn.Module):
Whether to apply QK LayerNorm in attention.
ffn_type : str
Feed-forward activation: ``"swiglu"`` or ``"gelu"``.
- use_te : bool
- Build the Transformer Engine fused LayerNorm modules.
"""
def __init__(
@@ -581,19 +562,16 @@ def __init__(
residue_scaling_factor: float = 1.0,
qk_layernorm: bool = True,
ffn_type: str = "swiglu",
- use_te: bool = False,
):
super().__init__()
attn_cls = (
EsmcFlashMultiHeadAttention if use_flash_attn else EsmcMultiHeadAttention
)
- self.attn = attn_cls(
- d_model, n_heads, bias=bias, qk_layernorm=qk_layernorm, use_te=use_te
- )
+ self.attn = attn_cls(d_model, n_heads, bias=bias, qk_layernorm=qk_layernorm)
if ffn_type == "swiglu":
- self.ffn = esmc_swiglu_ln_ffn(d_model, expansion_ratio, bias, use_te)
+ self.ffn = esmc_swiglu_ln_ffn(d_model, expansion_ratio, bias)
elif ffn_type == "gelu":
self.ffn = esmc_gelu_ln_ffn(d_model, expansion_ratio, bias)
else:
@@ -654,8 +632,6 @@ class EsmcTransformerStack(nn.Module):
FFN expansion ratio.
use_flash_attn : bool
Use the Flash Attention 2 kernel when available.
- use_te : bool
- Build the Transformer Engine fused LayerNorm modules.
"""
def __init__(
@@ -669,7 +645,6 @@ def __init__(
ffn_type: str = "swiglu",
expansion_ratio: float = 8 / 3,
use_flash_attn: bool = False,
- use_te: bool = False,
):
super().__init__()
self.blocks = nn.ModuleList(
@@ -685,7 +660,6 @@ def __init__(
bias=bias,
qk_layernorm=qk_layernorm,
ffn_type=ffn_type,
- use_te=use_te,
)
for _ in range(n_layers)
]
diff --git a/esm/models/esmc/model.py b/esm/models/esmc/model.py
index dea2ee95..9d647ae1 100644
--- a/esm/models/esmc/model.py
+++ b/esm/models/esmc/model.py
@@ -19,12 +19,7 @@
from esm.models.esmc.checkpoint_layout import native_to_published, published_to_native
from esm.models.esmc.config import EsmcConfig
-from esm.models.esmc.kernels import (
- FLASH_ATTN_INSTALLED,
- TE_INSTALLED,
- pad_input,
- unpad_input,
-)
+from esm.models.esmc.kernels import FLASH_ATTN_INSTALLED, pad_input, unpad_input
from esm.models.esmc.layers import EsmcRotaryEmbedding, EsmcTransformerStack
from esm.models.esmc.sae import EsmcSaeLayer
from esm.models.hub import HubPreTrainedModel, resolve_model_dir
@@ -102,37 +97,6 @@ class EsmcSequenceClassifierOutput:
attentions: tuple[torch.Tensor, ...] | None = None
-# ---------------------------------------------------------------------------
-# State-dict / weight loading helpers
-# ---------------------------------------------------------------------------
-
-
-def _drop_te_extra_state(
- module: nn.Module, state_dict: dict, prefix: str, local_metadata: dict
-) -> None:
- """Strip Transformer Engine's non-tensor ``_extra_state`` entries; TE
- returns these as ``io.BytesIO`` objects the checkpoint writer cannot
- dtype-infer."""
- for key in list(state_dict):
- if key.endswith("_extra_state"):
- del state_dict[key]
-
-
-def _ignore_te_extra_state(module: nn.Module, incompatible_keys) -> None:
- """The load-side counterpart of :func:`_drop_te_extra_state`.
-
- Without it a model cannot even reload its own ``state_dict``, since the keys
- stripped on the way out are still expected on the way in. They hold nothing
- but fp8 scaling metadata, which is empty until fp8 runs and is recomputed
- when it does.
- """
- incompatible_keys.missing_keys[:] = [
- key
- for key in incompatible_keys.missing_keys
- if not key.endswith("_extra_state")
- ]
-
-
# ---------------------------------------------------------------------------
# Base class
# ---------------------------------------------------------------------------
@@ -181,8 +145,8 @@ def _adapt_checkpoint_keys(
``lm_head.*``. Keys are first translated out of the published layout
(see :mod:`.checkpoint_layout`), then matched against the target tree:
this keeps keys that already match, strips the ``esmc.`` prefix when
- loading the bare encoder, drops Transformer-Engine ``_extra_state``
- entries, and ignores keys with no counterpart (e.g. ``lm_head.*`` when
+ loading the bare encoder, drops ``_extra_state`` entries left by
+ checkpoints saved with Transformer Engine, and ignores keys with no counterpart (e.g. ``lm_head.*`` when
loading :class:`EsmcModel`).
"""
adapted: dict[str, torch.Tensor] = {}
@@ -204,8 +168,8 @@ def _normalize_checkpoint_layout(
def save_pretrained(self, save_directory: str | os.PathLike) -> None:
"""Write ``config.json`` and ``model.safetensors`` in published layout.
- The in-memory state dict keeps Transformer Engine's fused parameters, so
- it is translated on the way out; see :mod:`.checkpoint_layout`.
+ The in-memory state dict keeps the projections fused, so it is
+ translated on the way out; see :mod:`.checkpoint_layout`.
"""
directory = Path(save_directory)
directory.mkdir(parents=True, exist_ok=True)
@@ -284,11 +248,10 @@ class EsmcModel(EsmcPreTrainedModel):
def __init__(self, config: EsmcConfig):
super().__init__(config)
- # Every fused kernel is CUDA-only, so all of them hang off one condition:
- # the device this model's parameters are being created on. Deriving them
- # separately is what let the LayerNorm and attention gates disagree.
+ # Fused kernels are CUDA-only; gate on the device the parameters are
+ # being created on.
on_cuda = torch.get_default_device().type == "cuda"
- self._use_flash_attn = (
+ self._flash_attn_requested = (
on_cuda
and FLASH_ATTN_INSTALLED
and config.attn_implementation == "flash_attention_2"
@@ -298,12 +261,14 @@ def __init__(self, config: EsmcConfig):
config.hidden_size,
config.num_attention_heads,
config.num_hidden_layers,
- use_flash_attn=self._use_flash_attn,
- use_te=on_cuda and TE_INSTALLED,
+ use_flash_attn=self._flash_attn_requested,
)
self._sae_models: nn.ModuleDict = nn.ModuleDict()
- self._register_state_dict_hook(_drop_te_extra_state)
- self.register_load_state_dict_post_hook(_ignore_te_extra_state)
+
+ @property
+ def _use_flash_attn(self) -> bool:
+ """Whether forward takes the varlen flash path; re-read so ``.to("cpu")`` falls back."""
+ return self._flash_attn_requested and self.embed.weight.is_cuda
def get_input_embeddings(self) -> nn.Embedding:
return self.embed
diff --git a/esm/models/esmc/tokenizer.py b/esm/models/esmc/tokenizer.py
index 6a4b03b5..a956a0af 100644
--- a/esm/models/esmc/tokenizer.py
+++ b/esm/models/esmc/tokenizer.py
@@ -3,7 +3,7 @@
from tokenizers import AddedToken, Tokenizer
from tokenizers.models import BPE
from tokenizers.processors import TemplateProcessing
-from transformers import PreTrainedTokenizerFast
+from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}
diff --git a/esm/models/esmfold2/__init__.py b/esm/models/esmfold2/__init__.py
index 7409ab38..ce27971f 100644
--- a/esm/models/esmfold2/__init__.py
+++ b/esm/models/esmfold2/__init__.py
@@ -17,6 +17,7 @@
DNAInput,
LigandInput,
Modification,
+ PocketConditioning,
ProteinInput,
RNAInput,
StructurePredictionInput,
@@ -46,6 +47,7 @@
"MolecularComplex",
"MolecularComplexMetadata",
"MolecularComplexResult",
+ "PocketConditioning",
"ProteinInput",
"RNAInput",
"StructurePredictionInput",
diff --git a/esm/models/esmfold2/config.py b/esm/models/esmfold2/config.py
index ffe39908..dd5e3bd0 100644
--- a/esm/models/esmfold2/config.py
+++ b/esm/models/esmfold2/config.py
@@ -1,16 +1,3 @@
-# Copyright 2026 Biohub. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
"""ESMFold2 model configuration."""
from __future__ import annotations
diff --git a/esm/models/esmfold2/distributed/esmc_tp.py b/esm/models/esmfold2/distributed/esmc_tp.py
index 8fdb7260..d8bc9e83 100644
--- a/esm/models/esmfold2/distributed/esmc_tp.py
+++ b/esm/models/esmfold2/distributed/esmc_tp.py
@@ -21,109 +21,109 @@
"""Tensor-parallel ESM-C over the CP ranks (inference-only).
-The replicated ESM-C 6B forward is the dominant per-rank floor (T0 re-hook:
-~19.8 GB live, identical on every rank). Tensor parallelism shards its weights
-AND intra-block activations across the CP ranks.
-
-**Phase 1 (this module): the SwiGLU MLP** — ~67% of ESM-C's params. Each
-``UnifiedTransformerBlock.ffn`` is a Transformer Engine ``LayerNormMLP``; TE has
-native TP (``set_parallel_mode=True`` → column-parallel fc1 with the SwiGLU
-gate/up split handled internally + row-parallel fc2 + an internal all-reduce), so
-the wrap is: construct a TP ``LayerNormMLP`` over the CP process group and copy the
-full weights in sharded. The MLP's ``ffn_hidden`` (6912) divides 4/9/16, so this
-works at every CP grid (unlike head-parallel attention, blocked by n_heads=40).
-
-Validated bit-exact (modulo bf16 rounding) under ``torch.inference_mode`` — the
-column/row shard + internal all-reduce reproduce the serial output; the ESM-C
-forward already runs under the model's ``@torch.inference_mode`` so TE-TP (plain
-sharded matmuls + NCCL, no autograd hooks) is inference-safe.
-
-Attention TP (qkv column + out_proj row + the full-width QK-LayerNorm all-reduce,
-constrained to TP=4 by the 40 heads) is a later phase — see ``fix_9x_ESMC_TP.md``.
+The replicated ESM-C 6B forward is the dominant per-rank floor (~19.8 GB live,
+identical on every rank). This module shards its SwiGLU MLP, ~67% of ESM-C's
+params, across the CP ranks.
+
+fc1 is column-parallel: each rank keeps its slice of both gate and up. fc2 is
+row-parallel: each rank keeps the matching input columns. One all-reduce after
+fc2 restores the full, replicated block output. ``ffn_hidden`` (6912) divides
+4/9/16, so this works at every CP grid.
+
+Attention TP (blocked to TP=4 by n_heads=40) is not implemented.
"""
import torch
+import torch.distributed as dist
import torch.nn as nn
+from torch import Tensor
+from esm.models.esmc.layers import EsmcLayerNormMLP
+from esm.models.esmfold2.fp8 import Fp8LayerNormMLP, Fp8Linear
+
+
+class AllReduced(nn.Module):
+ """Runs one rank's MLP shard and sums the partial outputs across ranks."""
+
+ def __init__(self, shard: nn.Module, tp_group: dist.ProcessGroup) -> None:
+ super().__init__()
+ self.shard = shard
+ self.tp_group = tp_group
+
+ def forward(self, x: Tensor) -> Tensor:
+ out = self.shard(x)
+ dist.all_reduce(out, group=self.tp_group)
+ return out
-def _tp_layernorm_mlp(old, tp_group, tp_size: int, rank: int):
- """Build a TE LayerNormMLP TP-shard of ``old`` (a serial te.LayerNormMLP),
- holding only this rank's slice. Bit-exact (bf16) vs the full module."""
- import transformer_engine.pytorch as te
-
- hidden = old.layer_norm_weight.shape[0]
- ffn = old.fc2_weight.shape[1] # fc2 = [hidden, ffn]
- if ffn % tp_size:
- raise ValueError(f"ESM-C ffn_hidden={ffn} not divisible by tp_size={tp_size}")
- dtype = old.fc1_weight.dtype
- device = old.fc1_weight.device
-
- new = te.LayerNormMLP(
- hidden,
- ffn,
- eps=getattr(old, "eps", 1e-5),
- normalization=getattr(old, "normalization", "LayerNorm"),
- activation=getattr(old, "activation", "swiglu"),
- bias=False, # ESM-C trained with bias=False
- zero_centered_gamma=getattr(old, "zero_centered_gamma", False),
- set_parallel_mode=True,
- tp_group=tp_group,
- tp_size=tp_size,
- sequence_parallel=False,
- params_dtype=dtype,
- ).to(device=device, dtype=dtype)
- new.eval()
-
- fl = ffn // tp_size
- with torch.no_grad():
- new.layer_norm_weight.copy_(old.layer_norm_weight) # ty:ignore[call-non-callable]
- new.layer_norm_bias.copy_(old.layer_norm_bias) # ty:ignore[unresolved-attribute]
- # fc1 full = [2*ffn, hidden] laid out [gate(ffn); up(ffn)]; column-parallel
- # SwiGLU shard = cat([gate[r], up[r]]) (confirmed bit-exact by the spike).
- gate, up = old.fc1_weight[:ffn], old.fc1_weight[ffn:]
- new.fc1_weight.copy_(
- torch.cat(
- [gate[rank * fl : (rank + 1) * fl], up[rank * fl : (rank + 1) * fl]],
- dim=0,
- )
- ) # ty:ignore[call-non-callable]
- # fc2 full = [hidden, ffn]; row-parallel shards the input (ffn) dim.
- new.fc2_weight.copy_(old.fc2_weight[:, rank * fl : (rank + 1) * fl]) # ty:ignore[call-non-callable]
- return new
+
+def _gate_up_rows(fc1: Tensor, ffn: int, lo: int, hi: int) -> Tensor:
+ """fc1 is packed ``[gate; up]``, so a shard takes the same rows from each half."""
+ return torch.cat([fc1[lo:hi], fc1[ffn + lo : ffn + hi]])
+
+
+def _shard_bf16(src: EsmcLayerNormMLP, lo: int, hi: int) -> EsmcLayerNormMLP:
+ with torch.device("meta"):
+ shard = EsmcLayerNormMLP(src.hidden_size, hi - lo, src.eps)
+ shard.layer_norm_weight = src.layer_norm_weight
+ shard.layer_norm_bias = src.layer_norm_bias
+ ffn = src.ffn_hidden_size
+ shard.fc1_weight = nn.Parameter(
+ _gate_up_rows(src.fc1_weight, ffn, lo, hi), requires_grad=False
+ )
+ shard.fc2_weight = nn.Parameter(
+ src.fc2_weight[:, lo:hi].contiguous(), requires_grad=False
+ )
+ return shard
+
+
+def _shard_fp8(src: Fp8LayerNormMLP, lo: int, hi: int) -> Fp8LayerNormMLP:
+ # Scales are per output row, so a column slice of fc2 keeps its scales.
+ ffn = src.fc2.in_features
+ fc1 = Fp8Linear(
+ _gate_up_rows(src.fc1.weight_fp8, ffn, lo, hi),
+ _gate_up_rows(src.fc1.weight_scale, ffn, lo, hi),
+ )
+ fc2 = Fp8Linear(src.fc2.weight_fp8[:, lo:hi].contiguous(), src.fc2.weight_scale)
+ return Fp8LayerNormMLP(
+ src.layer_norm_weight, src.layer_norm_bias, src.eps, fc1, fc2
+ )
def tp_shard_esmc_mlp(model: nn.Module, dist_manager, tp_group=None) -> int:
- """Replace every ESM-C block's SwiGLU MLP with a TE tensor-parallel shard
+ """Replace every ESM-C block's SwiGLU MLP with its tensor-parallel shard
over ``tp_group`` (default: the full CP process group). Returns the number of
blocks sharded. No-op (returns 0) if ESM-C isn't loaded.
-
- Each rank ends up holding ``1/tp_size`` of every MLP weight; the per-block
- forward all-reduces internally so the block output stays replicated — the
- ESM-C hidden states reach all CP ranks unchanged for the #14 lm_z builder.
"""
esmc = getattr(model, "esmc", None)
if esmc is None:
return 0
- blocks = esmc.transformer.blocks
if tp_group is None:
tp_group = dist_manager.group["cp"]
- tp_size = torch.distributed.get_world_size(tp_group)
- rank = torch.distributed.get_rank(tp_group)
+ tp_size = dist.get_world_size(tp_group)
+ rank = dist.get_rank(tp_group)
if tp_size == 1:
return 0
n = 0
- for block in blocks:
- ffn = getattr(block, "ffn", None)
- # Only the TE LayerNormMLP path is supported (the accelerated build);
- # the pure-PyTorch fallback would need a DTensor TP and isn't on this path.
- if ffn is None or type(ffn).__name__ != "LayerNormMLP":
- raise TypeError(
- f"expected block.ffn to be a TE LayerNormMLP, got "
- f"{type(ffn).__name__ if ffn is not None else None}"
+ for block in esmc.transformer.blocks:
+ ffn = block.ffn
+ width = (
+ ffn.fc2.in_features
+ if isinstance(ffn, Fp8LayerNormMLP)
+ else ffn.ffn_hidden_size
+ )
+ if width % tp_size:
+ raise ValueError(
+ f"ESM-C ffn_hidden={width} not divisible by tp_size={tp_size}"
)
- block.ffn = _tp_layernorm_mlp(ffn, tp_group, tp_size, rank)
- del ffn
+ lo, hi = rank * (width // tp_size), (rank + 1) * (width // tp_size)
+ if isinstance(ffn, Fp8LayerNormMLP):
+ shard: nn.Module = _shard_fp8(ffn, lo, hi)
+ elif isinstance(ffn, EsmcLayerNormMLP):
+ shard = _shard_bf16(ffn, lo, hi)
+ else:
+ raise TypeError(f"cannot shard block.ffn of type {type(ffn).__name__}")
+ block.ffn = AllReduced(shard, tp_group)
n += 1
return n
diff --git a/esm/models/esmfold2/distributed/msa_wrapper.py b/esm/models/esmfold2/distributed/msa_wrapper.py
index 16378899..260e8796 100644
--- a/esm/models/esmfold2/distributed/msa_wrapper.py
+++ b/esm/models/esmfold2/distributed/msa_wrapper.py
@@ -265,11 +265,11 @@ def wrap_model_with_cp(
``bf16`` (default True): run the distributed trunk in bf16 — quality-neutral,
~1.3-1.6x faster, lower peak VRAM (see ``wrap_model_with_cp_trunks``).
- ``offload_esmc`` (default True): offload the ESM-C LM (~12 GB) to CPU after
+ ``offload_esmc`` (default True): offload the ESM-C LM (~6 GB in fp8, ~12 GB in bf16) to CPU after
its one-shot use, freeing it for the trunk/diffusion. Sets
``model._offload_esmc``. Validated with bf16: 2.14x end-to-end + 18.5 GB lower
peak vs fp32, pLDDT unchanged. Set False for high-throughput tiny-fold or
- concurrent-fold use (the per-fold ~12 GB CPU<->GPU transfer isn't worth it).
+ concurrent-fold use (the per-fold ~6-12 GB CPU<->GPU transfer isn't worth it).
"""
from esm.models.esmfold2.distributed.utils import wrap_model_with_cp_trunks
@@ -343,7 +343,7 @@ def wrap_model_with_cp(
replaced += wrap_model_with_cp_structure_head(model, dist_manager)
# Tensor-parallel ESM-C (the dominant replicated floor). Phase 1: shard the
- # SwiGLU MLP (~67% of ESM-C weights) over the CP ranks via TE-native TP.
+ # SwiGLU MLP (~67% of ESM-C weights) over the CP ranks.
if tp_esmc:
from esm.models.esmfold2.distributed.esmc_tp import tp_shard_esmc_mlp
diff --git a/esm/models/esmfold2/experimental.py b/esm/models/esmfold2/experimental.py
index e73bbefe..d79583cf 100644
--- a/esm/models/esmfold2/experimental.py
+++ b/esm/models/esmfold2/experimental.py
@@ -1,17 +1,3 @@
-# coding=utf-8
-# Copyright 2026 Biohub. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
"""ESMFold2 experimental variant — older architecture from during the
development of ESMFold2.
@@ -428,6 +414,9 @@ def _maybe_mask(x: Tensor) -> Tensor:
pair_repr = pair_repr + _maybe_mask(
self.outer_product_mean(msa_repr, msa_attention_mask)
)
+ # Eager on purpose, including under "fast": _maybe_mask scales the
+ # update before the residual add, and the fused kernel adds its residual
+ # inside. model.py's MSAEncoderBlock has no such mask and does fuse.
pair_repr = pair_repr + _maybe_mask(
self.tri_mul_out(pair_repr, mask=pair_attention_mask)
)
@@ -621,7 +610,13 @@ def __init__(self, config: EsmFold2Config) -> None:
self.post_init()
def set_kernel_backend(self, backend: str | None) -> None:
- """Select kernel backend (None / "fused" / "cuequivariance")."""
+ """Select kernel backend (None / "fused" / "fast").
+
+ "fast" reaches the trunk here, but not this model's own MSA encoder: its
+ block multiplies each triangle update by the MSA track mask before the
+ residual add, which the kernel's fused residual cannot express. That
+ encoder keeps the eager engine under every backend.
+ """
self.folding_trunk.set_kernel_backend(backend)
if self.confidence_head is not None:
self.confidence_head.set_kernel_backend(backend)
diff --git a/esm/models/esmfold2/fp8.py b/esm/models/esmfold2/fp8.py
new file mode 100644
index 00000000..39c6cb91
--- /dev/null
+++ b/esm/models/esmfold2/fp8.py
@@ -0,0 +1,223 @@
+"""fp8 inference for the ESMC backbone, on PyTorch's native ``torch._scaled_mm``.
+
+Weights are stored as fp8 (e4m3) with one fp32 scale per output row.
+Activations are quantized per row on the fly. The fp8 GEMM needs compute
+capability >= 8.9 (Ada or Hopper). Elsewhere the weight is upcast and the matmul
+runs in the input dtype, so fp8 still halves weight memory.
+"""
+
+import functools
+import logging
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from esm.models.esmc.layers import EsmcLayerNormLinear, EsmcLayerNormMLP
+
+try:
+ from esm.models.esmfold2.kernels.fp8_quantize import quantize_rows_fp8
+except (ImportError, RuntimeError):
+ quantize_rows_fp8: Any = None
+
+logger = logging.getLogger(__name__)
+
+_FP8 = torch.float8_e4m3fn
+_FP8_MAX = torch.finfo(_FP8).max
+
+
+@functools.cache
+def _has_fp8_gemm(device: torch.device) -> bool:
+ """Whether ``device`` runs fp8 GEMMs. Warns once per device on a fallback."""
+ fp8 = device.type == "cuda" and torch.cuda.get_device_capability(device) >= (8, 9)
+ if not fp8:
+ logger.warning(
+ "fp8 ESMC on %s: no fp8 GEMM here (needs compute capability >= 8.9). "
+ "Weights stay fp8, but each matmul upcasts them and runs without the "
+ "fp8 speedup.",
+ device,
+ )
+ elif quantize_rows_fp8 is None:
+ logger.warning(
+ "fp8 ESMC on %s: Triton is unavailable, so activations are quantized "
+ "with slower PyTorch ops.",
+ device,
+ )
+ return fp8
+
+
+def _quantize_rows(
+ x: Tensor, *, swiglu: bool = False, norm: tuple[Tensor, Tensor, float] | None = None
+) -> tuple[Tensor, Tensor]:
+ """Quantize a 2D tensor to fp8 with one fp32 scale per row.
+
+ ``swiglu`` and ``norm`` fold the preceding op in, as in :func:`quantize_rows_fp8`.
+ """
+ # Triton only emits e4m3 on compute capability >= 8.9.
+ if quantize_rows_fp8 is not None and _has_fp8_gemm(x.device):
+ return quantize_rows_fp8(x, swiglu=swiglu, norm=norm)
+ return _quantize_rows_torch(x, swiglu=swiglu, norm=norm)
+
+
+def _quantize_rows_torch(
+ x: Tensor, *, swiglu: bool = False, norm: tuple[Tensor, Tensor, float] | None = None
+) -> tuple[Tensor, Tensor]:
+ """Reference for :func:`quantize_rows_fp8`, and the path without Triton."""
+ x = x.float()
+ if swiglu:
+ gate, up = x.chunk(2, dim=-1)
+ x = F.silu(gate) * up
+ if norm is not None:
+ weight, bias, eps = norm
+ x = F.layer_norm(x, x.shape[-1:], weight.float(), bias.float(), eps)
+ scale = x.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) / _FP8_MAX
+ return (x / scale).to(_FP8), scale
+
+
+class Fp8Linear(nn.Module):
+ """A frozen, bias-free Linear whose weight lives in fp8.
+
+ Weight and scale are stored as integer views so ``module.to(dtype)``
+ cannot cast them.
+ """
+
+ weight_fp8: Tensor
+ weight_scale: Tensor
+
+ def __init__(self, weight_fp8: Tensor, weight_scale: Tensor) -> None:
+ super().__init__()
+ self.out_features, self.in_features = weight_fp8.shape
+ self.register_buffer("weight_fp8", weight_fp8.view(torch.uint8))
+ self.register_buffer("weight_scale", weight_scale.view(torch.int32))
+
+ @classmethod
+ def from_weight(cls, weight: Tensor) -> "Fp8Linear":
+ return cls(*_quantize_rows(weight.detach()))
+
+ def forward(self, x: Tensor) -> Tensor:
+ if not _has_fp8_gemm(x.device):
+ weight = self.weight_fp8.view(_FP8).float() * self.weight_scale.view(
+ torch.float32
+ )
+ return F.linear(x, weight.to(x.dtype))
+ return self.fp8_gemm(x)
+
+ def fp8_gemm(
+ self,
+ x: Tensor,
+ *,
+ swiglu: bool = False,
+ norm: tuple[Tensor, Tensor, float] | None = None,
+ ) -> Tensor:
+ """Quantize ``x`` (after the optional fused op) and run the fp8 GEMM."""
+ lead = x.shape[:-1]
+ xq, x_scale = _quantize_rows(
+ x.reshape(-1, x.shape[-1]), swiglu=swiglu, norm=norm
+ )
+ out = torch._scaled_mm(
+ xq,
+ self.weight_fp8.view(_FP8).t(),
+ scale_a=x_scale,
+ scale_b=self.weight_scale.view(torch.float32).t(),
+ out_dtype=torch.bfloat16,
+ )
+ return out.reshape(*lead, self.out_features)
+
+
+class _LayerNormThen(nn.Module):
+ """LayerNorm (reusing the source module's parameters) before fp8 GEMMs."""
+
+ def __init__(self, layer_norm_weight: Tensor, layer_norm_bias: Tensor, eps: float):
+ super().__init__()
+ self.eps = eps
+ self.layer_norm_weight = layer_norm_weight
+ self.layer_norm_bias = layer_norm_bias
+
+ def _norm(self, x: Tensor) -> Tensor:
+ return F.layer_norm(
+ x, x.shape[-1:], self.layer_norm_weight, self.layer_norm_bias, self.eps
+ )
+
+ def _norm_args(self) -> tuple[Tensor, Tensor, float]:
+ return self.layer_norm_weight, self.layer_norm_bias, self.eps
+
+
+class Fp8LayerNormLinear(_LayerNormThen):
+ """fp8 counterpart of :class:`EsmcLayerNormLinear`."""
+
+ def __init__(self, src: EsmcLayerNormLinear) -> None:
+ super().__init__(src.layer_norm_weight, src.layer_norm_bias, src.eps)
+ self.linear = Fp8Linear.from_weight(src.weight)
+
+ def forward(self, x: Tensor) -> Tensor:
+ if not _has_fp8_gemm(x.device):
+ return self.linear(self._norm(x))
+ return self.linear.fp8_gemm(x, norm=self._norm_args())
+
+
+class Fp8LayerNormMLP(_LayerNormThen):
+ """fp8 counterpart of :class:`EsmcLayerNormMLP`. ``fc1`` packs ``[gate; up]``."""
+
+ def __init__(
+ self,
+ layer_norm_weight: Tensor,
+ layer_norm_bias: Tensor,
+ eps: float,
+ fc1: Fp8Linear,
+ fc2: Fp8Linear,
+ ) -> None:
+ super().__init__(layer_norm_weight, layer_norm_bias, eps)
+ self.fc1 = fc1
+ self.fc2 = fc2
+
+ @classmethod
+ def from_module(cls, src: EsmcLayerNormMLP) -> "Fp8LayerNormMLP":
+ return cls(
+ src.layer_norm_weight,
+ src.layer_norm_bias,
+ src.eps,
+ Fp8Linear.from_weight(src.fc1_weight),
+ Fp8Linear.from_weight(src.fc2_weight),
+ )
+
+ def forward(self, x: Tensor) -> Tensor:
+ if not _has_fp8_gemm(x.device):
+ x1, x2 = self.fc1(self._norm(x)).chunk(2, dim=-1)
+ return self.fc2(F.silu(x1) * x2)
+ hidden = self.fc1.fp8_gemm(x, norm=self._norm_args())
+ return self.fc2.fp8_gemm(hidden, swiglu=True)
+
+
+def _fp8_counterpart(module: nn.Module) -> nn.Module | None:
+ if isinstance(module, EsmcLayerNormMLP):
+ return Fp8LayerNormMLP.from_module(module)
+ if isinstance(module, EsmcLayerNormLinear):
+ return Fp8LayerNormLinear(module)
+ if isinstance(module, nn.Linear) and module.bias is None:
+ return Fp8Linear.from_weight(module.weight)
+ return None
+
+
+def _swap_children(parent: nn.Module) -> int:
+ swapped = 0
+ for name, child in list(parent.named_children()):
+ new = _fp8_counterpart(child)
+ if new is None:
+ swapped += _swap_children(child)
+ else:
+ setattr(parent, name, new)
+ swapped += 1
+ return swapped
+
+
+def convert_esmc_to_fp8_(esmc: nn.Module) -> int:
+ """Swap every projection in ``esmc``'s transformer blocks for its fp8
+ counterpart, in place. Returns the number of modules swapped.
+
+ Each source weight is freed as soon as its fp8 copy exists. Works on any
+ device, so a checkpoint can be converted on the CPU before it is moved.
+ """
+ with torch.no_grad():
+ return _swap_children(esmc.get_submodule("transformer.blocks"))
diff --git a/esm/models/esmfold2/hf_adapter.py b/esm/models/esmfold2/hf_adapter.py
index 6bb71927..3b787b72 100644
--- a/esm/models/esmfold2/hf_adapter.py
+++ b/esm/models/esmfold2/hf_adapter.py
@@ -182,8 +182,8 @@ def from_pretrained(
model = upstream.EsmFold2Model.from_pretrained(
pretrained_model_name_or_path, dtype=dtype, **kwargs
)
- # No "fp8": ours quantizes TransformerEngine modules in place, and the
- # port's ESMC is plain PyTorch with none for that walk to find.
+ # No "fp8": ours swaps the native ESMC's fused LayerNorm modules, and
+ # the port's ESMC has a different module tree.
esmc_dtypes = {"bf16": torch.bfloat16, "fp32": torch.float32}
if esmc_precision not in esmc_dtypes:
raise ValueError(
diff --git a/esm/models/esmfold2/kernels/__init__.py b/esm/models/esmfold2/kernels/__init__.py
index ef99caeb..a4b02864 100644
--- a/esm/models/esmfold2/kernels/__init__.py
+++ b/esm/models/esmfold2/kernels/__init__.py
@@ -1,11 +1,3 @@
-# coding=utf-8
-# Copyright 2026 Biohub. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
"""Triton inference kernels for ESMFold2."""
from esm.models.esmfold2.kernels.fused_attention_pair_bias import fused_pair_bias
diff --git a/esm/models/esmfold2/kernels/fp8_quantize.py b/esm/models/esmfold2/kernels/fp8_quantize.py
new file mode 100644
index 00000000..17cc29aa
--- /dev/null
+++ b/esm/models/esmfold2/kernels/fp8_quantize.py
@@ -0,0 +1,89 @@
+"""Row-wise fp8 quantization, optionally fused with the op that feeds the GEMM.
+
+One program per row: an optional LayerNorm or SwiGLU, then the row amax, the
+scale and the fp8 cast. Replaces five to seven separate PyTorch kernels per GEMM.
+"""
+
+import torch
+import triton
+import triton.language as tl
+from torch import Tensor
+from triton.language.extra import libdevice
+
+_FP8_MAX = torch.finfo(torch.float8_e4m3fn).max
+
+
+@triton.jit
+def _quantize_rows_kernel(
+ x_ptr,
+ w_ptr,
+ b_ptr,
+ q_ptr,
+ s_ptr,
+ K,
+ stride_x,
+ eps,
+ fp8_max,
+ NORM: tl.constexpr,
+ SWIGLU: tl.constexpr,
+ BLOCK: tl.constexpr,
+):
+ row = tl.program_id(0).to(tl.int64)
+ offs = tl.arange(0, BLOCK)
+ mask = offs < K
+ base = x_ptr + row * stride_x
+ x = tl.load(base + offs, mask=mask, other=0.0).to(tl.float32)
+ if SWIGLU:
+ up = tl.load(base + K + offs, mask=mask, other=0.0).to(tl.float32)
+ x = x * tl.sigmoid(x) * up
+ if NORM:
+ mean = tl.sum(x, axis=0) / K
+ centred = tl.where(mask, x - mean, 0.0)
+ rstd = tl.rsqrt(tl.sum(centred * centred, axis=0) / K + eps)
+ w = tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+ b = tl.load(b_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+ x = centred * rstd * w + b
+ # IEEE division, as torch does: Triton's default `/` is approximate and flips fp8 roundings.
+ scale = tl.math.div_rn(tl.maximum(tl.max(tl.abs(x), axis=0), 1e-12), fp8_max)
+ q = tl.math.div_rn(x, scale)
+ # Snap to the e4m3 grid here: sm89 casts via fp16, and that double rounding
+ # sends ~0.3% of values one fp8 step the wrong way. The cast below is then exact.
+ exponent = tl.maximum(((q.to(tl.int32, bitcast=True) >> 23) & 0xFF) - 127, -6)
+ step = ((exponent - 3 + 127) << 23).to(tl.float32, bitcast=True)
+ q = libdevice.rint(q / step) * step
+ tl.store(q_ptr + row * K + offs, q.to(tl.float8e4nv), mask=mask)
+ tl.store(s_ptr + row, scale)
+
+
+def quantize_rows_fp8(
+ x: Tensor, *, swiglu: bool = False, norm: tuple[Tensor, Tensor, float] | None = None
+) -> tuple[Tensor, Tensor]:
+ """Quantize a 2D CUDA tensor to fp8 e4m3 with one fp32 scale per row.
+
+ ``swiglu``: ``x`` is ``[gate, up]`` along the last dim; quantize
+ ``silu(gate) * up``. ``norm``: ``(weight, bias, eps)`` of a LayerNorm applied
+ first.
+ """
+ assert x.dim() == 2 and x.stride(-1) == 1
+ rows, width = x.shape
+ K = width // 2 if swiglu else width
+ q = torch.empty((rows, K), device=x.device, dtype=torch.float8_e4m3fn)
+ scale = torch.empty((rows, 1), device=x.device, dtype=torch.float32)
+ w, b, eps = norm if norm is not None else (x, x, 0.0)
+ block = triton.next_power_of_2(K)
+ _quantize_rows_kernel[(rows,)](
+ x,
+ w,
+ b,
+ q,
+ scale,
+ K,
+ x.stride(0),
+ eps,
+ _FP8_MAX,
+ NORM=norm is not None, # ty:ignore[invalid-argument-type]
+ SWIGLU=swiglu, # ty:ignore[invalid-argument-type]
+ BLOCK=block,
+ num_warps=8 if block >= 4096 else 4, # ty:ignore[unknown-argument]
+ )
+ return q, scale
diff --git a/esm/models/esmfold2/kernels/fused_attention_pair_bias.py b/esm/models/esmfold2/kernels/fused_attention_pair_bias.py
index df35e0ab..3f60adea 100644
--- a/esm/models/esmfold2/kernels/fused_attention_pair_bias.py
+++ b/esm/models/esmfold2/kernels/fused_attention_pair_bias.py
@@ -11,14 +11,11 @@
# ruff: noqa: E402
-import os
import warnings
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning)
-os.environ.setdefault("CUEQ_DEFAULT_CONFIG", "1")
-os.environ.setdefault("CUEQ_DISABLE_AOT_TUNING", "1")
import torch
import triton
@@ -87,10 +84,12 @@ def _pair_bias_kernel(
mask_k = offs_k < K
mask_h = offs_h < NUM_HEADS
+ # int64 before the scale, not after. Each term is added to the pointer on
+ # its own, so b*Q*K*DIM_Z wraps at 5 diffusion samples of 1,449 tokens and
+ # q*K*DIM_Z wraps at 2,897 tokens. In-tile offsets stay int32.
z_full_ptrs = (
z_ptr
- + pid_b * Q * K * DIM_Z
- + pid_q * K * DIM_Z
+ + (pid_b.to(tl.int64) * Q + pid_q) * K * DIM_Z
+ offs_k[:, None] * DIM_Z
+ offs_z_full[None, :]
)
@@ -109,9 +108,9 @@ def _pair_bias_kernel(
# head-blocks have the same value so this avoids racy duplicate writes.
if SAVE_STATS:
if pid_hblk == 0:
- stats_ptrs = mean_ptr + pid_b * Q * K + pid_q * K + offs_k
+ stats_ptrs = mean_ptr + (pid_b.to(tl.int64) * Q + pid_q) * K + offs_k
tl.store(stats_ptrs, mean, mask=mask_k)
- stats_ptrs2 = rstd_ptr + pid_b * Q * K + pid_q * K + offs_k
+ stats_ptrs2 = rstd_ptr + (pid_b.to(tl.int64) * Q + pid_q) * K + offs_k
tl.store(stats_ptrs2, rstd, mask=mask_k)
acc = tl.zeros([TILE_K, HEADS_PER_BLK], dtype=tl.float32)
@@ -122,8 +121,7 @@ def _pair_bias_kernel(
z_slice_ptrs = (
z_ptr
- + pid_b * Q * K * DIM_Z
- + pid_q * K * DIM_Z
+ + (pid_b.to(tl.int64) * Q + pid_q) * K * DIM_Z
+ offs_k[:, None] * DIM_Z
+ offs_c[None, :]
)
@@ -155,9 +153,7 @@ def _pair_bias_kernel(
out_ptrs = (
out_ptr
- + pid_b * NUM_HEADS * Q * K
- + offs_h[None, :] * Q * K
- + pid_q * K
+ + ((pid_b.to(tl.int64) * NUM_HEADS + offs_h[None, :]) * Q + pid_q) * K
+ offs_k[:, None]
)
tl.store(
@@ -229,18 +225,18 @@ def _pair_bias_backward_kernel(
mask_z = offs_z < DIM_Z
mask_h = offs_h < NUM_HEADS
+ # int64 before the scale, same as the forward kernel.
z_ptrs = (
z_ptr
- + pid_b * Q * K * DIM_Z
- + pid_q * K * DIM_Z
+ + (pid_b.to(tl.int64) * Q + pid_q) * K * DIM_Z
+ offs_k[:, None] * DIM_Z
+ offs_z[None, :]
)
z = tl.load(z_ptrs, mask=mask_k[:, None] & mask_z[None, :], other=0.0).to(
tl.float32
)
- mean_ptrs = mean_ptr + pid_b * Q * K + pid_q * K + offs_k
- rstd_ptrs = rstd_ptr + pid_b * Q * K + pid_q * K + offs_k
+ mean_ptrs = mean_ptr + (pid_b.to(tl.int64) * Q + pid_q) * K + offs_k
+ rstd_ptrs = rstd_ptr + (pid_b.to(tl.int64) * Q + pid_q) * K + offs_k
mean = tl.load(mean_ptrs, mask=mask_k, other=0.0)
rstd = tl.load(rstd_ptrs, mask=mask_k, other=0.0)
@@ -259,9 +255,7 @@ def _pair_bias_backward_kernel(
d_bias_ptrs = (
d_bias_ptr
- + pid_b * NUM_HEADS * Q * K
- + offs_h[None, :] * Q * K
- + pid_q * K
+ + ((pid_b.to(tl.int64) * NUM_HEADS + offs_h[None, :]) * Q + pid_q) * K
+ offs_k[:, None]
)
d_bias = tl.load(d_bias_ptrs, mask=mask_k[:, None] & mask_h[None, :], other=0.0).to(
@@ -298,8 +292,7 @@ def _pair_bias_backward_kernel(
d_z_ptrs = (
d_z_ptr
- + pid_b * Q * K * DIM_Z
- + pid_q * K * DIM_Z
+ + (pid_b.to(tl.int64) * Q + pid_q) * K * DIM_Z
+ offs_k[:, None] * DIM_Z
+ offs_z[None, :]
)
diff --git a/esm/models/esmfold2/kernels/fused_dual_gemm.py b/esm/models/esmfold2/kernels/fused_dual_gemm.py
index e642275b..5e26b48e 100644
--- a/esm/models/esmfold2/kernels/fused_dual_gemm.py
+++ b/esm/models/esmfold2/kernels/fused_dual_gemm.py
@@ -12,14 +12,11 @@
from __future__ import annotations
-import os
import warnings
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning)
-os.environ.setdefault("CUEQ_DEFAULT_CONFIG", "1")
-os.environ.setdefault("CUEQ_DISABLE_AOT_TUNING", "1")
import torch
import triton
diff --git a/esm/models/esmfold2/kernels/fused_lnlin_swiglu.py b/esm/models/esmfold2/kernels/fused_lnlin_swiglu.py
index b163c5aa..f5a9660b 100644
--- a/esm/models/esmfold2/kernels/fused_lnlin_swiglu.py
+++ b/esm/models/esmfold2/kernels/fused_lnlin_swiglu.py
@@ -85,6 +85,7 @@ def _lnlin_swiglu_fwd_kernel(
stride_out_m,
stride_out_n,
HAS_LN_BIAS: tl.constexpr,
+ SAVE_LIN: tl.constexpr,
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
@@ -157,18 +158,20 @@ def _lnlin_swiglu_fwd_kernel(
wa_ptrs += BLOCK_SIZE_K * stride_wk
wb_ptrs += BLOCK_SIZE_K * stride_wk
- a_bf = a_acc.to(Lin_ptr.type.element_ty)
- b_bf = b_acc.to(Lin_ptr.type.element_ty)
-
out_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
- lin_a_ptrs = (
- Lin_ptr + offs_m[:, None] * stride_lin_m + offs_n[None, :] * stride_lin_n
- )
- lin_b_ptrs = (
- Lin_ptr + offs_m[:, None] * stride_lin_m + (N + offs_n[None, :]) * stride_lin_n
- )
- tl.store(lin_a_ptrs, a_bf, mask=out_mask)
- tl.store(lin_b_ptrs, b_bf, mask=out_mask)
+ # The [M, 2N] linear output is read only by backward. A forward-only call
+ # skips the store, which is the single largest write this kernel makes.
+ if SAVE_LIN:
+ lin_a_ptrs = (
+ Lin_ptr + offs_m[:, None] * stride_lin_m + offs_n[None, :] * stride_lin_n
+ )
+ lin_b_ptrs = (
+ Lin_ptr
+ + offs_m[:, None] * stride_lin_m
+ + (N + offs_n[None, :]) * stride_lin_n
+ )
+ tl.store(lin_a_ptrs, a_acc.to(Lin_ptr.type.element_ty), mask=out_mask)
+ tl.store(lin_b_ptrs, b_acc.to(Lin_ptr.type.element_ty), mask=out_mask)
# SwiGLU = silu(a) * b. SiLU computed in fp32 for numerical accuracy.
sig = tl.sigmoid(a_acc)
@@ -253,7 +256,11 @@ def _ln_stats_settings(K: int) -> tuple[int, int]:
def _lnlin_swiglu_fwd(
- x_2d: torch.Tensor, W12: torch.Tensor, LN_W: torch.Tensor, LN_B: torch.Tensor | None
+ x_2d: torch.Tensor,
+ W12: torch.Tensor,
+ LN_W: torch.Tensor,
+ LN_B: torch.Tensor | None,
+ save_lin: bool = True,
):
assert x_2d.is_contiguous(), "X must be contiguous"
M, K = x_2d.shape
@@ -262,7 +269,11 @@ def _lnlin_swiglu_fwd(
N = two_N // 2
out = torch.empty((M, N), dtype=x_2d.dtype, device=x_2d.device)
- lin = torch.empty((M, two_N), dtype=x_2d.dtype, device=x_2d.device)
+ # Placeholder when the caller has no backward: the kernel still needs a
+ # valid pointer of the right dtype, but never dereferences it.
+ lin = torch.empty(
+ (M, two_N) if save_lin else (1,), dtype=x_2d.dtype, device=x_2d.device
+ )
Mean = torch.empty((M,), dtype=x_2d.dtype, device=x_2d.device)
Rstd = torch.empty((M,), dtype=x_2d.dtype, device=x_2d.device)
@@ -300,11 +311,12 @@ def _lnlin_swiglu_fwd(
x_2d.stride(1),
W12.stride(0),
W12.stride(1),
- lin.stride(0),
- lin.stride(1),
+ two_N if save_lin else 0,
+ 1,
out.stride(0),
out.stride(1),
HAS_LN_BIAS=(LN_B is not None), # ty:ignore[invalid-argument-type]
+ SAVE_LIN=save_lin, # ty:ignore[invalid-argument-type]
BLOCK_SIZE_M=cfg["BLOCK_SIZE_M"],
BLOCK_SIZE_N=cfg["BLOCK_SIZE_N"],
BLOCK_SIZE_K=cfg["BLOCK_SIZE_K"],
@@ -342,11 +354,14 @@ class FusedLNLinearSwiGLUFunction(torch.autograd.Function):
@staticmethod
@torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.bfloat16)
- def forward(ctx, X, W12, LN_W, LN_B):
+ def forward(ctx, X, W12, LN_W, LN_B, needs_backward: bool = True):
x_shape = X.shape
x_2d = X.contiguous().view(-1, x_shape[-1])
- out, lin, mean, rstd = _lnlin_swiglu_fwd(x_2d, W12, LN_W, LN_B)
- ctx.save_for_backward(x_2d, W12, LN_W, LN_B, mean, rstd, lin)
+ out, lin, mean, rstd = _lnlin_swiglu_fwd(
+ x_2d, W12, LN_W, LN_B, save_lin=needs_backward
+ )
+ if needs_backward:
+ ctx.save_for_backward(x_2d, W12, LN_W, LN_B, mean, rstd, lin)
ctx.x_shape = x_shape
ctx.has_ln_bias = LN_B is not None
return out.view(*x_shape[:-1], out.shape[-1])
@@ -375,7 +390,7 @@ def backward(ctx, dout):
dx_norm, x_2d, [K], mean.float(), rstd.float(), LN_W, LN_B, output_mask
)
# If LN_B was None (no bias), ATen returns dLN_B=None and we just pass it through.
- return dX.view(ctx.x_shape), dW12, dLN_W, dLN_B
+ return dX.view(ctx.x_shape), dW12, dLN_W, dLN_B, None
class FusedLNLinearSwiGLU(nn.Module):
@@ -410,4 +425,12 @@ def reset_parameters(self):
nn.init.uniform_(self.W12, -bound, bound)
def forward(self, X: torch.Tensor) -> torch.Tensor:
- return FusedLNLinearSwiGLUFunction.apply(X, self.W12, self.LN_W, self.LN_B)
+ # Decided here, not inside forward(): autograd runs Function.forward under
+ # no_grad, so torch.is_grad_enabled() is always False in there.
+ needs_backward = torch.is_grad_enabled() and any(
+ t is not None and t.requires_grad
+ for t in (X, self.W12, self.LN_W, self.LN_B)
+ )
+ return FusedLNLinearSwiGLUFunction.apply(
+ X, self.W12, self.LN_W, self.LN_B, needs_backward
+ )
diff --git a/esm/models/esmfold2/kernels/trimul_with_residual.py b/esm/models/esmfold2/kernels/trimul_with_residual.py
index 37b4b92c..5736c8f1 100644
--- a/esm/models/esmfold2/kernels/trimul_with_residual.py
+++ b/esm/models/esmfold2/kernels/trimul_with_residual.py
@@ -36,14 +36,11 @@
from __future__ import annotations
-import os
import warnings
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning)
-os.environ.setdefault("CUEQ_DEFAULT_CONFIG", "1")
-os.environ.setdefault("CUEQ_DISABLE_AOT_TUNING", "1")
import torch
import triton
diff --git a/esm/models/esmfold2/layers.py b/esm/models/esmfold2/layers.py
index d4a796e7..87425fb3 100644
--- a/esm/models/esmfold2/layers.py
+++ b/esm/models/esmfold2/layers.py
@@ -1,16 +1,9 @@
-# coding=utf-8
-# Copyright 2026 Biohub. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
"""Shared building blocks for ESMFold2 HuggingFace model variants."""
from __future__ import annotations
import random
+import warnings
from contextlib import contextmanager
from functools import partial
from typing import Any, cast
@@ -34,18 +27,6 @@
pad_input: Any = None
FLASH_ATTN_AVAILABLE = False
-try:
- from cuequivariance_torch import attention_pair_bias as _cue_attn_pair_bias
- from cuequivariance_torch.primitives.triangle import (
- triangle_multiplicative_update as _cue_tri_mul,
- )
-
- CUE_AVAILABLE = True
-except ImportError:
- _cue_attn_pair_bias: Any = None
- _cue_tri_mul: Any = None
- CUE_AVAILABLE = False
-
# Vendored inference-only Triton kernels.
try:
# `@triton.autotune` resolves `driver.active` when it decorates, so importing
@@ -73,8 +54,13 @@
from esm.models.esmfold2.config import EsmFold2Config
BACKEND_FUSED = "fused"
-BACKEND_CUEQ = "cuequivariance"
-_VALID_BACKENDS = (None, BACKEND_FUSED, BACKEND_CUEQ)
+BACKEND_FAST = "fast"
+_VALID_BACKENDS = (None, BACKEND_FUSED, BACKEND_FAST)
+
+# "fast" is a superset of "fused": every bit-exact kernel still runs, plus the
+# approximate ones. So a fused-path gate must accept it, and only the gates
+# named ``_fast_*`` below are exclusive to it.
+_KERNEL_BACKENDS = (BACKEND_FUSED, BACKEND_FAST)
# The vendored fused Triton kernels (LN+SwiGLU, trimul-with-residual) operate in
# bfloat16 only. Single-source that dtype here so every fused-path buffer and
@@ -86,7 +72,7 @@ def _fused_active(module: nn.Module, tensor: Tensor) -> bool:
"""Common preconditions for the vendored fused Triton inference kernels."""
return (
TRITON_KERNELS_AVAILABLE
- and getattr(module, "_kernel_backend", None) == BACKEND_FUSED
+ and getattr(module, "_kernel_backend", None) in _KERNEL_BACKENDS
and not torch.is_grad_enabled()
and tensor.is_cuda
)
@@ -96,7 +82,23 @@ def _fused_pair_stack_active(module: nn.Module, tensor: Tensor) -> bool:
"""Fused pair-stack kernels that support autograd."""
return (
TRITON_KERNELS_AVAILABLE
- and getattr(module, "_kernel_backend", None) == BACKEND_FUSED
+ and getattr(module, "_kernel_backend", None) in _KERNEL_BACKENDS
+ and tensor.is_cuda
+ )
+
+
+def _fast_pair_stack_active(module: nn.Module, tensor: Tensor) -> bool:
+ """Only under "fast": the approximate pair-stack kernels.
+
+ Inference only, unlike ``_fused_pair_stack_active``. That gate is permissive
+ because its kernel is bit-exact, so letting gradients through costs nothing.
+ This one changes values, and an opt-in inference accelerator must not
+ silently move training gradients.
+ """
+ return (
+ TRITON_KERNELS_AVAILABLE
+ and getattr(module, "_kernel_backend", None) == BACKEND_FAST
+ and not torch.is_grad_enabled()
and tensor.is_cuda
)
@@ -106,8 +108,36 @@ def _to_fused_kernel_dtype(t: Tensor) -> Tensor:
return t if t.dtype == _FUSED_KERNEL_DTYPE else t.to(_FUSED_KERNEL_DTYPE)
-def _cueq_active(module: nn.Module) -> bool:
- return CUE_AVAILABLE and getattr(module, "_kernel_backend", None) == BACKEND_CUEQ
+_fast_trimul_refusal_warned = False
+
+
+def fast_trimul_refusal_reason(pair: Tensor) -> str:
+ """Why the approximate triangle-multiplication kernel cannot serve a call."""
+ if not TRITON_KERNELS_AVAILABLE:
+ return "the vendored Triton kernels are unavailable"
+ if not pair.is_cuda:
+ return f"the pair tensor is on {pair.device.type}, not cuda"
+ if torch.is_grad_enabled():
+ return "gradients are enabled and this kernel is inference-only"
+ return "the kernel is not installed"
+
+
+def warn_fast_trimul_refused(where: str, reason: str) -> None:
+ """Warn once per process that the fast trimul fell back to the eager engine."""
+ global _fast_trimul_refusal_warned
+ if _fast_trimul_refusal_warned:
+ return
+ _fast_trimul_refusal_warned = True
+ warnings.warn(
+ f"ESMFold2: {where} is running the eager triangle multiplication because "
+ f"{reason}; set_kernel_backend('{BACKEND_FAST}') has no effect there.",
+ stacklevel=3,
+ )
+
+
+def _any_requires_grad(*tensors: Tensor | None) -> bool:
+ """True if any given tensor is not None and requires grad."""
+ return any(t is not None and t.requires_grad for t in tensors)
class DropoutResidual(nn.Module):
@@ -1077,16 +1107,57 @@ def _can_use_fused_pair_bias(
and hasattr(self, "pair_norm")
)
- def _can_use_cueq_pair_bias(
- self, z: Tensor, n_queries: int, beta: Tensor | float
- ) -> bool:
- return (
- _cueq_active(self)
- and n_queries > 750
- and z.dim() == 4
- and self._is_zero_beta(beta)
- and hasattr(self, "pair_bias_proj")
+ def _pair_bias_cache_slot(
+ self,
+ inference_cache: dict[str, Any] | None,
+ z: Tensor,
+ attention_mask: Tensor | None,
+ bsz: int,
+ num_diffusion_samples: int,
+ ) -> dict[str, Any] | None:
+ """This block's slot in the per-fold pair-bias cache, or ``None``.
+
+ The pair bias is a projection of ``z``. Within one fold ``z``, the mask
+ and the projection weights are all fixed, so every diffusion step of
+ every block recomputes the identical tensor. The slot lives in the
+ caller's per-fold ``inference_cache``, keyed by this module instance, so
+ two blocks never share and nothing survives ``fold()``.
+
+ The slot is invalidated unless ``z`` and the mask are the *same tensor
+ objects* as last time. In-place mutation of a cached ``z`` would go
+ unnoticed; nothing in the diffusion path does that.
+ """
+ if inference_cache is None:
+ return None
+ if not (hasattr(self, "pair_norm") and hasattr(self, "pair_bias_proj")):
+ return None
+ if torch.is_grad_enabled() or _any_requires_grad(z, attention_mask):
+ warnings.warn(
+ "ESMFold2: per-fold pair-bias cache not engaged because autograd "
+ "is active; the bias is recomputed on every diffusion step.",
+ stacklevel=3,
+ )
+ return None
+ slot: dict[str, Any] = inference_cache.setdefault("pair_bias", {}).setdefault(
+ self, {}
)
+ key = (tuple(z.shape), z.dtype, bsz, num_diffusion_samples)
+ if (
+ slot.get("z") is not z
+ or slot.get("mask") is not attention_mask
+ or slot.get("key") != key
+ ):
+ slot.clear()
+ slot["z"] = z
+ slot["mask"] = attention_mask
+ slot["key"] = key
+ return slot
+
+ @staticmethod
+ def _expand_for_samples(t: Tensor, bsz: int, num_diffusion_samples: int) -> Tensor:
+ if t.shape[0] != bsz and num_diffusion_samples > 1:
+ return t.repeat_interleave(num_diffusion_samples, dim=0)
+ return t
def forward(
self,
@@ -1096,6 +1167,7 @@ def forward(
beta: Tensor | float = 0.0,
attention_mask: Tensor | None = None,
num_diffusion_samples: int = 1,
+ inference_cache: dict[str, Any] | None = None,
) -> Tensor:
bsz, n_queries, d_model = a.shape
@@ -1111,9 +1183,13 @@ def forward(
k = k.view(bsz, n_keys, self.num_heads, self.head_dim)
v = v.view(bsz, n_keys, self.num_heads, self.head_dim)
- # Expand z for num_diffusion_samples
- if z.dim() == 4 and z.shape[0] != bsz and num_diffusion_samples > 1:
- z = z.repeat_interleave(num_diffusion_samples, dim=0)
+ # Keyed on the pre-expansion z and mask. z is expanded per diffusion
+ # sample at the two bias sites below rather than here, so a cache hit
+ # skips the [B*S, Q, K, C] copy as well as the projection.
+ bias_slot = self._pair_bias_cache_slot(
+ inference_cache, z, attention_mask, bsz, num_diffusion_samples
+ )
+
if (
attention_mask is not None
and attention_mask.shape[0] != bsz
@@ -1124,31 +1200,44 @@ def forward(
)
if self._can_use_fused_pair_bias(z, n_queries, beta):
- kernel_mask = (
- attention_mask
- if attention_mask is not None
- else torch.ones(bsz, n_queries, device=a.device, dtype=torch.bool)
- )
- pair_norm_w = self.pair_norm.weight
- pair_norm_b = (
- self.pair_norm.bias
- if self.pair_norm.bias is not None
- else torch.zeros_like(pair_norm_w)
- )
- z_bf = z if z.dtype == torch.bfloat16 else z.to(torch.bfloat16)
- bias = _fused_pair_bias(
- z_bf,
- kernel_mask,
- self.pair_bias_proj.weight,
- num_heads=self.num_heads,
- pair_norm_w=pair_norm_w,
- pair_norm_b=pair_norm_b,
- ) # (B, H, Q, K)
q_bhqd = q.transpose(1, 2)
k_bhqd = k.transpose(1, 2)
v_bhqd = v.transpose(1, 2)
+ attn_bias = None
+ if bias_slot is not None and bias_slot.get("fused_dtype") == q_bhqd.dtype:
+ attn_bias = bias_slot.get("fused")
+ if attn_bias is None:
+ z_exp = self._expand_for_samples(z, bsz, num_diffusion_samples)
+ kernel_mask = (
+ attention_mask
+ if attention_mask is not None
+ else torch.ones(bsz, n_queries, device=a.device, dtype=torch.bool)
+ )
+ pair_norm_w = self.pair_norm.weight
+ pair_norm_b = (
+ self.pair_norm.bias
+ if self.pair_norm.bias is not None
+ else torch.zeros_like(pair_norm_w)
+ )
+ z_bf = (
+ z_exp if z_exp.dtype == torch.bfloat16 else z_exp.to(torch.bfloat16)
+ )
+ bias = _fused_pair_bias(
+ z_bf,
+ kernel_mask,
+ self.pair_bias_proj.weight,
+ num_heads=self.num_heads,
+ pair_norm_w=pair_norm_w,
+ pair_norm_b=pair_norm_b,
+ ) # (B, H, Q, K)
+ # Cache the SDPA-ready tensor: the cast is the same size as the
+ # kernel output, so storing both would double the cost.
+ attn_bias = bias.to(q_bhqd.dtype)
+ if bias_slot is not None:
+ bias_slot["fused"] = attn_bias
+ bias_slot["fused_dtype"] = q_bhqd.dtype
attn_out = F.scaled_dot_product_attention(
- q_bhqd, k_bhqd, v_bhqd, attn_mask=bias.to(q_bhqd.dtype)
+ q_bhqd, k_bhqd, v_bhqd, attn_mask=attn_bias
)
g = torch.sigmoid(self.g_proj(x)).view(
bsz, n_queries, self.num_heads, self.head_dim
@@ -1159,55 +1248,36 @@ def forward(
out = torch.sigmoid(self.out_gate(s)) * out
return out
- if self._can_use_cueq_pair_bias(z, n_queries, beta):
- kernel_mask = (
- attention_mask
- if attention_mask is not None
- else torch.ones(bsz, n_queries, device=a.device, dtype=torch.bool)
- )
- out, _ = _cue_attn_pair_bias(
- s=x,
- q=q.transpose(1, 2),
- k=k.transpose(1, 2),
- v=v.transpose(1, 2),
- z=z,
- mask=kernel_mask,
- num_heads=self.num_heads,
- w_proj_z=self.pair_bias_proj.weight,
- w_proj_g=self.g_proj.weight,
- w_proj_o=self.out_proj.weight,
- w_ln_z=self.pair_norm.weight,
- b_ln_z=self.pair_norm.bias,
- return_z_proj=False,
- is_cached_z_proj=False,
- )
- else:
- # Standard attention with pair bias
- g = torch.sigmoid(self.g_proj(x)).view(
- bsz, n_queries, self.num_heads, self.head_dim
- )
-
- logits = (
- torch.einsum("... i h d, ... j h d -> ... i j h", q, k) * self.scale
- )
+ # Standard attention with pair bias
+ g = torch.sigmoid(self.g_proj(x)).view(
+ bsz, n_queries, self.num_heads, self.head_dim
+ )
- if z.dim() == 4:
- pair_bias = self.pair_bias_proj(self.pair_norm(z))
- else:
- pair_bias = z.unsqueeze(-1)
- logits = logits + pair_bias.to(dtype=logits.dtype)
+ logits = torch.einsum("... i h d, ... j h d -> ... i j h", q, k) * self.scale
- if attention_mask is not None:
- min_val = torch.finfo(logits.dtype).min
- mask_bias = torch.where(
- attention_mask.bool()[:, None, :, None], 0.0, min_val
- )
- logits = logits + mask_bias.to(dtype=logits.dtype)
+ if z.dim() == 4:
+ pair_bias = None if bias_slot is None else bias_slot.get("eager")
+ if pair_bias is None:
+ z_exp = self._expand_for_samples(z, bsz, num_diffusion_samples)
+ pair_bias = self.pair_bias_proj(self.pair_norm(z_exp))
+ if bias_slot is not None:
+ bias_slot["eager"] = pair_bias
+ else:
+ # A [B, Q, K] bias passed straight through; the unsqueeze is a view.
+ pair_bias = z.unsqueeze(-1)
+ logits = logits + pair_bias.to(dtype=logits.dtype)
+
+ if attention_mask is not None:
+ min_val = torch.finfo(logits.dtype).min
+ mask_bias = torch.where(
+ attention_mask.bool()[:, None, :, None], 0.0, min_val
+ )
+ logits = logits + mask_bias.to(dtype=logits.dtype)
- attn = torch.softmax(logits, dim=-2).to(dtype=v.dtype)
- ctx = torch.einsum("... i j h, ... j h d -> ... i h d", attn, v)
- ctx = g * ctx
- out = self.out_proj(ctx.reshape(bsz, n_queries, d_model))
+ attn = torch.softmax(logits, dim=-2).to(dtype=v.dtype)
+ ctx = torch.einsum("... i j h, ... j h d -> ... i h d", attn, v)
+ ctx = g * ctx
+ out = self.out_proj(ctx.reshape(bsz, n_queries, d_model))
if s is not None:
out = torch.sigmoid(self.out_gate(s)) * out
@@ -1317,6 +1387,7 @@ def forward(
attention_mask: Tensor | None = None,
num_diffusion_samples: int = 1,
return_intermediates: bool = False,
+ inference_cache: dict[str, Any] | None = None,
) -> tuple[Tensor, list[Tensor]]:
intermediates: list[Tensor] = []
x = a
@@ -1328,6 +1399,7 @@ def forward(
beta,
attention_mask=attention_mask,
num_diffusion_samples=num_diffusion_samples,
+ inference_cache=inference_cache,
)
x = x + transition(x, s)
if return_intermediates:
@@ -1389,7 +1461,7 @@ def forward(
relative_position_encoding: Tensor,
sigma_data: float | None = None,
num_diffusion_samples: int = 1,
- inference_cache: dict[str, Tensor] | None = None,
+ inference_cache: dict[str, Any] | None = None,
) -> tuple[Tensor, Tensor]:
sigma = self.sigma_data if sigma_data is None else float(sigma_data)
base_batch = z_trunk.shape[0]
@@ -1545,7 +1617,7 @@ def forward(
num_diffusion_samples: int = 1,
return_token_repr: bool = False,
return_atom_repr: bool = False,
- inference_cache: dict[str, Tensor] | None = None,
+ inference_cache: dict[str, Any] | None = None,
) -> dict[str, Tensor | None]:
bsz = x_noisy.shape[0]
sigma = self.sigma_data if sigma_data is None else float(sigma_data)
@@ -1598,6 +1670,7 @@ def forward(
beta=0.0,
attention_mask=token_attention_mask,
num_diffusion_samples=num_diffusion_samples,
+ inference_cache=inference_cache,
)
# Step 6: token norm
@@ -1827,7 +1900,7 @@ def sample(
device = s_inputs.device
target_batch = s_inputs.shape[0] * num_diffusion_samples
- inference_cache: dict[str, Tensor] | None = {} if use_inference_cache else None
+ inference_cache: dict[str, Any] | None = {} if use_inference_cache else None
steps = (
self.inference_num_steps
@@ -2223,7 +2296,6 @@ def compute_lm_hidden_states(
residue_index: Tensor,
mol_type: Tensor,
token_mask: Tensor,
- pad_to_multiple: int | None = None,
lm_mask_pct: float = 0.0,
mask_token_id: int = 32,
) -> Tensor:
@@ -2297,11 +2369,8 @@ def compute_lm_hidden_states(
expand_map[prot_pos_b] = per_token_lm_pos[inverse_ordered]
expand_maps.append(expand_map)
- # Pad to longest LM input; round to ``pad_to_multiple`` when fp8 is on
- # (TE fp8 kernels assert prod(shape[:-1]) % 8 == 0).
+ # Pad to longest LM input.
max_len = max(lm_lengths)
- if pad_to_multiple is not None and pad_to_multiple > 1:
- max_len = ((max_len + pad_to_multiple - 1) // pad_to_multiple) * pad_to_multiple
lm_input_ids = torch.full(
(B, max_len),
1,
@@ -2372,7 +2441,6 @@ def __init__(self, input_channels: int, latent_channels: int, flow: str) -> None
)
self.proj_gate = nn.Linear(self.input_channels, self.input_channels, bias=False)
- self._use_kernels: bool = False
# Default chunked for memory on long sequences; tests override with
# ``set_chunk_size(None)`` for the unchunked path under bit-exact bf16
# parity checks.
@@ -2387,9 +2455,6 @@ def split_kernel_weights(self) -> tuple[Tensor, Tensor]:
self.proj_bundle.weight[2 * self.latent_channels :, :],
)
- def _kernel_flow_direction(self) -> str:
- return self.flow
-
def _triangular_contract(self, left_stream: Tensor, right_stream: Tensor) -> Tensor:
return torch.einsum(self._einsum_equation, left_stream, right_stream)
@@ -2416,37 +2481,6 @@ def forward(self, pair_grid: Tensor, visibility: Tensor | None = None) -> Tensor
if visibility is None:
visibility = pair_grid.new_ones(pair_grid.shape[:-1])
- if self._use_kernels:
- p_in_weight, g_in_weight = self.split_kernel_weights()
-
- try:
- return _cue_tri_mul(
- pair_grid,
- direction=self._kernel_flow_direction(),
- mask=visibility,
- norm_in_weight=self.norm_start.weight,
- norm_in_bias=self.norm_start.bias,
- p_in_weight=p_in_weight,
- g_in_weight=g_in_weight,
- norm_out_weight=self.norm_mix.weight,
- norm_out_bias=self.norm_mix.bias,
- p_out_weight=self.proj_emit.weight,
- g_out_weight=self.proj_gate.weight,
- eps=_EPS,
- )
- except Exception as e:
- import logging as _logging
-
- _logging.getLogger(__name__).warning(
- "cuequivariance triangle_multiplicative_update kernel failed "
- "(flow=%s, shape=%s, dtype=%s); falling back to chunked einsum. "
- "Error: %s",
- self.flow,
- tuple(pair_grid.shape),
- pair_grid.dtype,
- e,
- )
-
normalized_grid = self.norm_start(pair_grid)
bundled = self.proj_bundle(normalized_grid)
signal, gate_logits = bundled.split(2 * self.latent_channels, dim=-1)
@@ -2471,22 +2505,88 @@ class TriangleMultiplicativeUpdate(nn.Module):
def __init__(self, dim: int = 128, _outgoing: bool = True) -> None:
super().__init__()
flow = "outgoing" if _outgoing else "incoming"
+ self._flow = flow
self._engine = TriangleMultiplicativeBlock(
input_channels=dim, latent_channels=dim, flow=flow
)
+ self._kernel_backend: str | None = None
def set_kernel_backend(self, backend: str | None) -> None:
- # Engine uses cueq when backend=="cuequivariance"; the "fused" backend
- # routes through the parent PairUpdateBlock's fused path (bypassing this).
- self._engine._use_kernels = backend == BACKEND_CUEQ
- if backend == BACKEND_CUEQ and not CUE_AVAILABLE:
- raise RuntimeError(
- "backend='cuequivariance' but cuequivariance_torch is not installed."
+ if backend not in _VALID_BACKENDS:
+ raise ValueError(
+ f"backend must be one of {_VALID_BACKENDS}, got {backend!r}"
)
+ self._kernel_backend = backend
def set_chunk_size(self, chunk_size: int | None) -> None:
self._engine.set_chunk_size(chunk_size)
+ def fast_backend_requested(self) -> bool:
+ """``set_kernel_backend("fast")`` was called, regardless of runtime fit."""
+ return self._kernel_backend == BACKEND_FAST
+
+ def fused_residual_ready(self, pair: Tensor) -> bool:
+ """Trunk gate: fused preconditions including dtype."""
+ return (
+ _fused_pair_stack_active(self, pair)
+ and _fused_trimul_with_residual is not None
+ and pair.dtype == _FUSED_KERNEL_DTYPE
+ )
+
+ def fast_residual_available(self, pair: Tensor) -> bool:
+ """MSA-encoder gate: "fast" only, dtype excluded so the caller can cast."""
+ return (
+ _fast_pair_stack_active(self, pair)
+ and _fused_trimul_with_residual is not None
+ )
+
+ def fused_residual_update(
+ self,
+ pair: Tensor,
+ *,
+ residual: Tensor,
+ drop_mask: Tensor | None,
+ mask: Tensor | None,
+ ) -> Tensor:
+ """``residual + drop_mask * TriMul(pair)`` through the fused Triton kernel.
+
+ The single call site for the kernel. The trunk's ``PairUpdateBlock`` and
+ the MSA encoder's ``MSAEncoderBlock`` both route here, so the weight
+ marshalling lives in one place. Callers check their own gate first.
+
+ Parameters
+ ----------
+ pair
+ ``(B, L, L, c_z)`` bf16 CUDA pair tensor.
+ residual
+ Added in-kernel at the final GEMM; pass ``pair`` for a plain residual.
+ drop_mask
+ Row-shared dropout mask, or ``None`` for inference.
+ mask
+ ``(B, L, L)`` pair mask, or ``None``. Pass fp32: a bool mask also
+ runs, but a second dtype means a second Triton specialization with
+ its own autotune result, so the two are not guaranteed bit-equal.
+ """
+ assert _fused_trimul_with_residual is not None
+ engine = self._engine
+ p_in_weight, g_in_weight = engine.split_kernel_weights()
+ return _fused_trimul_with_residual(
+ pair,
+ self._flow,
+ residual=residual,
+ drop_mask=drop_mask,
+ norm_in_weight=_to_fused_kernel_dtype(engine.norm_start.weight),
+ norm_in_bias=_to_fused_kernel_dtype(engine.norm_start.bias),
+ p_in_weight=_to_fused_kernel_dtype(p_in_weight),
+ g_in_weight=_to_fused_kernel_dtype(g_in_weight),
+ norm_out_weight=_to_fused_kernel_dtype(engine.norm_mix.weight),
+ norm_out_bias=_to_fused_kernel_dtype(engine.norm_mix.bias),
+ p_out_weight=_to_fused_kernel_dtype(engine.proj_emit.weight),
+ g_out_weight=_to_fused_kernel_dtype(engine.proj_gate.weight),
+ mask=mask,
+ eps=_EPS,
+ )
+
def forward(self, z: Tensor, mask: Tensor | None = None) -> Tensor:
return self._engine(z, visibility=mask)
@@ -2512,13 +2612,13 @@ def set_chunk_size(self, chunk_size: int | None) -> None:
self._chunk_size = chunk_size
def set_kernel_backend(self, backend: str | None) -> None:
- """Install / uninstall FusedLNLinearSwiGLU (no cueq equivalent)."""
+ """Install / uninstall FusedLNLinearSwiGLU."""
if backend not in _VALID_BACKENDS:
raise ValueError(
f"backend must be one of {_VALID_BACKENDS}, got {backend!r}"
)
self._kernel_backend = backend
- if backend == BACKEND_FUSED and TRITON_KERNELS_AVAILABLE:
+ if backend in _KERNEL_BACKENDS and TRITON_KERNELS_AVAILABLE:
assert _FusedLNLinearSwiGLU is not None
d_model = self.norm.normalized_shape[0]
d_inner = self.ffn.hidden_features
@@ -2617,7 +2717,7 @@ def set_kernel_backend(self, backend: str | None) -> None:
self.pair_transition.set_kernel_backend(backend)
self._kernel_backend = backend
self.row_drop = DropoutResidual(
- 0.0, batch_dim=1, use_fused_kernels=(backend == BACKEND_FUSED)
+ 0.0, batch_dim=1, use_fused_kernels=(backend in _KERNEL_BACKENDS)
)
def set_chunk_size(self, chunk_size: int | None) -> None:
@@ -2626,35 +2726,18 @@ def set_chunk_size(self, chunk_size: int | None) -> None:
self.pair_transition.set_chunk_size(chunk_size)
def _can_use_fused_trimul_with_residual(self, pair: Tensor) -> bool:
- return (
- _fused_pair_stack_active(self, pair)
- and pair.dtype == _FUSED_KERNEL_DTYPE
- and _fused_trimul_with_residual is not None
- )
+ return self.tri_mul_out.fused_residual_ready(pair)
def _fused_trimul_with_residual(
self, pair: Tensor, direction: str, pair_attention_mask: Tensor | None
) -> Tensor:
"""Fused TriMul+residual call; weights from the corresponding engine."""
tri = self.tri_mul_out if direction == "outgoing" else self.tri_mul_in
- engine: TriangleMultiplicativeBlock = tri._engine
- p_in_weight, g_in_weight = engine.split_kernel_weights()
-
- return _fused_trimul_with_residual(
+ return tri.fused_residual_update(
pair,
- direction,
residual=pair,
drop_mask=None, # inference: no dropout, matches internal's eval path
- norm_in_weight=_to_fused_kernel_dtype(engine.norm_start.weight),
- norm_in_bias=_to_fused_kernel_dtype(engine.norm_start.bias),
- p_in_weight=_to_fused_kernel_dtype(p_in_weight),
- g_in_weight=_to_fused_kernel_dtype(g_in_weight),
- norm_out_weight=_to_fused_kernel_dtype(engine.norm_mix.weight),
- norm_out_bias=_to_fused_kernel_dtype(engine.norm_mix.bias),
- p_out_weight=_to_fused_kernel_dtype(engine.proj_emit.weight),
- g_out_weight=_to_fused_kernel_dtype(engine.proj_gate.weight),
mask=pair_attention_mask,
- eps=_EPS,
)
def forward(
@@ -2704,7 +2787,7 @@ def forward(
orig_dtype = pair.dtype
fused_on = (
len(self.blocks) > 0
- and getattr(self.blocks[0], "_kernel_backend", None) == BACKEND_FUSED
+ and getattr(self.blocks[0], "_kernel_backend", None) in _KERNEL_BACKENDS
)
if pair.is_cuda and fused_on and orig_dtype != torch.bfloat16:
pair = pair.to(torch.bfloat16)
diff --git a/esm/models/esmfold2/model.py b/esm/models/esmfold2/model.py
index b0051f06..09d9ddab 100644
--- a/esm/models/esmfold2/model.py
+++ b/esm/models/esmfold2/model.py
@@ -12,29 +12,18 @@
"""
import math
-from contextlib import contextmanager
-from typing import Any, cast
+from typing import cast
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
-try:
- import transformer_engine.pytorch as te
- from transformer_engine.common.recipe import DelayedScaling, Format
-
- TE_AVAILABLE = True
-except ImportError:
- te: Any = None
- DelayedScaling: Any = None
- Format: Any = None
- TE_AVAILABLE = False
-
from esm.models.esmc import EsmcModel
from esm.models.esmc.checkpoint_layout import published_to_native_subtree
from esm.models.esmfold2.config import EsmFold2Config, default_module_flags
from esm.models.esmfold2.constants import MOL_TYPE_NONPOLYMER
+from esm.models.esmfold2.fp8 import convert_esmc_to_fp8_
from esm.models.esmfold2.hf_checkpoint import hf_state_dict_to_native, is_hf_layout
from esm.models.esmfold2.layers import (
CHAR_VOCAB_SIZE,
@@ -52,11 +41,14 @@
TriangleMultiplicativeUpdate,
_categorical_mean,
_compute_intra_token_idx,
+ _to_fused_kernel_dtype,
compute_lm_hidden_states,
+ fast_trimul_refusal_reason,
gather_rep_atom_coords,
gather_token_to_atom,
maybe_apply_msa_column_masking,
maybe_subsample_msa,
+ warn_fast_trimul_refused,
)
from esm.models.hub import HubPreTrainedModel, resolve_model_dir
@@ -407,136 +399,6 @@ def _inverse_softplus(value: float) -> float:
return value + math.log(-math.expm1(-value))
-def _convert_te_modules_to_fp8_inplace(module: nn.Module) -> None:
- """Re-init each TE module via quantized_model_init so weights live as fp8.
-
- Must be called inside torch.no_grad(); covers nn.Linear, te.Linear,
- te.LayerNormLinear, te.LayerNormMLP — the last two hold 99% of ESMC weight.
- """
- if not TE_AVAILABLE:
- raise RuntimeError("transformer_engine is not available; cannot use fp8.")
- from transformer_engine.pytorch import quantized_model_init
-
- def _walk(mod: nn.Module) -> None:
- for name, child in list(mod.named_children()):
- replaced = False
- if isinstance(child, nn.Linear):
- in_f, out_f = child.in_features, child.out_features
- has_bias = child.bias is not None
- device = child.weight.device
- dtype = child.weight.dtype
- w = child.weight.data
- b = child.bias.data if has_bias else None
- setattr(mod, name, nn.Identity())
- del child
- torch.cuda.empty_cache()
- with quantized_model_init(enabled=True):
- new_mod = te.Linear(
- in_f, out_f, bias=has_bias, params_dtype=dtype
- ).to(device)
- new_mod.weight.quantize_(w) # ty:ignore[call-non-callable, unresolved-attribute]
- if has_bias:
- assert b is not None
- new_mod.bias.data.copy_(b) # ty:ignore[call-non-callable]
- del w, b
- replaced = True
- elif isinstance(child, te.Linear):
- # te.Linear with bf16 weight → re-init inside quantized_model_init for fp8.
- in_f, out_f = child.in_features, child.out_features
- has_bias = child.bias is not None
- device = child.weight.device
- dtype = (
- child.weight.dtype
- if not hasattr(child.weight, "_data")
- else torch.bfloat16
- )
- state = {k: v.detach().clone() for k, v in child.state_dict().items()}
- setattr(mod, name, nn.Identity())
- del child
- torch.cuda.empty_cache()
- with quantized_model_init(enabled=True):
- new_mod = te.Linear(
- in_f, out_f, bias=has_bias, params_dtype=dtype
- ).to(device) # ty:ignore[no-matching-overload]
- new_mod.load_state_dict(state, strict=False)
- replaced = True
- elif hasattr(te, "LayerNormLinear") and isinstance(
- child, te.LayerNormLinear
- ):
- state = {k: v.detach().clone() for k, v in child.state_dict().items()}
- hidden_size = child.in_features
- out_features = child.out_features
- has_bias = child.use_bias
- device = next(child.parameters()).device
- setattr(mod, name, nn.Identity())
- del child
- torch.cuda.empty_cache()
- with quantized_model_init(enabled=True):
- new_mod = te.LayerNormLinear(
- hidden_size,
- out_features,
- bias=has_bias,
- params_dtype=torch.bfloat16,
- ).to(device)
- new_mod.load_state_dict(state, strict=False)
- replaced = True
- elif hasattr(te, "LayerNormMLP") and isinstance(child, te.LayerNormMLP):
- state = {k: v.detach().clone() for k, v in child.state_dict().items()}
- fc1_weight: Tensor = child.fc1_weight
- hidden_size = int(fc1_weight.shape[1])
- # fc1 packed as (2*ffn_hidden_size, hidden_size) for swiglu.
- ffn_hidden_size = int(fc1_weight.shape[0]) // 2
- has_bias = (
- getattr(child, "fc1_bias", None) is not None
- and child.fc1_bias is not None
- )
- device = fc1_weight.device
- setattr(mod, name, nn.Identity())
- del child
- torch.cuda.empty_cache()
- with quantized_model_init(enabled=True):
- new_mod = te.LayerNormMLP(
- hidden_size=hidden_size,
- ffn_hidden_size=ffn_hidden_size,
- bias=has_bias,
- activation="swiglu",
- params_dtype=torch.bfloat16,
- ).to(device)
- new_mod.load_state_dict(state, strict=False)
- replaced = True
-
- if replaced:
- # Freeze via .eval()+.requires_grad_(False); per-param ops would unwrap Float8Tensor.
- new_mod.eval().requires_grad_(False)
- setattr(mod, name, new_mod)
- torch.cuda.empty_cache()
- else:
- _walk(child)
-
- _walk(module)
- torch.cuda.empty_cache()
-
-
-@contextmanager
-def _lm_precision_context(fp8: bool):
- """bf16 autocast (+ optional TE fp8 autocast) around the LM forward.
-
- te.autocast keeps te.Linear outputs bf16 instead of the fp32 default
- (~425 MB at L=1024 in the hidden-state cache).
- """
- with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
- if fp8 and TE_AVAILABLE:
- fp8_recipe = DelayedScaling(
- fp8_format=Format.HYBRID,
- amax_history_len=1,
- amax_compute_algo="most_recent",
- )
- with te.autocast(enabled=True, recipe=fp8_recipe):
- yield
- else:
- yield
-
-
class EsmFold2Model(HubPreTrainedModel):
"""ESMFold2 — all-atom structure prediction with an ESMC PLM backbone.
@@ -566,8 +428,9 @@ class EsmFold2Model(HubPreTrainedModel):
* ``model.set_chunk_size(int|None)``: caps L² ops (triangle / OPM /
pair transition) at this token-axis chunk. Default 64 — fits
L≈2k on an 80 GB GPU. Pass ``None`` for faster inference at L<600.
- * ``model.set_kernel_backend(None | "fused" | "cuequivariance")``:
- select kernel backend (None = reference path).
+ * ``model.set_kernel_backend(None | "fused" | "fast")``:
+ select kernel backend (None = reference path, "fused" = bit-exact
+ Triton kernels, "fast" = those plus approximate ones).
"""
config_class = EsmFold2Config
@@ -617,9 +480,8 @@ def __init__(self, config: EsmFold2Config) -> None:
self.esmc: nn.Module | None = (
EsmcModel(config.esmc_config) if config.esmc_config is not None else None
)
- self._esmc_fp8: bool = False # set by load_esmc(fp8=True)
# When True, ESM-C is offloaded to CPU after its one-shot hidden-state
- # computation (restored on the next call), freeing its ~12 GB for the
+ # computation (restored on the next call), freeing its ~6 GB (fp8) for the
# recycling trunk + diffusion — the freed blocks stay in the caching pool
# so the trunk reuses them without cudaMalloc, relieving the allocator
# pressure that drove rank desync / spin-wait. Opt-in (the CP entry point
@@ -679,7 +541,7 @@ def __init__(self, config: EsmFold2Config) -> None:
self.post_init()
- def load_esmc(self, esmc_model_path: str, precision: str = "bf16") -> None:
+ def load_esmc(self, esmc_model_path: str, precision: str = "fp8") -> None:
"""Fetch the ESMC LM from a separate repo and attach it.
Only needed when the backbone is *not* bundled into this checkpoint; see
@@ -694,17 +556,17 @@ def load_esmc(self, esmc_model_path: str, precision: str = "bf16") -> None:
)
self.set_esmc_precision(precision)
- def set_esmc_precision(self, precision: str = "bf16") -> None:
- """Cast the attached ESMC LM and freeze it.
+ def set_esmc_precision(self, precision: str = "fp8") -> None:
+ """Cast the attached ESMC LM in place and freeze it.
- ``precision``: ``"bf16"`` (default), ``"fp32"``, or ``"fp8"``.
- ``"fp8"`` requires H100 + TransformerEngine ≥ 2.x and quantizes
- every TE module's weights to fp8 storage.
+ ``precision``: ``"fp8"`` (default; ESMFold2 was trained on an fp8 LM),
+ ``"bf16"`` or ``"fp32"``. ``"fp8"`` stores every projection weight in fp8.
+ It runs fp8 GEMMs on compute capability >= 8.9 and upcasts elsewhere.
"""
dtype_map = {
"bf16": torch.bfloat16,
"fp32": torch.float32,
- "fp8": torch.bfloat16, # underlying weights stay bf16, TE re-quantizes to fp8
+ "fp8": torch.bfloat16, # LayerNorms and embeddings stay bf16
}
if precision not in dtype_map:
raise ValueError(
@@ -713,20 +575,13 @@ def set_esmc_precision(self, precision: str = "bf16") -> None:
if self.esmc is None:
raise RuntimeError("no ESMC LM is attached; nothing to cast.")
- esmc = self.esmc.to(device=self.device, dtype=dtype_map[precision]).eval()
+ # No device move: converting before the move keeps the fp32 checkpoint off the GPU.
+ esmc = self.esmc.to(dtype=dtype_map[precision]).eval()
for p in esmc.parameters():
p.requires_grad_(False)
if precision == "fp8":
- if not TE_AVAILABLE:
- raise RuntimeError(
- "transformer_engine is not available; cannot use fp8."
- )
- with torch.no_grad():
- _convert_te_modules_to_fp8_inplace(esmc)
- self._esmc_fp8 = True
- else:
- self._esmc_fp8 = False
+ convert_esmc_to_fp8_(esmc)
self.esmc = esmc
@@ -736,7 +591,7 @@ def from_pretrained(
pretrained_model_name_or_path,
*,
load_esmc: bool = True,
- esmc_precision: str = "bf16",
+ esmc_precision: str = "fp8",
config: EsmFold2Config | None = None,
device: torch.device | str = "cpu",
dtype: torch.dtype | None = None,
@@ -758,22 +613,26 @@ def from_pretrained(
device=device,
dtype=dtype,
)
- model = cls._load_pretrained(local_dir, config, device=device, dtype=dtype)
+ # Built on the CPU and cast there: the bundled ESMC is stored in fp32, and
+ # moving it first put the whole fp32 copy on the GPU.
+ model = cls._load_pretrained(local_dir, config, device="cpu", dtype=dtype)
# A bundled backbone came in with the trunk; only the separate-repo
# arrangement needs a second fetch.
if load_esmc and config.esmc_config is None:
model.load_esmc(model.config.esmc_id, precision=esmc_precision)
elif config.esmc_config is not None:
model.set_esmc_precision(esmc_precision)
- return model
+ return model.to(device)
def set_kernel_backend(self, backend: str | None) -> None:
"""Select kernel backend.
Args:
backend: ``None`` (reference path), ``"fused"`` (vendored Triton
- kernels), or ``"cuequivariance"`` (cuequivariance kernels
- where applicable; vanilla python fallback otherwise).
+ kernels, bit-exact against the reference path), or ``"fast"``
+ (those plus approximate kernels — currently the MSA encoder's
+ triangle multiplication, which is bf16 where the eager engine
+ upcasts to fp32).
"""
self.folding_trunk.set_kernel_backend(backend)
if self.lm_encoder is not None:
@@ -794,12 +653,12 @@ def apply_torch_compile(
"""
import torch._dynamo
- from esm.models.esmfold2.layers import BACKEND_FUSED
+ from esm.models.esmfold2.layers import _KERNEL_BACKENDS
stacked = [
name
for name, module in self.named_modules()
- if getattr(module, "_kernel_backend", None) == BACKEND_FUSED
+ if getattr(module, "_kernel_backend", None) in _KERNEL_BACKENDS
]
if stacked:
raise RuntimeError(
@@ -861,9 +720,7 @@ def _compute_lm_hidden_states(
# previous fold (see self._offload_esmc). No-op when already resident.
if self._offload_esmc:
self.esmc.to(self.device)
- # fp8 TE kernels require prod(shape[:-1]) % 8 == 0.
- pad_to = 8 if self._esmc_fp8 else None
- with _lm_precision_context(self._esmc_fp8):
+ with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
return compute_lm_hidden_states(
self.esmc,
input_ids,
@@ -871,7 +728,6 @@ def _compute_lm_hidden_states(
residue_index,
mol_type,
tok_mask,
- pad_to_multiple=pad_to,
lm_mask_pct=lm_mask_pct,
)
@@ -1197,7 +1053,7 @@ def forward(
else:
lm_z = self.language_model(lm_hidden_states.detach())
del lm_hidden_states
- # ESM-C is done for this forward — offload to free its ~12 GB for the
+ # ESM-C is done for this forward — offload to free its ~6 GB (fp8) for the
# trunk/diffusion. Return the freed blocks to the driver (not just the
# caching pool) so cuSOLVER's out-of-pool cudaMalloc — for its handle in
# the diffusion rigid-align SVD — can succeed near the OOM boundary.
@@ -1462,6 +1318,34 @@ def set_chunk_size(self, chunk_size: int | None) -> None:
self.msa_transition.set_chunk_size(chunk_size)
self.pair_transition.set_chunk_size(chunk_size)
+ def _triangle_updates(self, pair: Tensor, pair_attention_mask: Tensor) -> Tensor:
+ """Both triangle updates, fused under "fast", eager otherwise.
+
+ The eager engine upcasts the L^2 x c_z contraction to fp32, which makes
+ this the same operation as the trunk's at ~9x the cost. The kernel is
+ bf16, so this is not bitwise and sits behind "fast", not "fused".
+ Cast around the pair of updates the way ``FoldingTrunk`` casts around
+ the whole trunk; the kernel's own residual keeps the intermediate bf16.
+ """
+ if not self.tri_mul_out.fast_residual_available(pair):
+ if self.tri_mul_out.fast_backend_requested():
+ warn_fast_trimul_refused(
+ "the MSA encoder", fast_trimul_refusal_reason(pair)
+ )
+ pair = pair + self.tri_mul_out(pair, mask=pair_attention_mask)
+ return pair + self.tri_mul_in(pair, mask=pair_attention_mask)
+
+ orig_dtype = pair.dtype
+ fused_pair = _to_fused_kernel_dtype(pair)
+ # Bool runs too, but fp32 matches the trunk's mask, so the kernel keeps
+ # one Triton specialization instead of two.
+ mask = pair_attention_mask.to(torch.float32)
+ for tri in (self.tri_mul_out, self.tri_mul_in):
+ fused_pair = tri.fused_residual_update(
+ fused_pair, residual=fused_pair, drop_mask=None, mask=mask
+ )
+ return fused_pair.to(orig_dtype)
+
def forward(
self,
m: Tensor,
@@ -1473,8 +1357,7 @@ def forward(
if not self.is_final_block:
m = m + self.msa_pair_weighted_averaging(m, pair, pair_attention_mask)
m = m + self.msa_transition(m)
- pair = pair + self.tri_mul_out(pair, mask=pair_attention_mask)
- pair = pair + self.tri_mul_in(pair, mask=pair_attention_mask)
+ pair = self._triangle_updates(pair, pair_attention_mask)
pair = pair + self.pair_transition(pair)
return m, pair
diff --git a/esm/models/esmfold2/prepare_input.py b/esm/models/esmfold2/prepare_input.py
index b75fbc16..59849b27 100644
--- a/esm/models/esmfold2/prepare_input.py
+++ b/esm/models/esmfold2/prepare_input.py
@@ -1496,6 +1496,14 @@ def prepare_esmfold2_input(
all tensors for the model forward pass, and chain_infos contains
metadata for output processing.
"""
+ # Rejected here rather than dropped in ``build_feature_tensors``: no
+ # released checkpoint consumes pocket features, so a fold with ``pocket``
+ # set would otherwise run unconditioned with no signal to the caller.
+ if input.pocket is not None:
+ raise NotImplementedError(
+ "pocket conditioning is not implemented for ESMFold2; fold without it."
+ )
+
chains, tokens, atoms = build_chains_from_input(input, seed)
features = build_feature_tensors(chains, tokens, atoms, input)
return features, chains
diff --git a/esm/models/esmfold2/processor.py b/esm/models/esmfold2/processor.py
index 47ff27d0..e6389334 100644
--- a/esm/models/esmfold2/processor.py
+++ b/esm/models/esmfold2/processor.py
@@ -50,9 +50,10 @@ def _lm_dropout_context(model: Any, lm_dropout: float | None):
Applies dropout to ``lm_z`` at inference (``training=True``, fresh mask per
loop) so repeated folds give a diverse ensemble. Release models read
``config.lm_encoder.lm_dropout``, the experimental model a top-level
- ``config.lm_dropout`` — disambiguated by ``config.type``. ``None``/``0`` is a no-op.
+ ``config.lm_dropout`` — disambiguated by ``config.type``. ``None`` is a no-op;
+ ``0`` disables dropout for the call.
"""
- if not lm_dropout:
+ if lm_dropout is None:
yield
return
diff --git a/esm/models/esmfold2/protein_utils.py b/esm/models/esmfold2/protein_utils.py
index 4c0650c0..9c01c98f 100644
--- a/esm/models/esmfold2/protein_utils.py
+++ b/esm/models/esmfold2/protein_utils.py
@@ -1,17 +1,3 @@
-# coding=utf-8
-# Copyright 2026 Biohub. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
"""Self-contained protein featurization for ESMFold2 inference."""
from __future__ import annotations
diff --git a/esm/models/esmfold2/types.py b/esm/models/esmfold2/types.py
index 81be960b..715346d6 100644
--- a/esm/models/esmfold2/types.py
+++ b/esm/models/esmfold2/types.py
@@ -14,6 +14,7 @@
DNAInput,
LigandInput,
Modification,
+ PocketConditioning,
ProteinInput,
RNAInput,
StructurePredictionInput,
@@ -28,6 +29,7 @@
"DNAInput",
"LigandInput",
"DistogramConditioning",
+ "PocketConditioning",
"CovalentBond",
"StructurePredictionInput",
]
diff --git a/esm/models/hub.py b/esm/models/hub.py
index 7452ac0b..a7787d1b 100644
--- a/esm/models/hub.py
+++ b/esm/models/hub.py
@@ -23,16 +23,6 @@
_SAFETENSORS_SINGLE = "model.safetensors"
-def drop_te_extra_state(
- module: nn.Module, state_dict: dict, prefix: str, local_metadata: dict
-) -> None:
- """Strip Transformer Engine's non-tensor ``_extra_state`` entries; TE returns
- these as ``io.BytesIO`` objects the checkpoint writer cannot dtype-infer."""
- for key in list(state_dict):
- if key.endswith("_extra_state"):
- del state_dict[key]
-
-
def read_safetensors_dir(directory: str | os.PathLike) -> dict[str, torch.Tensor]:
"""Read a (possibly sharded) safetensors checkpoint into a single dict."""
directory = Path(directory)
@@ -140,7 +130,7 @@ def _normalize_checkpoint_layout(
Runs before any key accounting, so a family whose in-memory parameters
are packed differently from its published ones (ESMC fuses q/k/v and
- gate/up for Transformer Engine) is still held to the same
+ gate/up) is still held to the same
nothing-dropped guarantee. The default is the identity.
"""
return raw
diff --git a/esm/sdk/api.py b/esm/sdk/api.py
index df1c28e1..68eea750 100644
--- a/esm/sdk/api.py
+++ b/esm/sdk/api.py
@@ -14,6 +14,7 @@
from esm.utils import encoding
from esm.utils.constants.models import ESM3_OPEN_SMALL
from esm.utils.misc import get_chainbreak_boundaries_from_sequence
+from esm.utils.structure.input_builder import MSAInput, StructurePredictionInput
from esm.utils.structure.protein_chain import ProteinChain
from esm.utils.structure.protein_complex import SINGLE_LETTER_CHAIN_IDS, ProteinComplex
from esm.utils.types import FunctionAnnotation, PathOrBuffer
@@ -342,6 +343,8 @@ class GenerationConfig:
tokens as generation conditioning.
only_compute_backbone_rmsd (bool): Only compute the RMSD of the backbone atoms.
Affects the returned crmsd.
+ include_pae (bool): Whether to include Predicted Aligned Error (PAE) matrix in
+ the response. Local inference ignores this flag.
"""
track: str = ""
@@ -358,6 +361,7 @@ class GenerationConfig:
top_p: float = 1.0
condition_on_coordinates_only: bool = True
only_compute_backbone_rmsd: bool = False
+ include_pae: bool = False
def use_entropy_based_unmasking_strategy(self):
"""Use entropy based unmasking strategy during generation."""
@@ -439,6 +443,37 @@ class FoldMaxAccuracyConfig:
include_embeddings: bool = False
+@define
+class FoldMaxAccuracyRequest:
+ """One maximum-accuracy folding task.
+
+ Parameters
+ ----------
+ all_atom_input
+ Molecular complex to fold.
+ config
+ Per-task configuration. When omitted, use the call's shared configuration
+ or the defaults in :class:`FoldMaxAccuracyConfig`.
+ """
+
+ all_atom_input: StructurePredictionInput
+ config: FoldMaxAccuracyConfig | None = None
+
+
+@define
+class FoldRequest:
+ sequence: str
+ config: FoldingConfig | None = None
+ msa: MSAInput = None
+ target_structure: ProteinComplex | None = None
+
+
+@define
+class FoldAllAtomRequest:
+ all_atom_input: StructurePredictionInput
+ config: FoldingConfig | None = None
+
+
## Low Level Endpoint Types
@define
class SamplingTrackConfig:
diff --git a/esm/sdk/base_forge_client.py b/esm/sdk/base_forge_client.py
index 714f7243..35e73789 100644
--- a/esm/sdk/base_forge_client.py
+++ b/esm/sdk/base_forge_client.py
@@ -1,24 +1,49 @@
import asyncio
import time
+import warnings
from abc import ABC, abstractmethod
-from collections.abc import Iterator
-from contextlib import suppress
-from typing import Any, Generic, Literal, TypeVar, overload
+from collections import deque
+from collections.abc import (
+ AsyncGenerator,
+ AsyncIterator,
+ Awaitable,
+ Callable,
+ Generator,
+ Iterable,
+ Iterator,
+)
+from concurrent.futures import ThreadPoolExecutor
+from contextvars import copy_context
+from typing import Any, Generic, Literal, TypeVar, cast, overload
from urllib.parse import urljoin
+import attr
import httpx
+from esm import __version__
from esm.sdk.api import ESMProteinError
+from esm.sdk.batch import BatchError, BatchSubmission, BatchTimeoutError
from esm.sdk.retry import retry_decorator
from esm.utils.decoding import assemble_message
POLL_BACKOFF_FACTOR = 2.0
POLLS_PER_INTERVAL = 3
DEFAULT_POLL_INTERVAL = 2
+TRequest = TypeVar("TRequest", bound=attr.AttrsInstance)
+TResponse = TypeVar("TResponse")
+
+
+def _check_result_options(page_size: int, max_concurrency: int) -> None:
+ if not 1 <= page_size <= 1000:
+ raise ValueError("page_size must be between 1 and 1000")
+ if max_concurrency < 1:
+ raise ValueError("max_concurrency must be positive")
def _check_polling(interval: int | None, ceiling: int | None) -> None:
"""Reject a pair one caller supplied that cannot both be honored."""
+ if any(value is not None and value <= 0 for value in (interval, ceiling)):
+ raise ValueError("Polling intervals must be positive")
if interval is not None and ceiling is not None and ceiling < interval:
raise ValueError(
f"poll_max_interval ({ceiling}) is below poll_interval ({interval}); the "
@@ -49,6 +74,23 @@ def _poll_intervals(initial: float, maximum: float | None = None) -> Iterator[fl
interval = min(interval * POLL_BACKOFF_FACTOR, ceiling)
+def _sdk_version() -> str:
+ return __version__
+
+
+def merge_headers(base: dict[str, str], overrides: dict[str, str]) -> dict[str, str]:
+ """Merge headers so an override replaces a base key of any casing."""
+ overridden = {name.lower() for name in overrides}
+ merged = {k: v for k, v in base.items() if k.lower() not in overridden}
+ merged.update(overrides)
+ return merged
+
+
+def build_biohub_component() -> str:
+ """Build the ``X-Biohub-Component`` sent on Forge/Biohub Platform requests."""
+ return f"biohub/forge-sdk-python/{_sdk_version()}"
+
+
class _BaseForgeInferenceClient:
def __init__(
self,
@@ -67,7 +109,10 @@ def __init__(
self.model = model # Name of the model to run.
self.url = url
self.token = token
- self.headers = {"Authorization": f"Bearer {self.token}"}
+ self.headers = {
+ "Authorization": f"Bearer {self.token}",
+ "X-Biohub-Component": build_biohub_component(),
+ }
self.request_timeout = request_timeout
self.min_retry_wait = min_retry_wait
self.max_retry_wait = max_retry_wait
@@ -93,8 +138,8 @@ def close(self):
self._client.close()
async def aclose(self):
- if self.async_client is not None:
- await self.async_client.aclose()
+ if self._async_client is not None:
+ await self._async_client.aclose()
async def __aenter__(self):
return self
@@ -114,14 +159,17 @@ def prepare_request(
request: dict[str, Any],
potential_sequence_of_concern: bool | None = None,
return_bytes: bool = False,
+ disable_cache: bool = False,
headers: dict[str, str] = {},
) -> tuple[dict[str, Any], dict[str, str]]:
if potential_sequence_of_concern is not None:
request["potential_sequence_of_concern"] = potential_sequence_of_concern
- headers = {**self.headers, **headers}
+ headers = merge_headers(self.headers, headers)
if return_bytes:
headers["return-bytes"] = "true"
+ if disable_cache:
+ headers["X-Disable-Cache"] = "true"
return request, headers
def prepare_data(self, response, endpoint: str) -> dict[str, Any]:
@@ -150,12 +198,17 @@ async def _async_post(
potential_sequence_of_concern: bool | None = None,
params: dict[str, Any] = {},
headers: dict[str, str] = {},
+ disable_cache: bool = False,
return_bytes: bool = False,
timeout: int | None = None,
):
try:
request, headers = self.prepare_request(
- request, potential_sequence_of_concern, return_bytes, headers
+ request,
+ potential_sequence_of_concern,
+ return_bytes,
+ disable_cache,
+ headers,
)
response = await self.async_client.post(
url=urljoin(self.url, f"/api/v1/{endpoint}"),
@@ -189,11 +242,16 @@ def _post(
params: dict[str, Any] = {},
headers: dict[str, str] = {},
return_bytes: bool = False,
+ disable_cache: bool = False,
timeout: int | None = None,
):
try:
request, headers = self.prepare_request(
- request, potential_sequence_of_concern, return_bytes, headers
+ request,
+ potential_sequence_of_concern,
+ return_bytes,
+ disable_cache,
+ headers,
)
response = self.client.post(
url=urljoin(self.url, f"/api/v1/{endpoint}"),
@@ -232,7 +290,7 @@ def __init__(
request_timeout: int | None = 120,
min_retry_wait: int = 1,
max_retry_wait: int = 10,
- max_retry_attempts: int = 5,
+ max_retry_attempts: int = 3,
poll_interval: int | None = None,
poll_max_interval: int | None = None,
transfer_timeout: int | None = 60,
@@ -253,39 +311,88 @@ def __init__(
# S3 result download)
self.transfer_timeout = transfer_timeout
- @retry_decorator
- def submit(self, endpoint: str, payload: list[dict[str, Any]]) -> str:
- response_data = self._post(
- "batch/submit",
- {"endpoint": endpoint, "payload": payload},
- timeout=self.transfer_timeout,
- )
- task_id = response_data.get("task_id")
- if not task_id:
- raise ESMProteinError(
- error_code=500, error_msg="API did not return a valid task_id."
+ def submit(
+ self, endpoint: str, payload: list[dict[str, Any]], disable_cache: bool = False
+ ) -> BatchSubmission:
+ try:
+ response_data = self._post(
+ "batch/submit",
+ {"endpoint": endpoint, "payload": payload},
+ disable_cache=disable_cache,
+ timeout=self.transfer_timeout,
)
- return task_id
+ except ESMProteinError as error:
+ raise BatchError.from_error(error) from error
+ return self._submission_response(endpoint, len(payload), response_data)
- @retry_decorator
- async def async_submit(self, endpoint: str, payload: list[dict[str, Any]]) -> str:
- response_data = await self._async_post(
- "batch/submit",
- {"endpoint": endpoint, "payload": payload},
- timeout=self.transfer_timeout,
- )
- task_id = response_data.get("task_id")
- if not task_id:
- raise ESMProteinError(
- error_code=500, error_msg="API did not return a valid task_id."
+ async def async_submit(
+ self, endpoint: str, payload: list[dict[str, Any]], disable_cache: bool = False
+ ) -> BatchSubmission:
+ try:
+ response_data = await self._async_post(
+ "batch/submit",
+ {"endpoint": endpoint, "payload": payload},
+ disable_cache=disable_cache,
+ timeout=self.transfer_timeout,
)
- return task_id
+ except ESMProteinError as error:
+ raise BatchError.from_error(error) from error
+ return self._submission_response(endpoint, len(payload), response_data)
- def cancel(self, task_id: str) -> dict[str, Any]:
- return self._post("batch/cancel", {"task_id": task_id})
+ @staticmethod
+ def _submission_response(
+ endpoint: str, task_count: int, data: dict[str, Any]
+ ) -> BatchSubmission:
+ batch_id = data.get("batch_id")
+ task_id = data.get("task_id")
+ valid_batch_id = isinstance(batch_id, str) and bool(batch_id)
+ valid_task_id = isinstance(task_id, str) and bool(task_id)
+ total = data.get("total_items")
+ valid_total = type(total) is int and total > 0
+ submission = None
+ if valid_task_id:
+ submission = BatchSubmission(
+ endpoint=endpoint,
+ batch_id=batch_id if valid_batch_id else None,
+ task_ids=(task_id,),
+ )
+ elif valid_batch_id:
+ submission = BatchSubmission(endpoint=endpoint, batch_id=batch_id)
+
+ # Legacy support
+ if (
+ task_count == 1
+ and valid_task_id
+ and "batch_id" not in data
+ and "total_items" not in data
+ ):
+ return cast(BatchSubmission, submission)
- async def async_cancel(self, task_id: str) -> dict[str, Any]:
- return await self._async_post("batch/cancel", {"task_id": task_id})
+ if not valid_batch_id or not valid_total or total != task_count:
+ raise BatchError(
+ 502, f"Unexpected submission response: {data}", submission=submission
+ )
+ return cast(BatchSubmission, submission)
+
+ @staticmethod
+ def _cancel_request(task_id: str | None, batch_id: str | None) -> dict[str, str]:
+ if bool(task_id) == bool(batch_id):
+ raise ValueError("Provide exactly one of task_id or batch_id")
+ if task_id:
+ return {"task_id": task_id}
+ return {"batch_id": cast(str, batch_id)}
+
+ def cancel(
+ self, task_id: str | None = None, *, batch_id: str | None = None
+ ) -> dict[str, Any]:
+ return self._post("batch/cancel", self._cancel_request(task_id, batch_id))
+
+ async def async_cancel(
+ self, task_id: str | None = None, *, batch_id: str | None = None
+ ) -> dict[str, Any]:
+ return await self._async_post(
+ "batch/cancel", self._cancel_request(task_id, batch_id)
+ )
@retry_decorator
def get_status(self, task_id: str) -> dict[str, Any]:
@@ -295,13 +402,68 @@ def get_status(self, task_id: str) -> dict[str, Any]:
async def async_get_status(self, task_id: str) -> dict[str, Any]:
return await self._async_post("batch/status", {"task_id": task_id})
- def wait_for_completion(
+ @staticmethod
+ def _list_batch_request(
+ batch_id: str,
+ status: list[str] | None,
+ include_responses: bool,
+ limit: int | None,
+ cursor: str | None,
+ ) -> dict[str, Any]:
+ request: dict[str, Any] = {
+ "batch_id": batch_id,
+ "include_responses": include_responses,
+ }
+ if status is not None:
+ request["status"] = status
+ if limit is not None:
+ request["limit"] = limit
+ if cursor is not None:
+ request["cursor"] = cursor
+ return request
+
+ @retry_decorator
+ def list_batch(
self,
- task_id: str,
+ batch_id: str,
+ *,
+ status: list[str] | None = None,
+ include_responses: bool = False,
+ limit: int | None = None,
+ cursor: str | None = None,
+ ) -> dict[str, Any]:
+ return self._post(
+ "batch/list",
+ self._list_batch_request(
+ batch_id, status, include_responses, limit, cursor
+ ),
+ )
+
+ @retry_decorator
+ async def async_list_batch(
+ self,
+ batch_id: str,
+ *,
+ status: list[str] | None = None,
+ include_responses: bool = False,
+ limit: int | None = None,
+ cursor: str | None = None,
+ ) -> dict[str, Any]:
+ return await self._async_post(
+ "batch/list",
+ self._list_batch_request(
+ batch_id, status, include_responses, limit, cursor
+ ),
+ )
+
+ def wait(
+ self,
+ submission: BatchSubmission,
timeout: int,
poll_interval: int | None = None,
poll_max_interval: int | None = None,
- ) -> dict:
+ ) -> dict[str, Any] | list[dict[str, Any]]:
+ """Return a terminal batch summary, or task states for a task-only handle."""
_check_polling(poll_interval, poll_max_interval)
deadline = time.monotonic() + timeout
intervals = _poll_intervals(
@@ -311,37 +473,49 @@ def wait_for_completion(
)
)
+ completed: dict[str, dict[str, Any]] = {}
while True:
- response = self.get_status(task_id)
- job_status = response.get("status")
- if job_status == "done":
- return response
- elif job_status == "cancelled":
- raise ESMProteinError(
- error_code=500, error_msg=f"Job {task_id} cancelled."
- )
- elif job_status == "failed":
- raise ESMProteinError(
- error_code=500,
- error_msg=f"Job {task_id} failed with error: '{response.get('error')}'.",
- )
+ if submission.batch_id is not None:
+ try:
+ summary = self.list_batch(submission.batch_id)
+ except ESMProteinError as error:
+ raise BatchError.from_error(error, submission) from error
+ if self._batch_is_complete(summary):
+ return summary
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ break
+ time.sleep(min(next(intervals), remaining))
+ continue
+ polled = False
+ for task_id in submission.task_ids:
+ if task_id in completed:
+ continue
+ if polled and time.monotonic() >= deadline:
+ raise BatchTimeoutError(submission, timeout)
+ try:
+ response = self.get_status(task_id)
+ except ESMProteinError as error:
+ raise BatchError.from_error(error, submission, task_id) from error
+ if response["status"] in {"done", "failed", "cancelled"}:
+ completed[task_id] = response
+ polled = True
+ if len(completed) == len(submission.task_ids):
+ return [completed[task_id] for task_id in submission.task_ids]
remaining = deadline - time.monotonic()
if remaining <= 0:
break
time.sleep(min(next(intervals), remaining))
- raise ESMProteinError(
- error_code=500,
- error_msg=f"Job {task_id} timed out after {timeout} seconds.",
- )
+ raise BatchTimeoutError(submission, timeout)
- async def async_wait_for_completion(
+ async def async_wait(
self,
- task_id: str,
+ submission: BatchSubmission,
timeout: int,
poll_interval: int | None = None,
poll_max_interval: int | None = None,
- ) -> dict:
+ ) -> dict[str, Any] | list[dict[str, Any]]:
_check_polling(poll_interval, poll_max_interval)
deadline = time.monotonic() + timeout
intervals = _poll_intervals(
@@ -351,29 +525,363 @@ async def async_wait_for_completion(
)
)
+ completed: dict[str, dict[str, Any]] = {}
while True:
- response = await self.async_get_status(task_id)
- job_status = response.get("status")
- if job_status == "done":
- return response
- elif job_status == "cancelled":
- raise ESMProteinError(
- error_code=500, error_msg=f"Job {task_id} cancelled."
- )
- elif job_status == "failed":
- raise ESMProteinError(
- error_code=500,
- error_msg=f"Job {task_id} failed with error: '{response.get('error')}'.",
- )
+ if submission.batch_id is not None:
+ try:
+ summary = await self.async_list_batch(submission.batch_id)
+ except ESMProteinError as error:
+ raise BatchError.from_error(error, submission) from error
+ if self._batch_is_complete(summary):
+ return summary
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ break
+ await asyncio.sleep(min(next(intervals), remaining))
+ continue
+ polled = False
+ for task_id in submission.task_ids:
+ if task_id in completed:
+ continue
+ if polled and time.monotonic() >= deadline:
+ raise BatchTimeoutError(submission, timeout)
+ try:
+ response = await self.async_get_status(task_id)
+ except ESMProteinError as error:
+ raise BatchError.from_error(error, submission, task_id) from error
+ if response["status"] in {"done", "failed", "cancelled"}:
+ completed[task_id] = response
+ polled = True
+ if len(completed) == len(submission.task_ids):
+ return [completed[task_id] for task_id in submission.task_ids]
remaining = deadline - time.monotonic()
if remaining <= 0:
break
await asyncio.sleep(min(next(intervals), remaining))
- raise ESMProteinError(
- error_code=500,
- error_msg=f"Job {task_id} timed out after {timeout} seconds.",
+ raise BatchTimeoutError(submission, timeout)
+
+ @staticmethod
+ def _batch_is_complete(summary: dict[str, Any]) -> bool:
+ counts = summary["counts"]
+ if counts.get("queued", 0) + counts.get("in_progress", 0) > 0:
+ return False
+ # Cancelled/failed dispatch can leave tasks uncreated, but running work must finish.
+ if summary["status"] in {"failed", "cancelled"}:
+ return True
+ terminal = sum(
+ counts.get(state, 0) for state in ("done", "failed", "cancelled")
)
+ return summary["status"] == "done" and terminal == summary["total"]
+
+ @staticmethod
+ def _require_done(task_id: str, status: dict[str, Any]) -> None:
+ state = status["status"]
+ if state == "done":
+ return
+ if state == "failed":
+ raise BatchError(
+ 500,
+ f"Job {task_id} failed with error: '{status.get('error')}'.",
+ task_id=task_id,
+ task_status=state,
+ )
+ message = f"Job {task_id} is {state}."
+ if state == "cancelled" and status.get("error"):
+ message += f" {status['error']}"
+ raise BatchError(409, message, task_id=task_id, task_status=state)
+
+ def wait_for_completion(
+ self,
+ task_id: str,
+ timeout: int,
+ poll_interval: int | None = None,
+ poll_max_interval: int | None = None,
+ ) -> dict[str, Any]:
+ submission = BatchSubmission(endpoint="", task_ids=(task_id,))
+ statuses = self.wait(submission, timeout, poll_interval, poll_max_interval)
+ status = cast(list[dict[str, Any]], statuses)[0]
+ self._require_done(task_id, status)
+ return status
+
+ async def async_wait_for_completion(
+ self,
+ task_id: str,
+ timeout: int,
+ poll_interval: int | None = None,
+ poll_max_interval: int | None = None,
+ ) -> dict[str, Any]:
+ submission = BatchSubmission(endpoint="", task_ids=(task_id,))
+ statuses = await self.async_wait(
+ submission, timeout, poll_interval, poll_max_interval
+ )
+ status = cast(list[dict[str, Any]], statuses)[0]
+ self._require_done(task_id, status)
+ return status
+
+ @staticmethod
+ def _check_listing_complete(
+ submission: BatchSubmission, seen: set[str], summary: dict[str, Any]
+ ) -> None:
+ total = summary["total"]
+ if len(seen) != total:
+ raise BatchError(
+ 409,
+ f"Batch {submission.batch_id} lists {len(seen)} of {total} tasks "
+ f"(status: {summary['status']}); not all tasks are available.",
+ submission=submission,
+ )
+
+ @staticmethod
+ def _record_listed_tasks(
+ submission: BatchSubmission, seen: set[str], items: list[dict[str, Any]]
+ ) -> None:
+ for item in items:
+ task_id = item["task_id"]
+ if task_id in seen:
+ raise BatchError(
+ 502, f"Batch listing repeated task {task_id}", submission
+ )
+ seen.add(task_id)
+
+ def _result_pages(
+ self, submission: BatchSubmission, page_size: int
+ ) -> Iterator[list[dict[str, Any]]]:
+ if submission.batch_id is None:
+ for start in range(0, len(submission.task_ids), page_size):
+ yield [
+ {"task_id": task_id}
+ for task_id in submission.task_ids[start : start + page_size]
+ ]
+ return
+ cursor = None
+ cursors: set[str] = set()
+ seen: set[str] = set()
+ while True:
+ try:
+ page = self.list_batch(
+ submission.batch_id,
+ include_responses=True,
+ limit=page_size,
+ cursor=cursor,
+ )
+ except ESMProteinError as error:
+ raise BatchError.from_error(error, submission) from error
+ self._record_listed_tasks(submission, seen, page["items"])
+ yield page["items"]
+ cursor = page["cursor"]
+ if cursor is None:
+ break
+ if cursor in cursors:
+ raise BatchError(
+ 502, "Batch listing repeated a pagination cursor", submission
+ )
+ cursors.add(cursor)
+ self._check_listing_complete(submission, seen, page)
+
+ async def _async_result_pages(
+ self, submission: BatchSubmission, page_size: int
+ ) -> AsyncIterator[list[dict[str, Any]]]:
+ if submission.batch_id is None:
+ for page in self._result_pages(submission, page_size):
+ yield page
+ return
+ cursor = None
+ cursors: set[str] = set()
+ seen: set[str] = set()
+ while True:
+ try:
+ page = await self.async_list_batch(
+ submission.batch_id,
+ include_responses=True,
+ limit=page_size,
+ cursor=cursor,
+ )
+ except ESMProteinError as error:
+ raise BatchError.from_error(error, submission) from error
+ self._record_listed_tasks(submission, seen, page["items"])
+ yield page["items"]
+ cursor = page["cursor"]
+ if cursor is None:
+ break
+ if cursor in cursors:
+ raise BatchError(
+ 502, "Batch listing repeated a pagination cursor", submission
+ )
+ cursors.add(cursor)
+ self._check_listing_complete(submission, seen, page)
+
+ def _retrieve_result(
+ self,
+ item: dict[str, Any],
+ process_response: Callable[[dict[str, Any]], TResponse],
+ submission: BatchSubmission | None = None,
+ ) -> TResponse | BatchError:
+ task_id = item["task_id"]
+ try:
+ if "status" not in item or (
+ item["status"] == "done" and "response" not in item
+ ):
+ item = self.get_status(task_id)
+ self._require_done(task_id, item)
+ return process_response(item)
+ except ESMProteinError as error:
+ return BatchError.from_error(error, submission, task_id)
+ except (KeyError, TypeError, ValueError) as error:
+ return BatchError(
+ 502, f"Cannot decode task {task_id}: {error}", submission, task_id
+ )
+
+ async def _async_retrieve_result(
+ self,
+ item: dict[str, Any],
+ process_response: Callable[[dict[str, Any]], Awaitable[TResponse]],
+ submission: BatchSubmission | None = None,
+ ) -> TResponse | BatchError:
+ task_id = item["task_id"]
+ try:
+ if "status" not in item or (
+ item["status"] == "done" and "response" not in item
+ ):
+ item = await self.async_get_status(task_id)
+ self._require_done(task_id, item)
+ return await process_response(item)
+ except ESMProteinError as error:
+ return BatchError.from_error(error, submission, task_id)
+ except (KeyError, TypeError, ValueError) as error:
+ return BatchError(
+ 502, f"Cannot decode task {task_id}: {error}", submission, task_id
+ )
+
+ def _iter_results(
+ self,
+ submission: BatchSubmission,
+ process_response: Callable[[dict[str, Any]], TResponse],
+ *,
+ page_size: int,
+ max_concurrency: int,
+ ) -> Generator[tuple[str, TResponse | ESMProteinError], None, None]:
+ # Initialize once before worker threads share the HTTP connection pool.
+ _ = self.client
+ with ThreadPoolExecutor(
+ max_workers=min(max_concurrency, page_size)
+ ) as executor:
+ for items in self._result_pages(submission, page_size):
+ # Bound prefetched outputs as well as simultaneous downloads.
+ for start in range(0, len(items), max_concurrency):
+ chunk = items[start : start + max_concurrency]
+ pending = deque(
+ executor.submit(
+ copy_context().run,
+ self._retrieve_result,
+ item,
+ process_response,
+ submission,
+ )
+ for item in chunk
+ )
+ for item in chunk:
+ output = cast(
+ TResponse | ESMProteinError, pending.popleft().result()
+ )
+ yield item["task_id"], output
+ del output
+
+ async def _async_iter_results(
+ self,
+ submission: BatchSubmission,
+ process_response: Callable[[dict[str, Any]], Awaitable[TResponse]],
+ *,
+ page_size: int,
+ max_concurrency: int,
+ ) -> AsyncGenerator[tuple[str, TResponse | ESMProteinError], None]:
+ async for items in self._async_result_pages(submission, page_size):
+ for start in range(0, len(items), max_concurrency):
+ chunk = items[start : start + max_concurrency]
+ results: dict[int, TResponse | ESMProteinError] = {}
+
+ async def retrieve_into(index: int, item: dict[str, Any]) -> None:
+ # Keep outputs outside Task results, which can outlive aclose().
+ results[index] = await self._async_retrieve_result(
+ item, process_response, submission
+ )
+
+ tasks = [
+ asyncio.create_task(retrieve_into(index, item))
+ for index, item in enumerate(chunk)
+ ]
+ try:
+ await asyncio.gather(*tasks)
+ for index, item in enumerate(chunk):
+ yield item["task_id"], results.pop(index)
+ finally:
+ for task in tasks:
+ task.cancel()
+ await asyncio.gather(*tasks, return_exceptions=True)
+ results.clear()
+
+ def _collect_results(
+ self,
+ submission: BatchSubmission,
+ process_response: Callable[[dict[str, Any]], TResponse],
+ *,
+ page_size: int,
+ max_concurrency: int,
+ ) -> list[TResponse | ESMProteinError]:
+ results: list[tuple[str, TResponse | ESMProteinError]] = []
+ try:
+ results.extend(
+ self._iter_results(
+ submission,
+ process_response,
+ page_size=page_size,
+ max_concurrency=max_concurrency,
+ )
+ )
+ except BatchError as error:
+ error.partial_results = results
+ raise
+ return [output for _, output in results]
+
+ async def _async_collect_results(
+ self,
+ submission: BatchSubmission,
+ process_response: Callable[[dict[str, Any]], Awaitable[TResponse]],
+ *,
+ page_size: int,
+ max_concurrency: int,
+ ) -> list[TResponse | ESMProteinError]:
+ results: list[tuple[str, TResponse | ESMProteinError]] = []
+ try:
+ async for result in self._async_iter_results(
+ submission,
+ process_response,
+ page_size=page_size,
+ max_concurrency=max_concurrency,
+ ):
+ results.append(result)
+ except BatchError as error:
+ error.partial_results = results
+ raise
+ return [output for _, output in results]
+
+ def _read_response(
+ self, status: dict[str, Any], *, return_bytes: bool = False
+ ) -> Any:
+ response = status["response"]
+ if isinstance(response, str) and response.startswith(("https://", "http://")):
+ return self.get_result_from_s3(response, return_bytes=return_bytes)
+ return response
+
+ async def _async_read_response(
+ self, status: dict[str, Any], *, return_bytes: bool = False
+ ) -> Any:
+ response = status["response"]
+ if isinstance(response, str) and response.startswith(("https://", "http://")):
+ return await self.async_get_result_from_s3(
+ response, return_bytes=return_bytes
+ )
+ return response
@overload
def get_result_from_s3(
@@ -432,10 +940,8 @@ async def async_get_result_from_s3(
)
-TResponse = TypeVar("TResponse")
-
-
-class EndpointHandler(ABC, Generic[TResponse]):
+class EndpointHandler(ABC, Generic[TRequest, TResponse]):
+ request_type: type[TRequest]
poll_interval: int | None = None
poll_max_interval: int = 30
default_timeout: int = 30 * 60 # 30 minutes
@@ -465,106 +971,476 @@ def endpoint_name(self) -> str:
pass
@abstractmethod
- def _prepare_request(self, **kwargs) -> list[dict[str, Any]]:
+ def _prepare_request(
+ self, request: TRequest, *, model_name: str | None = None
+ ) -> dict[str, Any]:
pass
@abstractmethod
- def _process_response(self, response: dict, **kwargs) -> TResponse:
+ def _process_response(self, response: dict) -> TResponse:
pass
@abstractmethod
- async def _async_process_response(self, response: dict, **kwargs) -> TResponse:
+ async def _async_process_response(self, response: dict) -> TResponse:
pass
+ def _prepare_requests(
+ self,
+ requests: Iterable[TRequest] | None,
+ shared: dict[str, Any],
+ *,
+ model_name: str | None = None,
+ ) -> list[dict[str, Any]]:
+ if requests is None:
+ # attrs rejects unknown and missing fields before any HTTP request.
+ return [
+ self._prepare_request(
+ self.request_type(**shared), model_name=model_name
+ )
+ ]
+ requests = list(requests)
+ if not requests:
+ raise ValueError("requests must contain at least one task")
+ name = self.request_type.__name__
+ fields = attr.fields_dict(self.request_type)
+ for key in shared:
+ if key != "config" or key not in fields:
+ raise TypeError(
+ f"{key}= cannot be shared across requests; set it on each {name}"
+ )
+ for index, request in enumerate(requests):
+ if not isinstance(request, self.request_type):
+ raise TypeError(
+ f"requests[{index}] must be a {name}, got {type(request).__name__}"
+ )
+ return [
+ self._prepare_request(
+ attr.evolve(
+ request,
+ **{k: v for k, v in shared.items() if getattr(request, k) is None},
+ ),
+ model_name=model_name,
+ )
+ for request in requests
+ ]
+
+ def submit(
+ self,
+ requests: Iterable[TRequest] | None = None,
+ *,
+ model_name: str | None = None,
+ disable_cache: bool = False,
+ **kwargs: Any,
+ ) -> BatchSubmission:
+ """Submit one task from request fields as keywords, or an iterable of requests.
+
+ Pass the iterable positionally or as ``requests=``.
+ Iterables, including generators, are materialized before validation.
+ With ``requests``, only ``config`` supplies a default for requests whose
+ config is ``None``. Other inputs must be set on each request.
+ """
+ payloads = self._prepare_requests(requests, kwargs, model_name=model_name)
+ return self._batch_client.submit(
+ self.endpoint_name, payloads, disable_cache=disable_cache
+ )
+
+ async def async_submit(
+ self,
+ requests: Iterable[TRequest] | None = None,
+ *,
+ model_name: str | None = None,
+ disable_cache: bool = False,
+ **kwargs: Any,
+ ) -> BatchSubmission:
+ payloads = self._prepare_requests(requests, kwargs, model_name=model_name)
+ return await self._batch_client.async_submit(
+ self.endpoint_name, payloads, disable_cache=disable_cache
+ )
+
+ def _check_submission(self, submission: BatchSubmission) -> None:
+ if submission.endpoint != self.endpoint_name:
+ raise ValueError(
+ f"Expected a {self.endpoint_name} submission, got {submission.endpoint}"
+ )
+
+ def wait(
+ self,
+ submission: BatchSubmission,
+ timeout: int | None = None,
+ *,
+ poll_interval: int | None = None,
+ poll_max_interval: int | None = None,
+ ) -> dict[str, Any] | list[dict[str, Any]]:
+ """Return a batch summary (task states without a batch ID); never cancel."""
+ self._check_submission(submission)
+ interval, ceiling = self._resolved_polling(poll_interval, poll_max_interval)
+ return self._batch_client.wait(
+ submission,
+ self.default_timeout if timeout is None else timeout,
+ poll_interval=interval,
+ poll_max_interval=ceiling,
+ )
+
+ async def async_wait(
+ self,
+ submission: BatchSubmission,
+ timeout: int | None = None,
+ *,
+ poll_interval: int | None = None,
+ poll_max_interval: int | None = None,
+ ) -> dict[str, Any] | list[dict[str, Any]]:
+ self._check_submission(submission)
+ interval, ceiling = self._resolved_polling(poll_interval, poll_max_interval)
+ return await self._batch_client.async_wait(
+ submission,
+ self.default_timeout if timeout is None else timeout,
+ poll_interval=interval,
+ poll_max_interval=ceiling,
+ )
+
+ def result(self, task_id: str) -> TResponse:
+ """Retrieve a completed result; otherwise raise with its ``task_status``."""
+ output = self._batch_client._retrieve_result(
+ {"task_id": task_id}, self._process_response
+ )
+ if isinstance(output, BatchError):
+ raise output
+ return output
+
+ async def async_result(self, task_id: str) -> TResponse:
+ output = await self._batch_client._async_retrieve_result(
+ {"task_id": task_id}, self._async_process_response
+ )
+ if isinstance(output, BatchError):
+ raise output
+ return output
+
+ def iter_results(
+ self,
+ submission: BatchSubmission,
+ *,
+ page_size: int = 100,
+ max_concurrency: int = 8,
+ ) -> Generator[tuple[str, TResponse | ESMProteinError], None, None]:
+ """Yield authoritative task IDs and decoded outputs in server listing order.
+
+ Parameters
+ ----------
+ submission
+ Handle returned by ``submit()``, containing the endpoint and job IDs.
+ page_size
+ Entries per listing request, between 1 and 1000.
+ max_concurrency
+ Maximum number of prefetched task outputs and simultaneous downloads.
+
+ Notes
+ -----
+ Incomplete dispatch raises ``BatchError`` after available results are yielded.
+ Task-only handles use the supplied task ID order.
+ """
+ self._check_submission(submission)
+ _check_result_options(page_size, max_concurrency)
+ return self._batch_client._iter_results(
+ submission,
+ self._process_response,
+ page_size=page_size,
+ max_concurrency=max_concurrency,
+ )
+
+ def async_iter_results(
+ self,
+ submission: BatchSubmission,
+ *,
+ page_size: int = 100,
+ max_concurrency: int = 8,
+ ) -> AsyncGenerator[tuple[str, TResponse | ESMProteinError], None]:
+ """Asynchronous equivalent of :meth:`iter_results`; use ``async for``."""
+ self._check_submission(submission)
+ _check_result_options(page_size, max_concurrency)
+ return self._batch_client._async_iter_results(
+ submission,
+ self._async_process_response,
+ page_size=page_size,
+ max_concurrency=max_concurrency,
+ )
+
+ def results(
+ self,
+ submission: BatchSubmission,
+ *,
+ page_size: int = 100,
+ max_concurrency: int = 8,
+ ) -> list[TResponse | ESMProteinError]:
+ """Collect decoded results in server listing order, retaining per-task errors.
+
+ Parameters
+ ----------
+ submission
+ Handle returned by ``submit()``, containing the endpoint and batch or task IDs.
+ page_size
+ Entries per listing request, between 1 and 1000. All pages are read;
+ this does not limit the number of returned results or their memory.
+ max_concurrency
+ Maximum number of simultaneous task result retrievals.
+
+ Notes
+ -----
+ Inline results need no status read. Stored outputs use fresh status
+ reads and downloads. Failed, cancelled, and unfinished tasks return
+ errors with ``task_status`` set.
+ Listing failures and incomplete dispatch raise ``BatchError`` with the
+ submission and collected ``(task_id, output)`` pairs in ``partial_results``.
+ Use :meth:`iter_results` to consume outputs without collecting the whole batch.
+ """
+ self._check_submission(submission)
+ _check_result_options(page_size, max_concurrency)
+ return self._batch_client._collect_results(
+ submission,
+ self._process_response,
+ page_size=page_size,
+ max_concurrency=max_concurrency,
+ )
+
+ async def async_results(
+ self,
+ submission: BatchSubmission,
+ *,
+ page_size: int = 100,
+ max_concurrency: int = 8,
+ ) -> list[TResponse | ESMProteinError]:
+ """Asynchronous equivalent of :meth:`results`, with the same limits."""
+ self._check_submission(submission)
+ _check_result_options(page_size, max_concurrency)
+ return await self._batch_client._async_collect_results(
+ submission,
+ self._async_process_response,
+ page_size=page_size,
+ max_concurrency=max_concurrency,
+ )
+
+ @staticmethod
+ def _cancellation_target(submission: BatchSubmission) -> dict[str, str]:
+ if submission.batch_id is not None:
+ return {"batch_id": submission.batch_id}
+ if len(submission.task_ids) == 1:
+ return {"task_id": submission.task_id}
+ raise ValueError("Cancelling multiple tasks requires their batch_id")
+
+ def _cancel_after_error(
+ self, submission: BatchSubmission, error: BatchError
+ ) -> None:
+ try:
+ error.cancellation = self._batch_client.cancel(
+ **self._cancellation_target(submission)
+ )
+ except ESMProteinError as cancellation_error:
+ error.cancellation_error = cancellation_error
+ warnings.warn(
+ f"Could not cancel job: {cancellation_error.error_msg}",
+ RuntimeWarning,
+ stacklevel=3,
+ )
+
+ async def _async_cancel_after_error(
+ self, submission: BatchSubmission, error: BatchError
+ ) -> None:
+ try:
+ error.cancellation = await self._batch_client.async_cancel(
+ **self._cancellation_target(submission)
+ )
+ except ESMProteinError as cancellation_error:
+ error.cancellation_error = cancellation_error
+ warnings.warn(
+ f"Could not cancel job: {cancellation_error.error_msg}",
+ RuntimeWarning,
+ stacklevel=3,
+ )
+
+ @overload
+ def run(
+ self,
+ requests: Iterable[TRequest],
+ *,
+ timeout: int | None = None,
+ model_name: str | None = None,
+ **kwargs: Any,
+ ) -> list[TResponse | ESMProteinError] | ESMProteinError: ...
+
+ @overload
+ def run(
+ self,
+ requests: None = None,
+ *,
+ timeout: int | None = None,
+ model_name: str | None = None,
+ **kwargs: Any,
+ ) -> TResponse | ESMProteinError: ...
+
def run(
self,
+ requests: Iterable[TRequest] | None = None,
+ *,
timeout: int | None = None,
+ model_name: str | None = None,
+ disable_cache: bool = False,
cancel_on_timeout: bool = True,
poll_interval: int | None = None,
poll_max_interval: int | None = None,
+ page_size: int = 100,
+ max_concurrency: int = 8,
**kwargs,
- ) -> TResponse | ESMProteinError:
- """
- Submit and execute a batch job, waiting for completion by polling the status of the job.
- Args:
- timeout: Maximum time to wait for job completion, in seconds.
- poll_interval: Seconds between status polls
- poll_max_interval: Ceiling the backoff may grow `poll_interval` to.
- cancel_on_timeout: If True, cancels the batch job if it times out or is interrupted.
- **kwargs: Arguments to pass to the batch job.
- Returns:
- The response from the batch job or an ESMProteinError if the job fails.
+ ) -> TResponse | list[TResponse | ESMProteinError] | ESMProteinError:
+ """Submit, wait, and decode tasks in server listing order.
+
+ ``requests`` (positional or named) selects list output, even for one item; keywords
+ combine with it as in :meth:`submit`.
+ Invalid SDK arguments raise ``TypeError`` or ``ValueError`` before submission.
+ Server submission and waiting failures retain accepted job IDs in
+ ``BatchError.submission``.
+ Retrieval failures return ``BatchError`` with collected ``(task_id, output)``
+ pairs in ``partial_results``.
+ Timeout or interruption attempts to cancel queued work by default;
+ running tasks can still finish and be billed. Set ``cancel_on_timeout=False``
+ to leave work running and recover through ``BatchError.submission``.
+ Use ``submit()`` then ``wait()`` to save a handle before waiting;
+ ``wait()`` never cancels work. ``page_size`` and ``max_concurrency``
+ configure result retrieval, as on ``results()``.
"""
timeout = self.default_timeout if timeout is None else timeout
- task_id = None
- task_timed_out = False
- keyboard_interrupted = False
+ _check_result_options(page_size, max_concurrency)
+ submission = None
interval, max_interval = self._resolved_polling(
poll_interval, poll_max_interval
)
try:
- request = self._prepare_request(**kwargs)
- task_id = self._batch_client.submit(self.endpoint_name, request)
- response = self._batch_client.wait_for_completion(
- task_id, timeout, poll_interval=interval, poll_max_interval=max_interval
- )
- return self._process_response(response, **kwargs)
- except KeyboardInterrupt:
- keyboard_interrupted = True
+ submission = self.submit(
+ requests=requests,
+ disable_cache=disable_cache,
+ model_name=model_name,
+ **kwargs,
+ )
+ self.wait(
+ submission,
+ timeout,
+ poll_interval=interval,
+ poll_max_interval=max_interval,
+ )
+ outputs = self._batch_client._collect_results(
+ submission,
+ self._process_response,
+ page_size=page_size,
+ max_concurrency=max_concurrency,
+ )
+ return outputs if requests is not None else outputs[0]
+ except KeyboardInterrupt as interruption:
+ if submission is not None:
+ interruption.add_note(
+ f"Submitted batch {submission.batch_id}; tasks {submission.task_ids}"
+ )
+ setattr(interruption, "submission", submission)
+ if cancel_on_timeout:
+ error = BatchError(499, "Interrupted", submission)
+ self._cancel_after_error(submission, error)
+ setattr(interruption, "cancellation", error.cancellation)
+ setattr(
+ interruption, "cancellation_error", error.cancellation_error
+ )
raise
- except ESMProteinError as e:
- if "timed out" in e.error_msg:
- task_timed_out = True
- return e
- finally:
+ except ESMProteinError as original:
+ error = BatchError.from_error(original, submission)
if (
cancel_on_timeout
- and task_id
- and (task_timed_out or keyboard_interrupted)
+ and submission is not None
+ and isinstance(error, BatchTimeoutError)
):
- with suppress(
- ESMProteinError
- ): # Don't surface errors from canceling the task
- with suppress(KeyboardInterrupt):
- self._batch_client.cancel(task_id)
+ self._cancel_after_error(submission, error)
+ return error
+ @overload
async def async_run(
self,
+ requests: Iterable[TRequest],
+ *,
timeout: int | None = None,
+ model_name: str | None = None,
+ **kwargs: Any,
+ ) -> list[TResponse | ESMProteinError] | ESMProteinError: ...
+
+ @overload
+ async def async_run(
+ self,
+ requests: None = None,
+ *,
+ timeout: int | None = None,
+ model_name: str | None = None,
+ **kwargs: Any,
+ ) -> TResponse | ESMProteinError: ...
+
+ async def async_run(
+ self,
+ requests: Iterable[TRequest] | None = None,
+ *,
+ timeout: int | None = None,
+ model_name: str | None = None,
+ disable_cache: bool = False,
cancel_on_timeout: bool = True,
poll_interval: int | None = None,
poll_max_interval: int | None = None,
+ page_size: int = 100,
+ max_concurrency: int = 8,
**kwargs,
- ) -> TResponse | ESMProteinError:
+ ) -> TResponse | list[TResponse | ESMProteinError] | ESMProteinError:
+ """Async equivalent of ``run()``, including cancellation on timeout by default.
+
+ Cancelling the awaiting coroutine also attempts to cancel queued work.
+ Set ``cancel_on_timeout=False`` to leave work running.
+ """
timeout = self.default_timeout if timeout is None else timeout
- task_id = None
- task_timed_out = False
- keyboard_interrupted = False
+ _check_result_options(page_size, max_concurrency)
+ submission = None
interval, max_interval = self._resolved_polling(
poll_interval, poll_max_interval
)
try:
- request = self._prepare_request(**kwargs)
- task_id = await self._batch_client.async_submit(self.endpoint_name, request)
- response = await self._batch_client.async_wait_for_completion(
- task_id, timeout, poll_interval=interval, poll_max_interval=max_interval
- )
- return await self._async_process_response(response, **kwargs)
- except KeyboardInterrupt:
- keyboard_interrupted = True
+ submission = await self.async_submit(
+ requests=requests,
+ disable_cache=disable_cache,
+ model_name=model_name,
+ **kwargs,
+ )
+ await self.async_wait(
+ submission,
+ timeout,
+ poll_interval=interval,
+ poll_max_interval=max_interval,
+ )
+ outputs = await self._batch_client._async_collect_results(
+ submission,
+ self._async_process_response,
+ page_size=page_size,
+ max_concurrency=max_concurrency,
+ )
+ return outputs if requests is not None else outputs[0]
+ except (KeyboardInterrupt, asyncio.CancelledError) as interruption:
+ if submission is not None:
+ interruption.add_note(
+ f"Submitted batch {submission.batch_id}; tasks {submission.task_ids}"
+ )
+ setattr(interruption, "submission", submission)
+ if cancel_on_timeout:
+ error = BatchError(499, "Interrupted", submission)
+ await self._async_cancel_after_error(submission, error)
+ setattr(interruption, "cancellation", error.cancellation)
+ setattr(
+ interruption, "cancellation_error", error.cancellation_error
+ )
raise
- except ESMProteinError as e:
- if "timed out" in e.error_msg:
- task_timed_out = True
- return e
- finally:
+ except ESMProteinError as original:
+ error = BatchError.from_error(original, submission)
if (
cancel_on_timeout
- and task_id
- and (task_timed_out or keyboard_interrupted)
+ and submission is not None
+ and isinstance(error, BatchTimeoutError)
):
- with suppress(
- ESMProteinError
- ): # Don't surface errors from canceling the task
- with suppress(KeyboardInterrupt):
- await self._batch_client.async_cancel(task_id)
+ await self._async_cancel_after_error(submission, error)
+ return error
diff --git a/esm/sdk/base_forge_client_test.py b/esm/sdk/base_forge_client_test.py
new file mode 100644
index 00000000..27aa5f26
--- /dev/null
+++ b/esm/sdk/base_forge_client_test.py
@@ -0,0 +1,73 @@
+import httpx
+import pytest
+
+from esm.sdk.base_forge_client import _BaseForgeInferenceClient, _sdk_version
+
+BIOHUB_COMPONENT = f"biohub/forge-sdk-python/{_sdk_version()}"
+
+
+def _client() -> _BaseForgeInferenceClient:
+ return _BaseForgeInferenceClient(
+ model="esmc-300m-2024-12",
+ url="https://biohub.ai",
+ token="test-token",
+ request_timeout=1,
+ min_retry_wait=1,
+ max_retry_wait=1,
+ max_retry_attempts=1,
+ )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("user_agent", [None, "custom-transport/1.0"])
+async def test_posts_component_and_preserves_transport_user_agent(
+ asynchronous: bool, user_agent: str | None
+) -> None:
+ client = _client()
+ requests: list[httpx.Request] = []
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ requests.append(request)
+ return httpx.Response(200, json={"outputs": {}})
+
+ transport = httpx.MockTransport(handle_request)
+ transport_headers = {} if user_agent is None else {"User-Agent": user_agent}
+ override = {"x-biohub-component": f"biohub/forge-sdk-python-test/{_sdk_version()}"}
+ request_headers = [{}, override, {}]
+
+ if asynchronous:
+ async with httpx.AsyncClient(
+ transport=transport, headers=transport_headers
+ ) as http_client:
+ client._async_client = http_client
+ for headers in request_headers:
+ assert await client._async_post("fold", {}, headers=headers) == {
+ "outputs": {}
+ }
+ else:
+ with httpx.Client(
+ transport=transport, headers=transport_headers
+ ) as http_client:
+ client._client = http_client
+ for headers in request_headers:
+ assert client._post("fold", {}, headers=headers) == {"outputs": {}}
+
+ assert [request.headers.get_list("x-biohub-component") for request in requests] == [
+ [BIOHUB_COMPONENT],
+ [override["x-biohub-component"]],
+ [BIOHUB_COMPONENT],
+ ]
+ for request in requests:
+ assert request.headers["authorization"] == "Bearer test-token"
+ assert request.headers.get_list("user-agent") == [
+ user_agent or f"python-httpx/{httpx.__version__}"
+ ]
+
+ assert client.headers == {
+ "Authorization": "Bearer test-token",
+ "X-Biohub-Component": BIOHUB_COMPONENT,
+ }
+ assert override == {
+ "x-biohub-component": f"biohub/forge-sdk-python-test/{_sdk_version()}"
+ }
diff --git a/esm/sdk/batch.py b/esm/sdk/batch.py
new file mode 100644
index 00000000..53b4d2e4
--- /dev/null
+++ b/esm/sdk/batch.py
@@ -0,0 +1,126 @@
+"""Durable identifiers and errors for the Batch API."""
+
+from dataclasses import dataclass
+from typing import Any
+
+from esm.sdk.api import ESMProteinError
+
+
+@dataclass(frozen=True)
+class BatchSubmission:
+ """A durable handle; task IDs contain only identifiers returned by the server."""
+
+ endpoint: str
+ task_ids: tuple[str, ...] = ()
+ batch_id: str | None = None
+
+ def __post_init__(self) -> None:
+ if not self.batch_id and not self.task_ids:
+ raise ValueError("A submission must contain a batch ID or task IDs")
+ if len(set(self.task_ids)) != len(self.task_ids):
+ raise ValueError("Task IDs must be unique")
+
+ @property
+ def task_id(self) -> str:
+ """Return the task ID of a single-task submission."""
+ if len(self.task_ids) != 1:
+ raise ValueError(
+ "This submission does not contain exactly one known task ID; "
+ "use list_batch() to discover batch tasks"
+ )
+ return self.task_ids[0]
+
+ def to_dict(self) -> dict[str, Any]:
+ """Return JSON-serializable identifiers and decoding context."""
+ return {
+ "endpoint": self.endpoint,
+ "batch_id": self.batch_id,
+ "task_ids": list(self.task_ids),
+ }
+
+ @classmethod
+ def from_dict(cls, data: dict[str, Any]) -> "BatchSubmission":
+ """Restore a submission saved by :meth:`to_dict`."""
+ return cls(
+ endpoint=data["endpoint"],
+ batch_id=data["batch_id"],
+ task_ids=tuple(data["task_ids"]),
+ )
+
+
+class BatchError(ESMProteinError):
+ """An error retaining the accepted submission and any collected partial results.
+
+ ``partial_results`` contains ``(task_id, output)`` pairs collected before
+ ``results()`` or ``run()`` failed. Outputs include per-task errors.
+ Streaming methods leave this list empty; callers retain the yielded outputs.
+ """
+
+ def __init__(
+ self,
+ error_code: int,
+ error_msg: str,
+ submission: BatchSubmission | None = None,
+ task_id: str | None = None,
+ *,
+ task_status: str | None = None,
+ ) -> None:
+ super().__init__(error_code, error_msg)
+ self.submission = submission
+ self.task_id = task_id
+ self.task_status = task_status
+ self.partial_results: list[tuple[str, Any]] = []
+ self.cancellation: dict[str, Any] | None = None
+ self.cancellation_error: ESMProteinError | None = None
+
+ def __reduce__(self) -> tuple[Any, ...]:
+ # Timeout constructors differ from Exception.args; attrs also omits our metadata.
+ return (
+ _restore_batch_error,
+ (type(self), self.error_code, self.error_msg, self.__dict__),
+ )
+
+ @property
+ def batch_id(self) -> str | None:
+ return self.submission.batch_id if self.submission else None
+
+ @property
+ def task_ids(self) -> tuple[str, ...]:
+ return self.submission.task_ids if self.submission else ()
+
+ @classmethod
+ def from_error(
+ cls,
+ error: ESMProteinError,
+ submission: BatchSubmission | None = None,
+ task_id: str | None = None,
+ ) -> "BatchError":
+ if isinstance(error, BatchError):
+ if submission is not None:
+ error.submission = submission
+ if task_id is not None:
+ error.task_id = task_id
+ return error
+ return cls(error.error_code, error.error_msg, submission, task_id)
+
+
+class BatchTimeoutError(BatchError):
+ """The caller stopped waiting; the saved submission remains recoverable."""
+
+ def __init__(self, submission: BatchSubmission, timeout: float) -> None:
+ identifier = submission.batch_id or ", ".join(submission.task_ids)
+ super().__init__(
+ 408,
+ f"Job {identifier} timed out after {timeout} seconds.",
+ submission=submission,
+ task_id=submission.task_ids[0] if len(submission.task_ids) == 1 else None,
+ )
+
+
+def _restore_batch_error(
+ error_type: type[BatchError], error_code: int, error_msg: str, state: dict[str, Any]
+) -> BatchError:
+ error = error_type.__new__(error_type)
+ ESMProteinError.__init__(error, error_code, error_msg)
+ error.__dict__.update(state)
+ return error
diff --git a/esm/sdk/batch_test.py b/esm/sdk/batch_test.py
new file mode 100644
index 00000000..aa33b517
--- /dev/null
+++ b/esm/sdk/batch_test.py
@@ -0,0 +1,874 @@
+"""Public batch lifecycle contracts against a mocked HTTP server."""
+
+import asyncio
+import json
+from collections.abc import AsyncIterator
+from typing import Any
+
+import httpx
+import pytest
+import pytest_asyncio
+
+from esm.sdk import base_forge_client
+from esm.sdk.api import (
+ ESMProtein,
+ ESMProteinError,
+ FoldAllAtomRequest,
+ FoldMaxAccuracyConfig,
+ FoldMaxAccuracyRequest,
+ FoldRequest,
+)
+from esm.sdk.batch import BatchError, BatchSubmission, BatchTimeoutError
+from esm.sdk.forge import ForgeBatchClient
+from esm.utils.constants.models import ESMFOLD2_CUTOFF_2025
+from esm.utils.residue_constants import restype_1to3
+from esm.utils.structure.input_builder import ProteinInput, StructurePredictionInput
+from esm.utils.structure.molecular_complex import MolecularComplexResult
+
+
+def _input(sequence: str) -> StructurePredictionInput:
+ return StructurePredictionInput(sequences=[ProteinInput(id="A", sequence=sequence)])
+
+
+def _fold_result(sequence: str = "ACDEF") -> dict[str, Any]:
+ return {
+ "complex": {
+ "id": sequence,
+ "sequence": [restype_1to3[residue] for residue in sequence],
+ "atom_positions": [[0.0, 0.0, 0.0] for _ in sequence],
+ "atom_elements": ["C" for _ in sequence],
+ "token_to_atoms": [[index, index + 1] for index in range(len(sequence))],
+ "chain_id": [0 for _ in sequence],
+ "plddt": [90.0 for _ in sequence],
+ "metadata": {"entity_lookup": {1: "protein"}, "chain_lookup": {0: "A"}},
+ }
+ }
+
+
+class _BatchAPI:
+ def __init__(self) -> None:
+ self.requests: list[httpx.Request] = []
+ self.replies: list[httpx.Response | Exception] = []
+ self.clients: list[ForgeBatchClient] = []
+
+ def respond(self, request: httpx.Request) -> httpx.Response:
+ self.requests.append(request)
+ reply = self.replies.pop(0)
+ if isinstance(reply, Exception):
+ raise reply
+ return reply
+
+ def client(self, model: str | None = None) -> ForgeBatchClient:
+ client = ForgeBatchClient(
+ token="test-token",
+ model=model,
+ min_retry_wait=0,
+ max_retry_wait=0,
+ max_retry_attempts=3,
+ )
+ transport = httpx.MockTransport(self.respond)
+ client._batch_client._client = httpx.Client(transport=transport)
+ client._batch_client._async_client = httpx.AsyncClient(transport=transport)
+ self.clients.append(client)
+ return client
+
+ def data(self, value: dict[str, Any]) -> None:
+ self.replies.append(
+ httpx.Response(200, json={"status": "success", "data": value})
+ )
+
+ def accepted(self, count: int = 1) -> None:
+ reply: dict[str, Any] = {"batch_id": "batch-A", "total_items": count}
+ if count == 1:
+ reply["task_id"] = "task-1"
+ self.data(reply)
+
+ def completed(self, task_id: str = "task-1") -> None:
+ self.data({"status": "done", "response": f"https://results.example/{task_id}"})
+
+ def download(self, sequence: str = "ACDEF") -> None:
+ self.replies.append(httpx.Response(200, json=_fold_result(sequence)))
+
+ def page(
+ self,
+ items: list[dict[str, Any]],
+ cursor: str | None = None,
+ total: int | None = None,
+ ) -> None:
+ self.data(
+ {
+ "status": "done",
+ "items": items,
+ "cursor": cursor,
+ "total": len(items) if total is None else total,
+ }
+ )
+
+
+class _Clock:
+ def __init__(self) -> None:
+ self.now = 1000.0
+
+ def monotonic(self) -> float:
+ return self.now
+
+ def sleep(self, seconds: float) -> None:
+ self.now += seconds
+
+
+@pytest_asyncio.fixture
+async def api() -> AsyncIterator[_BatchAPI]:
+ api = _BatchAPI()
+ yield api
+ for client in api.clients:
+ client._batch_client.client.close()
+ await client._batch_client.async_client.aclose()
+
+
+async def _call(target: Any, asynchronous: bool, method: str, *args, **kwargs) -> Any:
+ if asynchronous:
+ return await getattr(target, f"async_{method}")(*args, **kwargs)
+ return getattr(target, method)(*args, **kwargs)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize(
+ "count,request_shape,disable_cache",
+ [(1, "single", False), (3, "list", True), (2, "positional_generator", False)],
+)
+async def test_submit_retains_only_explicit_ids(
+ api, asynchronous, count, request_shape, disable_cache
+):
+ client = api.client(model="unrelated-client-model")
+ sequences = ["ACDEF", "GHIKL", "MNPQR"][:count]
+ api.accepted(count)
+ requests = (
+ FoldMaxAccuracyRequest(all_atom_input=_input(sequence))
+ for sequence in sequences
+ )
+ kwargs = (
+ {"requests": list(requests) if request_shape == "list" else requests}
+ if request_shape != "single"
+ else {"all_atom_input": _input(sequences[0])}
+ )
+ args = (kwargs.pop("requests"),) if request_shape == "positional_generator" else ()
+
+ submission = await _call(
+ client.fold_max_accuracy,
+ asynchronous,
+ "submit",
+ *args,
+ disable_cache=disable_cache,
+ **kwargs,
+ )
+
+ assert submission.batch_id == "batch-A"
+ saved = json.loads(json.dumps(submission.to_dict()))
+ assert saved == {
+ "endpoint": "fold_max_accuracy",
+ "batch_id": "batch-A",
+ "task_ids": list(submission.task_ids),
+ }
+ assert BatchSubmission.from_dict(saved) == submission
+ assert submission.task_ids == (("task-1",) if count == 1 else ())
+ assert len(api.requests) == 1
+ assert api.requests[0].headers.get("X-Disable-Cache") == (
+ "true" if disable_cache else None
+ )
+ request = json.loads(api.requests[0].content)
+ assert request["endpoint"] == "fold_max_accuracy"
+ assert set(request) == {"endpoint", "payload"}
+ assert [
+ payload["all_atom_input"]["sequences"][0]["sequence"]
+ for payload in request["payload"]
+ ] == sequences
+ assert all(
+ payload["model"] == ESMFOLD2_CUTOFF_2025 for payload in request["payload"]
+ )
+ if count == 1:
+ assert submission.task_id == "task-1"
+ else:
+ with pytest.raises(ValueError, match="exactly one known task ID"):
+ _ = submission.task_id
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_fold_max_accuracy_typed_requests_resolve_configs_independently(
+ api, asynchronous
+):
+ client = api.client()
+ item_config = FoldMaxAccuracyConfig(include_embeddings=True)
+ requests = [
+ FoldMaxAccuracyRequest(all_atom_input=_input("ACDEF")),
+ FoldMaxAccuracyRequest(all_atom_input=_input("GHIKL"), config=item_config),
+ ]
+ config = FoldMaxAccuracyConfig(include_pae=True)
+ api.accepted(len(requests))
+ api.data({"status": "done", "counts": {"done": 2}, "total": 2})
+ api.page(
+ [
+ {"task_id": str(index), "status": "done", "response": _fold_result()}
+ for index in range(len(requests))
+ ]
+ )
+
+ await _call(
+ client.fold_max_accuracy,
+ asynchronous,
+ "run",
+ requests=iter(requests),
+ config=config,
+ )
+
+ payload = json.loads(api.requests[0].content)["payload"]
+ assert [(item["include_pae"], item["include_embeddings"]) for item in payload] == [
+ (True, False),
+ (False, True),
+ ]
+ assert requests[0].config is None
+ assert requests[1].config is item_config
+ assert not api.replies
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize(
+ "invalid_input", ["empty_generator", "ambiguous", "model_override", "dict"]
+)
+async def test_fold_max_accuracy_rejects_invalid_request_inputs_before_http(
+ api, asynchronous, invalid_input
+):
+ client = api.client()
+ typed = [FoldMaxAccuracyRequest(all_atom_input=_input("GHIKL"))]
+ kwargs: dict[str, Any] = {
+ "empty_generator": {
+ "requests": (FoldMaxAccuracyRequest(all_atom_input=item) for item in [])
+ },
+ "ambiguous": {"all_atom_input": _input("ACDEF"), "requests": typed},
+ "model_override": {"model_name": "esmfold2", "requests": typed},
+ "dict": {"requests": [{"all_atom_input": _input("ACDEF")}]},
+ }[invalid_input]
+
+ expected_error = ValueError if invalid_input == "empty_generator" else TypeError
+ expected_message = (
+ "fold_max_accuracy selects its own model"
+ if invalid_input == "model_override"
+ else None
+ )
+ with pytest.raises(expected_error, match=expected_message):
+ await _call(client.fold_max_accuracy, asynchronous, "run", **kwargs)
+
+ assert not api.requests
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_run_accepts_single_task_response(api, asynchronous):
+ client = api.client()
+ api.data({"message": "Batch request submitted", "task_id": "legacy-task"})
+ api.completed("legacy-task")
+ api.completed("legacy-task")
+ api.download()
+
+ result = await _call(
+ client.fold_all_atom,
+ asynchronous,
+ "run",
+ all_atom_input=_input("ACDEF"),
+ model_name="esmfold2-test",
+ )
+
+ assert isinstance(result, MolecularComplexResult)
+ assert [request.url.path for request in api.requests] == [
+ "/api/v1/batch/submit",
+ "/api/v1/batch/status",
+ "/api/v1/batch/status",
+ "/legacy-task",
+ ]
+ assert not api.replies
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("model", ["esmfold2-2026-05", "esmfold2-fast-2026-05"])
+@pytest.mark.parametrize("endpoint", ["fold", "fold_all_atom"])
+async def test_public_folding_submit_and_decode(api, asynchronous, model, endpoint):
+ client = api.client(model=model)
+ sequences = ["ACDEF", "GHIKL"]
+ api.accepted(count=2)
+ api.data({"status": "done", "counts": {"done": 2}, "total": 2})
+ api.page(
+ [
+ {
+ "task_id": f"batch-A-{index:05d}",
+ "status": "done",
+ "response_available": True,
+ }
+ for index in range(2)
+ ]
+ )
+ for index, sequence in enumerate(sequences):
+ api.completed(f"batch-A-{index:05d}")
+ response = (
+ {
+ "sequence": sequence,
+ "coordinates": [[[0.0, 0.0, 0.0]] * 37] * len(sequence),
+ }
+ if endpoint == "fold"
+ else _fold_result(sequence)
+ )
+ api.replies.append(httpx.Response(200, json=response))
+ requests = [
+ FoldRequest(sequence=sequence)
+ if endpoint == "fold"
+ else FoldAllAtomRequest(all_atom_input=_input(sequence))
+ for sequence in sequences
+ ]
+ results = await _call(
+ getattr(client, endpoint),
+ asynchronous,
+ "run",
+ requests=requests,
+ max_concurrency=1,
+ )
+ assert isinstance(results, list)
+ if endpoint == "fold":
+ assert all(isinstance(result, ESMProtein) for result in results)
+ assert [result.sequence for result in results] == sequences
+ assert all(result.coordinates.shape == (5, 37, 3) for result in results)
+ else:
+ assert all(isinstance(result, MolecularComplexResult) for result in results)
+ assert [result.complex.id for result in results] == sequences
+ submitted = json.loads(api.requests[0].content)
+ assert submitted["endpoint"] == endpoint
+ assert [item["model"] for item in submitted["payload"]] == [model, model]
+ assert (
+ len(
+ [
+ request
+ for request in api.requests
+ if request.url.path.endswith("/submit")
+ ]
+ )
+ == 1
+ )
+ assert not api.replies
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("method", ["submit", "run"])
+@pytest.mark.parametrize(
+ "failure, error_code",
+ [
+ (lambda: httpx.ReadTimeout("Server accepted the request, response lost"), 500),
+ (lambda: httpx.Response(503, json={"message": "Unavailable"}), 503),
+ (lambda: httpx.Response(200, content=b"invalid JSON"), 400),
+ ],
+ ids=["lost_response", "unavailable", "invalid_json"],
+)
+async def test_submission_failures_are_not_resent(
+ api, asynchronous, method, failure, error_code
+):
+ client = api.client()
+ api.replies.append(failure())
+ api.accepted()
+ options = {"all_atom_input": _input("ACDEF")}
+
+ if method == "submit":
+ with pytest.raises(BatchError) as raised:
+ await _call(client.fold_max_accuracy, asynchronous, method, **options)
+ error = raised.value
+ else:
+ error = await _call(client.fold_max_accuracy, asynchronous, method, **options)
+
+ assert isinstance(error, BatchError)
+ assert error.error_code == error_code
+ assert error.submission is None
+ assert error.task_status is None
+ assert len(api.requests) == 1
+ assert len(api.replies) == 1
+ assert api.requests[0].url.path == "/api/v1/batch/submit"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize(
+ "request_shape", ["single", "list_one", "positional_generator_many"]
+)
+async def test_run_composes_submit_wait_and_ordered_typed_results(
+ api, asynchronous, request_shape
+):
+ client = api.client()
+ sequences = ["ACDEF", "GHIKL"] if request_shape.endswith("many") else ["ACDEF"]
+ api.accepted(len(sequences))
+ task_ids = ["opaque-z", "opaque-a"][: len(sequences)]
+ api.data(
+ {"status": "done", "counts": {"done": len(sequences)}, "total": len(sequences)}
+ )
+ api.page(
+ [
+ {"task_id": task_id, "status": "done", "response": _fold_result(sequence)}
+ for task_id, sequence in reversed(list(zip(task_ids, sequences)))
+ ]
+ )
+ requests = (
+ FoldMaxAccuracyRequest(all_atom_input=_input(sequence))
+ for sequence in sequences
+ )
+ kwargs = (
+ {"all_atom_input": _input(sequences[0])}
+ if request_shape == "single"
+ else {
+ "requests": list(requests) if request_shape.startswith("list") else requests
+ }
+ )
+
+ args = (
+ (kwargs.pop("requests"),)
+ if request_shape == "positional_generator_many"
+ else ()
+ )
+ output = await _call(
+ client.fold_max_accuracy,
+ asynchronous,
+ "run",
+ *args,
+ disable_cache=True,
+ page_size=2,
+ max_concurrency=1,
+ **kwargs,
+ )
+
+ if request_shape == "single":
+ assert isinstance(output, MolecularComplexResult)
+ results = [output]
+ else:
+ assert isinstance(output, list)
+ results = output
+ assert [result.complex.id for result in results] == list(reversed(sequences))
+ assert api.requests[0].headers.get("X-Disable-Cache") == "true"
+ assert [request.url.path for request in api.requests] == [
+ "/api/v1/batch/submit",
+ "/api/v1/batch/list",
+ "/api/v1/batch/list",
+ ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_timeout_without_cancellation_keeps_a_recoverable_handle(
+ api, asynchronous
+):
+ client = api.client()
+ api.accepted()
+ api.data({"status": "in_progress", "counts": {"queued": 1}, "total": 1})
+
+ timeout = await _call(
+ client.fold_max_accuracy,
+ asynchronous,
+ "run",
+ all_atom_input=_input("ACDEF"),
+ timeout=0,
+ cancel_on_timeout=False,
+ )
+
+ assert isinstance(timeout, BatchTimeoutError)
+ assert timeout.error_code == 408
+ assert timeout.task_status is None
+ assert timeout.batch_id == "batch-A"
+ assert timeout.task_ids == ("task-1",)
+ assert timeout.submission is not None
+ restored = BatchSubmission.from_dict(
+ json.loads(json.dumps(timeout.submission.to_dict()))
+ )
+ api.page([{"task_id": "task-1", "status": "done", "response": _fold_result()}])
+ results = await _call(
+ api.client().fold_max_accuracy, asynchronous, "results", restored
+ )
+ assert results[0].complex.id == "ACDEF"
+ assert sum(request.url.path.endswith("/submit") for request in api.requests) == 1
+ assert not any(request.url.path.endswith("/cancel") for request in api.requests)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_multi_task_timeout_cancels_and_retains_batch_only_handle(
+ api, monkeypatch, asynchronous
+):
+ clock = _Clock()
+ monkeypatch.setattr(base_forge_client, "time", clock)
+ respond = api.respond
+
+ def slow_pending_task(request: httpx.Request) -> httpx.Response:
+ response = respond(request)
+ if request.url.path.endswith("/list"):
+ clock.now += 2
+ return response
+
+ monkeypatch.setattr(api, "respond", slow_pending_task)
+ client = api.client()
+ api.accepted(3)
+ api.data(
+ {
+ "status": "in_progress",
+ "counts": {"done": 1, "in_progress": 1, "queued": 1},
+ "total": 3,
+ }
+ )
+ cancellation = {"cancelled": 1, "already_running": 1}
+ api.data(cancellation)
+
+ timeout = await _call(
+ client.fold_max_accuracy,
+ asynchronous,
+ "run",
+ requests=[
+ FoldMaxAccuracyRequest(all_atom_input=_input("ACDEF")),
+ FoldMaxAccuracyRequest(all_atom_input=_input("GHIKL")),
+ FoldMaxAccuracyRequest(all_atom_input=_input("MNPQR")),
+ ],
+ timeout=1,
+ )
+
+ assert isinstance(timeout, BatchTimeoutError)
+ assert timeout.batch_id == "batch-A"
+ assert timeout.task_ids == ()
+ assert timeout.cancellation == cancellation
+ assert timeout.cancellation_error is None
+ assert [request.url.path for request in api.requests] == [
+ "/api/v1/batch/submit",
+ "/api/v1/batch/list",
+ "/api/v1/batch/cancel",
+ ]
+ assert json.loads(api.requests[-1].content) == {"batch_id": "batch-A"}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("dispatch_status", ["in_progress", "cancelled", "failed"])
+async def test_wait_polls_one_aggregate_per_cycle_until_no_active_tasks(
+ api, monkeypatch, asynchronous, dispatch_status
+):
+ client = api.client()
+ clock = _Clock()
+ monkeypatch.setattr(base_forge_client, "time", clock)
+
+ async def sleep(seconds: float) -> None:
+ clock.sleep(seconds)
+
+ monkeypatch.setattr(base_forge_client.asyncio, "sleep", sleep)
+ submission = BatchSubmission(
+ endpoint="fold_max_accuracy", batch_id="batch-A", task_ids=("a", "b", "c", "d")
+ )
+ api.data(
+ {"status": dispatch_status, "counts": {"in_progress": 1, "done": 1}, "total": 4}
+ )
+ terminal = {
+ "status": "done" if dispatch_status == "in_progress" else dispatch_status,
+ "counts": {"done": 2, "failed": 1, "cancelled": 1}
+ if dispatch_status == "in_progress"
+ else {"done": 2},
+ "total": 4,
+ }
+ api.data(terminal)
+
+ assert (
+ await _call(
+ client, asynchronous, "wait", submission, timeout=5, poll_interval=1
+ )
+ == terminal
+ )
+
+ assert len(api.requests) == 2
+ assert all(request.url.path.endswith("/list") for request in api.requests)
+ assert all(
+ json.loads(request.content)
+ == {"batch_id": "batch-A", "include_responses": False}
+ for request in api.requests
+ )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("iterate", [False, True], ids=["collect", "iterate"])
+async def test_results_paginate_and_only_refresh_stored_outputs(
+ api, asynchronous, iterate
+):
+ client = api.client()
+ task_ids = ("stored", "inline", "failed", "cancelled", "queued", "in_progress")
+ submission = BatchSubmission(
+ endpoint="fold_max_accuracy", batch_id="batch-A", task_ids=task_ids
+ )
+ api.page(
+ [
+ {"task_id": "inline", "status": "done", "response": _fold_result()},
+ {"task_id": "failed", "status": "failed", "error": "CUDA out of memory"},
+ ],
+ cursor="page-2",
+ total=6,
+ )
+ api.page(
+ [
+ {
+ "task_id": "cancelled",
+ "status": "cancelled",
+ "error": "Task expired after waiting too long in the queue.",
+ },
+ {"task_id": "queued", "status": "queued"},
+ ],
+ cursor="page-3",
+ total=6,
+ )
+ api.page(
+ [
+ {"task_id": "stored", "status": "done", "response_available": True},
+ {"task_id": "in_progress", "status": "in_progress"},
+ ],
+ total=6,
+ )
+ api.completed("fresh-url")
+ api.download()
+
+ if iterate:
+ if asynchronous:
+ entries = [
+ entry
+ async for entry in client.fold_max_accuracy.async_iter_results(
+ submission, page_size=2
+ )
+ ]
+ else:
+ entries = list(
+ client.fold_max_accuracy.iter_results(submission, page_size=2)
+ )
+ assert [task_id for task_id, _ in entries] == [
+ "inline",
+ "failed",
+ "cancelled",
+ "queued",
+ "stored",
+ "in_progress",
+ ]
+ results = [result for _, result in entries]
+ else:
+ results = await _call(
+ client.fold_max_accuracy, asynchronous, "results", submission, page_size=2
+ )
+
+ assert [results[index].complex.id for index in (0, 4)] == ["ACDEF", "ACDEF"]
+ failures = [results[index] for index in (1, 2, 3, 5)]
+ assert [result.error_code for result in failures] == [500, 409, 409, 409]
+ assert [result.task_status for result in failures] == [
+ "failed",
+ "cancelled",
+ "queued",
+ "in_progress",
+ ]
+ assert [result.task_id for result in failures] == [
+ "failed",
+ "cancelled",
+ "queued",
+ "in_progress",
+ ]
+ assert all(result.submission == submission for result in failures)
+ assert "CUDA out of memory" in failures[0].error_msg
+ assert "Task expired after waiting too long in the queue." in failures[1].error_msg
+ list_requests = [
+ json.loads(request.content)
+ for request in api.requests
+ if request.url.path.endswith("/list")
+ ]
+ assert list_requests == [
+ {"batch_id": "batch-A", "include_responses": True, "limit": 2},
+ {
+ "batch_id": "batch-A",
+ "include_responses": True,
+ "limit": 2,
+ "cursor": "page-2",
+ },
+ {
+ "batch_id": "batch-A",
+ "include_responses": True,
+ "limit": 2,
+ "cursor": "page-3",
+ },
+ ]
+ assert [
+ json.loads(request.content)
+ for request in api.requests
+ if request.url.path.endswith("/status")
+ ] == [{"task_id": "stored"}]
+ assert [
+ str(request.url) for request in api.requests if request.method == "GET"
+ ] == ["https://results.example/fresh-url"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_result_decode_error_is_not_a_task_failure(api, asynchronous):
+ client = api.client()
+ api.data({"status": "done", "response": {"complex": {}}})
+
+ with pytest.raises(BatchError) as error:
+ await _call(client.fold_max_accuracy, asynchronous, "result", "saved-task")
+
+ assert error.value.error_code == 502
+ assert error.value.task_status is None
+ assert error.value.task_id == "saved-task"
+ assert len(api.requests) == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("cancel_on_timeout", [None, False])
+async def test_async_cancellation_reraises_with_submission(api, cancel_on_timeout):
+ client = api.client()
+ polling = asyncio.Event()
+ await client._batch_client.async_client.aclose()
+
+ async def respond(request: httpx.Request) -> httpx.Response:
+ api.requests.append(request)
+ if request.url.path.endswith("/submit"):
+ data = {"batch_id": "batch-A", "task_id": "task-1", "total_items": 1}
+ elif request.url.path.endswith("/list"):
+ polling.set()
+ await asyncio.Event().wait()
+ raise AssertionError("Cancelled poll resumed")
+ else:
+ assert request.url.path.endswith("/cancel")
+ data = {"cancelled": 1}
+ return httpx.Response(200, json={"status": "success", "data": data})
+
+ client._batch_client._async_client = httpx.AsyncClient(
+ transport=httpx.MockTransport(respond)
+ )
+ options = (
+ {} if cancel_on_timeout is None else {"cancel_on_timeout": cancel_on_timeout}
+ )
+ task = asyncio.create_task(
+ client.fold_max_accuracy.async_run(all_atom_input=_input("ACDEF"), **options)
+ )
+ await asyncio.wait_for(polling.wait(), timeout=5)
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError) as cancelled:
+ await task
+
+ assert getattr(cancelled.value, "submission").task_ids == ("task-1",)
+ cancellations = [
+ request for request in api.requests if request.url.path.endswith("/cancel")
+ ]
+ assert len(cancellations) == int(cancel_on_timeout is not False)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("method", ["results", "run"])
+@pytest.mark.parametrize("failure", ["incomplete", "http"])
+async def test_collect_failure_retains_mixed_partial_results(
+ api, asynchronous, method, failure
+):
+ client = api.client()
+ submission = BatchSubmission(endpoint="fold_max_accuracy", batch_id="batch-A")
+ if method == "run":
+ api.accepted(3)
+ api.data({"status": "failed", "counts": {"done": 1, "failed": 1}, "total": 3})
+ api.page(
+ [
+ {"task_id": "opaque-success", "status": "done", "response_available": True},
+ {"task_id": "opaque-failed", "status": "failed", "error": "Model failed"},
+ ],
+ cursor="later" if failure == "http" else None,
+ total=3,
+ )
+ api.completed("stored-output")
+ api.download("GHIKL")
+ if failure == "http":
+ api.replies.append(httpx.Response(404, json={"message": "Batch not found"}))
+
+ if method == "run":
+ error = await _call(
+ client.fold_max_accuracy,
+ asynchronous,
+ method,
+ requests=[
+ FoldMaxAccuracyRequest(all_atom_input=_input("ACDEF")) for _ in range(3)
+ ],
+ page_size=2,
+ max_concurrency=1,
+ )
+ else:
+ with pytest.raises(BatchError) as raised:
+ await _call(
+ client.fold_max_accuracy,
+ asynchronous,
+ method,
+ submission,
+ page_size=2,
+ max_concurrency=1,
+ )
+ error = raised.value
+
+ assert isinstance(error, BatchError)
+ assert error.error_code == (404 if failure == "http" else 409)
+ assert error.submission == submission
+ assert error.task_id is None
+ assert error.task_status is None
+ assert [task_id for task_id, _ in error.partial_results] == [
+ "opaque-success",
+ "opaque-failed",
+ ]
+ successful, failed = [result for _, result in error.partial_results]
+ assert isinstance(successful, MolecularComplexResult)
+ assert successful.complex.id == "GHIKL"
+ assert isinstance(failed, BatchError)
+ assert failed.task_id == "opaque-failed"
+ assert failed.task_status == "failed"
+ assert failed.error_code == 500
+ assert failed.submission == submission
+ assert failed.partial_results == []
+ if failure == "http":
+ assert isinstance(error.__cause__, ESMProteinError)
+ assert error.__cause__.error_code == 404
+ assert sum(request.url.path.endswith("/status") for request in api.requests) == 1
+ assert [
+ str(request.url) for request in api.requests if request.method == "GET"
+ ] == ["https://results.example/stored-output"]
+ assert not any(request.url.path.endswith("/cancel") for request in api.requests)
+ assert api.replies == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("method", ["result", "results"])
+async def test_expired_result_preserves_http_410_without_retries(
+ api, asynchronous, method
+):
+ client = api.client()
+ submission = BatchSubmission(endpoint="fold_max_accuracy", batch_id="batch-A")
+ reason = "Results expired after the retention period."
+ if method == "results":
+ api.page([{"task_id": "task-1", "status": "done", "response_available": False}])
+ api.replies.append(httpx.Response(410, json={"message": reason}))
+
+ if method == "result":
+ with pytest.raises(BatchError) as raised:
+ await _call(client.fold_max_accuracy, asynchronous, method, "task-1")
+ result = raised.value
+ else:
+ results = await _call(
+ client.fold_max_accuracy, asynchronous, method, submission
+ )
+ assert len(results) == 1
+ result = results[0]
+ assert result.submission == submission
+
+ assert isinstance(result, BatchError)
+ assert result.error_code == 410
+ assert result.task_id == "task-1"
+ assert reason in result.error_msg
+ assert [request.url.path for request in api.requests] == [
+ *(["/api/v1/batch/list"] if method == "results" else []),
+ "/api/v1/batch/status",
+ ]
diff --git a/esm/sdk/forge.py b/esm/sdk/forge.py
index 21ca420e..bc8551ec 100644
--- a/esm/sdk/forge.py
+++ b/esm/sdk/forge.py
@@ -2,7 +2,6 @@
import asyncio
import base64
-import pickle
import warnings
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Sequence, TypeAlias, Union, cast
@@ -16,8 +15,11 @@
ESMProtein,
ESMProteinError,
ESMProteinTensor,
+ FoldAllAtomRequest,
FoldingConfig,
FoldMaxAccuracyConfig,
+ FoldMaxAccuracyRequest,
+ FoldRequest,
ForwardAndSampleOutput,
ForwardTrackData,
GenerationConfig,
@@ -34,10 +36,11 @@
_BaseForgeBatchClient,
_BaseForgeInferenceClient,
)
+from esm.sdk.batch import BatchSubmission
from esm.sdk.retry import retry_decorator
from esm.sdk.validation import validate_fold_max_accuracy_input
from esm.utils.compression import compress_state_dict
-from esm.utils.constants.api import MIMETYPE_ES_PICKLE
+from esm.utils.constants.api import ACCEPT_ES_PICKLE
from esm.utils.constants.models import (
DEFAULT_ESMFOLD2_FAST_LM_MASK_PCT,
ESMFOLD2,
@@ -114,10 +117,10 @@ def __init__(
url: str = "https://biohub.ai",
model: str | None = None,
token: str = "",
- request_timeout: int | None = 120,
+ request_timeout: int | None = 125,
min_retry_wait: int = 1,
max_retry_wait: int = 10,
- max_retry_attempts: int = 5,
+ max_retry_attempts: int = 3,
):
"""
Forge/Biohub Platform client for folding and inverse folding between sequence and structure spaces.
@@ -126,7 +129,8 @@ def __init__(
url: URL of the Forge/Biohub Platform server.
model: Name of the model to be used for folding / inv folding.
token: API token.
- request_timeout: Override the system default request timeout (120s), in seconds.
+ request_timeout: Override the system default request timeout (125s), in seconds.
+ Note the API has a 120s timeout
"""
super().__init__(
model=model or "",
@@ -195,9 +199,9 @@ def _process_fold_request(
return request
@staticmethod
- def _process_fold_response(data: dict[str, Any], sequence: str) -> ESMProtein:
+ def _process_fold_response(data: dict[str, Any]) -> ESMProtein:
return ESMProtein(
- sequence=sequence,
+ sequence=data["sequence"],
coordinates=maybe_tensor(data["coordinates"], convert_none_to_nan=True),
ptm=maybe_tensor(data.get("ptm", None)),
plddt=maybe_tensor(data.get("plddt", None), convert_none_to_nan=True),
@@ -248,6 +252,7 @@ async def async_fold(
msa: MSAInput = None,
config: FoldingConfig = FoldingConfig(),
target_structure: ProteinComplex | None = None,
+ disable_cache: bool = False,
) -> ESMProtein | ESMProteinError:
"""Predict coordinates for a protein sequence.
@@ -257,6 +262,7 @@ async def async_fold(
msa: Optional multi-sequence alignment data that may help some models fold the sequence.
config: Optional configuration for folding parameters.
target_structure: Optional target structure to use for distogram conditioning.
+ disable_cache: Optional flag to disable caching
Deprecated:
potential_sequence_of_concern: this parameter is largely deprecated
@@ -273,11 +279,16 @@ async def async_fold(
)
try:
- data = await self._async_post("fold", request)
+ data = await self._async_post(
+ "fold",
+ request,
+ headers={"Accept": ACCEPT_ES_PICKLE},
+ disable_cache=disable_cache,
+ )
except ESMProteinError as e:
return e
- return self._process_fold_response(data, sequence)
+ return self._process_fold_response(data)
@retry_decorator
def fold(
@@ -288,6 +299,7 @@ def fold(
msa: MSAInput = None,
config: FoldingConfig = FoldingConfig(),
target_structure: ProteinComplex | None = None,
+ disable_cache: bool = False,
) -> ESMProtein | ESMProteinError:
"""Predict coordinates for a protein sequence.
@@ -297,6 +309,7 @@ def fold(
msa: Optional multi-sequence alignment data that may help some models fold the sequence.
config: Optional configuration for folding parameters.
target_structure: Optional target structure to use for distogram conditioning.
+ disable_cache: Optional flag to disable caching
Deprecated:
potential_sequence_of_concern: this parameter is largely deprecated
@@ -314,11 +327,16 @@ def fold(
)
try:
- data = self._post("fold", request)
+ data = self._post(
+ "fold",
+ request,
+ headers={"Accept": ACCEPT_ES_PICKLE},
+ disable_cache=disable_cache,
+ )
except ESMProteinError as e:
return e
- return self._process_fold_response(data, sequence)
+ return self._process_fold_response(data)
@retry_decorator
async def async_fold_all_atom(
@@ -326,6 +344,7 @@ async def async_fold_all_atom(
all_atom_input: StructurePredictionInput,
config: FoldingConfig = FoldingConfig(),
model_name: str | None = None,
+ disable_cache: bool = False,
) -> MolecularComplexResult | list[MolecularComplexResult] | ESMProteinError:
"""Fold a molecular complex containing proteins, nucleic acids, and/or ligands.
@@ -333,13 +352,16 @@ async def async_fold_all_atom(
all_atom_input: StructurePredictionInput containing sequences for different molecule types
config: Optional configuration for folding parameters.
model_name: Override the client level model name if needed
+ disable_cache: Optional flag to disable caching
"""
request = self._process_fold_all_atom_request(
all_atom_input, config, model_name if model_name is not None else self.model
)
try:
- data = await self._async_post("fold_all_atom", request)
+ data = await self._async_post(
+ "fold_all_atom", request, disable_cache=disable_cache
+ )
except ESMProteinError as e:
return e
@@ -351,6 +373,7 @@ def fold_all_atom(
all_atom_input: StructurePredictionInput,
model_name: str | None = None,
config: FoldingConfig = FoldingConfig(),
+ disable_cache: bool = False,
) -> MolecularComplexResult | list[MolecularComplexResult] | ESMProteinError:
"""Predict coordinates for a molecular complex containing proteins, dna, rna, and/or ligands.
@@ -358,13 +381,14 @@ def fold_all_atom(
all_atom_input: StructurePredictionInput containing sequences for different molecule types
model_name: Override the client level model name if needed
config: Optional configuration for folding parameters.
+ disable_cache: Optional flag to disable caching
"""
request = self._process_fold_all_atom_request(
all_atom_input, config, model_name if model_name is not None else self.model
)
try:
- data = self._post("fold_all_atom", request)
+ data = self._post("fold_all_atom", request, disable_cache=disable_cache)
except ESMProteinError as e:
return e
@@ -477,7 +501,10 @@ async def async_inverse_fold(
try:
data = await self._async_post(
- "inverse_fold", request, potential_sequence_of_concern
+ "inverse_fold",
+ request,
+ potential_sequence_of_concern,
+ headers={"Accept": ACCEPT_ES_PICKLE},
)
except ESMProteinError as e:
return e
@@ -512,7 +539,12 @@ def inverse_fold(
)
try:
- data = self._post("inverse_fold", request, potential_sequence_of_concern)
+ data = self._post(
+ "inverse_fold",
+ request,
+ potential_sequence_of_concern,
+ headers={"Accept": ACCEPT_ES_PICKLE},
+ )
except ESMProteinError as e:
return e
@@ -525,10 +557,10 @@ def __init__(
model: str,
url: str = "https://biohub.ai",
token: str = "",
- request_timeout: int | None = 120,
+ request_timeout: int | None = 125,
min_retry_wait: int = 1,
max_retry_wait: int = 10,
- max_retry_attempts: int = 5,
+ max_retry_attempts: int = 3,
):
ESM3InferenceClient.__init__(self)
_BaseForgeInferenceClient.__init__(
@@ -567,6 +599,7 @@ def _process_generate_protein_request(
"strategy": config.strategy,
"temperature_annealing": config.temperature_annealing,
"only_compute_backbone_rmsd": config.only_compute_backbone_rmsd,
+ "include_pae": config.include_pae,
}
return request
@@ -971,7 +1004,10 @@ async def __async_generate_protein(
request = self._process_generate_protein_request(input, config, self.model)
try:
data = await self._async_post(
- "generate", request, input.potential_sequence_of_concern
+ "generate",
+ request,
+ input.potential_sequence_of_concern,
+ headers={"Accept": ACCEPT_ES_PICKLE},
)
except ESMProteinError as e:
return e
@@ -987,7 +1023,10 @@ async def __async_generate_protein_tensor(
try:
data = await self._async_post(
- "generate_tensor", request, input.potential_sequence_of_concern
+ "generate_tensor",
+ request,
+ input.potential_sequence_of_concern,
+ headers={"Accept": ACCEPT_ES_PICKLE},
)
except ESMProteinError as e:
return e
@@ -999,7 +1038,12 @@ def __generate_protein(
) -> ESMProtein | ESMProteinError:
request = self._process_generate_protein_request(input, config, self.model)
try:
- data = self._post("generate", request, input.potential_sequence_of_concern)
+ data = self._post(
+ "generate",
+ request,
+ input.potential_sequence_of_concern,
+ headers={"Accept": ACCEPT_ES_PICKLE},
+ )
except ESMProteinError as e:
return e
@@ -1014,7 +1058,10 @@ def __generate_protein_tensor(
try:
data = self._post(
- "generate_tensor", request, input.potential_sequence_of_concern
+ "generate_tensor",
+ request,
+ input.potential_sequence_of_concern,
+ headers={"Accept": ACCEPT_ES_PICKLE},
)
except ESMProteinError as e:
return e
@@ -1033,9 +1080,7 @@ async def async_forward_and_sample(
"forward_and_sample",
request,
input.potential_sequence_of_concern,
- headers={
- "Accept": f"{MIMETYPE_ES_PICKLE};protocol={pickle.HIGHEST_PROTOCOL}, application/json"
- },
+ headers={"Accept": ACCEPT_ES_PICKLE},
)
except ESMProteinError as e:
return e
@@ -1054,9 +1099,7 @@ def forward_and_sample(
"forward_and_sample",
request,
input.potential_sequence_of_concern,
- headers={
- "Accept": f"{MIMETYPE_ES_PICKLE};protocol={pickle.HIGHEST_PROTOCOL}, application/json"
- },
+ headers={"Accept": ACCEPT_ES_PICKLE},
)
except ESMProteinError as e:
return e
@@ -1071,7 +1114,10 @@ async def async_encode(
try:
data = await self._async_post(
- "encode", request, input.potential_sequence_of_concern
+ "encode",
+ request,
+ input.potential_sequence_of_concern,
+ headers={"Accept": ACCEPT_ES_PICKLE},
)
except ESMProteinError as e:
return e
@@ -1083,7 +1129,12 @@ def encode(self, input: ESMProtein) -> ESMProteinTensor | ESMProteinError:
request = self._process_encode_request(input, self.model)
try:
- data = self._post("encode", request, input.potential_sequence_of_concern)
+ data = self._post(
+ "encode",
+ request,
+ input.potential_sequence_of_concern,
+ headers={"Accept": ACCEPT_ES_PICKLE},
+ )
except ESMProteinError as e:
return e
@@ -1097,7 +1148,10 @@ async def async_decode(
try:
data = await self._async_post(
- "decode", request, input.potential_sequence_of_concern
+ "decode",
+ request,
+ input.potential_sequence_of_concern,
+ headers={"Accept": ACCEPT_ES_PICKLE},
)
except ESMProteinError as e:
return e
@@ -1109,7 +1163,12 @@ def decode(self, input: ESMProteinTensor) -> ESMProtein | ESMProteinError:
request = self._process_decode_request(input, self.model)
try:
- data = self._post("decode", request, input.potential_sequence_of_concern)
+ data = self._post(
+ "decode",
+ request,
+ input.potential_sequence_of_concern,
+ headers={"Accept": ACCEPT_ES_PICKLE},
+ )
except ESMProteinError as e:
return e
@@ -1130,9 +1189,7 @@ async def async_logits(
request,
input.potential_sequence_of_concern,
return_bytes=return_bytes,
- headers={
- "Accept": f"{MIMETYPE_ES_PICKLE};protocol={pickle.HIGHEST_PROTOCOL}, application/json"
- },
+ headers={"Accept": ACCEPT_ES_PICKLE},
)
except ESMProteinError as e:
return e
@@ -1154,9 +1211,7 @@ def logits(
request,
input.potential_sequence_of_concern,
return_bytes=return_bytes,
- headers={
- "Accept": f"{MIMETYPE_ES_PICKLE};protocol={pickle.HIGHEST_PROTOCOL}, application/json"
- },
+ headers={"Accept": ACCEPT_ES_PICKLE},
)
except ESMProteinError as e:
return e
@@ -1176,10 +1231,10 @@ def __init__(
model: str,
url: str = "https://biohub.ai",
token: str = "",
- request_timeout: int | None = 60,
+ request_timeout: int | None = 125,
min_retry_wait: int = 1,
max_retry_wait: int = 10,
- max_retry_attempts: int = 5,
+ max_retry_attempts: int = 3,
):
ESMCInferenceClient.__init__(self)
_BaseForgeInferenceClient.__init__(
@@ -1269,7 +1324,10 @@ async def async_encode(
try:
data = await self._async_post(
- "encode", request, input.potential_sequence_of_concern
+ "encode",
+ request,
+ input.potential_sequence_of_concern,
+ headers={"Accept": ACCEPT_ES_PICKLE},
)
except ESMProteinError as e:
return e
@@ -1287,7 +1345,12 @@ def encode(self, input: ESMProtein) -> ESMProteinTensor | ESMProteinError:
request = {"inputs": tracks, "model": self.model}
try:
- data = self._post("encode", request, input.potential_sequence_of_concern)
+ data = self._post(
+ "encode",
+ request,
+ input.potential_sequence_of_concern,
+ headers={"Accept": ACCEPT_ES_PICKLE},
+ )
except ESMProteinError as e:
return e
@@ -1309,7 +1372,10 @@ async def async_decode(
try:
data = await self._async_post(
- "decode", request, input.potential_sequence_of_concern
+ "decode",
+ request,
+ input.potential_sequence_of_concern,
+ headers={"Accept": ACCEPT_ES_PICKLE},
)
except ESMProteinError as e:
return e
@@ -1326,7 +1392,12 @@ def decode(self, input: ESMProteinTensor) -> ESMProtein | ESMProteinError:
request = {"model": self.model, "inputs": tokens}
try:
- data = self._post("decode", request, input.potential_sequence_of_concern)
+ data = self._post(
+ "decode",
+ request,
+ input.potential_sequence_of_concern,
+ headers={"Accept": ACCEPT_ES_PICKLE},
+ )
except ESMProteinError as e:
return e
@@ -1346,9 +1417,7 @@ async def async_logits(
request,
input.potential_sequence_of_concern,
return_bytes=return_bytes,
- headers={
- "Accept": f"{MIMETYPE_ES_PICKLE};protocol={pickle.HIGHEST_PROTOCOL}, application/json"
- },
+ headers={"Accept": ACCEPT_ES_PICKLE},
)
except ESMProteinError as e:
return e
@@ -1369,9 +1438,7 @@ def logits(
request,
input.potential_sequence_of_concern,
return_bytes=return_bytes,
- headers={
- "Accept": f"{MIMETYPE_ES_PICKLE};protocol={pickle.HIGHEST_PROTOCOL}, application/json"
- },
+ headers={"Accept": ACCEPT_ES_PICKLE},
)
except ESMProteinError as e:
return e
@@ -1385,9 +1452,12 @@ def raw_model(self):
)
-class FoldMaxAccuracyHandler(EndpointHandler[MolecularComplexResult]):
+class FoldMaxAccuracyHandler(
+ EndpointHandler[FoldMaxAccuracyRequest, MolecularComplexResult]
+):
"""Fold a molecular complex at the highest-accuracy settings."""
+ request_type = FoldMaxAccuracyRequest
poll_interval = 10
poll_max_interval = 60
default_timeout = 2 * 60 * 60 # 2 hours
@@ -1400,37 +1470,123 @@ def endpoint_name(self) -> str:
return "fold_max_accuracy"
def _prepare_request(
- self,
- all_atom_input: StructurePredictionInput,
- config: FoldMaxAccuracyConfig = FoldMaxAccuracyConfig(),
- ) -> list[dict[str, Any]]:
- validate_fold_max_accuracy_input(all_atom_input)
- return [
- {
- "all_atom_input": serialize_structure_prediction_input(all_atom_input),
- "model": ESMFOLD2_CUTOFF_2025,
- **attr.asdict(config),
- }
- ]
+ self, request: FoldMaxAccuracyRequest, *, model_name: str | None = None
+ ) -> dict[str, Any]:
+ if model_name is not None:
+ raise TypeError(
+ "fold_max_accuracy selects its own model; model_name= is not supported"
+ )
+ config = (
+ request.config if request.config is not None else FoldMaxAccuracyConfig()
+ )
+ validate_fold_max_accuracy_input(request.all_atom_input)
+ return {
+ "all_atom_input": serialize_structure_prediction_input(
+ request.all_atom_input
+ ),
+ "model": ESMFOLD2_CUTOFF_2025,
+ **attr.asdict(config),
+ }
+
+ def _process_response(self, response: dict) -> MolecularComplexResult:
+ result = self._batch_client._read_response(response)
+ return SequenceStructureForgeInferenceClient._process_fold_all_atom_response(
+ result
+ )
+
+ async def _async_process_response(self, response: dict) -> MolecularComplexResult:
+ result = await self._batch_client._async_read_response(response)
+ return SequenceStructureForgeInferenceClient._process_fold_all_atom_response(
+ result
+ )
+
+
+def _require_model(endpoint: str, model: str | None) -> str:
+ # A null model fails the submit route's body parse with an unhelpful message.
+ if model is None or not model.strip():
+ raise ValueError(
+ f"{endpoint} requires a non-empty model. Pass model= when constructing "
+ "ForgeBatchClient, or model_name= to submit()/run()."
+ )
+ return model
+
+
+class FoldHandler(EndpointHandler[FoldRequest, ESMProtein]):
+ request_type = FoldRequest
+
+ def __init__(self, batch_client: _BaseForgeBatchClient, model: str | None = None):
+ super().__init__(batch_client)
+ self.model = model
+
+ @property
+ def endpoint_name(self) -> str:
+ return "fold"
+
+ def _prepare_request(
+ self, request: FoldRequest, *, model_name: str | None = None
+ ) -> dict[str, Any]:
+ model = model_name if model_name is not None else self.model
+ return SequenceStructureForgeInferenceClient._process_fold_request(
+ request.sequence,
+ request.msa,
+ request.config if request.config is not None else FoldingConfig(),
+ request.target_structure,
+ _require_model(self.endpoint_name, model),
+ )
+
+ def _process_response(self, response: dict) -> ESMProtein:
+ result = self._batch_client._read_response(response)
+ return SequenceStructureForgeInferenceClient._process_fold_response(result)
+
+ async def _async_process_response(self, response: dict) -> ESMProtein:
+ result = await self._batch_client._async_read_response(response)
+ return SequenceStructureForgeInferenceClient._process_fold_response(result)
+
+
+class FoldAllAtomHandler(EndpointHandler[FoldAllAtomRequest, MolecularComplexResult]):
+ request_type = FoldAllAtomRequest
- def _process_response(self, response: dict, **kwargs) -> MolecularComplexResult:
- s3_url = response["response"]
- result = self._batch_client.get_result_from_s3(s3_url)
+ def __init__(self, batch_client: _BaseForgeBatchClient, model: str | None = None):
+ super().__init__(batch_client)
+ self.model = model
+
+ @property
+ def endpoint_name(self) -> str:
+ return "fold_all_atom"
+
+ def _prepare_request(
+ self, request: FoldAllAtomRequest, *, model_name: str | None = None
+ ) -> dict[str, Any]:
+ model = model_name if model_name is not None else self.model
+ return SequenceStructureForgeInferenceClient._process_fold_all_atom_request(
+ request.all_atom_input,
+ request.config if request.config is not None else FoldingConfig(),
+ _require_model(self.endpoint_name, model),
+ )
+
+ def _process_response(self, response: dict) -> MolecularComplexResult:
+ result = self._batch_client._read_response(response)
+ # Use the same logic as _process_fold_all_atom_response
return SequenceStructureForgeInferenceClient._process_fold_all_atom_response(
result
)
- async def _async_process_response(
- self, response: dict, **kwargs
- ) -> MolecularComplexResult:
- s3_url = response["response"]
- result = await self._batch_client.async_get_result_from_s3(s3_url)
+ async def _async_process_response(self, response: dict) -> MolecularComplexResult:
+ result = await self._batch_client._async_read_response(response)
+ # Use the same logic as _process_fold_all_atom_response
return SequenceStructureForgeInferenceClient._process_fold_all_atom_response(
result
)
class ForgeBatchClient:
+ """Submit one or many tasks and recover them through saved job identifiers.
+
+ Endpoint handlers accept SDK inputs and return decoded Python objects.
+ This client provides shared status, waiting, listing, and cancellation
+ operations. Each method has an ``async_`` equivalent.
+ """
+
def __init__(
self,
url: str = "https://biohub.ai",
@@ -1441,7 +1597,7 @@ def __init__(
poll_max_interval: int | None = None,
min_retry_wait: int = 2,
max_retry_wait: int = 2,
- max_retry_attempts: int = 5,
+ max_retry_attempts: int = 3,
transfer_timeout: int | None = 60,
):
self._batch_client = _BaseForgeBatchClient(
@@ -1458,9 +1614,154 @@ def __init__(
self.model = model
self._fold_max_accuracy: FoldMaxAccuracyHandler | None = None
+ self._fold: FoldHandler | None = None
+ self._fold_all_atom: FoldAllAtomHandler | None = None
+
+ def close(self) -> None:
+ self._batch_client.close()
+
+ async def aclose(self) -> None:
+ await self._batch_client.aclose()
+ self._batch_client.close()
+
+ def __enter__(self) -> ForgeBatchClient:
+ return self
+
+ def __exit__(self, exc_type, exc_value, traceback) -> None:
+ self.close()
+
+ async def __aenter__(self) -> ForgeBatchClient:
+ return self
+
+ async def __aexit__(self, exc_type, exc_value, traceback) -> None:
+ await self.aclose()
+
+ def get_status(self, task_id: str) -> dict[str, Any]:
+ """Read task state, result location, errors, billing, and available artifacts."""
+ return self._batch_client.get_status(task_id)
+
+ async def async_get_status(self, task_id: str) -> dict[str, Any]:
+ return await self._batch_client.async_get_status(task_id)
+
+ def wait(
+ self,
+ submission: BatchSubmission,
+ timeout: int = 1800,
+ *,
+ poll_interval: int | None = None,
+ poll_max_interval: int | None = None,
+ ) -> dict[str, Any] | list[dict[str, Any]]:
+ """Wait for a batch, returning its aggregate summary.
+
+ A polling timeout raises ``BatchTimeoutError`` with the submission and leaves
+ work running. The budget excludes submission and result downloads.
+ """
+ return self._batch_client.wait(
+ submission, timeout, poll_interval, poll_max_interval
+ )
+
+ async def async_wait(
+ self,
+ submission: BatchSubmission,
+ timeout: int = 1800,
+ *,
+ poll_interval: int | None = None,
+ poll_max_interval: int | None = None,
+ ) -> dict[str, Any] | list[dict[str, Any]]:
+ return await self._batch_client.async_wait(
+ submission, timeout, poll_interval, poll_max_interval
+ )
+
+ def cancel(
+ self, task_id: str | None = None, *, batch_id: str | None = None
+ ) -> dict[str, Any]:
+ """Cancel one task or queued work in a batch; running tasks may still finish."""
+ return self._batch_client.cancel(task_id, batch_id=batch_id)
+
+ async def async_cancel(
+ self, task_id: str | None = None, *, batch_id: str | None = None
+ ) -> dict[str, Any]:
+ return await self._batch_client.async_cancel(task_id, batch_id=batch_id)
+
+ def list_batch(
+ self,
+ batch_id: str,
+ *,
+ status: list[str] | None = None,
+ include_responses: bool = False,
+ limit: int | None = None,
+ cursor: str | None = None,
+ ) -> dict[str, Any]:
+ """Read one batch's progress and, optionally, a page of its tasks.
+
+ Parameters
+ ----------
+ batch_id
+ Saved batch ID returned by submission.
+ status
+ Task states to include, such as ``["done", "failed"]``. Only filters
+ the item page; aggregate counts always describe the whole batch.
+ include_responses
+ Include task entries and pagination information. S3 results are
+ marked ``response_available``; this call does not download them.
+ limit
+ Page size. Omit to use the server default.
+ cursor
+ Opaque cursor from the previous page. Use the same filters when
+ requesting subsequent pages.
+
+ Returns
+ -------
+ dict[str, Any]
+ ``batch_id``, ``endpoint``, ``status``, ``counts``, and ``total``.
+ With ``include_responses=True``, also ``items`` and ``cursor``;
+ a null cursor marks the last page. Item entries contain task IDs,
+ states, and available inline responses or errors.
+
+ Notes
+ -----
+ HTTP errors are raised as ``ESMProteinError``.
+ """
+ return self._batch_client.list_batch(
+ batch_id,
+ status=status,
+ include_responses=include_responses,
+ limit=limit,
+ cursor=cursor,
+ )
+
+ async def async_list_batch(
+ self,
+ batch_id: str,
+ *,
+ status: list[str] | None = None,
+ include_responses: bool = False,
+ limit: int | None = None,
+ cursor: str | None = None,
+ ) -> dict[str, Any]:
+ """Asynchronous equivalent of :meth:`list_batch`, with the same options."""
+ return await self._batch_client.async_list_batch(
+ batch_id,
+ status=status,
+ include_responses=include_responses,
+ limit=limit,
+ cursor=cursor,
+ )
@property
def fold_max_accuracy(self) -> FoldMaxAccuracyHandler:
if self._fold_max_accuracy is None:
self._fold_max_accuracy = FoldMaxAccuracyHandler(self._batch_client)
return self._fold_max_accuracy
+
+ @property
+ def fold(self) -> FoldHandler:
+ if self._fold is None:
+ self._fold = FoldHandler(self._batch_client, self.model)
+ return self._fold
+
+ @property
+ def fold_all_atom(self) -> FoldAllAtomHandler:
+ if self._fold_all_atom is None:
+ self._fold_all_atom = FoldAllAtomHandler(self._batch_client, self.model)
+ return self._fold_all_atom
diff --git a/esm/sdk/retry.py b/esm/sdk/retry.py
index aa2fe287..5b726d44 100644
--- a/esm/sdk/retry.py
+++ b/esm/sdk/retry.py
@@ -11,14 +11,15 @@
)
from esm.sdk.api import ESMProteinError
+from esm.sdk.batch import BatchError
skip_retries_var = ContextVar("skip_retries", default=False)
-def retry_if_specific_error(exception):
- """
- We only retry on specific errors.
- """
+def retry_if_specific_error(exception: object) -> bool:
+ """Retry transient request errors"""
+ if isinstance(exception, BatchError) and exception.task_status is not None:
+ return False
return isinstance(exception, ESMProteinError) and exception.error_code in {
429, # Too Many Requests
500, # Internal Server Error
diff --git a/esm/tokenization/sequence_tokenizer.py b/esm/tokenization/sequence_tokenizer.py
index 5ed59345..0166004e 100644
--- a/esm/tokenization/sequence_tokenizer.py
+++ b/esm/tokenization/sequence_tokenizer.py
@@ -1,7 +1,7 @@
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.processors import TemplateProcessing
-from transformers import PreTrainedTokenizerFast
+from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
from esm.tokenization.tokenizer_base import EsmTokenizerBase
from esm.utils.constants import esm3 as C
diff --git a/esm/utils/constants/api.py b/esm/utils/constants/api.py
index 6a6f0c34..a0714ffd 100644
--- a/esm/utils/constants/api.py
+++ b/esm/utils/constants/api.py
@@ -1,6 +1,11 @@
+import pickle
+
MAX_TOPK_SEQUENCE = 32
MAX_TOPK_STRUCTURE = MAX_TOPK_SEQUENCE
MAX_TOPK_SECONDARY_STRUCTURE = MAX_TOPK_SEQUENCE
MAX_TOPK_SASA = MAX_TOPK_SEQUENCE
MAX_TOPK_FUNCTION = MAX_TOPK_SEQUENCE
MIMETYPE_ES_PICKLE = "application/vnd.evolutionaryscale.pickle"
+ACCEPT_ES_PICKLE = (
+ f"{MIMETYPE_ES_PICKLE};protocol={pickle.HIGHEST_PROTOCOL}, application/json"
+)
diff --git a/esm/utils/generation.py b/esm/utils/generation.py
index 68b690ac..9ff799a4 100644
--- a/esm/utils/generation.py
+++ b/esm/utils/generation.py
@@ -3,6 +3,7 @@
from warnings import warn
import attr
+import numpy as np
import torch
from tqdm import tqdm
@@ -123,6 +124,20 @@ def iterative_sampling_raw(
# Function and residue annotation encoding/decoding is lossy
# There is no guarantee that decoding encoded tokens will yield the same input
raw_protein.function_annotations = input_protein.function_annotations
+ if input_protein.coordinates is not None:
+ input_protein_chain = input_protein.to_protein_chain()
+ raw_protein_chain = raw_protein.to_protein_chain()
+ has_coords = np.where(
+ input_protein.coordinates[:, :3].isfinite().all(-1).any(-1)
+ )[0]
+ raw_protein.crmsd = torch.tensor(
+ input_protein_chain.rmsd(
+ raw_protein_chain,
+ only_compute_backbone_rmsd=config.only_compute_backbone_rmsd,
+ mobile_inds=has_coords,
+ target_inds=has_coords,
+ )
+ )
return raw_proteins
diff --git a/esm/utils/structure/input_builder.py b/esm/utils/structure/input_builder.py
index 214f2981..96c21649 100644
--- a/esm/utils/structure/input_builder.py
+++ b/esm/utils/structure/input_builder.py
@@ -73,7 +73,8 @@ class PocketConditioning:
@dataclass
class CovalentBond:
"""``res_idx`` and ``atom_idx`` are 0-based. Use care when working with PDB
- files residue numbers that are 1-based.
+ files residue numbers that are 1-based. ``atom_idx`` counts the residue's heavy
+ atoms without CCD leaving atoms (a polymer's last modified residue keeps them).
"""
chain_id1: str
@@ -300,9 +301,10 @@ def _msa(chain: dict[str, Any]) -> MSAInput:
for b in bonds_blk
]
- return StructurePredictionInput(
+ result = StructurePredictionInput(
sequences=sequences,
pocket=pocket,
distogram_conditioning=distogram_conditioning,
covalent_bonds=covalent_bonds,
)
+ return result
diff --git a/esm/utils/structure/molecular_complex.py b/esm/utils/structure/molecular_complex.py
index 610939cb..8520edf7 100644
--- a/esm/utils/structure/molecular_complex.py
+++ b/esm/utils/structure/molecular_complex.py
@@ -1,8 +1,10 @@
from __future__ import annotations
import io
+import itertools
import os
import re
+from collections.abc import Iterator
from dataclasses import asdict, dataclass
from pathlib import Path
from subprocess import check_output
@@ -50,6 +52,18 @@ class MolecularComplexResult:
num_tokens: int | None = None
+def _unused_chain_labels(taken: set[str]) -> Iterator[str]:
+ """Bijective base-26 A..Z, AA, AB, skipping *taken*; ``chr(65 + n)`` emits
+ punctuation past Z, and reusing a taken label would merge two chains."""
+ for index in itertools.count():
+ label, remaining = "", index
+ while remaining >= 0:
+ label = chr(65 + remaining % 26) + label
+ remaining = remaining // 26 - 1
+ if label not in taken:
+ yield label
+
+
@dataclass
class MolecularComplexMetadata:
"""Metadata for MolecularComplex objects."""
@@ -418,10 +432,8 @@ def to_protein_complex(self) -> ProteinComplex:
# Convert MolecularComplex metadata to ProteinComplex format
unique_chain_ids = np.unique(protein_chain_ids)
entity_lookup = {int(cid): int(cid) for cid in unique_chain_ids}
- chain_lookup = {
- int(cid): self.metadata.chain_lookup.get(int(cid), chr(65 + int(cid)))
- for cid in unique_chain_ids
- }
+ chain_names = self._resolved_chain_names(name_null_chains=False)
+ chain_lookup = {int(cid): chain_names[int(cid)] for cid in unique_chain_ids}
protein_metadata = ProteinComplexMetadata(
entity_lookup=entity_lookup,
@@ -476,61 +488,44 @@ def from_cif_file(
id: str = "complex_from_string",
) -> "MolecularComplex":
"""Read a parsed CIF file using the same semantics as :meth:`from_mmcif`."""
- # Get structure - handle missing model information gracefully
- try:
- structure = pdbx.get_structure(
- mmcif_file, model=1, extra_fields=["b_factor"]
- )
- except (KeyError, ValueError):
- # Fallback for mmCIF files without model information
- try:
- structure = pdbx.get_structure(mmcif_file)
- except Exception:
- # Last resort: use the first available model or all atoms
- structure = pdbx.get_structure(mmcif_file, model=None)
+ block = mmcif_file.block
+ atom_site = block.get("atom_site")
+ # Load every alt-loc and apply biotite's first-alt-loc filter ourselves,
+ # so the identical mask can be applied to label_asym_id below. A
+ # residue may have only alt-loc "B" (or "L", "1", ...); any separate
+ # alt-loc rule would desynchronise the two arrays. altloc="all"
+ # requires the label_alt_id column.
+ has_altlocs = atom_site is not None and "label_alt_id" in atom_site
+ structure = pdbx.get_structure(
+ mmcif_file,
+ model=1,
+ extra_fields=["b_factor"],
+ altloc="all" if has_altlocs else "first",
+ )
# Type hint for pyright - structure is an AtomArray which is iterable
if TYPE_CHECKING:
structure: Any = structure
- # Read label_asym_id from the raw CIF atom_site category.
# Biotite's atom.chain_id uses auth_asym_id, which collapses ligands
# onto their parent protein chain. label_asym_id gives each entity a
# distinct chain identifier.
- block = mmcif_file.block
- label_asym_ids: list[str] | None = None
- if "atom_site" in block:
- atom_site = block["atom_site"]
- if "label_asym_id" in atom_site:
- _col = atom_site["label_asym_id"]
- _raw = (
- _col.as_array(str)
- if hasattr(_col, "as_array")
- else np.array(list(_col), dtype=str)
+ label_asym_ids: np.ndarray | None = None
+ if atom_site is not None and "label_asym_id" in atom_site:
+ # get_structure(model=1) keeps the first model's rows.
+ models = atom_site["pdbx_PDB_model_num"].as_array(str)
+ label_asym_ids = atom_site["label_asym_id"].as_array(str)[
+ models == models[0]
+ ]
+ if len(label_asym_ids) != len(structure):
+ raise ValueError(
+ f"{id}: {len(label_asym_ids)} model-1 label_asym_id rows but "
+ f"biotite loaded {len(structure)} atoms; cannot align chains."
)
- # biotite's get_structure(model=1) filters to model 1 AND
- # removes alternate conformations. We must apply the same
- # filters to label_asym_id to keep arrays aligned.
- keep = np.ones(len(_raw), dtype=bool)
- if "pdbx_PDB_model_num" in atom_site:
- _mc = atom_site["pdbx_PDB_model_num"]
- _models = (
- _mc.as_array(str)
- if hasattr(_mc, "as_array")
- else np.array(list(_mc), dtype=str)
- )
- keep &= _models == "1"
- if "label_alt_id" in atom_site:
- _ac = atom_site["label_alt_id"]
- _alts = (
- _ac.as_array(str)
- if hasattr(_ac, "as_array")
- else np.array(list(_ac), dtype=str)
- )
- keep &= np.isin(_alts, [".", "?", "", "A"])
- filtered = _raw[keep]
- if len(filtered) == len(structure):
- label_asym_ids = filtered.tolist()
- # If lengths still don't match, fall back to atom.chain_id
+ if has_altlocs:
+ first_altloc = bs.filter_first_altloc(structure, structure.altloc_id)
+ structure = structure[first_altloc]
+ if label_asym_ids is not None:
+ label_asym_ids = label_asym_ids[first_altloc]
# Get entity information from mmCIF
entity_info = {}
@@ -553,7 +548,9 @@ def from_cif_file(
pass
chain_labels = (
- label_asym_ids if label_asym_ids is not None else structure.chain_id
+ label_asym_ids.tolist()
+ if label_asym_ids is not None
+ else structure.chain_id
)
# Include water-only chains so numeric chain IDs remain unchanged.
_, chain_first, chain_numbers = np.unique(
@@ -561,16 +558,22 @@ def from_cif_file(
)
chain_lookup = {i: chain_labels[first] for i, first in enumerate(chain_first)}
names = structure.res_name
- residues = structure.res_id
+ # Residues are ordered by first appearance in the file rather than by
+ # res_id, so insertion-code residues (52, 52A, 52B, ...) keep file order.
+ residue_keys = np.rec.fromarrays(
+ (chain_numbers, structure.res_id, structure.ins_code)
+ )
+ _, residue_first, residue_inverse = np.unique(
+ residue_keys, return_index=True, return_inverse=True
+ )
+ residue_rank = residue_first[residue_inverse]
# Stable sorting preserves atom order within each (chain, residue, name).
- order = np.lexsort((names, residues, chain_numbers))
+ order = np.lexsort((names, residue_rank, chain_numbers))
order = order[names[order] != "HOH"]
if len(order):
boundary = np.ones(len(order), dtype=bool)
- boundary[1:] = (
- (chain_numbers[order[1:]] != chain_numbers[order[:-1]])
- | (residues[order[1:]] != residues[order[:-1]])
- | (names[order[1:]] != names[order[:-1]])
+ boundary[1:] = (residue_rank[order[1:]] != residue_rank[order[:-1]]) | (
+ names[order[1:]] != names[order[:-1]]
)
starts = np.flatnonzero(boundary)
first_atoms = order[starts]
@@ -628,7 +631,8 @@ def from_atomarray(
This is the inverse of :meth:`to_mmcif` (which expands the flat
token/atom arrays into an ``AtomArray``): it groups the AtomArray's atoms
- into tokens by ``(chain_id, res_id)`` in array order and rebuilds the flat
+ into tokens by ``(chain_id, res_id, ins_code)`` in array order and
+ rebuilds the flat
``atom_positions`` / ``atom_elements`` / ``atom_names`` / ``atom_hetero``
arrays plus the token-level ``token_to_atoms`` / ``chain_id`` /
``sequence`` / ``plddt`` arrays. Unlike ``ProteinChain.from_atomarray``
@@ -647,6 +651,7 @@ def from_atomarray(
coords = np.asarray(atom_arr.coord, dtype=np.float32)
chain_labels = np.asarray(atom_arr.chain_id, dtype=object)
res_ids = np.asarray(atom_arr.res_id)
+ ins_codes = np.asarray(atom_arr.ins_code, dtype=object)
res_names = np.asarray(atom_arr.res_name, dtype=object)
atom_names = np.asarray(atom_arr.atom_name, dtype=object)
elements = np.asarray(atom_arr.element, dtype=object)
@@ -682,9 +687,9 @@ def _close_token(start: int, end: int) -> None:
token_plddt.append(float(b_factor[start:end].mean()) / PLDDT_B_FACTOR_SCALE)
token_start = 0
- prev_key: tuple[Any, int] | None = None
+ prev_key: tuple[Any, int, str] | None = None
for i in range(n_atoms):
- key = (chain_labels[i], int(res_ids[i]))
+ key = (chain_labels[i], int(res_ids[i]), ins_codes[i])
if i > 0 and key != prev_key:
_close_token(token_start, i)
token_start = i
@@ -710,20 +715,34 @@ def _close_token(start: int, end: int) -> None:
atom_hetero=hetero,
)
- def _get_entity_mapping(
- self,
- ) -> tuple[dict[str, list[str]], dict[str, int], dict[int, tuple[str, ...]]]:
- """Compute chain→sequence, chain→entity_id, and entity_id→sequence mappings.
+ def _resolved_chain_names(self, *, name_null_chains: bool) -> dict[int, str]:
+ """Chain id → name, naming unnamed chains from the labels no chain already uses.
- Returns:
- (chain_sequences, chain_to_entity, entity_sequences)
+ Fallback labels are positional, so they are not stable across subsets of the
+ same complex. ``name_null_chains`` counts a blank name as unnamed.
"""
+ lookup = self.metadata.chain_lookup
+
+ def is_named(name: str) -> bool:
+ return not (name_null_chains and not name.strip())
+
+ # Reserve names of chains absent from this file too, so a fallback cannot
+ # collide with one; only present chains consume a label.
+ fallbacks = _unused_chain_labels({n for n in lookup.values() if is_named(n)})
+ return {
+ cid: lookup[cid]
+ if cid in lookup and is_named(lookup[cid])
+ else next(fallbacks)
+ for cid in (int(c) for c in np.unique(self.chain_id))
+ }
+
+ def _get_entity_mapping(
+ self, chain_names: dict[int, str]
+ ) -> tuple[dict[str, int], dict[int, tuple[str, ...]]]:
+ """Compute chain→entity_id and entity_id→sequence mappings."""
chain_sequences: dict[str, list[str]] = {}
for token_idx in range(len(self.token_to_atoms)):
- chain_id_numeric = self.chain_id[token_idx]
- chain_id_str = self.metadata.chain_lookup.get(
- int(chain_id_numeric), chr(65 + int(chain_id_numeric))
- )
+ chain_id_str = chain_names[int(self.chain_id[token_idx])]
if chain_id_str not in chain_sequences:
chain_sequences[chain_id_str] = []
chain_sequences[chain_id_str].append(self.sequence[token_idx])
@@ -740,10 +759,13 @@ def _get_entity_mapping(
entity_id_counter += 1
chain_to_entity[chain_id_str] = sequence_to_entity[seq_tuple]
- return chain_sequences, chain_to_entity, entity_sequences
+ return chain_to_entity, entity_sequences
def _add_entity_information(
- self, cif_file: CIFFile, entity_sequences: dict[int, tuple[str, ...]]
+ self,
+ cif_file: CIFFile,
+ entity_sequences: dict[int, tuple[str, ...]],
+ chain_to_entity: dict[str, int],
) -> None:
"""Add _entity category to CIF file so OST can identify ligands vs polymers."""
@@ -784,7 +806,6 @@ def _add_entity_information(
)
# Add _struct_asym to map chain IDs to entity IDs
- _, chain_to_entity, _ = self._get_entity_mapping()
if chain_to_entity:
asym_ids = sorted(chain_to_entity.keys())
asym_entity_ids = [str(chain_to_entity[c]) for c in asym_ids]
@@ -905,9 +926,12 @@ def _add_entity_information(
},
)
- def to_mmcif(self) -> str:
+ def to_mmcif(self, *, name_null_chains: bool = False) -> str:
"""Write MolecularComplex to mmcif string using biotite.
+ ``name_null_chains`` is opt-in because naming blank chains changes which
+ rows join downstream in the OST evals.
+
Returns:
String representation of the complex in mmCIF format
"""
@@ -927,7 +951,8 @@ def to_mmcif(self) -> str:
atom_names = np.empty(n_atoms, dtype=object)
# Build entity mappings: chains with identical sequences share entity ID
- _, chain_to_entity, entity_sequences = self._get_entity_mapping()
+ chain_names = self._resolved_chain_names(name_null_chains=name_null_chains)
+ chain_to_entity, entity_sequences = self._get_entity_mapping(chain_names)
atom_entity_ids = np.zeros(n_atoms, dtype=np.int32)
@@ -938,9 +963,7 @@ def to_mmcif(self) -> str:
for token_idx, (start, end) in enumerate(self.token_to_atoms):
token = self.sequence[token_idx]
chain_id_numeric = self.chain_id[token_idx]
- chain_id_str = self.metadata.chain_lookup.get(
- int(chain_id_numeric), chr(65 + int(chain_id_numeric))
- )
+ chain_id_str = chain_names[int(chain_id_numeric)]
# Track residue numbering per chain
if chain_id_numeric not in chain_res_counters:
@@ -1031,7 +1054,7 @@ def to_mmcif(self) -> str:
)
# Add _entity category for OST compatibility
- self._add_entity_information(cif_file, entity_sequences)
+ self._add_entity_information(cif_file, entity_sequences, chain_to_entity)
# biotite echoes unmasked float columns at full precision
# so we round every float column to conventional mmCIF precision
diff --git a/esm/utils/structure/protein_chain.py b/esm/utils/structure/protein_chain.py
index 83b0435d..48520db3 100644
--- a/esm/utils/structure/protein_chain.py
+++ b/esm/utils/structure/protein_chain.py
@@ -323,6 +323,19 @@ def to_pdb_string(self, include_insertions: bool = True) -> str:
def to_mmcif(self, path: PathOrBuffer):
f = CIFFile()
set_structure_pdbx(f, self.atom_array, data_block=self.id)
+ entity_id = str(self.entity_id if self.entity_id is not None else 1)
+ atom_site = f.block["atom_site"]
+ atom_site["label_entity_id"] = CIFColumn(
+ np.full(atom_site.row_count, entity_id)
+ )
+ asym_ids = np.unique(atom_site["label_asym_id"].as_array(str))
+ f.block["struct_asym"] = CIFCategory(
+ {"id": asym_ids, "entity_id": np.full(len(asym_ids), entity_id)}
+ )
+ f.block["entity"] = CIFCategory({"id": [entity_id], "type": ["polymer"]})
+ f.block["entity_poly"] = CIFCategory(
+ {"entity_id": [entity_id], "type": ["polypeptide(L)"]}
+ )
# incantations molstar needs to render pLDDT / confidence onto
# the structure with "alphafold-view"
@@ -486,9 +499,32 @@ def sasa(self, by_residue: bool = True):
return sasa_per_atom
def sap_score(self, aggregation: str = "atom") -> np.ndarray:
- """Computes per-atom SAP score.
- Can optionally aggregate by residue (by averaging over atoms. NOTE: this returns values only for residues that have coordinates!)
- or full-protein (sum of SAP score for atoms with SAP > 0, as in Lauer et al. 2011)."""
+ """Compute static spatial aggregation propensity (SAP) with a 5-Angstrom radius.
+
+ Each atom is a neighborhood center. Only side-chain SASA contributes,
+ normalized by the residue's fully exposed side-chain area in Ala-X-Ala
+ and weighted by its Black-Mould hydrophobicity from residue_constants.
+ This evaluates one conformation, without simulation averaging.
+
+ Parameters
+ ----------
+ aggregation : {"atom", "residue", "protein"}, default="atom"
+ Return per-atom scores, residue means (NaN for missing residues),
+ or the sum of positive atom scores over the entire chain. The latter
+ extends Lauer et al.'s CDR-only aggregation to all supplied atoms.
+
+ Notes
+ -----
+ The hydrophobicity table subtracts 0.5, leaving glycine at 0.001 rather
+ than exactly zero. SASA is computed using Biotite's default parameters.
+ Cached scores computed with the former reversed SASA mask must be
+ recomputed; both atom scores and their aggregations change.
+
+ References
+ ----------
+ Chennamsetty et al. (2009), doi:10.1073/pnas.0904191106.
+ Lauer et al. (2012), doi:10.1002/jps.22758.
+ """
sap_radius = 5.0
arr = self.atom_array_no_insertions
@@ -523,8 +559,9 @@ def sap_score(self, aggregation: str = "atom") -> np.ndarray:
assert len(res_hydrophobicity) == len(self)
# compute SAP score
- is_side_chain = ~bs.filter_peptide_backbone(arr)
- sasa_per_atom[is_side_chain] = 0
+ # Include backbone oxygens omitted by Biotite's peptide-backbone filter; retain all neighborhood centers.
+ is_backbone = np.isin(arr.atom_name, ["N", "CA", "C", "O", "OXT"])
+ sasa_per_atom[is_backbone] = 0
kdtree = KDTree(arr.coord)
neighbors = kdtree.query_ball_tree(kdtree, sap_radius, p=2.0)
sap_by_atom = np.zeros_like(sasa_per_atom)
diff --git a/esm/utils/structure/protein_chain_test.py b/esm/utils/structure/protein_chain_test.py
new file mode 100644
index 00000000..e7092243
--- /dev/null
+++ b/esm/utils/structure/protein_chain_test.py
@@ -0,0 +1,72 @@
+import biotite.structure as bs
+import numpy as np
+import pytest
+
+from esm.utils import residue_constants
+from esm.utils.structure.protein_chain import ProteinChain
+
+
+@pytest.fixture
+def sap_chain() -> ProteinChain:
+ positions = np.full((5, 37, 3), np.nan)
+ # Missing terminal/interior residues exercise residue aggregation locally.
+ for residue, atoms in {
+ 1: {"N": -5, "CA": -4, "C": -3, "O": -2, "OXT": -1, "CB": 0},
+ 3: {"N": 5, "CA": 6, "C": 7, "O": 8, "OXT": 9, "CB": 10, "CG": 11},
+ }.items():
+ for name, x in atoms.items():
+ positions[residue, residue_constants.atom_order[name]] = [x, 0, 0]
+ return ProteinChain.from_atom37(positions, sequence="AAADA")
+
+
+@pytest.mark.parametrize("atom_name", ["N", "CA", "C", "O", "OXT"])
+def test_sap_ignores_backbone_sasa(
+ sap_chain: ProteinChain, monkeypatch: pytest.MonkeyPatch, atom_name: str
+) -> None:
+ arr = sap_chain.atom_array_no_insertions
+ sasa = np.where(arr.atom_name == atom_name, 100.0, 0.0)
+ assert np.count_nonzero(sasa) == 2
+ monkeypatch.setattr(bs, "sasa", lambda _: sasa.copy())
+
+ np.testing.assert_array_equal(sap_chain.sap_score(), np.zeros(len(arr)))
+ assert sap_chain.sap_score(aggregation="protein") == 0.0
+
+
+@pytest.mark.parametrize("backbone_sasa", [0.0, 1000.0])
+def test_sap_sidechain_neighborhoods_and_aggregation(
+ sap_chain: ProteinChain, monkeypatch: pytest.MonkeyPatch, backbone_sasa: float
+) -> None:
+ arr = sap_chain.atom_array_no_insertions
+ sasa = np.full(len(arr), backbone_sasa)
+ sasa[(arr.res_name == "ALA") & (arr.atom_name == "CB")] = 32.39045
+ sasa[(arr.res_name == "ASP") & (arr.atom_name == "CB")] = 55.1045
+ sasa[(arr.res_name == "ASP") & (arr.atom_name == "CG")] = 22.0418
+ monkeypatch.setattr(bs, "sasa", lambda _: sasa.copy())
+
+ ala_score = 0.5 * 0.116
+ asp_cb_score = 0.5 * -0.472
+ asp_score = 0.7 * -0.472
+ expected = np.where(arr.res_name == "ALA", ala_score, asp_score)
+ # ASP N is exactly 5 A from both CB atoms, but outside the ASP CG neighborhood.
+ expected[(arr.res_name == "ASP") & (arr.atom_name == "N")] = (
+ ala_score + asp_cb_score
+ )
+ atom_scores = sap_chain.sap_score(aggregation="atom")
+ np.testing.assert_allclose(atom_scores, expected)
+ assert len(atom_scores) == len(arr)
+ assert np.any(atom_scores > 0) and np.any(atom_scores < 0)
+
+ np.testing.assert_allclose(
+ sap_chain.sap_score(aggregation="residue"),
+ [
+ np.nan,
+ ala_score,
+ np.nan,
+ (ala_score + asp_cb_score + 6 * asp_score) / 7,
+ np.nan,
+ ],
+ equal_nan=True,
+ )
+ protein_score = sap_chain.sap_score(aggregation="protein")
+ assert protein_score == pytest.approx(atom_scores[atom_scores > 0].sum())
+ assert protein_score == pytest.approx(6 * ala_score)
diff --git a/pixi.lock b/pixi.lock
index ab87343a..8fba0ff6 100644
--- a/pixi.lock
+++ b/pixi.lock
@@ -143,7 +143,6 @@ environments:
- conda: https://conda.anaconda.org/conda-forge/noarch/munkres-1.1.4-pyhd8ed1ab_1.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/packaging-26.3-pyhc364b38_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pip-26.2.1-pyh8b19718_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.11.3-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pycparser-3.0-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pygments-2.21.0-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pyparsing-3.3.2-pyhcf101f3_0.conda
@@ -165,18 +164,16 @@ environments:
- conda: https://conda.anaconda.org/conda-forge/noarch/wheel-0.48.0-pyhd8ed1ab_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/zipp-4.1.0-pyhcf101f3_0.conda
- pypi: .
+ - pypi: direct+https://github.com/evolutionaryscale/wheels/releases/download/py312-pt211-cu13-sm80-90/flash_attn-2.7.4.post1-cp312-cp312-linux_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/00/78/9cbcc1c073b9d4918e925af1a059762265dc65004e020511b2a06fbfd020/pydssp-0.9.1-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/03/b8/e484ef633af3887baeeb4b6ad12743363af7cce68ae51e938e00aaa0529d/transformers-4.57.6-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/17/66/1ef1960ffa751df4f6ed05a76ba677e331ce451902f319d14c3f8ae81215/rdkit-2026.3.5-cp312-cp312-manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/17/c2/613ca5f729b55b2f7a9c17a3cd59d3f053ee2816da2d19933acc6b8741ec/biopython-1.88-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/19/f2/5ba170c6195e95e42ab9e2fade98a70a078a3663af14440a67b457208041/cuequivariance_ops_torch_cu13-0.11.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/1a/c9/82638ef24d7877510f83baf821f5619a61b45568ce21c0a87a91576510aa/torch-2.11.0-cp312-cp312-manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/1d/04/cb4e8ba80ba81790d02dcc69d7529961140fc114abde927c99e9ba9c9457/biotite-1.7.1-cp312-cp312-manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/22/a6/858897256d0deac81a172289110f31629fc4cee19b6f01283303e18c8db3/ptyprocess-0.7.0-py2.py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/23/45/caa600acfab94560807a20a64b5830d2cd3c3202b7f1328644d70b7d6bd8/nvidia_ml_py-13.610.43-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/28/50/f203ff3a3ddfe19308efc83c5a3a29ed02bf786732ec35e68bf9162f3365/safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/2a/09/f8d8f8f31e4483c10a906437b4ce31bdf3d6d417b73fe33f1a8b59e34228/einops-0.8.2-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl
@@ -193,7 +190,6 @@ environments:
- pypi: https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/41/09/5b161152e2d90f7b87f781c2e1267494aef9c32498df793f73ad0a0a494a/matplotlib_inline-0.2.2-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/4e/d0/4b9da3d9c5749bf2c93f2818af4dc306f3b3ccc671f68f6820635fd53b74/cuequivariance_torch-0.11.1-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/52/89/d0baf1269e57ac548387c11e5d9ba78d1e118ceb43f1e5e5114bdd2e8391/boto3-1.43.75-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/52/f2/9e726e4489f947557290605e3ee33b6ba09fe5dac9213263e0ada75200cf/dna_features_viewer-3.1.5-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/54/6f/84908cad2d6aa5144abcf7b42709fe4fdb459bc640ec7ac5786e7693dabc/prompt_toolkit-3.0.53-py3-none-any.whl
@@ -206,7 +202,6 @@ environments:
- pypi: https://files.pythonhosted.org/packages/67/4e/a28359bf1c1ecf11eba22123168c138698f7cb576ac678f5a2e16cd5da08/hf_xet-1.6.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/6a/fd/6adabd4f6d5e686f97dd02ce7fce3fe4cf672cbac36b8f67ff4040e8ad8b/msgpack-1.2.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/6c/c2/63fdda36c56437eeb44aaf9493c8bcd62ce230ab1598924fc626ffbfa943/scikit_learn-1.9.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/72/f1/b1f191b91d049780eb4db2605b81af28a697710c4ac219708bcba192f80c/cuequivariance_ops_cu13-0.11.1-py3-none-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/7e/85/a5bfaebfd305ac18b57b0854d74e37e586809061a91fda62f0bd50c8518e/narwhals-2.24.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl
@@ -244,7 +239,6 @@ environments:
- pypi: https://files.pythonhosted.org/packages/da/35/f2287558c17e29fafc8ef3daf819bb9834061cfa43bff8014f7df7f63bdc/anyio-4.14.2-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/e3/1e/2942a0d6d52240d439f44384a40c65d359d9ca70325b65b265c020113121/cuda_pathfinder-1.6.1-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/e7/44/423ac00af4dd95a5aeb27207e2c0d9b7118702149bf4704c3ddb55bb7429/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/e7/a8/7f99cfa83f61d2d8cb29f31e1f90a6233873a6038b4cdfb30179ef90e5a6/cuequivariance-0.11.1-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/ec/cd/bd196b2cf014afb1009de8b0f05ecd54011d881944e62763f3c1b1e8ef37/pygtrie-2.5.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/f1/7b/ce1eafaf1a76852e2ec9b22edecf1daa58175c090266e9f6c64afcd81d91/stack_data-0.6.3-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/f9/1c/01bfd571a64e7f270e6bab5e33777debe0edc56759233ce84f27dec92d14/tqdm-4.70.0-py3-none-any.whl
@@ -607,6 +601,7 @@ environments:
- conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_7.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/_python_abi3_support-1.0-hd8ed1ab_3.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/backports-1.0-pyhd8ed1ab_5.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/backports.asyncio.runner-1.2.0-pyh5ded981_2.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/backports.tarfile-1.2.0-pyhcf101f3_2.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.7.22-hbd8a1cb_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.7.22-pyhd8ed1ab_0.conda
@@ -656,6 +651,7 @@ environments:
- conda: https://conda.anaconda.org/conda-forge/noarch/pyparsing-3.3.2-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pysocks-1.7.1-pyha55dd90_7.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-9.1.1-pyhc364b38_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-asyncio-1.4.0-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-cov-7.1.0-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-xdist-3.8.0-pyhd8ed1ab_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
@@ -680,6 +676,7 @@ environments:
- conda: https://conda.anaconda.org/conda-forge/noarch/wheel-0.48.0-pyhd8ed1ab_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/zipp-4.1.0-pyhcf101f3_0.conda
- pypi: .
+ - pypi: direct+https://github.com/evolutionaryscale/wheels/releases/download/py312-pt211-cu13-sm80-90/flash_attn-2.7.4.post1-cp312-cp312-linux_x86_64.whl
- pypi: git+https://github.com/nrontsis/DockQ.git?rev=ba4df5adaad7c77fd60851d0b7b05f2b77061ba2#ba4df5adaad7c77fd60851d0b7b05f2b77061ba2
- pypi: https://files.pythonhosted.org/packages/00/78/9cbcc1c073b9d4918e925af1a059762265dc65004e020511b2a06fbfd020/pydssp-0.9.1-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/03/b8/e484ef633af3887baeeb4b6ad12743363af7cce68ae51e938e00aaa0529d/transformers-4.57.6-py3-none-any.whl
@@ -687,12 +684,9 @@ environments:
- pypi: https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/17/66/1ef1960ffa751df4f6ed05a76ba677e331ce451902f319d14c3f8ae81215/rdkit-2026.3.5-cp312-cp312-manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/17/c2/613ca5f729b55b2f7a9c17a3cd59d3f053ee2816da2d19933acc6b8741ec/biopython-1.88-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/19/f2/5ba170c6195e95e42ab9e2fade98a70a078a3663af14440a67b457208041/cuequivariance_ops_torch_cu13-0.11.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/1a/c9/82638ef24d7877510f83baf821f5619a61b45568ce21c0a87a91576510aa/torch-2.11.0-cp312-cp312-manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/1d/04/cb4e8ba80ba81790d02dcc69d7529961140fc114abde927c99e9ba9c9457/biotite-1.7.1-cp312-cp312-manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/22/a6/858897256d0deac81a172289110f31629fc4cee19b6f01283303e18c8db3/ptyprocess-0.7.0-py2.py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/23/45/caa600acfab94560807a20a64b5830d2cd3c3202b7f1328644d70b7d6bd8/nvidia_ml_py-13.610.43-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/28/50/f203ff3a3ddfe19308efc83c5a3a29ed02bf786732ec35e68bf9162f3365/safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/2a/09/f8d8f8f31e4483c10a906437b4ce31bdf3d6d417b73fe33f1a8b59e34228/einops-0.8.2-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl
@@ -709,7 +703,6 @@ environments:
- pypi: https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/41/09/5b161152e2d90f7b87f781c2e1267494aef9c32498df793f73ad0a0a494a/matplotlib_inline-0.2.2-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/4e/d0/4b9da3d9c5749bf2c93f2818af4dc306f3b3ccc671f68f6820635fd53b74/cuequivariance_torch-0.11.1-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/52/89/d0baf1269e57ac548387c11e5d9ba78d1e118ceb43f1e5e5114bdd2e8391/boto3-1.43.75-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/52/f2/9e726e4489f947557290605e3ee33b6ba09fe5dac9213263e0ada75200cf/dna_features_viewer-3.1.5-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/54/6f/84908cad2d6aa5144abcf7b42709fe4fdb459bc640ec7ac5786e7693dabc/prompt_toolkit-3.0.53-py3-none-any.whl
@@ -722,7 +715,6 @@ environments:
- pypi: https://files.pythonhosted.org/packages/67/4e/a28359bf1c1ecf11eba22123168c138698f7cb576ac678f5a2e16cd5da08/hf_xet-1.6.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/6a/fd/6adabd4f6d5e686f97dd02ce7fce3fe4cf672cbac36b8f67ff4040e8ad8b/msgpack-1.2.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/6c/c2/63fdda36c56437eeb44aaf9493c8bcd62ce230ab1598924fc626ffbfa943/scikit_learn-1.9.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/72/f1/b1f191b91d049780eb4db2605b81af28a697710c4ac219708bcba192f80c/cuequivariance_ops_cu13-0.11.1-py3-none-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/7e/85/a5bfaebfd305ac18b57b0854d74e37e586809061a91fda62f0bd50c8518e/narwhals-2.24.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl
@@ -761,7 +753,6 @@ environments:
- pypi: https://files.pythonhosted.org/packages/da/35/f2287558c17e29fafc8ef3daf819bb9834061cfa43bff8014f7df7f63bdc/anyio-4.14.2-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/e3/1e/2942a0d6d52240d439f44384a40c65d359d9ca70325b65b265c020113121/cuda_pathfinder-1.6.1-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/e7/44/423ac00af4dd95a5aeb27207e2c0d9b7118702149bf4704c3ddb55bb7429/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/e7/a8/7f99cfa83f61d2d8cb29f31e1f90a6233873a6038b4cdfb30179ef90e5a6/cuequivariance-0.11.1-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/ec/cd/bd196b2cf014afb1009de8b0f05ecd54011d881944e62763f3c1b1e8ef37/pygtrie-2.5.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/f1/7b/ce1eafaf1a76852e2ec9b22edecf1daa58175c090266e9f6c64afcd81d91/stack_data-0.6.3-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/f9/1c/01bfd571a64e7f270e6bab5e33777debe0edc56759233ce84f27dec92d14/tqdm-4.70.0-py3-none-any.whl
@@ -771,6 +762,7 @@ environments:
osx-arm64:
- conda: https://conda.anaconda.org/conda-forge/noarch/_python_abi3_support-1.0-hd8ed1ab_3.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/backports-1.0-pyhd8ed1ab_5.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/backports.asyncio.runner-1.2.0-pyh5ded981_2.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/backports.tarfile-1.2.0-pyhcf101f3_2.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.7.22-hbd8a1cb_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.7.22-pyhd8ed1ab_0.conda
@@ -819,6 +811,7 @@ environments:
- conda: https://conda.anaconda.org/conda-forge/noarch/pyparsing-3.3.2-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pysocks-1.7.1-pyha55dd90_7.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-9.1.1-pyhc364b38_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-asyncio-1.4.0-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-cov-7.1.0-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-xdist-3.8.0-pyhd8ed1ab_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
@@ -1152,6 +1145,7 @@ environments:
- conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_7.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/_python_abi3_support-1.0-hd8ed1ab_3.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/backports-1.0-pyhd8ed1ab_5.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/backports.asyncio.runner-1.2.0-pyh5ded981_2.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/backports.tarfile-1.2.0-pyhcf101f3_2.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.7.22-hbd8a1cb_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.7.22-pyhd8ed1ab_0.conda
@@ -1201,6 +1195,7 @@ environments:
- conda: https://conda.anaconda.org/conda-forge/noarch/pyparsing-3.3.2-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pysocks-1.7.1-pyha55dd90_7.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-9.1.1-pyhc364b38_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-asyncio-1.4.0-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-cov-7.1.0-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-xdist-3.8.0-pyhd8ed1ab_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
@@ -1231,8 +1226,6 @@ environments:
- pypi: https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/1f/88/8dab5de10c638c083772a6be83a3d8106ced492f74a928c8693638e5bb50/scikit_learn-1.9.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/22/a6/858897256d0deac81a172289110f31629fc4cee19b6f01283303e18c8db3/ptyprocess-0.7.0-py2.py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/23/45/caa600acfab94560807a20a64b5830d2cd3c3202b7f1328644d70b7d6bd8/nvidia_ml_py-13.610.43-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/28/50/f203ff3a3ddfe19308efc83c5a3a29ed02bf786732ec35e68bf9162f3365/safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/2a/09/f8d8f8f31e4483c10a906437b4ce31bdf3d6d417b73fe33f1a8b59e34228/einops-0.8.2-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl
@@ -1253,7 +1246,6 @@ environments:
- pypi: https://files.pythonhosted.org/packages/41/09/5b161152e2d90f7b87f781c2e1267494aef9c32498df793f73ad0a0a494a/matplotlib_inline-0.2.2-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/49/d5/f34b9313ca9db131ad7d74156e3459694ced17f81233a248bd00eb6d67b1/botocore-1.43.79-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/4e/d0/4b9da3d9c5749bf2c93f2818af4dc306f3b3ccc671f68f6820635fd53b74/cuequivariance_torch-0.11.1-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/52/f2/9e726e4489f947557290605e3ee33b6ba09fe5dac9213263e0ada75200cf/dna_features_viewer-3.1.5-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/54/6f/84908cad2d6aa5144abcf7b42709fe4fdb459bc640ec7ac5786e7693dabc/prompt_toolkit-3.0.53-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/56/7a/123e033aaff487c77107195fa5a2b8686795ca537935a24efae476c41f05/nvidia_nvjitlink-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl
@@ -1265,7 +1257,6 @@ environments:
- pypi: https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/67/4e/a28359bf1c1ecf11eba22123168c138698f7cb576ac678f5a2e16cd5da08/hf_xet-1.6.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/6f/c6/9e51ca59bbe53c44299f8f750f07c5e39b16f990c9de999d6454d1300503/rdkit-2026.3.5-cp313-cp313-manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/72/f1/b1f191b91d049780eb4db2605b81af28a697710c4ac219708bcba192f80c/cuequivariance_ops_cu13-0.11.1-py3-none-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/79/d3/36a46a8ed992b781acbc05928bd5bee3c810cb0c3563bf81a7b0c04a1a76/msgpack-1.2.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/7a/90/eef59b9b64442b4c966b027e8db407f59d99200b82b71316fc4466562a53/boto3-1.43.79-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl
@@ -1289,7 +1280,6 @@ environments:
- pypi: https://files.pythonhosted.org/packages/a9/21/9b05698b46f218fc0e118e1f8168395c65c8a2c750ae2bab54fc4bd4e0e8/markupsafe-3.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/ad/66/0d785f0bc5e4315a96c989bb476d0fc07ea4f85132550c7b156ca2035d52/traitlets-5.16.1-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/b0/b4/878fefaad5b2bcc6fcf8d474a25e3e3774bc5133e4b58adff4d0bca238bc/nvidia_nccl_cu13-2.28.9-py3-none-manylinux_2_18_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/b3/a1/42c6e78a654c28ba03932e7c7bfb2918ae5828488aae30322ced83379a19/cuequivariance_ops_torch_cu13-0.11.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/bc/8e/1239df488393d61076653bfb29f759d0f60cab8e030abdf7c17c31539b51/ipython-9.16.1-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/bc/e7/5c595c75e9f41a44f30e526eda465ea0b4eec93470e074e4a111b253f13a/s3transfer-0.19.2-py3-none-any.whl
@@ -1302,7 +1292,6 @@ environments:
- pypi: https://files.pythonhosted.org/packages/d9/33/1f075bf72b0b747cb3288d011319aaf64083cf2efef8354174e3ed4540e2/ipython_pygments_lexers-1.1.1-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/da/35/f2287558c17e29fafc8ef3daf819bb9834061cfa43bff8014f7df7f63bdc/anyio-4.14.2-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/e7/44/423ac00af4dd95a5aeb27207e2c0d9b7118702149bf4704c3ddb55bb7429/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/e7/a8/7f99cfa83f61d2d8cb29f31e1f90a6233873a6038b4cdfb30179ef90e5a6/cuequivariance-0.11.1-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/eb/dc/55481808fd70ef1567cf13540ffd4702af3f74b112e35427564b03f79c2d/narwhals-2.25.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/ec/cd/bd196b2cf014afb1009de8b0f05ecd54011d881944e62763f3c1b1e8ef37/pygtrie-2.5.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/f0/6f/ab71525583a7824c70f71de387a1c5ceb27ddcb3fda2dacb734e5b875f14/biotraj-1.2.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
@@ -1315,6 +1304,7 @@ environments:
osx-arm64:
- conda: https://conda.anaconda.org/conda-forge/noarch/_python_abi3_support-1.0-hd8ed1ab_3.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/backports-1.0-pyhd8ed1ab_5.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/backports.asyncio.runner-1.2.0-pyh5ded981_2.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/backports.tarfile-1.2.0-pyhcf101f3_2.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.7.22-hbd8a1cb_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.7.22-pyhd8ed1ab_0.conda
@@ -1363,6 +1353,7 @@ environments:
- conda: https://conda.anaconda.org/conda-forge/noarch/pyparsing-3.3.2-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pysocks-1.7.1-pyha55dd90_7.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-9.1.1-pyhc364b38_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-asyncio-1.4.0-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-cov-7.1.0-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-xdist-3.8.0-pyhd8ed1ab_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
@@ -1694,6 +1685,7 @@ environments:
- conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_7.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/_python_abi3_support-1.0-hd8ed1ab_3.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/backports-1.0-pyhd8ed1ab_5.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/backports.asyncio.runner-1.2.0-pyh5ded981_2.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/backports.tarfile-1.2.0-pyhcf101f3_2.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/backports.zstd-1.7.0-py314h680f03e_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.7.22-hbd8a1cb_0.conda
@@ -1745,6 +1737,7 @@ environments:
- conda: https://conda.anaconda.org/conda-forge/noarch/pyparsing-3.3.2-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pysocks-1.7.1-pyha55dd90_7.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-9.1.1-pyhc364b38_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-asyncio-1.4.0-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-cov-7.1.0-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-xdist-3.8.0-pyhd8ed1ab_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
@@ -1777,8 +1770,6 @@ environments:
- pypi: https://files.pythonhosted.org/packages/14/71/986ceea9aa3da548bf1357cad89b63915ec6d21ec957c8113b29ece567df/regex-2026.7.19-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/19/03/8c63e8cf52958534ef688625965ab04c269a6cadd8caef16758b380a821a/msgpack-1.2.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/22/a6/858897256d0deac81a172289110f31629fc4cee19b6f01283303e18c8db3/ptyprocess-0.7.0-py2.py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/23/45/caa600acfab94560807a20a64b5830d2cd3c3202b7f1328644d70b7d6bd8/nvidia_ml_py-13.610.43-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/28/50/f203ff3a3ddfe19308efc83c5a3a29ed02bf786732ec35e68bf9162f3365/safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/2a/09/f8d8f8f31e4483c10a906437b4ce31bdf3d6d417b73fe33f1a8b59e34228/einops-0.8.2-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl
@@ -1797,7 +1788,6 @@ environments:
- pypi: https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/47/e8/b98ca2d39b2e0e4730c0ee52537e488e7008025bc77ca89552ff91021f7c/torch-2.11.0-cp314-cp314-manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/49/d5/f34b9313ca9db131ad7d74156e3459694ced17f81233a248bd00eb6d67b1/botocore-1.43.79-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/4e/d0/4b9da3d9c5749bf2c93f2818af4dc306f3b3ccc671f68f6820635fd53b74/cuequivariance_torch-0.11.1-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/52/f2/9e726e4489f947557290605e3ee33b6ba09fe5dac9213263e0ada75200cf/dna_features_viewer-3.1.5-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/54/6f/84908cad2d6aa5144abcf7b42709fe4fdb459bc640ec7ac5786e7693dabc/prompt_toolkit-3.0.53-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/56/7a/123e033aaff487c77107195fa5a2b8686795ca537935a24efae476c41f05/nvidia_nvjitlink-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl
@@ -1807,9 +1797,7 @@ environments:
- pypi: https://files.pythonhosted.org/packages/62/01/a7171c5e2e8755597bd8f1c1eb228a0876f502afdf25f936061f5dbe2880/cuda_pathfinder-1.7.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/66/2f/e762d2a929d8e22d96b937ab94ccd39dd2e20b2cf9dbd7a4fc29feb90bd1/cuequivariance_ops_torch_cu13-0.11.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/67/4e/a28359bf1c1ecf11eba22123168c138698f7cb576ac678f5a2e16cd5da08/hf_xet-1.6.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/72/f1/b1f191b91d049780eb4db2605b81af28a697710c4ac219708bcba192f80c/cuequivariance_ops_cu13-0.11.1-py3-none-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/7a/90/eef59b9b64442b4c966b027e8db407f59d99200b82b71316fc4466562a53/boto3-1.43.79-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl
@@ -1847,7 +1835,6 @@ environments:
- pypi: https://files.pythonhosted.org/packages/da/35/f2287558c17e29fafc8ef3daf819bb9834061cfa43bff8014f7df7f63bdc/anyio-4.14.2-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/df/3d/9e7eee57b37c80cec63322c0231bb6da3cfe535a91d7a4d64896fcb89357/triton-3.6.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
- pypi: https://files.pythonhosted.org/packages/e7/44/423ac00af4dd95a5aeb27207e2c0d9b7118702149bf4704c3ddb55bb7429/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/e7/a8/7f99cfa83f61d2d8cb29f31e1f90a6233873a6038b4cdfb30179ef90e5a6/cuequivariance-0.11.1-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/eb/dc/55481808fd70ef1567cf13540ffd4702af3f74b112e35427564b03f79c2d/narwhals-2.25.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/ec/cd/bd196b2cf014afb1009de8b0f05ecd54011d881944e62763f3c1b1e8ef37/pygtrie-2.5.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/f0/af/4d72d9e475ac83719160c662619e4bf7b95c19507cd582e7d0167a3c3dae/scikit_learn-1.9.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
@@ -1859,6 +1846,7 @@ environments:
osx-arm64:
- conda: https://conda.anaconda.org/conda-forge/noarch/_python_abi3_support-1.0-hd8ed1ab_3.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/backports-1.0-pyhd8ed1ab_5.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/backports.asyncio.runner-1.2.0-pyh5ded981_2.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/backports.tarfile-1.2.0-pyhcf101f3_2.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/backports.zstd-1.7.0-py314h680f03e_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.7.22-hbd8a1cb_0.conda
@@ -1909,6 +1897,7 @@ environments:
- conda: https://conda.anaconda.org/conda-forge/noarch/pyparsing-3.3.2-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pysocks-1.7.1-pyha55dd90_7.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-9.1.1-pyhc364b38_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-asyncio-1.4.0-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-cov-7.1.0-pyhcf101f3_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-xdist-3.8.0-pyhd8ed1ab_0.conda
- conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
@@ -2084,326 +2073,6 @@ environments:
- pypi: https://files.pythonhosted.org/packages/f5/b1/fa7c600e7dceae12e9606c7578cbc9ff1e1ed55844883ee5c92205e86226/safetensors-0.8.0-cp310-abi3-macosx_11_0_arm64.whl
- pypi: https://files.pythonhosted.org/packages/f9/1c/01bfd571a64e7f270e6bab5e33777debe0edc56759233ce84f27dec92d14/tqdm-4.70.0-py3-none-any.whl
- pypi: https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl
- gpu:
- channels:
- - url: https://conda.anaconda.org/conda-forge/
- indexes:
- - https://pypi.org/simple
- packages:
- linux-64:
- - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-20_gnu.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/alsa-lib-1.2.16.1-hb03c661_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/backports.zstd-1.7.0-py312h3f22e6b_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-1.2.0-h505cf86_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-bin-1.2.0-h9908984_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-python-1.2.0-py312he9c40d5_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/bzip2-1.0.8-hda65f42_10.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/c-ares-1.34.8-hebe6cf0_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/cairo-1.18.4-he90730b_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/cffi-2.1.1-py312h703531f_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/cmake-4.4.2-hff7ac6b_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/comrak-0.54.0-hb17b654_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.3-py312h0a2e395_4.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/coverage-7.15.4-py312h589e9ff_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/cryptography-50.0.0-py312ha4b625e_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/cyrus-sasl-2.1.28-hac629b4_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/dbus-1.16.2-h24cb091_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/double-conversion-3.4.0-hecca717_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/fontconfig-2.18.3-h4db4eae_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/fonttools-4.63.0-py312h8a5da7c_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/freetype-2.14.3-ha770c72_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/fribidi-1.0.16-hb03c661_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/graphite2-1.3.15-h54a6638_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/icu-78.3-py310h44b86e0_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/keyutils-1.6.3-h7cc23a3_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.5.0-py312h0a2e395_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/krb5-1.22.2-hbc21106_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/lcms2-2.19.1-h0c24ade_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/ld_impl_linux-64-2.46.1-default_hbd61a6d_102.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/lerc-4.2.0-hdb68285_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libblas-3.11.0-9_h4a7cf45_openblas.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlicommon-1.2.0-h39a168f_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlidec-1.2.0-ha411449_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlienc-1.2.0-h018ffa1_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libcblas-3.11.0-9_h0358290_openblas.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libclang-cpp22.1-22.1.8-default_h12ff6da_7.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libclang13-22.1.8-default_h96bbd2b_7.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libcups-2.3.3-h7a8fb5f_6.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libcurl-8.21.0-ha042cf0_5.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libdeflate-1.25-hd45a770_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libdrm-2.4.129-h7cc23a3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libedit-3.1.20250104-pl5321h373387f_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libegl-1.7.0-ha4b6fd6_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libegl-devel-1.7.0-ha4b6fd6_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libev-4.33-h280c20c_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libexpat-2.8.1-hecca717_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libffi-3.7.0-h3435931_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libfreetype-2.14.3-ha770c72_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libfreetype6-2.14.3-h5e6c136_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libgcc-16.1.0-ha9f2e26_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libgcc-ng-16.1.0-h69a702a_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran-16.1.0-h69a702a_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran5-16.1.0-h79bb938_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libgl-1.7.0-ha4b6fd6_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libgl-devel-1.7.0-ha4b6fd6_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libglib-2.88.3-h45c3219_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libglvnd-1.7.0-ha4b6fd6_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libglx-1.7.0-ha4b6fd6_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libglx-devel-1.7.0-ha4b6fd6_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libgomp-16.1.0-he0feb66_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libharfbuzz-14.3.1-h23af247_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libiconv-1.18-h0cb94f2_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libjpeg-turbo-3.2.0-hb03c661_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/liblapack-3.11.0-9_h47877c9_openblas.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libllvm22-22.1.8-hf7376ad_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/liblzma-5.8.3-hb03c661_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libnghttp2-1.68.1-h74cf4be_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libnsl-2.0.1-hb9d3cd8_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libntlm-1.8-hb9d3cd8_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libopenblas-0.3.34-pthreads_h94d23a6_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libopengl-1.7.0-ha4b6fd6_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libpciaccess-0.19-hb03c661_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libpng-1.6.58-h922cc85_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libpq-18.6-h9d76c99_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libpsl-0.23.1-hd9e3e90_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libraqm-0.11.0-h6406941_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libsqlite-3.53.4-h13e7031_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libssh2-1.11.1-h6154650_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libstdcxx-16.1.0-h934c35e_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libstdcxx-ng-16.1.0-hdf11a46_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libtiff-4.7.2-h9d88235_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libuuid-2.42.2-h5347b49_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libuv-1.52.1-h280c20c_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libvulkan-loader-1.4.357.0-h5279c79_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libwebp-base-1.6.0-hd42ef1d_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libxcb-1.17.0-hb83e432_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libxcrypt-4.4.38-h280c20c_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libxkbcommon-1.13.2-h51789e4_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libxml2-16-2.15.3-hca6bf5a_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libxml2-2.15.3-h49c6c72_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libxslt-1.1.43-h711ed8c_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/libzlib-1.3.2-h25fd6f3_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-3.11.1-py312h6a12a82_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.11.1-py312h4c94fcb_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/ncurses-6.6-hdb14827_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/nh3-0.3.6-py310hd8a072f_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-2.3.5-py312h33ff503_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/openjpeg-2.5.4-h55fea9a_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/openldap-2.6.13-hbde042b_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/openssl-3.6.3-h35e630c_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/pandas-3.0.5-py312h8ecdadd_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/pcre2-10.47-haa7fec5_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-12.3.0-py312h50c33e8_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/pixman-0.46.4-h54a6638_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/pkg-config-0.29.2-h7c397b8_1011.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb03c661_1003.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/pyside6-6.11.2-py312hf143aa3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.12.13-h8ab3286_1_cpython.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/pyyaml-6.0.3-py312h8a5da7c_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/qhull-2020.2-h434a139_5.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/qt6-main-6.11.2-pl5321h9df5c37_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/readline-8.3-hd6e31c0_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/rhash-1.4.6-hb03c661_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.18.0-py312h54fa4ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/secretstorage-3.5.0-py312h7900ff3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/statsmodels-0.14.6-py312h4f23490_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_hd70dff1_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.5.8-py312h4c3975b_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/ty-0.0.49-h4e94fc0_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/ukkonen-1.1.0-py312hd9148b4_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/unicodedata2-17.0.1-py312h4c3975b_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/wayland-1.26.0-h1964d1d_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xcb-util-0.4.1-h4f16b4b_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xcb-util-cursor-0.1.6-hb03c661_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xcb-util-image-0.4.0-hb711507_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xcb-util-keysyms-0.4.1-hb711507_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xcb-util-renderutil-0.3.10-hb711507_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xcb-util-wm-0.4.2-hb711507_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xkeyboard-config-2.48-h280c20c_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libice-1.1.2-h280c20c_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libsm-1.2.6-h0d788c3_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libx11-1.8.13-he1eb515_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxau-1.0.12-hb03c661_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxcomposite-0.4.7-hb03c661_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxcursor-1.2.3-hb9d3cd8_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxdamage-1.1.6-hb9d3cd8_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxdmcp-1.1.5-hb03c661_2.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxext-1.3.7-h7cc23a3_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxfixes-6.0.2-hb03c661_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxi-1.8.3-hb03c661_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxrandr-1.5.5-hb03c661_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxrender-0.9.12-hb03c661_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxtst-1.2.5-hb9d3cd8_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxxf86vm-1.1.7-hb03c661_0.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-xorgproto-2025.1-h280c20c_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/yaml-0.2.5-hebe6cf0_3.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-ng-2.3.3-hce19668_1.conda
- - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_7.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/_python_abi3_support-1.0-hd8ed1ab_3.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/backports-1.0-pyhd8ed1ab_5.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/backports.tarfile-1.2.0-pyhcf101f3_2.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.7.22-hbd8a1cb_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.7.22-pyhd8ed1ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/cfgv-3.5.0-pyhd8ed1ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/charset-normalizer-3.5.1-pyhd8ed1ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.12.13-py312hd8ed1ab_1.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhcf101f3_2.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/distlib-0.4.3-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/docutils-0.23-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/exceptiongroup-1.3.1-pyhd8ed1ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/execnet-2.1.2-pyhd8ed1ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/filelock-3.32.3-pyhd8ed1ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/font-ttf-dejavu-sans-mono-2.37-hab24e00_0.tar.bz2
- - conda: https://conda.anaconda.org/conda-forge/noarch/font-ttf-inconsolata-3.000-h77eed37_0.tar.bz2
- - conda: https://conda.anaconda.org/conda-forge/noarch/font-ttf-source-code-pro-2.038-h77eed37_0.tar.bz2
- - conda: https://conda.anaconda.org/conda-forge/noarch/font-ttf-ubuntu-0.83-h77eed37_3.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/fonts-conda-ecosystem-1-0.tar.bz2
- - conda: https://conda.anaconda.org/conda-forge/noarch/fonts-conda-forge-1-hc364b38_1.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/h2-4.4.1-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/hpack-4.2.0-pyhd8ed1ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/hyperframe-6.1.0-pyhd8ed1ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/id-1.6.1-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/identify-2.6.19-pyhd8ed1ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/idna-3.19-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/importlib-metadata-9.0.0-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/importlib_resources-7.1.0-pyhd8ed1ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/iniconfig-2.3.0-pyhd8ed1ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/jaraco.classes-3.4.0-pyhcf101f3_3.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/jaraco.context-6.1.2-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/jaraco.functools-4.6.0-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/jeepney-0.9.0-pyhd8ed1ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/keyring-25.7.0-pyha804496_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/markdown-it-py-4.2.0-pyhd8ed1ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/mdurl-0.1.2-pyhd8ed1ab_1.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/more-itertools-11.1.0-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/munkres-1.1.4-pyhd8ed1ab_1.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/nodeenv-1.10.0-pyhd8ed1ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/packaging-26.3-pyhc364b38_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/patsy-1.0.2-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/pip-26.2.1-pyh8b19718_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.11.3-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/pluggy-1.6.0-pyhf9edf01_1.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/pre-commit-4.6.2-pyha770c72_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/pycparser-3.0-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/pygments-2.21.0-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/pyparsing-3.3.2-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/pysocks-1.7.1-pyha55dd90_7.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-9.1.1-pyhc364b38_2.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-cov-7.1.0-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-xdist-3.8.0-pyhd8ed1ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/python-discovery-1.5.2-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.12.13-hd8ed1ab_1.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.12-8_cp312.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/readme_renderer-45.0-pyhcf101f3_1.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/requests-2.34.2-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/requests-toolbelt-1.0.0-pyhd8ed1ab_1.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3986-2.0.0-pyhd8ed1ab_1.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/rich-15.0.0-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/seaborn-0.13.2-hd8ed1ab_3.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/seaborn-base-0.13.2-pyhd8ed1ab_3.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/setuptools-81.0.0-pyh332efcf_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhe01879c_1.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/tomli-2.4.1-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/twine-7.0.0-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.16.0-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2026c-h151e31d_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/urllib3-2.7.0-pyhd8ed1ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-21.7.4-pyhcf101f3_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/wheel-0.48.0-pyhd8ed1ab_0.conda
- - conda: https://conda.anaconda.org/conda-forge/noarch/zipp-4.1.0-pyhcf101f3_0.conda
- - pypi: .
- - pypi: direct+https://download.pytorch.org/whl/cu130/xformers-0.0.35-py39-none-manylinux_2_28_x86_64.whl
- - pypi: direct+https://github.com/evolutionaryscale/wheels/releases/download/py312-pt211-cu13-sm80-90/flash_attn-2.7.4.post1-cp312-cp312-linux_x86_64.whl
- - pypi: git+https://github.com/nrontsis/DockQ.git?rev=ba4df5adaad7c77fd60851d0b7b05f2b77061ba2#ba4df5adaad7c77fd60851d0b7b05f2b77061ba2
- - pypi: https://files.pythonhosted.org/packages/00/78/9cbcc1c073b9d4918e925af1a059762265dc65004e020511b2a06fbfd020/pydssp-0.9.1-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/03/b8/e484ef633af3887baeeb4b6ad12743363af7cce68ae51e938e00aaa0529d/transformers-4.57.6-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/17/66/1ef1960ffa751df4f6ed05a76ba677e331ce451902f319d14c3f8ae81215/rdkit-2026.3.5-cp312-cp312-manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/17/c2/613ca5f729b55b2f7a9c17a3cd59d3f053ee2816da2d19933acc6b8741ec/biopython-1.88-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/19/f2/5ba170c6195e95e42ab9e2fade98a70a078a3663af14440a67b457208041/cuequivariance_ops_torch_cu13-0.11.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/1a/c9/82638ef24d7877510f83baf821f5619a61b45568ce21c0a87a91576510aa/torch-2.11.0-cp312-cp312-manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/1d/04/cb4e8ba80ba81790d02dcc69d7529961140fc114abde927c99e9ba9c9457/biotite-1.7.1-cp312-cp312-manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/22/a6/858897256d0deac81a172289110f31629fc4cee19b6f01283303e18c8db3/ptyprocess-0.7.0-py2.py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/23/45/caa600acfab94560807a20a64b5830d2cd3c3202b7f1328644d70b7d6bd8/nvidia_ml_py-13.610.43-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/28/50/f203ff3a3ddfe19308efc83c5a3a29ed02bf786732ec35e68bf9162f3365/safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/2a/09/f8d8f8f31e4483c10a906437b4ce31bdf3d6d417b73fe33f1a8b59e34228/einops-0.8.2-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/2a/8e/096d00c7c480ef2ff4265349b14e2261d4ab787ba1f74e2e80d1c58079c3/regex-2026.7.19-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/2e/76/932be4b50ef6ccedf9d3c6639b056a967a86258c6d9200643f01269211ca/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/32/d5/f9a850d79b0851d1d4ef6456097579a9005b31fea68726a4ae5f2d82ddd9/threadpoolctl-3.6.0-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/33/ef/6d27fc118f58cb24886da413545a7efb0853d405fddbfd8b2d9ac09fbed4/jupyterlab_widgets-3.0.17-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/34/95/40e17e20046b7bc820d29d09ae84ec157ec8dd6e6f6cd722626292c31b2e/widgetsnbextension-4.0.16-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/39/2a/6d2e9047d1fb243dbaa364b01e0297534b9ed7fd27dba1c9f361519cf69b/cuda_bindings-13.3.1-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/3c/35/a9bf80a609e74e3b000fef598933235c908fcefcef9026042b8e6dfde2a9/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/41/09/5b161152e2d90f7b87f781c2e1267494aef9c32498df793f73ad0a0a494a/matplotlib_inline-0.2.2-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/4e/d0/4b9da3d9c5749bf2c93f2818af4dc306f3b3ccc671f68f6820635fd53b74/cuequivariance_torch-0.11.1-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/52/89/d0baf1269e57ac548387c11e5d9ba78d1e118ceb43f1e5e5114bdd2e8391/boto3-1.43.75-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/52/f2/9e726e4489f947557290605e3ee33b6ba09fe5dac9213263e0ada75200cf/dna_features_viewer-3.1.5-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/54/6f/84908cad2d6aa5144abcf7b42709fe4fdb459bc640ec7ac5786e7693dabc/prompt_toolkit-3.0.53-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/56/7a/123e033aaff487c77107195fa5a2b8686795ca537935a24efae476c41f05/nvidia_nvjitlink-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/57/b2/453099f5f3b698d7d0eab38916aac44c7f76229f451709e2eb9db6615dcd/cuda_toolkit-13.0.2-py2.py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/5f/67/cba3777620cdacb99102da4042883709c41c709f4b6323c10781a9c3aa34/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/60/97/891a0971e1e4a8c5d2b20bbe0e524dc04548d2307fee33cdeba148fd4fc7/comm-0.2.3-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/67/4e/a28359bf1c1ecf11eba22123168c138698f7cb576ac678f5a2e16cd5da08/hf_xet-1.6.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/6a/fd/6adabd4f6d5e686f97dd02ce7fce3fe4cf672cbac36b8f67ff4040e8ad8b/msgpack-1.2.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/6c/c2/63fdda36c56437eeb44aaf9493c8bcd62ce230ab1598924fc626ffbfa943/scikit_learn-1.9.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/72/f1/b1f191b91d049780eb4db2605b81af28a697710c4ac219708bcba192f80c/cuequivariance_ops_cu13-0.11.1-py3-none-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/7e/85/a5bfaebfd305ac18b57b0854d74e37e586809061a91fda62f0bd50c8518e/narwhals-2.24.0-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/80/cd/c607b6337ddcf55cc0249527819b727ef28481d7c56ace54cca69400c2b9/biotraj-1.2.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/81/b6/100d5811cb60bfc13626910cb48872b8ca53f1dcaf2127aac71185575a7f/parallelbar-2.5-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/8e/37/efad0257dc6e593a18957422533ff0f87ede7c9c6ea010a2177d738fb82f/pure_eval-0.2.3-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/8e/5a/7fd1b784a87e96e0078f49a0a13a98b4c5f644ba5597a4a3b70a2ba3e613/py3dmol-2.5.5-py2.py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/94/46/19d3178e70f37fda251f3d9560f29654f09c6eedbe1c8b790adeee84da49/botocore-1.43.75-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/96/42/3e5985a0a7e57de470b320c6d6a1a67c844f6737a587f3d44dd13d1819e7/wcwidth-0.8.2-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/99/5d/8268b644392ee874ee82a635cd0df1773de230bde356c38de28e298392cc/parso-0.8.7-py2.py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/9a/93/242e2eab5fe682ffcb8b0084bde703a41d51e17ee0f3a31ff0d9d813620a/jedi-0.20.0-py2.py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/9b/5d/f25ac7d4fb77cbd53ddc6d05d833c6bf52b12770a44fa9a447eed470ca9a/msgpack_numpy-0.4.8-py2.py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/9e/c3/059298687310d527a58bb01f3b1965787ee3b40dce76752eda8b44e9a2c5/pexpect-4.9.0-py2.py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/a3/22/0b4b932655d17a6da1b92fa92ab12844b053bb2ac2475e179ba6f043da1e/nvidia_cudnn_cu13-9.19.0.56-py3-none-manylinux_2_27_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/a5/9f/be0a41ca4a4917abf5cb9ae0daff1a6060cc5de950aec0396de9f3b52bc5/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/a8/2f/7b57e29836ea8714f81e9898409196f47d772d5ddedddf1592eadb8ab743/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/a8/af/48ac8483240de756d2438c380746e7130d1c6f75802ef22f3c6d49982787/huggingface_hub-0.36.2-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/a8/db/253133d7e7cb40d3af384bb2f5c0b4a2b7fdcffbc95c688cc67a20a3c103/accelerate-1.14.0-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/ad/66/0d785f0bc5e4315a96c989bb476d0fc07ea4f85132550c7b156ca2035d52/traitlets-5.16.1-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/b0/b4/878fefaad5b2bcc6fcf8d474a25e3e3774bc5133e4b58adff4d0bca238bc/nvidia_nccl_cu13-2.28.9-py3-none-manylinux_2_18_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/b7/80/6252de3a70cfd7767718ad476893f1c7dc129f942cc7ed0322e3137c03d9/zstd-1.5.7.2-cp312-cp312-manylinux_2_14_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/bc/8e/1239df488393d61076653bfb29f759d0f60cab8e030abdf7c17c31539b51/ipython-9.16.1-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/bc/e7/5c595c75e9f41a44f30e526eda465ea0b4eec93470e074e4a111b253f13a/s3transfer-0.19.2-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/c1/ea/53f2148663b321f21b5a606bd5f191517cf40b7072c0497d3c92c4a13b1e/executing-2.2.1-py2.py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/c2/5b/ba933f896d9b0b07608d575a8501e2b4e32166b60d84c430a4a7285ebe64/cloudpathlib-0.24.0-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/c2/f3/d86c845465a2723ad7e1e5c36dcd75ddb82898b3f53be47ebd429fb2fa5d/nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/c3/55/298e9b3b864a198234997e87a1471c1b17d7f3546ace6d18fb5cf1ce24b2/ipywidgets-8.1.9-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/d4/2b/04b8a15f3a1c77bc79ddf5c73875327f34b4fa75982df2b76e45e402d364/asttokens-3.0.2-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/d7/c1/eb8f9debc45d3b7918a32ab756658a0904732f75e555402972246b0b8e71/tenacity-9.1.4-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/d9/33/1f075bf72b0b747cb3288d011319aaf64083cf2efef8354174e3ed4540e2/ipython_pygments_lexers-1.1.1-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/da/35/f2287558c17e29fafc8ef3daf819bb9834061cfa43bff8014f7df7f63bdc/anyio-4.14.2-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/e3/1e/2942a0d6d52240d439f44384a40c65d359d9ca70325b65b265c020113121/cuda_pathfinder-1.6.1-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/e7/44/423ac00af4dd95a5aeb27207e2c0d9b7118702149bf4704c3ddb55bb7429/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/e7/a8/7f99cfa83f61d2d8cb29f31e1f90a6233873a6038b4cdfb30179ef90e5a6/cuequivariance-0.11.1-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/ec/cd/bd196b2cf014afb1009de8b0f05ecd54011d881944e62763f3c1b1e8ef37/pygtrie-2.5.0-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/f1/7b/ce1eafaf1a76852e2ec9b22edecf1daa58175c090266e9f6c64afcd81d91/stack_data-0.6.3-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/f9/1c/01bfd571a64e7f270e6bab5e33777debe0edc56759233ce84f27dec92d14/tqdm-4.70.0-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/fa/18/623c77619c31d62efd55302939756966f3ecc8d724a14dab2b75f1508850/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
- - pypi: https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl
- - pypi: https://files.pythonhosted.org/packages/fd/53/43b0d71f4e702fa9733f8b4571fdca50a8813f1e450b656c239beff12315/nvidia_cusparselt_cu13-0.8.0-py3-none-manylinux2014_x86_64.whl
packages:
- conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-20_gnu.conda
build_number: 20
@@ -6215,6 +5884,20 @@ packages:
run_exports: {}
size: 7069
timestamp: 1733218168786
+- conda: https://conda.anaconda.org/conda-forge/noarch/backports.asyncio.runner-1.2.0-pyh5ded981_2.conda
+ sha256: 2ade43752e8494f110a2cfb9e4d5b1ea29e3dcb037fba63395442d00371e8bf9
+ md5: 0fd7e45c862b3305226a992f9f7b204a
+ depends:
+ - python >=3.11
+ - python
+ constrains:
+ - python >=3.11
+ license: PSF-2.0
+ license_family: PSF
+ purls: []
+ run_exports: {}
+ size: 10186
+ timestamp: 1753456386827
- conda: https://conda.anaconda.org/conda-forge/noarch/backports.tarfile-1.2.0-pyhcf101f3_2.conda
sha256: 25abdb37e186f0d6ac3b774a63c81c5bc4bf554b5096b51343fa5e7c381193b1
md5: bea46844deb274b2cc2a3a941745fa73
@@ -6952,6 +6635,22 @@ packages:
run_exports: {}
size: 306724
timestamp: 1782127176429
+- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-asyncio-1.4.0-pyhcf101f3_0.conda
+ sha256: 1eddccab8f91f718f890a2aabfc25e68c37a6171fb120e4ec4df7a5a96a877e9
+ md5: 787319c0ace4c39b7e883a86332db6a4
+ depends:
+ - pytest >=8.4,<10
+ - python >=3.10
+ - typing_extensions >=4.12
+ - backports.asyncio.runner >=1.1,<2
+ - python
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/pytest-asyncio?source=hash-mapping
+ run_exports: {}
+ size: 43101
+ timestamp: 1779802443053
- conda: https://conda.anaconda.org/conda-forge/noarch/pytest-cov-7.1.0-pyhcf101f3_0.conda
sha256: 44e42919397bd00bfaa47358a6ca93d4c21493a8c18600176212ec21a8d25ca5
md5: 67d1790eefa81ed305b89d8e314c7923
@@ -9669,22 +9368,8 @@ packages:
- pygtrie
- dna-features-viewer
- accelerate
- - cuequivariance-torch>=0.8.1 ; platform_machine == 'x86_64' and sys_platform == 'linux'
- - cuequivariance-ops-torch-cu13>=0.8.1 ; platform_machine == 'x86_64' and sys_platform == 'linux'
- - triton>=3,<4 ; extra == 'fused'
- - transformer-engine[pytorch]>=2,<3 ; extra == 'fold-cp'
- - cuequivariance-torch ; extra == 'cueq12'
- - cuequivariance-ops-torch-cu12 ; extra == 'cueq12'
- - cuequivariance-torch ; extra == 'cueq13'
- - cuequivariance-ops-torch-cu13 ; extra == 'cueq13'
+ - triton==3.6.* ; platform_machine == 'x86_64' and sys_platform == 'linux'
requires_python: '>=3.12'
-- pypi: direct+https://download.pytorch.org/whl/cu130/xformers-0.0.35-py39-none-manylinux_2_28_x86_64.whl
- name: xformers
- version: 0.0.35
- requires_dist:
- - torch>=2.10
- - numpy
- requires_python: '>=3.9'
- pypi: direct+https://github.com/evolutionaryscale/wheels/releases/download/py312-pt211-cu13-sm80-90/flash_attn-2.7.4.post1-cp312-cp312-linux_x86_64.whl
name: flash-attn
version: 2.7.4.post1
@@ -10234,25 +9919,6 @@ packages:
version: 1.2.1
sha256: 0e2bf9280bceb5efca998435904b5d3e9fdbcc11d90dc9df30aec7973252b720
requires_python: '>=3.10'
-- pypi: https://files.pythonhosted.org/packages/19/f2/5ba170c6195e95e42ab9e2fade98a70a078a3663af14440a67b457208041/cuequivariance_ops_torch_cu13-0.11.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
- name: cuequivariance-ops-torch-cu13
- version: 0.11.1
- sha256: 57fa7d057ea3cfbc8c8f85ce8d24e6ff12cc49073759bf8ff789608d6cd7752d
- requires_dist:
- - cuequivariance-ops-cu13==0.11.1
- - scipy
- - torch>=2.11
- - numpy>=1.23,<2.0a0 ; extra == 'test'
- - pytest ; extra == 'test'
- - onnx ; extra == 'plugins'
- - onnxscript ; extra == 'plugins'
- - onnxruntime ; extra == 'plugins'
- - onnxruntime-extensions ; extra == 'plugins'
- - onnx-graphsurgeon ; extra == 'plugins'
- - tensorrt>=10 ; extra == 'plugins'
- - polygraphy ; extra == 'plugins'
- - cuda-bindings ; extra == 'plugins'
- requires_python: '>=3.10'
- pypi: https://files.pythonhosted.org/packages/1a/c9/82638ef24d7877510f83baf821f5619a61b45568ce21c0a87a91576510aa/torch-2.11.0-cp312-cp312-manylinux_2_28_x86_64.whl
name: torch
version: 2.11.0
@@ -10362,15 +10028,6 @@ packages:
name: ptyprocess
version: 0.7.0
sha256: 4b41f3967fce3af57cc7e94b888626c18bf37a083e3651ca8feeb66d492fef35
-- pypi: https://files.pythonhosted.org/packages/23/45/caa600acfab94560807a20a64b5830d2cd3c3202b7f1328644d70b7d6bd8/nvidia_ml_py-13.610.43-py3-none-any.whl
- name: nvidia-ml-py
- version: 13.610.43
- sha256: f13c72698edef492f985cc225f14faafe68ae065a2e407f45bdf6f4b9b43fde8
-- pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
- name: opt-einsum
- version: 3.4.0
- sha256: 69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd
- requires_python: '>=3.8'
- pypi: https://files.pythonhosted.org/packages/26/0d/8603382f61abd0db35841148ddc1ffd607bf3100b11c6e1dab6d2fc44e72/torch-2.11.0-cp314-cp314-macosx_11_0_arm64.whl
name: torch
version: 2.11.0
@@ -10832,13 +10489,6 @@ packages:
requires_dist:
- pytest ; extra == 'tests'
requires_python: '>=3.8'
-- pypi: https://files.pythonhosted.org/packages/4e/d0/4b9da3d9c5749bf2c93f2818af4dc306f3b3ccc671f68f6820635fd53b74/cuequivariance_torch-0.11.1-py3-none-any.whl
- name: cuequivariance-torch
- version: 0.11.1
- sha256: 7a649eb59e623ef90f4862d84972a9aa12142892814f7a00dd652a1206628391
- requires_dist:
- - cuequivariance
- requires_python: '>=3.10'
- pypi: https://files.pythonhosted.org/packages/50/2e/ebdb85a8da151397a2790363676b7ed7c125924fe618e4c6d8befb0cc62c/msgpack-1.2.1-cp312-cp312-macosx_11_0_arm64.whl
name: msgpack
version: 1.2.1
@@ -10980,25 +10630,6 @@ packages:
version: 26.1.0
sha256: c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309
requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/66/2f/e762d2a929d8e22d96b937ab94ccd39dd2e20b2cf9dbd7a4fc29feb90bd1/cuequivariance_ops_torch_cu13-0.11.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
- name: cuequivariance-ops-torch-cu13
- version: 0.11.1
- sha256: c1219828f61e93ab72dadfba617227148a27c360f30f960310d347767f45d30d
- requires_dist:
- - cuequivariance-ops-cu13==0.11.1
- - scipy
- - torch>=2.11
- - numpy>=1.23,<2.0a0 ; extra == 'test'
- - pytest ; extra == 'test'
- - onnx ; extra == 'plugins'
- - onnxscript ; extra == 'plugins'
- - onnxruntime ; extra == 'plugins'
- - onnxruntime-extensions ; extra == 'plugins'
- - onnx-graphsurgeon ; extra == 'plugins'
- - tensorrt>=10 ; extra == 'plugins'
- - polygraphy ; extra == 'plugins'
- - cuda-bindings ; extra == 'plugins'
- requires_python: '>=3.10'
- pypi: https://files.pythonhosted.org/packages/67/4e/a28359bf1c1ecf11eba22123168c138698f7cb576ac678f5a2e16cd5da08/hf_xet-1.6.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
name: hf-xet
version: 1.6.0
@@ -11105,18 +10736,6 @@ packages:
requires_dist:
- numpy
- pillow
-- pypi: https://files.pythonhosted.org/packages/72/f1/b1f191b91d049780eb4db2605b81af28a697710c4ac219708bcba192f80c/cuequivariance_ops_cu13-0.11.1-py3-none-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
- name: cuequivariance-ops-cu13
- version: 0.11.1
- sha256: 2dc36684116ac1eb7ee835b7c58c7000df992f996d91566d033f2aab71263396
- requires_dist:
- - nvidia-cublas>=12.5.0
- - tqdm
- - nvidia-ml-py
- - platformdirs
- - numpy ; extra == 'test'
- - pytest ; extra == 'test'
- requires_python: '>=3.10'
- pypi: https://files.pythonhosted.org/packages/79/d3/36a46a8ed992b781acbc05928bd5bee3c810cb0c3563bf81a7b0c04a1a76/msgpack-1.2.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
name: msgpack
version: 1.2.1
@@ -11768,25 +11387,6 @@ packages:
version: 2026.7.19
sha256: d7da47a0f248977f08e2cb659ff3c17ddc13a4d39b3a7baa0a81bf5b415430f6
requires_python: '>=3.10'
-- pypi: https://files.pythonhosted.org/packages/b3/a1/42c6e78a654c28ba03932e7c7bfb2918ae5828488aae30322ced83379a19/cuequivariance_ops_torch_cu13-0.11.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
- name: cuequivariance-ops-torch-cu13
- version: 0.11.1
- sha256: 589d92f136437b295f65819c1d2af11ecbed1db19558ed9ee482291d1c837aae
- requires_dist:
- - cuequivariance-ops-cu13==0.11.1
- - scipy
- - torch>=2.11
- - numpy>=1.23,<2.0a0 ; extra == 'test'
- - pytest ; extra == 'test'
- - onnx ; extra == 'plugins'
- - onnxscript ; extra == 'plugins'
- - onnxruntime ; extra == 'plugins'
- - onnxruntime-extensions ; extra == 'plugins'
- - onnx-graphsurgeon ; extra == 'plugins'
- - tensorrt>=10 ; extra == 'plugins'
- - polygraphy ; extra == 'plugins'
- - cuda-bindings ; extra == 'plugins'
- requires_python: '>=3.10'
- pypi: https://files.pythonhosted.org/packages/b5/64/7660f8a4a8e53c924d0fa05dc3a55c9cee10bbd82b11c5afb27d44b096ce/markupsafe-3.0.3-cp314-cp314-macosx_11_0_arm64.whl
name: markupsafe
version: 3.0.3
@@ -12143,19 +11743,6 @@ packages:
version: 13.1.0.3
sha256: ee8722c1f0145ab246bccb9e452153b5e0515fd094c3678df50b2a0888b8b171
requires_python: '>=3'
-- pypi: https://files.pythonhosted.org/packages/e7/a8/7f99cfa83f61d2d8cb29f31e1f90a6233873a6038b4cdfb30179ef90e5a6/cuequivariance-0.11.1-py3-none-any.whl
- name: cuequivariance
- version: 0.11.1
- sha256: b570d4dc6f798a06c7333283a4b36e580272985abcb65843339984a84544395c
- requires_dist:
- - networkx
- - numpy
- - opt-einsum
- - scipy
- - sympy
- - graphviz ; extra == 'dev'
- - pytest ; extra == 'dev'
- requires_python: '>=3.10'
- pypi: https://files.pythonhosted.org/packages/eb/dc/55481808fd70ef1567cf13540ffd4702af3f74b112e35427564b03f79c2d/narwhals-2.25.0-py3-none-any.whl
name: narwhals
version: 2.25.0
diff --git a/pyproject.toml b/pyproject.toml
index 6a14c6bf..6ae059fe 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
[project]
name = "esm"
-version = "3.4.1.post1"
+version = "3.4.2"
description = "EvolutionaryScale open model repository"
readme = "README.md"
requires-python = ">=3.12"
@@ -51,20 +51,10 @@ dependencies = [
"pygtrie",
"dna_features_viewer",
"accelerate",
- "cuequivariance-torch>=0.8.1; sys_platform == 'linux' and platform_machine == 'x86_64'",
- "cuequivariance-ops-torch-cu13>=0.8.1; sys_platform == 'linux' and platform_machine == 'x86_64'",
+ # Pinned to what torch 2.11 requires, so CPU-index torch gets the same one.
+ "triton==3.6.*; sys_platform == 'linux' and platform_machine == 'x86_64'",
]
-[project.optional-dependencies]
-# ESMFold2 "fused" kernel backend (set_kernel_backend("fused"))
-# inference kernels for tri-mul / LN+SwiGLU / dropout-residual.
-fused = ["triton>=3,<4"]
-# ESMFold2 context-parallel ESM-C tensor parallelism (wrap_model_with_cp(tp_esmc=True)).
-fold-cp = ["transformer-engine[pytorch]>=2,<3"]
-# "cuequivariance" kernel backend — pick the build matching your CUDA toolkit.
-cueq12 = ["cuequivariance-torch", "cuequivariance-ops-torch-cu12"]
-cueq13 = ["cuequivariance-torch", "cuequivariance-ops-torch-cu13"]
-
# Pytest
[tool.pytest.ini_options]
addopts = """
@@ -75,9 +65,9 @@ addopts = """
--ignore=tests/oss_pytests/test_oss_client.py
"""
markers = [
- "gpu: requires a CUDA device; run in the monorepo's gpu-ci",
- "manual: too heavy for a hosted runner; run in the monorepo's research-ci",
- "merge_only: needs a real published checkpoint; run in the monorepo's research-ci",
+ "gpu: requires a CUDA device",
+ "manual: run by hand, never by CI",
+ "merge_only: needs a real published checkpoint (network or staged weights)",
]
[tool.setuptools]
@@ -135,6 +125,7 @@ matplotlib = "*"
numpy = ">=2.0.0,<2.4"
pre-commit = "*"
pytest = "*"
+pytest-asyncio = "*"
pytest-cov = "*"
pytest-xdist = "*"
seaborn = "*"
@@ -149,20 +140,14 @@ dockq = { git = "https://github.com/nrontsis/DockQ.git", rev = "ba4df5adaad7c77f
lint-all = "pre-commit run --all-files --show-diff-on-failure"
cov-test = "pytest -v --junitxml=pytest.xml --cov=esm"
-[tool.pixi.feature.fused]
-platforms = ["linux-64"]
-
-[tool.pixi.feature.fused.pypi-dependencies]
+# PyPI rejects URL requirements, so pip users get this command from a warning in
+# esm/models/esmc/kernels.py instead. The wheel is cp312-only.
+[tool.pixi.feature.py312.target.linux-64.pypi-dependencies]
flash_attn = { url = "https://github.com/evolutionaryscale/wheels/releases/download/py312-pt211-cu13-sm80-90/flash_attn-2.7.4.post1-cp312-cp312-linux_x86_64.whl" }
-xformers = { url = "https://download.pytorch.org/whl/cu130/xformers-0.0.35-py39-none-manylinux_2_28_x86_64.whl" }
-cuequivariance-torch = ">=0.8.1"
-cuequivariance-ops-torch-cu13 = ">=0.8.1"
-# `gpu` stays on 3.12: the prebuilt flash_attn wheel in `fused` is cp312-only.
[tool.pixi.environments]
default = {features = ["py312"], solve-group = "default"}
dev = {features = ["dev", "py312"], solve-group = "default"}
-gpu = {features = ["dev", "fused", "py312"], solve-group = "default"}
dev-py313 = {features = ["dev", "py313"], solve-group = "py313"}
dev-py314 = {features = ["dev", "py314"], solve-group = "py314"}
@@ -232,13 +217,6 @@ replace-imports-with-any = [
"biotite.sequence.align",
"biotite.sequence.align.**",
"scipy.spatial",
- "transformer_engine",
- "transformer_engine.**",
- "transformer_engine_torch",
- "xformers",
- "xformers.**",
- "cuequivariance_torch",
- "cuequivariance_torch.**",
"flash_attn",
"flash_attn.**",
"zstd",
diff --git a/tests/Makefile b/tests/Makefile
index 9e01e5a3..0c3f1b76 100644
--- a/tests/Makefile
+++ b/tests/Makefile
@@ -3,9 +3,19 @@ DOCKER_TAG ?= dev
DOCKER_IMAGE_OSS=oss_pytests:${DOCKER_TAG}
INFRA_PROVIDER ?= AWS
+CACHE_REF ?=
+CACHE_BRANCH ?= main
+CACHE_NAME ?= oss-sdk
+comma := ,
+_CACHE_FROM = $(if $(CACHE_REF),--cache-from type=registry$(comma)ref=$(CACHE_REF):cache-$(CACHE_NAME)-$(CACHE_BRANCH) --cache-from type=registry$(comma)ref=$(CACHE_REF):cache-$(CACHE_NAME)-main)
+_CACHE_TO = $(if $(CACHE_REF),--cache-to type=registry$(comma)ref=$(CACHE_REF):cache-$(CACHE_NAME)-$(CACHE_BRANCH)$(comma)mode=max$(comma)compression=zstd$(comma)compression-level=3$(comma)force-compression=true)
+
build-oss-ci:
- docker build \
+ docker buildx build \
--output=type=docker \
+ --provenance=false \
+ $(_CACHE_FROM) \
+ $(_CACHE_TO) \
-f oss_pytests/Dockerfile \
-t $(DOCKER_IMAGE_OSS) \
oss_pytests
@@ -15,7 +25,8 @@ start-docker-oss:
--rm \
-e INFRA_PROVIDER=${INFRA_PROVIDER} \
-e URL=${URL} \
- -e ESM_API_KEY=${ESM_API_KEY} \
+ -e ESM_API_KEY \
+ -e CUDA_VISIBLE_DEVICES="" \
--name=$(USER)-oss_pytests \
--network=host \
${DOCKER_IMAGE_OSS}
diff --git a/tests/compatibility/esmfold2_hf_adapter_test.py b/tests/compatibility/esmfold2_hf_adapter_test.py
index 86a2c07b..57640269 100644
--- a/tests/compatibility/esmfold2_hf_adapter_test.py
+++ b/tests/compatibility/esmfold2_hf_adapter_test.py
@@ -452,31 +452,3 @@ def test_without_the_port_installed_the_error_says_why(call):
"""The module still imports; only the two entry points needing their classes fail."""
with pytest.raises(ImportError, match="transformers 5.16"):
call()
-
-
-@pytest.mark.skipif(
- not upstream_available(),
- reason="needs transformers >= 5.16 (the ESMFold2 port); this repo pins 4.x",
-)
-def test_against_the_real_port(tiny_esmfold2_config):
- """The one thing the stub cannot cover: their classes construct and their
- ``fold`` accepts what the adapter sends. Weights are out of scope."""
- from transformers.models.esmfold2.modeling_esmfold2 import (
- EsmFold2Model as UpstreamEsmFold2Model,
- )
-
- adapter = EsmFold2HFAdapter(UpstreamEsmFold2Model(tiny_esmfold2_config).eval())
- features = prepare_protein_features(TINY_SEQUENCE)
- lm_hidden_states = torch.zeros(
- 1,
- features["res_type"].shape[-1],
- adapter.config.esmc_config.num_hidden_layers + 1,
- adapter.config.esmc_config.hidden_size,
- )
-
- with torch.no_grad():
- output = adapter(**features, lm_hidden_states=lm_hidden_states, **FAST_FOLD)
-
- assert "sample_atom_coords" in output
- for key in UPSTREAM_ABSENT_OUTPUT_KEYS:
- assert key in output
diff --git a/tests/compatibility/esmfold2_hf_checkpoint_test.py b/tests/compatibility/esmfold2_hf_checkpoint_test.py
index b0c4a586..00f82d10 100644
--- a/tests/compatibility/esmfold2_hf_checkpoint_test.py
+++ b/tests/compatibility/esmfold2_hf_checkpoint_test.py
@@ -15,7 +15,6 @@
from esm.models.esmfold2 import EsmFold2Config, EsmFold2Model
from esm.models.esmfold2.config import default_module_flags
-from esm.models.esmfold2.hf_adapter import upstream_available
from esm.models.esmfold2.hf_checkpoint import (
hf_state_dict_to_native,
is_hf_layout,
@@ -203,55 +202,3 @@ def test_an_unknown_keyword_to_the_config_reader_raises(hf_config):
def test_an_unknown_keyword_to_the_model_reader_raises(hf_config):
with pytest.raises(TypeError, match="some_knob"):
EsmFold2Model.from_pretrained(hf_config, some_knob=3)
-
-
-# ---------------------------------------------------------------------------
-# Against the port itself
-# ---------------------------------------------------------------------------
-
-
-@pytest.mark.skipif(
- not upstream_available(),
- reason="needs transformers >= 5.16 (the ESMFold2 port); this repo pins 4.x",
-)
-def test_rules_invert_their_converter(tiny_esmfold2_config):
- """The rules have to reproduce their converter key for key.
-
- Their ``rename_trunk_keys`` is the only source of truth for the mapping, so
- it generates the expectation here rather than a table being maintained by
- hand. When upstream moves a name, this fails naming the key that moved.
- """
- from transformers.models.esmfold2.convert_esmfold2_checkpoint import (
- rename_trunk_keys,
- )
-
- tiny_esmfold2_config.msa_encoder.enabled = True
- with torch.device("meta"):
- ours = set(EsmFold2Model(tiny_esmfold2_config).state_dict())
- theirs = rename_trunk_keys(
- dict(EsmFold2Model(tiny_esmfold2_config).state_dict())
- )
-
- got = set()
- for key in theirs:
- native = to_native_key(key)
- got.add(native)
- if native.endswith(".a_proj.weight"):
- got.add(native.replace(".a_proj.", ".b_proj."))
- assert got == ours
-
-
-@pytest.mark.manual
-def test_a_re_exported_port_checkpoint_loads(hf_port_checkpoint_dir):
- """End to end on the real artifact: ~26 GB, so not part of any suite.
-
- Its ``config.json`` states no module flags, so this also exercises the
- warn-and-default path against a checkpoint that really does carry the
- weights those defaults assume. ``_load_pretrained`` is strict, so reaching
- this line at all means every tensor found a parameter and every parameter
- was filled.
- """
- model = EsmFold2Model.from_pretrained(hf_port_checkpoint_dir, esmc_precision="fp32")
-
- assert model.esmc is not None
- assert json.loads(json.dumps(model.config.to_dict()))["hidden_size"] > 0
diff --git a/tests/conftest.py b/tests/conftest.py
index 834d19c8..73806f24 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,14 +1,16 @@
"""Shared fixtures for the ESMC and ESMFold2 test suites.
Tiny randomly-initialised models cover the structural contracts on CPU in
-milliseconds. The published-weight fixtures download once per session and skip
-when the Hub is unreachable, so an offline runner degrades to the structural
-tests instead of failing.
+milliseconds. Anything fetching from the Hub is marked ``merge_only`` and runs
+in the lane that has network; it warns when unreachable, and fails outright
+under ``ESM_TESTS_STRICT_ARTIFACTS``.
"""
import os
import random
+import warnings
from pathlib import Path
+from typing import NoReturn
import pytest
import torch
@@ -27,10 +29,12 @@
)
TINY_SEQUENCES = ["MQIFVKTLTGKT", "MKV"]
-
ESMC_300M_REPO = "biohub/ESMC-300M"
ESMFOLD2_REPO = "biohub/ESMFold2"
ESMFOLD2_EXPERIMENTAL_REPO = "biohub/ESMFold2-Experimental"
+#: ``PROTEIN_REF_POS`` is a frozen copy of this snapshot. Move both or neither,
+#: and never track ``main``: the release changes per-residue atom ordering.
+ESMFOLD2_CCD_REVISION = "69869f737beffec5294845ede23db5fc0b4f509e"
# The ESMFold2 weights in the upstream transformers layout. A staging repo:
# biohub/ESMFold2 has not been republished through their converter yet.
ESMFOLD2_HF_PORT_REPO = "Rocketknight1/ESMFold2-merged-temp"
@@ -117,6 +121,50 @@ def staged_model_dir(subdirectory: str) -> str | None:
return str(directory) if (directory / "config.json").exists() else None
+# Skip reasons carry their own classification: xdist ships reports back as text,
+# so a module-level registry would only ever fill up inside the worker.
+# A missing GPU or optional package stays a plain `pytest.skip`.
+ARTIFACT_UNAVAILABLE = "ARTIFACT UNAVAILABLE"
+KNOWN_GAP = "KNOWN GAP"
+
+#: Set by CI, where every artifact is reachable by construction.
+STRICT_ARTIFACTS = os.environ.get("ESM_TESTS_STRICT_ARTIFACTS") == "1"
+
+
+def artifact_unavailable(reason: str) -> NoReturn:
+ """An input could not be fetched. Warns always, fails under CI."""
+ message = f"{ARTIFACT_UNAVAILABLE}: {reason}"
+ if STRICT_ARTIFACTS:
+ pytest.fail(message, pytrace=False) # ty:ignore[invalid-argument-type,parameter-already-assigned]
+ warnings.warn(message, stacklevel=2)
+ pytest.skip(message) # ty:ignore[too-many-positional-arguments]
+
+
+def known_gap(reason: str) -> NoReturn:
+ """A gap we have accepted. Listed in the summary, never fatal."""
+ pytest.skip(f"{KNOWN_GAP}: {reason}") # ty:ignore[too-many-positional-arguments]
+
+
+def pytest_terminal_summary(terminalreporter, exitstatus, config):
+ """Name and count these skips at the end; a status-line tally is easy to miss."""
+ buckets: dict[str, list[str]] = {ARTIFACT_UNAVAILABLE: [], KNOWN_GAP: []}
+ for report in terminalreporter.stats.get("skipped", []):
+ longrepr = getattr(report, "longrepr", None)
+ reason = longrepr[2] if isinstance(longrepr, tuple) else ""
+ for label in buckets:
+ if label in reason:
+ buckets[label].append(f"{report.nodeid}\n {reason}")
+
+ for label, entries in buckets.items():
+ if not entries:
+ continue
+ terminalreporter.section(
+ f"{label} ({len(entries)} skipped)", sep="!", red=True, bold=True
+ )
+ for entry in entries:
+ terminalreporter.line(entry)
+
+
def _skip_if_unreachable(repo: str) -> str:
"""Resolve a Hub snapshot, or skip when the runner genuinely has no access.
@@ -138,18 +186,47 @@ def _skip_if_unreachable(repo: str) -> str:
RequestsConnectionError,
Timeout,
) as exc:
- pytest.skip(f"{repo} unreachable ({type(exc).__name__}); skipping") # ty:ignore[too-many-positional-arguments]
+ artifact_unavailable(f"{repo} ({type(exc).__name__})")
-#: Fixtures that resolve a real published checkpoint. Anything depending on one
-#: needs network (or staged weights) and takes seconds rather than milliseconds.
+#: Fixtures that fetch from the Hub. Anything using one needs network, so it is
+#: marked ``merge_only`` below and runs in the lane that has it.
WEIGHTS_FIXTURES = frozenset(
- {"esmc_300m_dir", "esmfold2_dir", "esmfold2_experimental_dir"}
+ {"ccd_pickle", "esmc_300m_dir", "esmfold2_dir", "esmfold2_experimental_dir"}
)
+def resolve_ccd_path() -> str | None:
+ """The pinned CCD snapshot, or ``None`` when the runner cannot reach it.
+
+ Beats an ambient ``ESMCFOLD_CCD_PATH``, which named three different releases
+ across ``.bashrc``, ``.env`` and the old default - so how you launched pytest
+ decided the numbers. Ambient is the fallback only when the Hub is down.
+ """
+ from huggingface_hub import hf_hub_download
+ from huggingface_hub.errors import LocalEntryNotFoundError, OfflineModeIsEnabled
+ from requests.exceptions import ConnectionError as RequestsConnectionError
+ from requests.exceptions import Timeout
+
+ try:
+ return hf_hub_download(ESMFOLD2_REPO, "ccd.pkl", revision=ESMFOLD2_CCD_REVISION)
+ except (
+ LocalEntryNotFoundError,
+ OfflineModeIsEnabled,
+ RequestsConnectionError,
+ Timeout,
+ ):
+ return os.environ.get("ESMCFOLD_CCD_PATH") or None
+
+
def pytest_configure(config):
- """Point ``conformers.py`` at the mounted CCD before it is imported."""
+ """Point ``conformers.py`` at the pinned CCD before it is imported.
+
+ Not a fixture: ``conformers.py`` reads the variable at import time.
+ """
+ path = resolve_ccd_path()
+ if path:
+ os.environ["ESMCFOLD_CCD_PATH"] = path
def pytest_collection_modifyitems(items):
@@ -387,19 +464,13 @@ def esmfold2_complex_features(model, chains, seed: int = 0):
def ccd_pickle() -> str:
"""The CCD pickle every non-protein featurization path needs.
- ``load_ccd`` reaches for this on its own, so a test that looks offline is
- not; depending on this fixture makes that reach explicit and skips the test
- when the pickle is not mounted, which is the case on a CI runner.
-
- Deliberately no Hub fallback. ``biohub/ESMFold2:ccd.pkl`` is a *different*
- CCD release from the mounted one, and the release changes per-residue atom
- ordering - enough to move the reversed-atom controls in
- ``esmfold2_sampler_test.py`` by several Angstrom. Silently substituting it
- would make those numbers depend on which copy the runner happened to find.
+ ``load_ccd`` reaches for this on its own, so depending on the fixture makes
+ that reach explicit. Was a cluster mount no CI runner has, which skipped 46
+ items silently.
"""
- path = os.environ.get("ESMCFOLD_CCD_PATH", "")
+ path = resolve_ccd_path()
if not path or not Path(path).exists():
- pytest.skip(f"CCD pickle not mounted at {path}") # ty:ignore[too-many-positional-arguments]
+ artifact_unavailable(f"{ESMFOLD2_REPO}:ccd.pkl at {ESMFOLD2_CCD_REVISION[:12]}")
return path
diff --git a/tests/models/esmc_test.py b/tests/models/esmc_test.py
index 9e89188d..39875b1a 100644
--- a/tests/models/esmc_test.py
+++ b/tests/models/esmc_test.py
@@ -7,6 +7,7 @@
import functools
import gzip
+import itertools
import json
import math
import pickle
@@ -20,6 +21,7 @@
import torch.nn as nn
from safetensors.torch import save_file
+from esm.layers.rotary import TritonRotaryEmbedding, apply_rotary_emb_torch
from esm.models.esmc import (
ESMC,
EsmcConfig,
@@ -28,16 +30,22 @@
EsmcMaskedLMOutput,
EsmcModel,
EsmcOutput,
- EsmcTokenizer,
)
from esm.models.esmc import model as model_module
from esm.models.esmc.checkpoint_layout import native_to_published, published_to_native
-from esm.models.esmc.kernels import FLASH_ATTN_INSTALLED, TE_INSTALLED
+from esm.models.esmc.kernels import FLASH_ATTN_INSTALLED, FLASH_ATTN_ROTARY_INSTALLED
from esm.models.esmc.layers import (
EsmcFlashMultiHeadAttention,
EsmcLayerNormLinear,
EsmcLayerNormMLP,
EsmcRotaryEmbedding,
+ EsmcTritonRotaryEmbedding,
+)
+from esm.models.esmfold2.fp8 import (
+ Fp8LayerNormLinear,
+ Fp8LayerNormMLP,
+ Fp8Linear,
+ convert_esmc_to_fp8_,
)
from esm.models.hub import CONFIG_NAME, read_safetensors_dir
from tests.conftest import (
@@ -57,10 +65,6 @@
# O(L^2) tokens. 16 evenly spaced positions hold every length to a few seconds.
PPL_MAX_POSITIONS = 16
-# Which layers ``build_esmc`` puts in the model.
-FUSED_LAYERS = "fused"
-REFERENCE_LAYERS = "reference"
-
def _tokens(tokenizer, sequences):
return tokenizer(sequences, return_tensors="pt", padding=True)
@@ -140,21 +144,12 @@ def build_esmc(
device: str,
dtype: torch.dtype | None = None,
attn: str = "sdpa",
- layers: str = FUSED_LAYERS,
model_class=EsmcModel,
):
- """Load ESMC with an explicit choice of fused vs pure-PyTorch layers.
-
- ``EsmcModel.__init__`` reads ``torch.get_default_device()`` to pick its
- kernels, so building for the CPU and then moving is the only way to get the
- reference layers onto CUDA. That also disables the flash-attention block, so
- ``layers="reference"`` is only meaningful with ``attn="sdpa"`` / ``"eager"``.
- """
- build_device = "cpu" if layers == REFERENCE_LAYERS else device
+ """Load ESMC on ``device`` with the given attention backend."""
model = model_class.from_pretrained(
- directory, device=build_device, attn_implementation=attn
+ directory, device=device, attn_implementation=attn
)
- model = model.to(device)
if dtype is not None:
model = model.to(dtype)
return model.eval()
@@ -341,6 +336,147 @@ def test_gradients_reach_every_parameter(tiny_esmc_config, esmc_tokenizer):
]
+@pytest.mark.gpu
+def test_rotary_is_differentiable_on_cuda():
+ """The CUDA RoPE branch must be autograd-aware.
+
+ The raw Triton driver (``flash_attn.ops.triton.rotary.apply_rotary``) returns a
+ tensor with no ``grad_fn``, so a backward through ESMC silently drops the
+ attention-score path. Only the CUDA branch can regress; CPU is plain torch.
+ """
+ # Without flash-attn the module falls back to torch and the test passes without
+ # ever running the branch it exists to cover.
+ assert FLASH_ATTN_ROTARY_INSTALLED
+
+ torch.manual_seed(0)
+ rotary = EsmcRotaryEmbedding(64).cuda()
+ q = torch.randn(2, 16, 4, 64, device="cuda", requires_grad=True)
+ k = torch.randn(2, 16, 4, 64, device="cuda", requires_grad=True)
+
+ q_rot, k_rot = rotary(q, k)
+
+ assert q_rot.grad_fn is not None and k_rot.grad_fn is not None
+ q_grad, k_grad = torch.autograd.grad((q_rot.sum() + k_rot.sum()), (q, k))
+ assert q_grad.abs().sum() > 0 and k_grad.abs().sum() > 0
+
+ # The rotation must be in the backward too, not just an identity pass-through.
+ rotary._update_cos_sin_cache(16, device=q.device, dtype=q.dtype)
+ q_ref = q.detach().clone().requires_grad_(True)
+ expected = torch.autograd.grad(
+ apply_rotary_emb_torch(
+ q_ref, rotary._cos_cached, rotary._sin_cached, rotary.interleaved
+ )
+ .square()
+ .sum(),
+ q_ref,
+ )[0]
+ q_sq = q.detach().clone().requires_grad_(True)
+ actual = torch.autograd.grad(rotary(q_sq, q_sq)[0].square().sum(), q_sq)[0]
+ torch.testing.assert_close(actual, expected, atol=1e-4, rtol=1e-4)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+ "rotary_cls",
+ [EsmcTritonRotaryEmbedding, TritonRotaryEmbedding],
+ ids=["esmc", "legacy"],
+)
+def test_varlen_triton_rotary_gradients_match_per_sequence_torch(
+ rotary_cls: type[EsmcTritonRotaryEmbedding] | type[TritonRotaryEmbedding],
+):
+ """The varlen RoPE backward must carry the rotation, not just the forward.
+
+ The old raw-driver path produced a correct forward and an identity backward, so
+ only a gradient oracle separates the two: this test's error is ~6 before the fix
+ and ~1e-7 after. :func:`apply_rotary_emb_torch` is applied per sequence because
+ ``cu_seqlens`` restarts positions at every boundary.
+
+ Parameters
+ ----------
+ rotary_cls : type
+ Varlen Triton RoPE module under test. Both classes wrap the same kernel
+ behind the same signature, so they share one oracle.
+ """
+ assert FLASH_ATTN_ROTARY_INSTALLED
+
+ torch.manual_seed(0)
+ seqlens = [7, 13, 5]
+ n_heads, head_dim = 4, 64
+ bounds = list(itertools.pairwise(itertools.accumulate(seqlens, initial=0)))
+ total = sum(seqlens)
+ cu_seqlens = torch.tensor(
+ [0, *itertools.accumulate(seqlens)], dtype=torch.int32, device="cuda"
+ )
+
+ q = torch.randn(total, n_heads, head_dim, device="cuda", requires_grad=True)
+ k = torch.randn(total, n_heads, head_dim, device="cuda", requires_grad=True)
+ # Weighting by a fixed random covector makes the loss rotation-dependent. A
+ # rotation-invariant loss such as sum-of-squares cannot see an identity backward.
+ w_q = torch.randn_like(q.detach())
+ w_k = torch.randn_like(k.detach())
+
+ rotary = rotary_cls(head_dim).cuda()
+ q_rot, k_rot = rotary(q, k, cu_seqlens, max(seqlens))
+ q_grad, k_grad = torch.autograd.grad(
+ (q_rot * w_q).sum() + (k_rot * w_k).sum(), (q, k)
+ )
+
+ assert rotary._cos_cached is not None and rotary._sin_cached is not None
+ q_ref = q.detach().clone().requires_grad_(True)
+ k_ref = k.detach().clone().requires_grad_(True)
+ expected = [
+ torch.cat(
+ [
+ apply_rotary_emb_torch(
+ x[start:end].unsqueeze(0),
+ rotary._cos_cached,
+ rotary._sin_cached,
+ rotary.interleaved,
+ ).squeeze(0)
+ for start, end in bounds
+ ]
+ )
+ for x in (q_ref, k_ref)
+ ]
+ torch.testing.assert_close(q_rot, expected[0], atol=1e-5, rtol=1e-5)
+ torch.testing.assert_close(k_rot, expected[1], atol=1e-5, rtol=1e-5)
+
+ q_grad_ref, k_grad_ref = torch.autograd.grad(
+ (expected[0] * w_q).sum() + (expected[1] * w_k).sum(), (q_ref, k_ref)
+ )
+ torch.testing.assert_close(q_grad, q_grad_ref, atol=1e-5, rtol=1e-5)
+ torch.testing.assert_close(k_grad, k_grad_ref, atol=1e-5, rtol=1e-5)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("attn_implementation", ["sdpa", "flash_attention_2"])
+def test_gradients_reach_every_parameter_on_cuda(
+ tiny_esmc_config, esmc_tokenizer, attn_implementation
+):
+ """CPU-only grad coverage cannot see the CUDA RoPE branch, so repeat it here."""
+ if attn_implementation == "flash_attention_2" and not FLASH_ATTN_INSTALLED:
+ pytest.skip("flash-attn is not installed") # ty:ignore[too-many-positional-arguments]
+
+ torch.manual_seed(0)
+ with torch.device("cuda"):
+ model = EsmcForMaskedLM(
+ replace(tiny_esmc_config, attn_implementation=attn_implementation)
+ ).train()
+ if attn_implementation == "flash_attention_2":
+ # The fused attention kernel rejects fp32.
+ model = model.to(torch.bfloat16)
+ batch = {k: v.cuda() for k, v in _tokens(esmc_tokenizer, TINY_SEQUENCES).items()}
+
+ model(**batch).logits.square().mean().backward()
+
+ dead = [
+ name
+ for name, p in model.named_parameters()
+ if p.grad is None or not torch.isfinite(p.grad).all() or p.grad.eq(0).all()
+ ]
+ assert not dead
+
+
def test_sequence_classification_head(tiny_esmc_config, esmc_tokenizer):
"""The head reads ````, so it must emit one row per sequence."""
torch.manual_seed(0)
@@ -1133,7 +1269,6 @@ def test_attention_mask_folds_into_a_multi_chain_sequence_id(tiny_esmc, esmc_tok
def _pretend_installed(monkeypatch) -> None:
"""Patch where the flags are read, not where they are defined."""
- monkeypatch.setattr(model_module, "TE_INSTALLED", True)
monkeypatch.setattr(model_module, "FLASH_ATTN_INSTALLED", True)
@@ -1192,19 +1327,6 @@ def test_eager_and_sdpa_build_the_same_attention(tiny_esmc_config):
# ---------------------------------------------------------------------------
-@pytest.mark.gpu
-def test_state_dict_round_trips_under_transformer_engine(tiny_esmc_config):
- """TE adds ``_extra_state`` to ``state_dict`` and a hook strips it, so the
- loader has to be told not to expect it back."""
- if not TE_INSTALLED:
- pytest.skip("TransformerEngine unavailable") # ty:ignore[too-many-positional-arguments]
- with torch.device("cuda"):
- model = EsmcForMaskedLM(tiny_esmc_config).eval()
-
- assert not [k for k in model.state_dict() if k.endswith("_extra_state")]
- model.load_state_dict(model.state_dict())
-
-
@pytest.mark.gpu
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize("name", LENGTH_NAMES)
@@ -1222,28 +1344,50 @@ def test_gpu_dtypes(tiny_esmc_config, esmc_tokenizer, dtype, name):
assert torch.isfinite(out.float()).all()
-@pytest.mark.gpu
-def test_fp8_backbone(esmc_300m_dir):
- """ESMC is commonly run fp8 as a folding backbone."""
- if not TE_INSTALLED:
- pytest.skip("TransformerEngine unavailable") # ty:ignore[too-many-positional-arguments]
- if torch.cuda.get_device_capability()[0] < 9:
- pytest.skip("fp8 needs compute capability >= 9.0 (Hopper)") # ty:ignore[too-many-positional-arguments]
+def _tensor_bytes(model: nn.Module) -> int:
+ tensors = list(model.parameters()) + list(model.buffers())
+ return sum(t.numel() * t.element_size() for t in tensors)
+
- from esm.models.esmfold2.model import _convert_te_modules_to_fp8_inplace
+@pytest.mark.gpu
+def test_fp8_backbone(esmc_300m_dir, esmc_tokenizer):
+ """ESMC is commonly run fp8 as a folding backbone. Every projection must end
+ up in fp8, and the output must stay close to bf16."""
+ if torch.cuda.get_device_capability() < (8, 9):
+ pytest.skip("fp8 needs compute capability >= 8.9") # ty:ignore[too-many-positional-arguments]
+ reference = EsmcModel.from_pretrained(
+ esmc_300m_dir, device="cuda", dtype=torch.bfloat16
+ ).eval()
model = EsmcModel.from_pretrained(
esmc_300m_dir, device="cuda", dtype=torch.bfloat16
- )
- with torch.no_grad():
- _convert_te_modules_to_fp8_inplace(model)
- tokenizer = EsmcTokenizer()
+ ).eval()
+ bf16_bytes = _tensor_bytes(model)
+ n_blocks = len(model.transformer.blocks)
+ # qkv, out_proj and the MLP per block.
+ assert convert_esmc_to_fp8_(model) == 3 * n_blocks
+ for block in model.transformer.blocks:
+ assert isinstance(block.attn.layernorm_qkv, Fp8LayerNormLinear) # ty:ignore[unresolved-attribute]
+ assert isinstance(block.attn.out_proj, Fp8Linear) # ty:ignore[unresolved-attribute]
+ assert isinstance(block.ffn, Fp8LayerNormMLP)
+ assert not [
+ name
+ for name, p in model.named_parameters()
+ if p.dim() == 2 and name != "embed.weight"
+ ], "a 2D weight was left in bf16"
+ assert _tensor_bytes(model) < 0.6 * bf16_bytes
+
enc = {
- k: v.cuda() for k, v in tokenizer(SHORT_SEQUENCE, return_tensors="pt").items()
+ k: v.cuda()
+ for k, v in esmc_tokenizer(MEDIUM_SEQUENCE, return_tensors="pt").items()
}
- with torch.no_grad():
- out = model(**enc).last_hidden_state
- assert torch.isfinite(out.float()).all()
+ with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
+ expected = reference(**enc).last_hidden_state.float()
+ actual = model(**enc).last_hidden_state.float()
+ assert actual.dtype == torch.float32 and torch.isfinite(actual).all()
+ # Measured 0.13 on this input; TE's fp8 recipe had more error than this path.
+ rel = (actual - expected).norm() / expected.norm()
+ assert rel < 0.15, rel
@pytest.mark.gpu
@@ -1283,9 +1427,10 @@ def test_flash_attention_2_dispatches_and_agrees_with_sdpa(
with torch.no_grad():
expected = reference(**enc).last_hidden_state.float()
actual = candidate(**enc).last_hidden_state.float()
- # The flash path packs the batch and runs flash_attn_varlen where sdpa
- # runs xformers, so the reduction orders differ and the gap compounds
- # over 30 blocks. Held to a few bf16 ULP at activation scale.
+ # The flash path packs the batch and runs flash_attn_varlen where the
+ # "sdpa" build runs attention unpacked, so the reduction orders differ
+ # and the gap compounds over 30 blocks. Held to a few bf16 ULP at
+ # activation scale.
torch.testing.assert_close(
actual, expected, atol=5 * torch.finfo(torch.bfloat16).eps, rtol=1e-2
)
@@ -1293,9 +1438,8 @@ def test_flash_attention_2_dispatches_and_agrees_with_sdpa(
@pytest.mark.gpu
def test_cpu_inference_works_on_a_cuda_machine(esmc_300m_dir, esmc_tokenizer):
- """Transformer Engine is selected whenever it imports and CUDA exists, but its
- modules reject CPU inputs, so a CPU-targeted load has to opt out of them.
- """
+ """Fused kernels are CUDA-only, so a CPU-targeted load on a CUDA machine
+ has to opt out of them."""
model = EsmcModel.from_pretrained(esmc_300m_dir, device="cpu")
enc = esmc_tokenizer(SHORT_SEQUENCE, return_tensors="pt")
with torch.no_grad():
@@ -1305,13 +1449,35 @@ def test_cpu_inference_works_on_a_cuda_machine(esmc_300m_dir, esmc_tokenizer):
@pytest.mark.gpu
-def test_cpu_and_gpu_agree(esmc_300m_dir, esmc_tokenizer):
- """Same weights, same input, both devices, fp32 on each side.
-
- Also a cross-implementation check: the CPU model uses the reference layers
- while the CUDA model uses the fused Transformer Engine ones, which is where
- almost all of the difference below comes from.
+def test_flash_build_moved_to_cpu_matches_a_cpu_build(tiny_esmc_config, esmc_tokenizer):
+ """The 3.x README loads on CUDA then calls ``.to("cpu")``; the CUDA-built
+ flash blocks must then fall back instead of handing CPU tensors to Triton.
"""
+ if not FLASH_ATTN_INSTALLED:
+ pytest.skip("flash-attn is not installed") # ty:ignore[too-many-positional-arguments]
+
+ with torch.device("cuda"):
+ moved = EsmcModel(
+ replace(tiny_esmc_config, attn_implementation="flash_attention_2")
+ ).eval()
+ assert moved._use_flash_attn, "flash attention did not actually dispatch"
+ moved = moved.to("cpu")
+ assert not moved._use_flash_attn
+
+ reference = EsmcModel(replace(tiny_esmc_config, attn_implementation="sdpa"))
+ reference.load_state_dict(moved.state_dict())
+ reference.eval()
+
+ batch = _tokens(esmc_tokenizer, TINY_SEQUENCES)
+ with torch.no_grad():
+ actual = moved(**batch).last_hidden_state
+ expected = reference(**batch).last_hidden_state
+ torch.testing.assert_close(actual, expected, atol=0.0, rtol=0.0)
+
+
+@pytest.mark.gpu
+def test_cpu_and_gpu_agree(esmc_300m_dir, esmc_tokenizer):
+ """Same weights, same input, both devices, fp32 on each side."""
on_cpu = EsmcModel.from_pretrained(esmc_300m_dir, device="cpu")
on_gpu = EsmcModel.from_pretrained(esmc_300m_dir, device="cuda")
@@ -1328,19 +1494,15 @@ def test_cpu_and_gpu_agree(esmc_300m_dir, esmc_tokenizer):
torch.testing.assert_close(actual, expected, atol=3e-3, rtol=1e-4)
-# Every build reachable on a GPU: dtype x attention backend x fused/reference
-# layers. flash_attention_2 is only paired with fp16 / bf16 because the kernel
-# rejects fp32, and only with the fused layers because a CPU-built model never
-# gets the flash block at all (see ``build_esmc``).
+# Every build reachable on a GPU: dtype x attention backend. flash_attention_2
+# is only paired with fp16 / bf16 because the kernel rejects fp32.
GPU_BUILDS = [
- (torch.float32, "sdpa", FUSED_LAYERS),
- (torch.float32, "sdpa", REFERENCE_LAYERS),
- (torch.float32, "eager", FUSED_LAYERS),
- (torch.bfloat16, "sdpa", FUSED_LAYERS),
- (torch.bfloat16, "sdpa", REFERENCE_LAYERS),
- (torch.bfloat16, "flash_attention_2", FUSED_LAYERS),
- (torch.float16, "sdpa", FUSED_LAYERS),
- (torch.float16, "flash_attention_2", FUSED_LAYERS),
+ (torch.float32, "sdpa"),
+ (torch.float32, "eager"),
+ (torch.bfloat16, "sdpa"),
+ (torch.bfloat16, "flash_attention_2"),
+ (torch.float16, "sdpa"),
+ (torch.float16, "flash_attention_2"),
]
# What each dtype can be held to against the CPU / fp32 reference: max|Δ| on the
@@ -1350,8 +1512,6 @@ def test_cpu_and_gpu_agree(esmc_300m_dir, esmc_tokenizer):
# fp32 4.5e-4 0 / 0 / 0
# bf16 1.46e-2 1 / 3 / 3
# fp16 1.62e-3 0 / 1 / 0
-# fp32 with the reference layers reproduces the CPU reference to 1.7e-6; the fp32
-# bound has to accommodate Transformer Engine.
REFERENCE_NORM_TOLERANCE = {
torch.float32: 1e-3,
torch.bfloat16: 0.03,
@@ -1361,17 +1521,17 @@ def test_cpu_and_gpu_agree(esmc_300m_dir, esmc_tokenizer):
@pytest.mark.gpu
-@pytest.mark.parametrize("dtype,attn,layers", GPU_BUILDS)
+@pytest.mark.parametrize("dtype,attn", GPU_BUILDS)
def test_every_gpu_build_reaches_the_reference(
- esmc_300m_dir, esmc_tokenizer, dtype, attn, layers
+ esmc_300m_dir, esmc_tokenizer, dtype, attn
):
"""Correct, not merely self-consistent: every GPU build vs the CPU reference.
No GPU build can reach it bit-for-bit, so what is asserted is the decision
path (argmax over the vocabulary) plus a bound on the hidden-state norms.
- Without Transformer Engine or flash-attn installed the builds that would use
- them degenerate to the reference ones rather than failing, so this stays a
- superset of whatever a given runner can reach.
+ Without flash-attn installed the builds that would use it degenerate to
+ sdpa rather than failing, so this stays a superset of whatever a given
+ runner can reach.
"""
try:
model = build_esmc(
@@ -1379,14 +1539,13 @@ def test_every_gpu_build_reaches_the_reference(
device="cuda",
dtype=dtype,
attn=attn,
- layers=layers,
model_class=EsmcForMaskedLM,
)
# Only a genuinely absent package is a reason to skip. `from_pretrained`
# signals a strict-key failure with RuntimeError, so catching that here would
# turn a checkpoint regression into a green skip across every build.
except ImportError as exc:
- pytest.skip(f"build {dtype}/{attn}/{layers} unavailable: {exc}") # ty:ignore[too-many-positional-arguments]
+ pytest.skip(f"build {dtype}/{attn} unavailable: {exc}") # ty:ignore[too-many-positional-arguments]
for name, sequence in SEQUENCES.items():
case = reference_values()["cases"][name]
diff --git a/tests/models/esmfold2_api_test.py b/tests/models/esmfold2_api_test.py
index 4fde2696..d8e7bddf 100644
--- a/tests/models/esmfold2_api_test.py
+++ b/tests/models/esmfold2_api_test.py
@@ -15,7 +15,7 @@
from esm.models.esmfold2 import layers as _layers
from esm.models.esmfold2.prepare_input import prepare_esmfold2_input
-from esm.models.esmfold2.processor import ESMFold2InputBuilder
+from esm.models.esmfold2.processor import ESMFold2InputBuilder, _lm_dropout_context
from esm.models.esmfold2.protein_utils import (
OUTPUT_TO_PDB_FEATURE_KEYS,
prepare_protein_features,
@@ -459,3 +459,21 @@ def test_experimental_forward_supports_include_embeddings(
pooled = output["output_embedding_pair_pooled"]
assert pooled.shape == (1, len(TINY_SEQUENCE), _pair_width(tiny_esmfold2_config))
assert pooled.dtype == torch.float32
+
+
+def test_lm_dropout_zero_disables_and_restores(tiny_esmfold2):
+ lm_encoder = tiny_esmfold2.config.lm_encoder
+ before = (lm_encoder.lm_dropout, lm_encoder.per_loop_lm_dropout)
+
+ with _lm_dropout_context(tiny_esmfold2, 0):
+ assert lm_encoder.lm_dropout == 0
+
+ assert (lm_encoder.lm_dropout, lm_encoder.per_loop_lm_dropout) == before
+
+
+def test_lm_dropout_none_keeps_the_checkpoint_rate(tiny_esmfold2):
+ lm_encoder = tiny_esmfold2.config.lm_encoder
+ before = lm_encoder.lm_dropout
+
+ with _lm_dropout_context(tiny_esmfold2, None):
+ assert lm_encoder.lm_dropout == before
diff --git a/tests/models/esmfold2_builds_test.py b/tests/models/esmfold2_builds_test.py
index 9e494fe0..2bdddf55 100644
--- a/tests/models/esmfold2_builds_test.py
+++ b/tests/models/esmfold2_builds_test.py
@@ -266,9 +266,9 @@ def recorder(value, _name=name, _original=getattr(module, method_name)):
@pytest.mark.parametrize("variant", KERNEL_BACKEND_VARIANTS)
-@pytest.mark.parametrize("backend", [None, "fused", "cuequivariance"])
+@pytest.mark.parametrize("backend", [None, "fused"])
def test_set_kernel_backend_reaches_every_module(
- tiny_esmfold2_config, monkeypatch, variant, backend
+ tiny_esmfold2_config, variant, backend
):
"""Every module that *can* take a backend must be *given* one.
@@ -279,13 +279,7 @@ def test_set_kernel_backend_reaches_every_module(
the reference path, and every output comparison still passes because the
reference path is correct, merely slow. That is exactly the defect behind
``fix(esm): no-op set_kernel_backend on the CP diffusion wrapper``.
-
- ``CUE_AVAILABLE`` is forced on because ``TriangleMultiplicativeUpdate``
- refuses ``cuequivariance`` outright when the package is missing, which would
- reduce this case to a skip on every CPU runner. Nothing here executes a
- kernel, so the flag only unblocks the traversal.
"""
- monkeypatch.setattr(_layers, "CUE_AVAILABLE", True)
model = model_variant(variant, tiny_esmfold2_config).eval()
calls, expected = instrument_fanout(model, "set_kernel_backend")
@@ -484,9 +478,10 @@ def cpu_fp32_result(kernel_esmfold2_config):
# fp32, and that dominates.
#
# Observed at L=197, first the floor (two eager CUDA runs, so only CUDA's own
-# reduction-order nondeterminism separates them), then the worst case across all
-# seven GPU_BUILDS:
-# gpu-vs-gpu floor cpu-vs-gpu, all seven builds
+# reduction-order nondeterminism separates them), then the worst case across
+# every GPU_BUILDS row. Measured while the matrix still carried two
+# cuequivariance rows, so these are upper bounds for the five that remain:
+# gpu-vs-gpu floor cpu-vs-gpu, all builds
# sample_atom_coords 1.19e-03 A 9.96e-03 to 1.014e-02 A
# distogram_logits 2.72e-02 2.44
# pae 2.41e-02 2.28 to 2.30
@@ -522,8 +517,6 @@ def cpu_fp32_result(kernel_esmfold2_config):
(None, True, None),
("fused", False, None),
("fused", True, 64),
- ("cuequivariance", False, None),
- ("cuequivariance", True, 64),
]
@@ -558,9 +551,6 @@ def assert_backend_dispatched(model, backend):
for module in transitions:
assert module._fused_swiglu is None
- for tri_mul in tri_muls:
- assert tri_mul._engine._use_kernels == (backend == "cuequivariance")
-
def assert_chunking_dispatched(model, chunk_size):
engines = [
@@ -593,8 +583,6 @@ def skip_if_flash_missing(flash: bool):
def skip_if_backend_missing(backend):
if backend == "fused" and not _layers.TRITON_KERNELS_AVAILABLE:
pytest.skip("triton kernels unavailable") # ty:ignore[too-many-positional-arguments]
- if backend == "cuequivariance" and not _layers.CUE_AVAILABLE:
- pytest.skip("cuequivariance unavailable") # ty:ignore[too-many-positional-arguments]
@pytest.mark.gpu
diff --git a/tests/models/esmfold2_cpu_only_test.py b/tests/models/esmfold2_cpu_only_test.py
index ef3e6b53..ae91ee79 100644
--- a/tests/models/esmfold2_cpu_only_test.py
+++ b/tests/models/esmfold2_cpu_only_test.py
@@ -1,11 +1,3 @@
-# coding=utf-8
-# Copyright 2026 Biohub. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
"""ESMFold2 has to import and stay usable with no CUDA device.
``@triton.autotune`` resolves ``driver.active`` at *decoration* time, so merely
diff --git a/tests/models/esmfold2_fp8_test.py b/tests/models/esmfold2_fp8_test.py
new file mode 100644
index 00000000..e04beb8f
--- /dev/null
+++ b/tests/models/esmfold2_fp8_test.py
@@ -0,0 +1,111 @@
+"""fp8 ESMC backbone for ESMFold2, on native ``torch._scaled_mm``."""
+
+import pytest
+import torch
+
+from esm.models.esmfold2.fp8 import (
+ Fp8LayerNormLinear,
+ Fp8LayerNormMLP,
+ Fp8Linear,
+ _has_fp8_gemm,
+ _quantize_rows_torch,
+ convert_esmc_to_fp8_,
+)
+
+
+def _dequantized(module: Fp8Linear) -> torch.Tensor:
+ weight = module.weight_fp8.view(torch.float8_e4m3fn).float()
+ return weight * module.weight_scale.view(torch.float32)
+
+
+def test_fp8_weight_survives_dtype_casts():
+ """``model.to(bf16)`` runs after conversion in several load paths; it must not
+ upcast the fp8 weight or round its fp32 scale."""
+ torch.manual_seed(0)
+ weight = torch.randn(64, 32) * 0.02
+ module = Fp8Linear.from_weight(weight).to(torch.bfloat16)
+
+ assert module.weight_fp8.dtype == torch.uint8
+ assert module.weight_scale.dtype == torch.int32
+ # e4m3 keeps 3 mantissa bits, so one rounding step is at most 2**-4 relative.
+ torch.testing.assert_close(_dequantized(module), weight, atol=0, rtol=2**-4)
+
+
+def test_the_upcast_fallback_warns_once_per_device(caplog):
+ _has_fp8_gemm.cache_clear()
+ module = Fp8Linear.from_weight(torch.randn(16, 32))
+ with caplog.at_level("WARNING"), torch.no_grad():
+ module(torch.randn(3, 32))
+ module(torch.randn(3, 32))
+ warnings = [r for r in caplog.records if "no fp8 GEMM" in r.getMessage()]
+ assert len(warnings) == 1
+
+
+def test_cpu_conversion_runs_the_upcast_path(tiny_esmc):
+ """Conversion happens on the CPU before the move to GPU, and a CPU forward
+ upcasts the fp8 weight instead of failing."""
+ ids = torch.randint(4, 24, (2, 13))
+ with torch.no_grad():
+ expected = tiny_esmc(input_ids=ids).last_hidden_state
+ assert convert_esmc_to_fp8_(tiny_esmc) == 3 * len(tiny_esmc.transformer.blocks)
+ actual = tiny_esmc(input_ids=ids).last_hidden_state
+ rel = (actual - expected).norm() / expected.norm()
+ assert rel < 0.05, rel
+
+
+@pytest.mark.gpu
+def test_esmfold2_fp8_backbone_is_fully_fp8(tiny_esmfold2, tiny_esmc):
+ """Every ESMC projection must be fp8. The old path converted only
+ ``out_proj`` and left the LayerNorm-fused QKV and MLP weights in bf16."""
+ if torch.cuda.get_device_capability() < (8, 9):
+ pytest.skip("fp8 needs compute capability >= 8.9") # ty:ignore[too-many-positional-arguments]
+ tiny_esmfold2.esmc = tiny_esmc
+ model = tiny_esmfold2.cuda()
+ model.set_esmc_precision("fp8")
+
+ for block in model.esmc.transformer.blocks:
+ assert isinstance(block.attn.layernorm_qkv, Fp8LayerNormLinear)
+ assert isinstance(block.attn.out_proj, Fp8Linear)
+ assert isinstance(block.ffn, Fp8LayerNormMLP)
+ assert not [
+ name
+ for name, p in model.esmc.named_parameters()
+ if p.dim() == 2 and name != "embed.weight"
+ ], "a 2D weight was left in bf16"
+
+ ids = torch.randint(4, 24, (2, 37), device="cuda")
+ with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
+ out = model.esmc(input_ids=ids).last_hidden_state
+ assert out.shape == (2, 37, tiny_esmc.config.hidden_size)
+ assert torch.isfinite(out.float()).all()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("fused", ["plain", "swiglu", "norm"])
+def test_triton_quantization_matches_the_torch_reference(fused):
+ # Triton rejects fp8e4nv below compute capability 8.9 (e.g. A10G, A100).
+ if torch.cuda.get_device_capability() < (8, 9):
+ pytest.skip("fp8 needs compute capability >= 8.9") # ty:ignore[too-many-positional-arguments]
+ from esm.models.esmfold2.kernels.fp8_quantize import quantize_rows_fp8
+
+ torch.manual_seed(0)
+ # 6912 is ESMC-6B's MLP width: not a power of two, so the tail is masked.
+ x = torch.randn(37, 2 * 6912 if fused == "swiglu" else 6912, device="cuda")
+ x = x.to(torch.bfloat16)
+ norm = None
+ if fused == "norm":
+ norm = (
+ torch.randn(6912, device="cuda"),
+ torch.randn(6912, device="cuda"),
+ 1e-5,
+ )
+ kwargs = {"swiglu": fused == "swiglu", "norm": norm}
+
+ q, scale = quantize_rows_fp8(x, **kwargs)
+ q_ref, scale_ref = _quantize_rows_torch(x, **kwargs)
+ torch.testing.assert_close(scale, scale_ref, rtol=1e-5, atol=0)
+ # Different reduction and sigmoid orders may flip a rounding: allow one e4m3
+ # step (2**-3 relative) on a handful of elements.
+ deq, deq_ref = q.float() * scale, q_ref.float() * scale_ref
+ torch.testing.assert_close(deq, deq_ref, rtol=2**-3, atol=scale_ref.max().item())
+ assert (q.view(torch.uint8) != q_ref.view(torch.uint8)).float().mean() < 1e-3
diff --git a/tests/models/esmfold2_fused_transition_test.py b/tests/models/esmfold2_fused_transition_test.py
new file mode 100644
index 00000000..5b885e1c
--- /dev/null
+++ b/tests/models/esmfold2_fused_transition_test.py
@@ -0,0 +1,132 @@
+"""The forward-only pair ``Transition`` must be bitwise and must skip the backward buffer.
+
+``FusedLNLinearSwiGLU`` writes an ``[M, 2*d_inner]`` linear-output buffer that
+only backward reads. A forward-only call now skips that store. These tests pin
+the two things that could go wrong: the forward result changing, and the store
+being skipped on a call that does have a backward.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from esm.models.esmfold2.layers import Transition
+
+D_MODEL = 256
+D_INNER = Transition(D_MODEL).ffn.hidden_features
+
+
+def _transition(device: str = "cuda") -> Transition:
+ torch.manual_seed(0)
+ return Transition(D_MODEL).to(device=device, dtype=torch.bfloat16).eval()
+
+
+def _with_buffer(t: Transition, x: torch.Tensor) -> torch.Tensor:
+ """The pre-change path: force the [M, 2N] store, then the same addmm residual."""
+ # Local: the kernels package needs CUDA at import, and collection has to
+ # work on the CPU boxes (see esmfold2_cpu_only_test).
+ from esm.models.esmfold2.kernels.fused_lnlin_swiglu import (
+ FusedLNLinearSwiGLUFunction,
+ )
+
+ fused = t._fused_swiglu
+ assert fused is not None
+ hidden = FusedLNLinearSwiGLUFunction.apply(
+ x, fused.W12, fused.LN_W, fused.LN_B, True
+ )
+ return t._addmm_residual(x, hidden)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("shape", [(1, 64, 64), (1, 137, 137), (1, 200, 200)])
+def test_forward_only_transition_is_bitwise(shape):
+ """Skipping the store must not move a single bit of the output."""
+ t = _transition()
+ t.set_chunk_size(None)
+ t.set_kernel_backend("fused")
+ x = torch.randn(*shape, D_MODEL, device="cuda", dtype=torch.bfloat16)
+
+ with torch.inference_mode():
+ ref = _with_buffer(t, x)
+ got = t(x)
+
+ assert torch.equal(got, ref), (
+ f"max delta {(got.float() - ref.float()).abs().max().item():.3e}"
+ )
+
+
+@pytest.mark.gpu
+def test_frozen_weights_with_input_grad_keeps_the_buffer():
+ """The binder-design shape: params frozen, input requires grad.
+
+ ``_fused_pair_stack_active`` deliberately permits grad for this path, so the
+ no-backward detection must not key off the weights. Getting it wrong here
+ silently breaks the design gradient.
+ """
+ t = _transition()
+ t.set_chunk_size(None)
+ t.set_kernel_backend("fused")
+ t.requires_grad_(False) # frozen weights, as binder design runs it
+ rows = 128 * 128
+ buffer_bytes = rows * 2 * D_INNER * 2
+ x = torch.randn(1, 128, 128, D_MODEL, device="cuda", dtype=torch.bfloat16)
+
+ torch.cuda.synchronize()
+ torch.cuda.empty_cache()
+ torch.cuda.reset_peak_memory_stats()
+ base = torch.cuda.memory_allocated()
+ xr = x.clone().requires_grad_(True)
+ out = t(xr)
+ torch.cuda.synchronize()
+ peak = torch.cuda.max_memory_allocated() - base
+
+ assert out.requires_grad, "frozen weights must still yield a differentiable output"
+ assert peak >= buffer_bytes, f"peak {peak} < buffer {buffer_bytes}"
+
+ out.float().square().sum().backward()
+ assert xr.grad is not None
+
+ # And the gradient matches the path that always stored the buffer.
+ xr2 = x.clone().requires_grad_(True)
+ _with_buffer(t, xr2).float().square().sum().backward()
+ assert xr2.grad is not None
+ assert torch.equal(xr.grad, xr2.grad)
+
+
+@pytest.mark.gpu
+def test_no_grad_with_input_grad_still_skips_the_buffer():
+ """Grad globally off outranks a grad-requiring input: no backward can run."""
+ t = _transition()
+ t.set_chunk_size(None)
+ t.set_kernel_backend("fused")
+ rows = 128 * 128
+ buffer_bytes = rows * 2 * D_INNER * 2
+ x = torch.randn(1, 128, 128, D_MODEL, device="cuda", dtype=torch.bfloat16)
+
+ torch.cuda.synchronize()
+ torch.cuda.empty_cache()
+ torch.cuda.reset_peak_memory_stats()
+ base = torch.cuda.memory_allocated()
+ with torch.no_grad():
+ out = t(x.clone().requires_grad_(True))
+ torch.cuda.synchronize()
+ peak = torch.cuda.max_memory_allocated() - base
+ del out
+
+ assert peak < buffer_bytes, f"peak {peak} still fits the {buffer_bytes}-byte buffer"
+
+
+@pytest.mark.gpu
+def test_chunked_forward_is_bitwise_too():
+ """Chunking changes the row count per call; every chunk must still match."""
+ t = _transition()
+ t.set_kernel_backend("fused")
+ t.set_chunk_size(32)
+ x = torch.randn(1, 70, 70, D_MODEL, device="cuda", dtype=torch.bfloat16)
+ with torch.inference_mode():
+ got = t(x)
+ ref = torch.cat(
+ [_with_buffer(t, x[:, s : s + 32]) for s in range(0, 70, 32)], dim=1
+ )
+ assert torch.equal(got, ref)
diff --git a/tests/models/esmfold2_inputs_test.py b/tests/models/esmfold2_inputs_test.py
index 6142d228..c0b4c401 100644
--- a/tests/models/esmfold2_inputs_test.py
+++ b/tests/models/esmfold2_inputs_test.py
@@ -125,7 +125,9 @@ def run_model(
#: Why the ``ref_pos`` comparison below is the one thing in this file that can
#: skip: the CCD pickle is an environment-provided data artifact, like a
#: published checkpoint, and the fast path's table is pinned to one release of it.
-MISMATCHED_CCD = "the loaded CCD is not the snapshot PROTEIN_REF_POS was frozen from"
+MISMATCHED_CCD = (
+ "KNOWN GAP: the loaded CCD is not the snapshot PROTEIN_REF_POS was frozen from"
+)
C1_SEQUENCES = {
"single_residue": "M",
@@ -215,11 +217,12 @@ def test_fast_path_ref_pos_reproduces_the_ccd_conformers(name, ccd_pickle):
Split out from the shared-key test because it is the one feature that
depends on an external data artifact rather than on code: the table is a
- frozen copy of one CCD release, so pointing ``ESMCFOLD_CCD_PATH`` at a
- later one makes the two paths hand the model different geometry. Measured
- against ``biohub/ESMFold2:ccd.pkl``: bit-identical. Against the internal
- ``ccd_20251126`` snapshot: 6.5 A max, and 1.1 A on intra-residue distances,
- i.e. a re-generated conformer rather than a re-oriented one.
+ frozen copy of ``ESMFOLD2_CCD_REVISION``, so overriding
+ ``ESMCFOLD_CCD_PATH`` with a different release makes the two paths hand the
+ model different geometry. Regenerate with
+ ``python -m tests.regenerate_reference protein_ref_pos`` when the pin moves;
+ the guard below turns a stale table into a skip rather than a confusing
+ coordinate mismatch.
"""
if not ccd_matches_the_hardcoded_table():
pytest.skip(MISMATCHED_CCD) # ty:ignore[too-many-positional-arguments]
@@ -562,36 +565,6 @@ def test_lm_shim_wraps_each_chain_in_bos_and_eos(recording_esmc, ccd_pickle):
assert torch.isfinite(hidden).all()
-def test_lm_shim_marks_padding_with_sequence_id_minus_one(recording_esmc, ccd_pickle):
- """``pad_to_multiple`` is the fp8 path, and it is the one that pads at B=1.
-
- TE's fp8 kernels assert ``prod(shape[:-1]) % 8 == 0``, so ``_esmc_fp8``
- rounds the LM input up. Those extra columns must be PAD *and* carry
- ``sequence_id == -1``, or the last chain would grow a tail of phantom
- residues it attends to.
- """
- features = batched(prepare_esmfold2_input(protein_complex_input(C4_CHAINS))[0])
- unpadded = lm_hidden_states_for(recording_esmc, features)
- recording_esmc.calls.clear()
- padded = lm_hidden_states_for(recording_esmc, features, pad_to_multiple=8)
-
- seen_ids, seen_sequence_id = recording_esmc.calls[0]
- length = seen_ids.shape[1]
- assert length % 8 == 0
- # 6 + 4 + 5 residues + 6 special tokens = 21, rounded up to 24.
- assert length == 24
-
- is_pad = seen_ids[0] == PAD
- assert is_pad.sum() == 3
- assert torch.equal(seen_sequence_id[0] == -1, is_pad)
- assert (seen_sequence_id[0][~is_pad] >= 0).all()
-
- # Measured max|Δ| 2.4e-7 on the tiny backbone: the padded batch takes SDPA's
- # masked path and the unpadded one its unmasked path, so the reductions
- # differ in order but not in what they sum over.
- torch.testing.assert_close(unpadded, padded, atol=1e-6, rtol=0)
-
-
def test_lm_shim_gives_each_chain_its_own_sequence_id(recording_esmc, ccd_pickle):
"""``sequence_id`` increments once per chain and covers every chain.
diff --git a/tests/models/esmfold2_msa_trimul_test.py b/tests/models/esmfold2_msa_trimul_test.py
new file mode 100644
index 00000000..300c444d
--- /dev/null
+++ b/tests/models/esmfold2_msa_trimul_test.py
@@ -0,0 +1,307 @@
+"""``set_kernel_backend("fast")`` reaches the MSA encoder's triangle updates.
+
+The MSA encoder ran the eager fp32 contraction while the trunk ran the bf16
+Triton kernel on the identical operation. The kernel is NOT bitwise against the
+eager engine, so it sits behind "fast", never behind "fused".
+
+These tests pin what can go wrong: the kernel not firing under "fast", the
+kernel firing under "fused" (which must stay bit-exact), the trunk's own gate
+moving, and the CPU fallback losing its output.
+"""
+
+import contextlib
+
+import pytest
+import torch
+
+from esm.models.esmfold2.layers import PairUpdateBlock, TriangleMultiplicativeUpdate
+from esm.models.esmfold2.model import MSAEncoderBlock
+
+D_PAIR = 256
+
+
+def _msa_block(device: str) -> MSAEncoderBlock:
+ torch.manual_seed(0)
+ return (
+ MSAEncoderBlock(
+ d_msa=64,
+ d_pair=D_PAIR,
+ d_hidden=32,
+ n_heads_msa=8,
+ msa_head_width=16,
+ is_final_block=True,
+ )
+ .to(device)
+ .eval()
+ )
+
+
+def _pair_and_mask(L: int, device: str, dtype: torch.dtype):
+ torch.manual_seed(1)
+ pair = torch.randn(1, L, L, D_PAIR, device=device, dtype=dtype)
+ mask = torch.ones(1, L, L, device=device, dtype=torch.bool)
+ mask[:, :, L - 5 :] = False # padded tail, so the mask is load-bearing
+ mask[:, L - 5 :, :] = False
+ return pair, mask
+
+
+def _amp(device: str):
+ """The recycle loop runs the MSA encoder under bf16 autocast with fp32 weights.
+
+ Without it the eager engine's LayerNorm dtype-errors on a bf16 pair.
+ """
+ if device != "cuda":
+ return contextlib.nullcontext()
+ return torch.autocast("cuda", dtype=torch.bfloat16)
+
+
+def _eager_updates(blk: MSAEncoderBlock, pair, mask):
+ """The pre-change arithmetic: two eager updates, each with a plain residual."""
+ out = pair + blk.tri_mul_out(pair, mask=mask)
+ return out + blk.tri_mul_in(out, mask=mask)
+
+
+def test_invalid_backend_raises():
+ tri = TriangleMultiplicativeUpdate(dim=32)
+ with pytest.raises(ValueError, match="backend must be one of"):
+ tri.set_kernel_backend("bogus")
+
+
+def test_set_kernel_backend_is_recorded_on_the_wrapper():
+ """It used to be a no-op, so the MSA encoder could never see the request."""
+ tri = TriangleMultiplicativeUpdate(dim=32)
+ assert not tri.fast_backend_requested()
+ tri.set_kernel_backend("fast")
+ assert tri.fast_backend_requested()
+ tri.set_kernel_backend(None)
+ assert not tri.fast_backend_requested()
+
+
+def test_fused_is_not_fast_on_the_wrapper():
+ """ "fused" must not arm the approximate MSA kernel -- that is the whole split."""
+ tri = TriangleMultiplicativeUpdate(dim=32)
+ tri.set_kernel_backend("fused")
+ assert not tri.fast_backend_requested()
+
+
+def test_cpu_msa_block_warns_and_keeps_the_eager_result():
+ """No CUDA means no kernel: the fallback must warn and stay bit-exact."""
+ import esm.models.esmfold2.layers as layers
+
+ blk = _msa_block("cpu")
+ blk.set_chunk_size(None)
+ pair, mask = _pair_and_mask(24, "cpu", torch.float32)
+
+ with torch.inference_mode():
+ blk.set_kernel_backend(None)
+ ref = blk._triangle_updates(pair, mask)
+ blk.set_kernel_backend("fast")
+ layers._fast_trimul_refusal_warned = False
+ with pytest.warns(UserWarning, match="eager triangle multiplication"):
+ got = blk._triangle_updates(pair, mask)
+
+ assert torch.equal(got, ref)
+
+
+def test_cpu_fused_backend_does_not_warn():
+ """ "fused" never asks for this kernel, so a CPU fold must stay silent."""
+ import warnings as _warnings
+
+ import esm.models.esmfold2.layers as layers
+
+ blk = _msa_block("cpu")
+ blk.set_chunk_size(None)
+ pair, mask = _pair_and_mask(24, "cpu", torch.float32)
+
+ layers._fast_trimul_refusal_warned = False
+ blk.set_kernel_backend("fused")
+ with torch.inference_mode(), _warnings.catch_warnings():
+ _warnings.simplefilter("error") # any warning here is a regression
+ blk._triangle_updates(pair, mask)
+
+
+@pytest.mark.gpu
+def test_fused_msa_updates_fire_and_track_the_eager_path():
+ """The kernel must run twice per block and land within bf16 tolerance."""
+ from esm.models.esmfold2 import layers
+
+ blk = _msa_block("cuda")
+ blk.set_chunk_size(None)
+ blk.set_kernel_backend("fast")
+ pair, mask = _pair_and_mask(128, "cuda", torch.bfloat16)
+
+ calls = {"n": 0}
+ orig = layers._fused_trimul_with_residual
+
+ def counting(*a, **k):
+ calls["n"] += 1
+ return orig(*a, **k)
+
+ layers._fused_trimul_with_residual = counting
+ try:
+ with torch.inference_mode(), _amp("cuda"):
+ eager = _eager_updates(blk, pair, mask)
+ fused = blk._triangle_updates(pair, mask)
+ finally:
+ layers._fused_trimul_with_residual = orig
+
+ assert calls["n"] == 2, "one call per triangle update, or the gate refused"
+ assert fused.dtype == pair.dtype
+ assert fused.shape == eager.shape
+
+ # Not bitwise, on purpose: fp32 eager contraction vs bf16 kernel.
+ assert not torch.equal(fused, eager)
+ scale = eager.double().abs().max().clamp_min(1e-6)
+ rel = (fused.double() - eager.double()).abs().max() / scale
+ assert rel < 0.05, f"relative max delta {rel.item():.3e}"
+
+
+@pytest.mark.gpu
+def test_fused_msa_updates_ignore_chunk_size():
+ """Chunking is an eager-only knob; the fused path must not change with it."""
+ blk = _msa_block("cuda")
+ blk.set_kernel_backend("fast")
+ pair, mask = _pair_and_mask(128, "cuda", torch.bfloat16)
+
+ with torch.inference_mode(), _amp("cuda"):
+ blk.set_chunk_size(None)
+ unchunked = blk._triangle_updates(pair, mask)
+ blk.set_chunk_size(64)
+ chunked = blk._triangle_updates(pair, mask)
+
+ assert torch.equal(chunked, unchunked)
+
+
+@pytest.mark.gpu
+def test_fp32_pair_is_cast_rather_than_refused():
+ """The trunk casts at the stack level; the MSA block casts around the pair."""
+ blk = _msa_block("cuda")
+ blk.set_chunk_size(None)
+ blk.set_kernel_backend("fast")
+ pair, mask = _pair_and_mask(64, "cuda", torch.float32)
+
+ with torch.inference_mode(), _amp("cuda"):
+ out = blk._triangle_updates(pair, mask)
+
+ assert out.dtype == torch.float32
+
+
+@pytest.mark.gpu
+def test_trunk_gate_is_unchanged_by_the_delegation():
+ """PairUpdateBlock still demands bf16 on CUDA and nothing else."""
+ blk = PairUpdateBlock(d_pair=D_PAIR).cuda().eval()
+ blk.set_kernel_backend("fast")
+ bf16 = torch.zeros(1, 8, 8, D_PAIR, device="cuda", dtype=torch.bfloat16)
+
+ assert blk._can_use_fused_trimul_with_residual(bf16)
+ assert not blk._can_use_fused_trimul_with_residual(bf16.float())
+ assert not blk._can_use_fused_trimul_with_residual(bf16.cpu())
+
+ blk.set_kernel_backend(None)
+ assert not blk._can_use_fused_trimul_with_residual(bf16)
+
+
+@pytest.mark.gpu
+def test_fast_does_not_touch_gradients():
+ """ "fast" is an inference accelerator: it must not move training gradients.
+
+ The kernel is approximate, so arming it under autograd would silently change
+ what a training run optimises. The eager engine must serve instead, and the
+ backward must still work.
+ """
+ from esm.models.esmfold2 import layers
+
+ blk = _msa_block("cuda")
+ blk.set_chunk_size(None)
+ blk.set_kernel_backend("fast")
+ pair, mask = _pair_and_mask(64, "cuda", torch.bfloat16)
+ pair = pair.clone().requires_grad_(True)
+
+ assert not blk.tri_mul_out.fast_residual_available(pair), (
+ "the approximate kernel armed itself with gradients enabled"
+ )
+
+ calls = {"n": 0}
+ orig = layers._fused_trimul_with_residual
+
+ def counting(*a, **k):
+ calls["n"] += 1
+ return orig(*a, **k)
+
+ layers._fused_trimul_with_residual = counting
+ try:
+ with _amp("cuda"):
+ out = blk._triangle_updates(pair, mask)
+ out.float().pow(2).mean().backward()
+ finally:
+ layers._fused_trimul_with_residual = orig
+
+ assert calls["n"] == 0
+ assert pair.grad is not None
+ assert torch.isfinite(pair.grad).all()
+ assert pair.grad.abs().sum() > 0
+
+
+@pytest.mark.gpu
+def test_fast_arms_itself_under_inference_mode():
+ """The mirror of the gradient test: no grad means the kernel does serve."""
+ blk = _msa_block("cuda")
+ blk.set_kernel_backend("fast")
+ pair, _ = _pair_and_mask(64, "cuda", torch.bfloat16)
+
+ with torch.inference_mode():
+ assert blk.tri_mul_out.fast_residual_available(pair)
+ with torch.no_grad():
+ assert blk.tri_mul_out.fast_residual_available(pair)
+
+
+@pytest.mark.gpu
+def test_fused_backend_leaves_the_msa_encoder_bit_exact():
+ """The contract: "fused" is bit-exact, so it must not touch this path.
+
+ If "fused" ever arms the MSA kernel, the bit-exactness guarantee of the
+ default accelerated backend silently breaks for every MSA fold.
+ """
+ from esm.models.esmfold2 import layers
+
+ blk = _msa_block("cuda")
+ blk.set_chunk_size(None)
+ pair, mask = _pair_and_mask(128, "cuda", torch.bfloat16)
+
+ calls = {"n": 0}
+ orig = layers._fused_trimul_with_residual
+
+ def counting(*a, **k):
+ calls["n"] += 1
+ return orig(*a, **k)
+
+ layers._fused_trimul_with_residual = counting
+ try:
+ with torch.inference_mode(), _amp("cuda"):
+ blk.set_kernel_backend(None)
+ ref = blk._triangle_updates(pair, mask)
+ blk.set_kernel_backend("fused")
+ got = blk._triangle_updates(pair, mask)
+ finally:
+ layers._fused_trimul_with_residual = orig
+
+ assert calls["n"] == 0, "the approximate MSA kernel fired under 'fused'"
+ assert torch.equal(got, ref)
+
+
+@pytest.mark.gpu
+def test_fast_keeps_the_trunk_identical_to_fused():
+ """ "fast" is a superset: the trunk's bit-exact kernel must still run, unchanged."""
+ blk = PairUpdateBlock(d_pair=D_PAIR).cuda().eval()
+ pair, mask = _pair_and_mask(64, "cuda", torch.bfloat16)
+ mask_f32 = mask.to(torch.float32)
+
+ with torch.inference_mode():
+ blk.set_kernel_backend("fused")
+ assert blk._can_use_fused_trimul_with_residual(pair)
+ fused = blk(pair, pair_attention_mask=mask_f32)
+ blk.set_kernel_backend("fast")
+ assert blk._can_use_fused_trimul_with_residual(pair)
+ fast = blk(pair, pair_attention_mask=mask_f32)
+
+ assert torch.equal(fast, fused)
diff --git a/tests/models/esmfold2_pair_bias_cache_test.py b/tests/models/esmfold2_pair_bias_cache_test.py
new file mode 100644
index 00000000..816937d5
--- /dev/null
+++ b/tests/models/esmfold2_pair_bias_cache_test.py
@@ -0,0 +1,239 @@
+"""The per-fold pair-bias cache in ``AttentionPairBias``.
+
+The pair bias is a projection of ``z`` into ``[B, H, Q, K]``. Within one fold
+``z``, the mask and the projection weights are all fixed, so every diffusion
+step of every DiT block used to recompute the same tensor. The cache stores it
+in the caller's per-fold ``inference_cache``.
+
+This is a caching change only, so every test here is a ``torch.equal`` test, not
+an ``allclose`` one. The cache must also stay out of the way of autograd: the
+binder-design path takes input gradients through these blocks with frozen
+weights, and a detached cache hit would silently break it.
+"""
+
+import warnings
+
+import pytest
+import torch
+
+from esm.models.esmfold2.layers import AttentionPairBias
+
+_D_MODEL = 32
+_D_PAIR = 16
+_N_HEADS = 4
+_N_TOKENS = 7
+
+
+def _block(seed: int = 0) -> AttentionPairBias:
+ torch.manual_seed(seed)
+ block = AttentionPairBias(
+ d_model=_D_MODEL, d_pair=_D_PAIR, num_heads=_N_HEADS, use_conditioning=False
+ ).eval()
+ # Random weights: the default LayerNorm/Linear init would make several
+ # tensors identical and hide a keying mistake.
+ for p in block.parameters():
+ torch.nn.init.normal_(p, std=0.1)
+ return block
+
+
+def _inputs(batch: int = 1, seed: int = 1):
+ torch.manual_seed(seed)
+ a = torch.randn(batch, _N_TOKENS, _D_MODEL)
+ z = torch.randn(batch, _N_TOKENS, _N_TOKENS, _D_PAIR)
+ mask = torch.ones(batch, _N_TOKENS, dtype=torch.bool)
+ return a, z, mask
+
+
+class _ProjCounter:
+ """Counts pair-bias projections, i.e. cache misses."""
+
+ def __init__(self, block: AttentionPairBias) -> None:
+ self.calls = 0
+ self._handle = block.pair_bias_proj.register_forward_hook(self._hook)
+
+ def _hook(self, *_args) -> None:
+ self.calls += 1
+
+ def close(self) -> None:
+ self._handle.remove()
+
+
+def test_cache_is_bitwise_and_projects_once():
+ """Three cached calls give the uncached answer and one projection."""
+ block = _block()
+ a, z, mask = _inputs()
+
+ with torch.no_grad():
+ counter = _ProjCounter(block)
+ reference = block(a, None, z, attention_mask=mask)
+ assert counter.calls == 1
+
+ cache: dict = {}
+ outs = [
+ block(a, None, z, attention_mask=mask, inference_cache=cache)
+ for _ in range(3)
+ ]
+ counter.close()
+
+ assert counter.calls == 2, "the cached calls projected more than once"
+ for i, out in enumerate(outs):
+ assert torch.equal(out, reference), f"cached call {i} is not bitwise equal"
+
+
+def test_cache_holds_the_expanded_bias_for_several_diffusion_samples():
+ """``z`` is expanded per sample inside the block; the cache holds the result."""
+ block = _block()
+ samples = 3
+ a_base, z, mask = _inputs(batch=1)
+ a = a_base.repeat_interleave(samples, dim=0)
+
+ with torch.no_grad():
+ counter = _ProjCounter(block)
+ reference = block(
+ a, None, z, attention_mask=mask, num_diffusion_samples=samples
+ )
+ cache: dict = {}
+ first = block(
+ a,
+ None,
+ z,
+ attention_mask=mask,
+ num_diffusion_samples=samples,
+ inference_cache=cache,
+ )
+ second = block(
+ a,
+ None,
+ z,
+ attention_mask=mask,
+ num_diffusion_samples=samples,
+ inference_cache=cache,
+ )
+ counter.close()
+
+ assert counter.calls == 2
+ assert torch.equal(first, reference)
+ assert torch.equal(second, reference)
+ cached = cache["pair_bias"][block]["eager"]
+ assert cached.shape[0] == samples
+
+
+def test_two_blocks_sharing_one_cache_do_not_collide():
+ """The slot is keyed by module instance, so different weights stay apart."""
+ first, second = _block(seed=0), _block(seed=1)
+ a, z, mask = _inputs()
+
+ with torch.no_grad():
+ ref_first = first(a, None, z, attention_mask=mask)
+ ref_second = second(a, None, z, attention_mask=mask)
+ cache: dict = {}
+ got_first = first(a, None, z, attention_mask=mask, inference_cache=cache)
+ got_second = second(a, None, z, attention_mask=mask, inference_cache=cache)
+
+ assert not torch.equal(ref_first, ref_second), "the two blocks are not distinct"
+ assert torch.equal(got_first, ref_first)
+ assert torch.equal(got_second, ref_second)
+
+
+def test_grad_enabled_alone_refuses_the_cache_and_warns():
+ """First guard, in isolation: frozen weights, no grad-tracking input."""
+ block = _block().requires_grad_(False)
+ a, z, mask = _inputs()
+ cache: dict = {}
+
+ assert torch.is_grad_enabled()
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ block(a, None, z, attention_mask=mask, inference_cache=cache)
+
+ assert "pair_bias" not in cache, "the cache engaged under grad"
+ assert any("pair-bias cache not engaged" in str(w.message) for w in caught), (
+ "refusing to cache was silent"
+ )
+
+
+def test_a_grad_tracking_input_alone_refuses_the_cache():
+ """Second guard, in isolation: no grad context, grad-tracking ``z``."""
+ block = _block().requires_grad_(False)
+ a, z_base, mask = _inputs()
+ z = z_base.clone().requires_grad_(True)
+
+ with torch.no_grad():
+ assert block._pair_bias_cache_slot({}, z_base, mask, a.shape[0], 1) is not None
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ slot = block._pair_bias_cache_slot({}, z, mask, a.shape[0], 1)
+
+ assert slot is None, "the cache engaged on a grad-tracking z"
+ assert any("pair-bias cache not engaged" in str(w.message) for w in caught)
+
+
+def test_the_input_gradient_is_unchanged_when_a_cache_is_offered():
+ """The binder-design path: frozen weights, gradient w.r.t. the pair input."""
+ block = _block().requires_grad_(False)
+ a, z_base, mask = _inputs()
+
+ def input_grad(cache: dict | None) -> torch.Tensor:
+ z = z_base.clone().requires_grad_(True)
+ out = block(a, None, z, attention_mask=mask, inference_cache=cache)
+ return torch.autograd.grad(out.square().sum(), z)[0]
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ reference = input_grad(None)
+ cache: dict = {}
+ with_cache = input_grad(cache)
+ again = input_grad(cache)
+
+ assert "pair_bias" not in cache
+ assert torch.equal(with_cache, reference), "the input gradient changed"
+ assert torch.equal(again, reference), "a second grad call diverged"
+
+
+def test_nothing_survives_the_caller_dropping_the_cache():
+ """All state lives in the caller's dict; the module keeps none of it."""
+ block = _block()
+ a, z, mask = _inputs()
+ before = set(block.__dict__)
+
+ with torch.no_grad():
+ cache: dict = {}
+ block(a, None, z, attention_mask=mask, inference_cache=cache)
+
+ assert cache["pair_bias"][block]["eager"] is not None
+ assert set(block.__dict__) == before, "the block grew per-fold state"
+
+ with torch.no_grad():
+ fresh: dict = {}
+ counter = _ProjCounter(block)
+ block(a, None, z, attention_mask=mask, inference_cache=fresh)
+ counter.close()
+ assert counter.calls == 1, "a new cache dict reused an old bias"
+
+
+@pytest.mark.gpu
+def test_fused_path_cache_is_bitwise():
+ """Same check on the fused Triton kernel, which is the hot path."""
+ block = _block().cuda()
+ a, z, mask = _inputs()
+ a, z, mask = a.cuda(), z.cuda(), mask.cuda()
+ block.set_kernel_backend("fused")
+
+ with torch.no_grad():
+ assert block._can_use_fused_pair_bias(z, _N_TOKENS, 0.0), (
+ "the fused pair-bias path did not engage"
+ )
+ reference = block(a, None, z, attention_mask=mask)
+ cache: dict = {}
+ first = block(a, None, z, attention_mask=mask, inference_cache=cache)
+ second = block(a, None, z, attention_mask=mask, inference_cache=cache)
+
+ assert "fused" in cache["pair_bias"][block]
+ assert torch.equal(first, reference)
+ assert torch.equal(second, reference)
+ assert cache["pair_bias"][block]["fused"].shape == (
+ 1,
+ _N_HEADS,
+ _N_TOKENS,
+ _N_TOKENS,
+ )
diff --git a/tests/models/esmfold2_test.py b/tests/models/esmfold2_test.py
index 503a472b..0bb8c1fc 100644
--- a/tests/models/esmfold2_test.py
+++ b/tests/models/esmfold2_test.py
@@ -892,13 +892,11 @@ def test_an_unexpected_key_is_refused(tmp_path, tiny_esmfold2_config):
[EsmFold2Model, EsmFold2ExperimentalModel],
ids=["release", "experimental"],
)
-def test_only_transformer_engine_extra_state_is_excused(
- tmp_path, tiny_esmfold2_config, model_class
-):
+def test_only_extra_state_is_excused(tmp_path, tiny_esmfold2_config, model_class):
"""The ignore list must stay exactly one entry wide.
- ``_extra_state`` is Transformer Engine's non-tensor bookkeeping, which no
- module tree has a home for. Every other unexpected key has to raise, so the
+ ``_extra_state`` is Transformer Engine bookkeeping that published
+ checkpoints still carry, and no module tree has a home for it. Every other unexpected key has to raise, so the
list is pinned as a literal *and* exercised: an added ``_extra_state`` entry
loads, and the near-miss spelling next to it does not.
"""
@@ -1059,7 +1057,19 @@ def test_bundled_backbone_is_frozen_and_cast(
metadata={"format": "pt"},
)
+ # fp8 is the default: projections become fp8 buffers, the rest is bf16.
loaded = EsmFold2Model.from_pretrained(tmp_path, load_esmc=False)
+ assert {p.dtype for p in loaded.esmc.parameters()} == {torch.bfloat16}
+ assert not [
+ n
+ for n, p in loaded.esmc.named_parameters()
+ if p.dim() == 2 and n != "embed.weight"
+ ]
+ assert not any(p.requires_grad for p in loaded.esmc.parameters())
+
+ loaded = EsmFold2Model.from_pretrained(
+ tmp_path, load_esmc=False, esmc_precision="bf16"
+ )
assert {t.dtype for t in loaded.esmc.state_dict().values()} == {torch.bfloat16}
assert not any(p.requires_grad for p in loaded.esmc.parameters())
@@ -1176,8 +1186,7 @@ def test_loads_from_hub_with_full_key_coverage(repo, expected):
@pytest.mark.gpu
-@pytest.mark.parametrize("backend", ["fused", "cuequivariance"])
-def test_kernel_backends_agree(kernel_esmfold2_config, backend):
+def test_kernel_backends_agree(kernel_esmfold2_config):
"""Whatever accelerators a user has installed, the answer must match.
Compared on the final coordinates and pLDDT rather than intermediate
@@ -1188,12 +1197,10 @@ def test_kernel_backends_agree(kernel_esmfold2_config, backend):
compared two runs of one configuration and measured CUDA's own
nondeterminism rather than any kernel.
"""
- from esm.models.esmfold2.layers import CUE_AVAILABLE, TRITON_KERNELS_AVAILABLE
+ from esm.models.esmfold2.layers import TRITON_KERNELS_AVAILABLE
- if backend == "fused" and not TRITON_KERNELS_AVAILABLE:
+ if not TRITON_KERNELS_AVAILABLE:
pytest.skip("triton kernels unavailable") # ty:ignore[too-many-positional-arguments]
- if backend == "cuequivariance" and not CUE_AVAILABLE:
- pytest.skip("cuequivariance unavailable") # ty:ignore[too-many-positional-arguments]
torch.manual_seed(0)
model = EsmFold2Model(kernel_esmfold2_config).eval().cuda()
@@ -1217,13 +1224,13 @@ def run():
model.set_kernel_backend(None)
reference = run()
- model.set_kernel_backend(backend)
+ model.set_kernel_backend("fused")
candidate = run()
# Tolerances follow the measured noise, not a round number. On H100 the
# coordinates differ by 2e-4 between two runs of the *same* backend - CUDA
- # matmul selection is not deterministic - so that is the floor; the fused and
- # cuequivariance deltas were 3.4e-4 and 4.3e-4. pLDDT was bit-identical.
+ # matmul selection is not deterministic - so that is the floor; the fused
+ # delta was 3.4e-4. pLDDT was bit-identical.
# rtol is 0 for coordinates: they are Angstroms with |x| ~ 21, so a relative
# term would swamp the absolute bound.
#
diff --git a/tests/models/prepare_input_test.py b/tests/models/prepare_input_test.py
index 1fc5415f..980cc2e9 100644
--- a/tests/models/prepare_input_test.py
+++ b/tests/models/prepare_input_test.py
@@ -7,11 +7,13 @@
_entity_key,
build_chains_from_input,
compute_token_bonds,
+ prepare_esmfold2_input,
)
from esm.models.esmfold2.types import (
CovalentBond,
LigandInput,
Modification,
+ PocketConditioning,
ProteinInput,
StructurePredictionInput,
)
@@ -106,6 +108,32 @@ def test_build_chains_numbers_copies_across_both_id_forms():
assert [c.sym_id for c in chains] == [0, 1, 2]
+def test_pocket_conditioning_raises_not_implemented():
+ """``pocket`` must fail loudly, not fold unconditioned.
+
+ No released checkpoint consumes pocket features, so the featurizer rejects
+ the input rather than emitting the all-zeros ``pocket_feature`` that
+ ``forward`` would then drop.
+ """
+ spi = StructurePredictionInput(
+ sequences=[ProteinInput(id="A", sequence="MKT")],
+ pocket=PocketConditioning(binder_chain_id="A", contacts=[("A", 0)]),
+ )
+
+ with pytest.raises(NotImplementedError, match="pocket conditioning"):
+ prepare_esmfold2_input(spi)
+
+
+def test_pocket_unset_still_emits_zero_pocket_feature(ccd_pickle):
+ """The happy path is untouched: no ``pocket`` means no raise."""
+ spi = StructurePredictionInput(sequences=[ProteinInput(id="A", sequence="MKT")])
+
+ features, _ = prepare_esmfold2_input(spi)
+
+ assert "pocket_feature" in features
+ assert not features["pocket_feature"].any()
+
+
def _covalent_bond_input(bonds: list[CovalentBond] | None) -> StructurePredictionInput:
"""Protein with a lysine at residue 3 plus an ethanol ligand.
diff --git a/tests/oss_pytests/requirements.txt b/tests/oss_pytests/requirements.txt
index e725bcba..961750ad 100644
--- a/tests/oss_pytests/requirements.txt
+++ b/tests/oss_pytests/requirements.txt
@@ -1,2 +1,2 @@
-esm >=3.4.0,<4.0.0
+esm >=3.4.1.post1,<4.0.0
pytest
diff --git a/tests/oss_pytests/test_oss_client.py b/tests/oss_pytests/test_oss_client.py
index d534328a..bf17aec5 100644
--- a/tests/oss_pytests/test_oss_client.py
+++ b/tests/oss_pytests/test_oss_client.py
@@ -7,10 +7,12 @@
from esm.sdk.api import ( # pyright: ignore
ESMProtein,
ESMProteinTensor,
+ FoldingConfig,
ForwardAndSampleOutput,
GenerationConfig,
LogitsConfig,
LogitsOutput,
+ SAEConfig,
SamplingConfig,
SamplingTrackConfig,
)
@@ -18,12 +20,11 @@
API_TOKEN = os.environ.get("ESM_API_KEY", "")
URL = os.environ.get("URL")
+assert URL is not None, "URL must be set in the environment"
@pytest.mark.sdk
def test_oss_esm3_client():
- assert URL is not None
-
sequence = "MALWMRLLPLLALLAL___PDPAAA"
model = "esm3-small-2024-03"
esm3_client = client(model=model, url=URL, token=API_TOKEN)
@@ -54,32 +55,40 @@ def test_oss_esm3_client():
@pytest.mark.sdk
def test_oss_esmc_client():
- assert URL is not None
-
sequence = "MALWMRLLPLLALLALAVPDPAAA"
model = "esmc-300m-2024-12"
+ sae_model = "esmc-300m-2024-12-sae-layer23-k64-codebook65536"
esmc = esmc_client(model=model, url=URL, token=API_TOKEN)
protein = ESMProtein(sequence)
encoded_protein = esmc.encode(input=protein)
assert isinstance(encoded_protein, ESMProteinTensor)
+ assert encoded_protein.sequence is not None
decoded_protein = esmc.decode(input=encoded_protein)
assert isinstance(decoded_protein, ESMProtein)
logits_config = LogitsConfig(
- sequence=True, return_embeddings=True, return_hidden_states=True
+ sequence=True,
+ return_embeddings=True,
+ return_hidden_states=True,
+ sae_config=SAEConfig(models=[sae_model], normalize_features=False),
)
result = esmc.logits(input=encoded_protein, config=logits_config)
assert isinstance(result, LogitsOutput)
assert result.logits is not None
assert isinstance(result.logits.sequence, torch.Tensor)
+ assert result.sae_outputs is not None
+ assert sae_model in result.sae_outputs
+
+ # SAE features are returned sparse, one row per non-pad token.
+ features = result.sae_outputs[sae_model].to_dense()
+ assert features.shape[0] == encoded_protein.sequence.shape[0]
+ assert (features != 0).any()
@pytest.mark.sdk
def test_oss_sequence_structure_forge_inference_client():
- assert URL is not None
-
sequence = "MALWMRLLPLLALLALAVPDPAAA"
model = "esm3-small-2024-03"
client = SequenceStructureForgeInferenceClient(
@@ -88,3 +97,18 @@ def test_oss_sequence_structure_forge_inference_client():
encoded_protein = client.fold(sequence=sequence)
assert isinstance(encoded_protein, ESMProtein)
+
+
+@pytest.mark.sdk
+def test_oss_esmfold2_client():
+ sequence = "MALWMRLLPLLALLALAVPDPAAA"
+ model = "esmfold2-fast-2026-05"
+ fold_client = SequenceStructureForgeInferenceClient(
+ model=model, url=URL, token=API_TOKEN
+ )
+
+ config = FoldingConfig(num_sampling_steps=1, num_loops=1)
+ folded_protein = fold_client.fold(sequence=sequence, config=config)
+ assert isinstance(folded_protein, ESMProtein)
+ assert folded_protein.coordinates is not None
+ assert folded_protein.plddt is not None
diff --git a/tests/packaging_test.py b/tests/packaging_test.py
new file mode 100644
index 00000000..2b0b73e5
--- /dev/null
+++ b/tests/packaging_test.py
@@ -0,0 +1,113 @@
+"""Contracts on the published `esm` metadata.
+
+Parsed with ``tomllib``, never grepped: the bug these guard against is a table
+header going missing, and a grep for the keys under it still passes.
+"""
+
+import tomllib
+from pathlib import Path
+
+import pytest
+from packaging.requirements import Requirement
+from packaging.utils import canonicalize_name
+
+# In the monorepo the template lives under `opensource/root/`; the sync moves it
+# to the published repo root, beside `tests/`.
+_CANDIDATES = (
+ Path(__file__).parents[1] / "root" / "pyproject.toml",
+ Path(__file__).parents[1] / "pyproject.toml",
+)
+
+#: Everything a CPU-only install must not drag in.
+ACCELERATORS = frozenset(
+ {
+ "triton",
+ "xformers",
+ "cuequivariance-torch",
+ "cuequivariance-ops-torch-cu12",
+ "cuequivariance-ops-torch-cu13",
+ "flash-attn",
+ "transformer-engine",
+ }
+)
+
+LINUX_X86 = {"sys_platform": "linux", "platform_machine": "x86_64"}
+OTHER_PLATFORMS = {
+ "macos-arm64": {"sys_platform": "darwin", "platform_machine": "arm64"},
+ "linux-aarch64": {"sys_platform": "linux", "platform_machine": "aarch64"},
+ "windows": {"sys_platform": "win32", "platform_machine": "AMD64"},
+}
+
+# PEP 621. Anything else directly under `[project]` is a table header that was
+# dropped, not a field.
+PROJECT_FIELDS = frozenset(
+ {
+ "name",
+ "version",
+ "description",
+ "readme",
+ "requires-python",
+ "license",
+ "license-files",
+ "authors",
+ "maintainers",
+ "keywords",
+ "classifiers",
+ "urls",
+ "scripts",
+ "gui-scripts",
+ "entry-points",
+ "dependencies",
+ "optional-dependencies",
+ "dynamic",
+ }
+)
+
+
+@pytest.fixture(scope="module")
+def project() -> dict:
+ for path in _CANDIDATES:
+ if path.exists():
+ return tomllib.loads(path.read_text())["project"]
+ raise AssertionError(f"no pyproject.toml at any of {[str(p) for p in _CANDIDATES]}")
+
+
+@pytest.fixture(scope="module")
+def base_requirements(project) -> list[Requirement]:
+ return [Requirement(spec) for spec in project["dependencies"]]
+
+
+def test_there_are_no_extras(project):
+ assert "optional-dependencies" not in project
+
+
+def test_no_unrecognised_project_fields(project):
+ assert not set(project) - PROJECT_FIELDS
+
+
+def test_no_direct_references(project):
+ """PyPI rejects any release whose metadata carries a URL requirement."""
+ for spec in project["dependencies"]:
+ assert Requirement(spec).url is None, f"direct reference: {spec}"
+
+
+def test_triton_is_the_only_accelerator_and_linux_x86_only(base_requirements):
+ found = [r for r in base_requirements if canonicalize_name(r.name) in ACCELERATORS]
+ assert [canonicalize_name(r.name) for r in found] == ["triton"]
+ assert found[0].marker is not None
+ assert found[0].marker.evaluate(LINUX_X86)
+
+
+@pytest.mark.parametrize("environment", OTHER_PLATFORMS.values(), ids=OTHER_PLATFORMS)
+def test_base_install_is_accelerator_free_everywhere(base_requirements, environment):
+ """The real minimal-install contract; there is no `cpu` extra to express it.
+
+ An extra cannot remove a base dependency, so the only thing that can keep a
+ platform clean is the markers on the requirements themselves.
+ """
+ selected = {
+ canonicalize_name(r.name)
+ for r in base_requirements
+ if r.marker is None or r.marker.evaluate(environment)
+ }
+ assert not selected & ACCELERATORS
diff --git a/tests/regenerate_reference.py b/tests/regenerate_reference.py
index 49e72060..1e7a42ed 100644
--- a/tests/regenerate_reference.py
+++ b/tests/regenerate_reference.py
@@ -147,7 +147,50 @@ def esmfold2_tiny() -> None:
)
-REGENERATORS = {"esmc": esmc_300m, "esmfold2": esmfold2_tiny}
+def protein_ref_pos() -> None:
+ """Rewrite ``PROTEIN_REF_POS`` from the pinned CCD release.
+
+ The table is the fast path's literal copy of the CCD conformers, so it is
+ only correct for one release. Run this whenever ``ESMFOLD2_CCD_REVISION``
+ moves, then re-run ``esmfold2_inputs_test.py``: the fast-path parity tests
+ are what prove the two paths still agree.
+ """
+ from esm.models.esmfold2.conformers import get_idealized_atom_pos
+ from esm.models.esmfold2.protein_utils import (
+ PROTEIN_HEAVY_ATOMS,
+ PROTEIN_RESIDUE_TO_RES_TYPE,
+ PROTEIN_UNK_RES_TYPE,
+ )
+
+ header = "PROTEIN_REF_POS: dict[str, dict[str, tuple[float, float, float]]] = {"
+ lines = [header]
+ for res_3, atom_names in PROTEIN_HEAVY_ATOMS.items():
+ lines.append(f' "{res_3}": {{')
+ # ``.get`` with the UNK fallback mirrors ``prepare_protein_features``:
+ # UNK is in PROTEIN_HEAVY_ATOMS but not in the residue->res_type map.
+ res_type = PROTEIN_RESIDUE_TO_RES_TYPE.get(res_3, PROTEIN_UNK_RES_TYPE)
+ for name in atom_names:
+ pos = get_idealized_atom_pos(res_type, name)
+ if pos is None:
+ raise RuntimeError(f"{res_3}:{name} missing from the loaded CCD")
+ x, y, z = (float(v) for v in pos)
+ lines.append(f' "{name}": ({x}, {y}, {z}),')
+ lines.append(" },")
+ lines.append("}")
+
+ source = Path("esm/models/esmfold2/protein_utils.py")
+ text = source.read_text()
+ start = text.index(header)
+ end = text.index("\n}\n", start) + len("\n}")
+ source.write_text(text[:start] + "\n".join(lines) + text[end:])
+ print(f"rewrote PROTEIN_REF_POS in {source}")
+
+
+REGENERATORS = {
+ "esmc": esmc_300m,
+ "esmfold2": esmfold2_tiny,
+ "protein_ref_pos": protein_ref_pos,
+}
def main() -> None:
diff --git a/tests/sdk/base_forge_client_test.py b/tests/sdk/base_forge_client_test.py
index 31bcd1c8..b80086a0 100644
--- a/tests/sdk/base_forge_client_test.py
+++ b/tests/sdk/base_forge_client_test.py
@@ -3,6 +3,7 @@
import itertools
import pytest
+from attr import define
from esm.sdk import base_forge_client
from esm.sdk.api import ESMProteinError
@@ -11,23 +12,34 @@
EndpointHandler,
_poll_intervals,
)
+from esm.sdk.batch import BatchSubmission
from esm.sdk.forge import ForgeBatchClient
+from esm.utils.structure.input_builder import ProteinInput, StructurePredictionInput
-class _QuietHandler(EndpointHandler[dict]):
+@define
+class _QuietRequest:
+ pass
+
+
+class _QuietHandler(EndpointHandler[_QuietRequest, dict]):
"""An endpoint that states no pacing of its own, so the layers above it decide."""
+ request_type = _QuietRequest
+
@property
def endpoint_name(self) -> str:
return "quiet"
- def _prepare_request(self, **kwargs) -> list[dict]:
- return []
+ def _prepare_request(
+ self, request: _QuietRequest, *, model_name: str | None = None
+ ) -> dict:
+ return {}
- def _process_response(self, response: dict, **kwargs) -> dict:
+ def _process_response(self, response: dict) -> dict:
return response
- async def _async_process_response(self, response: dict, **kwargs) -> dict:
+ async def _async_process_response(self, response: dict) -> dict:
return response
@@ -182,8 +194,9 @@ def test_a_failed_job_raises_with_the_server_reason(clock):
def _stub_submission(
monkeypatch: pytest.MonkeyPatch, handler: EndpointHandler, task_id: str
) -> None:
- monkeypatch.setattr(handler, "_prepare_request", lambda **kwargs: [])
- monkeypatch.setattr(handler._batch_client, "submit", lambda *a, **kw: task_id)
+ monkeypatch.setattr(handler, "_prepare_request", lambda request, **kwargs: {})
+ submission = BatchSubmission(endpoint=handler.endpoint_name, task_ids=(task_id,))
+ monkeypatch.setattr(handler._batch_client, "submit", lambda *a, **kw: submission)
def test_the_users_interval_reaches_the_wait_loop(clock, monkeypatch):
@@ -192,7 +205,13 @@ def test_the_users_interval_reaches_the_wait_loop(clock, monkeypatch):
handler = client.fold_max_accuracy
_stub_submission(monkeypatch, handler, "task-7")
- result = handler.run(timeout=90, cancel_on_timeout=False)
+ result = handler.run(
+ all_atom_input=StructurePredictionInput(
+ sequences=[ProteinInput(id="A", sequence="ACDEF")]
+ ),
+ timeout=90,
+ cancel_on_timeout=False,
+ )
assert isinstance(result, ESMProteinError)
assert clock.sleeps == [45, 45]
@@ -204,7 +223,14 @@ def test_a_per_call_interval_reaches_the_wait_loop(clock, monkeypatch):
handler = client.fold_max_accuracy
_stub_submission(monkeypatch, handler, "task-8")
- result = handler.run(timeout=60, cancel_on_timeout=False, poll_interval=20)
+ result = handler.run(
+ all_atom_input=StructurePredictionInput(
+ sequences=[ProteinInput(id="A", sequence="ACDEF")]
+ ),
+ timeout=60,
+ cancel_on_timeout=False,
+ poll_interval=20,
+ )
assert isinstance(result, ESMProteinError)
assert clock.sleeps == [20, 20, 20]
diff --git a/tests/utils/molecular_complex_test.py b/tests/utils/molecular_complex_test.py
index bea335ce..179a33b1 100644
--- a/tests/utils/molecular_complex_test.py
+++ b/tests/utils/molecular_complex_test.py
@@ -6,13 +6,19 @@
3. Correct atom counts after roundtrip
"""
+import io
+import string
+
import biotite.structure as bs
import biotite.structure.io.pdbx as pdbx
import numpy as np
import pytest
from esm.utils import residue_constants
-from esm.utils.structure.molecular_complex import MolecularComplex
+from esm.utils.structure.molecular_complex import (
+ MolecularComplex,
+ MolecularComplexMetadata,
+)
# Minimal CIF with protein (chain A) + ligand (chain B in label_asym_id, chain A in auth_asym_id)
# This is the standard PDB convention that was previously broken.
@@ -126,6 +132,17 @@ def test_from_mmcif_ligand_on_separate_chain():
assert ligand_chain == "B"
+def test_from_mmcif_keeps_label_chains_when_residue_has_only_altloc_b():
+ """A residue whose only alt-loc is "B" must not collapse chains to auth IDs."""
+ cif = PROTEIN_LIGAND_CIF.replace(" . LIG B", " B LIG B")
+ assert cif.count(" B LIG B") == 3
+
+ mc = MolecularComplex.from_mmcif(cif)
+
+ assert mc.sequence == ["ALA", "GLY", "LIG"]
+ assert [mc.metadata.chain_lookup[int(c)] for c in mc.chain_id] == ["A", "A", "B"]
+
+
def test_to_mmcif_entity_categories():
"""to_mmcif should emit _entity category with polymer/non-polymer types."""
mc = MolecularComplex.from_mmcif(PROTEIN_LIGAND_CIF)
@@ -368,3 +385,108 @@ def test_from_mmcif_groups_residues_without_reordering_their_atoms(
np.testing.assert_array_equal(actual.chain_id, [0, 0, 2])
assert actual.metadata.chain_lookup == {0: "A", 1: "AA", 2: "Z"}
np.testing.assert_array_equal(actual.plddt, np.array([0.8, 0.2, 1], np.float32))
+
+
+def _poly_ala_chains(chain_lookup: dict[int, str]) -> MolecularComplex:
+ """Lengths differ so no two chains share a sequence, hence none share an entity."""
+ chain_of_token = [
+ cid for i, cid in enumerate(sorted(chain_lookup)) for _ in range(i + 1)
+ ]
+ n_tokens = len(chain_of_token)
+ return MolecularComplex(
+ id="synthetic",
+ sequence=["ALA"] * n_tokens,
+ atom_positions=np.arange(3 * n_tokens, dtype=np.float32).reshape(n_tokens, 3),
+ atom_elements=np.array(["C"] * n_tokens, dtype=object),
+ token_to_atoms=np.array([[i, i + 1] for i in range(n_tokens)], dtype=np.int32),
+ chain_id=np.array(chain_of_token, dtype=np.int64),
+ plddt=np.ones(n_tokens, dtype=np.float32),
+ metadata=MolecularComplexMetadata(entity_lookup={}, chain_lookup=chain_lookup),
+ atom_names=np.array(["CA"] * n_tokens, dtype=object),
+ atom_hetero=np.zeros(n_tokens, dtype=bool),
+ )
+
+
+def _asym_and_entity(cif: str) -> list[tuple[str, str]]:
+ """Each distinct (label_asym_id, label_entity_id) pair, in file order."""
+ site = pdbx.CIFFile.read(io.StringIO(cif)).block["atom_site"]
+ rows = zip(
+ site["label_asym_id"].as_array(str), site["label_entity_id"].as_array(str)
+ )
+ return list(dict.fromkeys(rows))
+
+
+def test_to_mmcif_names_null_chains_with_distinct_entities() -> None:
+ # 28 chains so labels run past Z, where chr(65 + n) emitted punctuation.
+ rows = _asym_and_entity(
+ _poly_ala_chains({i: "" for i in range(28)}).to_mmcif(name_null_chains=True)
+ )
+
+ assert [asym for asym, _ in rows] == [*string.ascii_uppercase, "AA", "AB"]
+ assert len({entity for _, entity in rows}) == 28
+
+
+def test_to_mmcif_null_chain_does_not_take_an_existing_name() -> None:
+ rows = _asym_and_entity(
+ _poly_ala_chains({0: "", 1: "A"}).to_mmcif(name_null_chains=True)
+ )
+
+ assert [asym for asym, _ in rows] == ["B", "A"]
+ assert len({entity for _, entity in rows}) == 2
+
+
+def test_to_mmcif_leaves_null_chains_blank_by_default() -> None:
+ rows = _asym_and_entity(_poly_ala_chains({i: "" for i in range(3)}).to_mmcif())
+
+ # All three chains collapse onto one unnamed asym id, hence onto one entity.
+ assert {asym for asym, _ in rows} == {""}
+ assert len({entity for _, entity in rows}) == 1
+
+
+_RES_NAMES = ["GLY", "GLY", "ALA", "SER", "CYS"]
+
+
+def _chain_a(residues: list[tuple[int, str]]) -> bs.AtomArray:
+ """Chain A of ``_RES_NAMES`` with the given ``(res_id, ins_code)``, 2 atoms each."""
+ atoms = bs.AtomArray(2 * len(residues))
+ atoms.coord = np.arange(3 * len(atoms), dtype=np.float32).reshape(-1, 3)
+ atoms.chain_id = ["A"] * len(atoms)
+ atoms.res_id = np.repeat([res_id for res_id, _ in residues], 2)
+ atoms.ins_code = np.repeat([ins for _, ins in residues], 2)
+ atoms.res_name = np.repeat(_RES_NAMES, 2)
+ atoms.atom_name = ["N", "CA"] * len(residues)
+ atoms.element = ["N", "C"] * len(residues)
+ atoms.hetero = np.zeros(len(atoms), dtype=bool)
+ atoms.set_annotation("b_factor", np.full(len(atoms), 50.0))
+ return atoms
+
+
+def _from_cif(atoms: bs.AtomArray) -> MolecularComplex:
+ cif = pdbx.CIFFile()
+ pdbx.set_structure(cif, atoms)
+ return MolecularComplex.from_cif_file(cif)
+
+
+@pytest.mark.parametrize(
+ "build", [MolecularComplex.from_atomarray, _from_cif], ids=["atomarray", "cif"]
+)
+@pytest.mark.parametrize(
+ "residues",
+ [
+ # Antibody-style insertion codes: 52, 52A, 52B, 52C, 53.
+ [(52, ""), (52, "A"), (52, "B"), (52, "C"), (53, "")],
+ # Author numbering that goes backwards (e.g. circular permutants).
+ [(10, ""), (11, ""), (1, ""), (2, ""), (3, "")],
+ ],
+ ids=["insertion_codes", "backwards_numbering"],
+)
+def test_residues_are_separate_tokens_in_file_order(build, residues) -> None:
+ atoms = _chain_a(residues)
+
+ actual = build(atoms)
+
+ assert actual.sequence == _RES_NAMES
+ np.testing.assert_array_equal(
+ actual.token_to_atoms, [[0, 2], [2, 4], [4, 6], [6, 8], [8, 10]]
+ )
+ np.testing.assert_array_equal(actual.atom_positions, atoms.coord)