Scale Batch Inference with Ray

In this tutorial, you scale Boltz-2 batch inference across every visible GPU with Ray. You create a worklist of independent protein requests, configure one complete model replica per GPU, run the pipeline, and inspect completion, throughput, and confidence summaries.

The complete runnable script is available at examples/quickstart/boltz2_ray.py. The sections below explain its worklist, Ray configuration, and output.

Ray replicas process independent requests concurrently. They do not split one structure prediction across several GPUs.

Before You Begin

Complete the installation on a system with at least one supported GPU. Run the serial Quickstart first to verify the wheel, checkpoint, and metadata.

If you do not have a repository checkout, create quickstart_ray.py with the following example. Both versions configure the Ray workers to import BioNeMo Inference Runtime (BioIR) from the installed wheel rather than an editable checkout.

Start at the top-level imports.

1
import json
2
import logging
3
import os
4
import time
5
from pathlib import Path
6
7
os.environ.setdefault("RAY_BACKEND_LOG_LEVEL", "fatal")
8
9
import ray
10
import torch
11
12
from bionemo_ir.data.schemas import InputRequest, MSARecord, Polymer
13
from bionemo_ir.pipeline.processor.engine_proc import (
14
    EngineProcessorConfig,
15
    build_processor,
16
)
17
from bionemo_ir.pipeline.stages.configs import (
18
    EngineStageConfig,
19
    FeatureGeneratorStageConfig,
20
    WriterStageConfig,
21
)

Understand the Worklist

The first part of quickstart_ray.py detects the visible GPUs and creates one prediction request for each GPU:

1
SEQUENCE = "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG"
2
3
replicas = torch.cuda.device_count()
4
if not replicas:
5
    raise RuntimeError("this example requires at least one visible GPU")
6
7
rows = []
8
for index in range(replicas):
9
    record_id = f"1UBQ-{index + 1}"
10
    request = InputRequest(
11
        input_id=record_id,
12
        polymers=[
13
            Polymer(
14
                chain_id=["A1"],
15
                sequence=SEQUENCE,
16
                msas=[MSARecord(content=f">{record_id}\n{SEQUENCE}\n")],
17
            )
18
        ],
19
    )
20
    rows.append({"record": request, "__record_id": record_id})

This block creates the Ray worklist:

• torch.cuda.device_count() determines how many complete model replicas the example can run and stops with an error when no GPU is visible.
• Each InputRequest contains protein chain A1, the 1UBQ ubiquitin sequence, and an inline, query-only unpaired MSA.
• Each request has a unique identifier, such as 1UBQ-1 or 1UBQ-2. __record_id becomes the CIF filename.
• The number of requests matches the number of replicas so the example verifies that each replica can complete one prediction.

Use a larger, representative worklist when measuring sustained throughput. Refer to Input Requests for other biomolecule types, MSAs, and templates.

Understand the Pipeline Configuration

The next block configures deterministic feature generation and one Ray engine actor per visible GPU:

1
SAMPLING_STEPS = 50
2
SEED = 42
3
4
5
def seed_worker() -> None:
6
    """Seed each Ray worker before it loads a pipeline stage."""
7
    torch.manual_seed(SEED)
8
9
10
launch_dir = Path.cwd()
11
config = EngineProcessorConfig(
12
    model_source="boltz-2",
13
    executor_backend="ray",
14
    runtime_args={"num_sampling_steps": SAMPLING_STEPS},
15
    feature_generator_stage=FeatureGeneratorStageConfig(
16
        init_context={"random_seed": SEED},
17
    ),
18
    engine_stage=EngineStageConfig(compute=replicas),
19
    writer_stage=WriterStageConfig(
20
        output_path=str(launch_dir / f"output/ray-{replicas}gpu"),
21
        format="cif",
22
    ),
23
    should_continue_on_error=False,
24
)

The configuration controls the distributed prediction pipeline:

• executor_backend="ray" sends rows through the Ray Data pipeline.
• EngineStageConfig(compute=replicas) creates one engine actor per visible GPU. Each actor reserves one GPU and loads one complete Boltz-2 model replica.
• num_sampling_steps=50 shortens the diffusion stage for this example.
• random_seed=42 seeds feature generation. The worker setup hook also seeds PyTorch before each worker loads its pipeline stage.
• WriterStageConfig uses an absolute path captured before the workers change directories and writes one CIF structure per request.
• should_continue_on_error=False stops the example if any request fails.

Refer to EngineProcessorConfig, Runtime Args, and Ray Multi-GPU Replicas for the available controls.

Understand Ray Execution

The next block starts Ray, materializes the worklist, waits for every result, and shuts Ray down:

1
# Keep workers outside a mounted source checkout so they use the installed wheel.
2
os.chdir("/tmp")
3
logging.getLogger().setLevel(logging.WARNING)
4
logging.getLogger("ray.data._internal").setLevel(logging.ERROR)
5
ray.init(
6
    include_dashboard=False,
7
    logging_level="ERROR",
8
    log_to_driver=False,
9
    runtime_env={"worker_process_setup_hook": seed_worker},
10
)
11
data_context = ray.data.DataContext.get_current()
12
data_context.enable_progress_bars = False
13
data_context.enable_operator_progress_bars = False
14
15
start = time.perf_counter()
16
try:
17
    dataset = ray.data.from_items(rows)
18
    outputs = list(build_processor(config)(dataset).materialize().iter_rows())
19
    resources = ray.cluster_resources()
20
finally:
21
    elapsed = time.perf_counter() - start
22
    ray.shutdown()

The script starts workers from /tmp so a mounted source checkout cannot shadow the installed wheel. ray.init() disables the dashboard and verbose driver logs, then runs seed_worker in each worker process.

ray.data.from_items(rows) creates the distributed dataset. build_processor(config) assembles the Ray pipeline, and materialize() waits for every request to complete. The finally block records elapsed worklist time and shuts Ray down even if inference fails.

Understand the Results

The final block calculates whole-worklist throughput and builds a compact summary from one representative prediction:

1
row = sorted(outputs, key=lambda output: output["__record_id"])[0]
2
scores = json.loads(row["scores"])
3
throughput = len(outputs) * 3600 / elapsed
4
summary = {
5
    "visible_gpus": replicas,
6
    "ray_gpus": resources.get("GPU", 0),
7
    "completed_structures": f"{len(outputs)}/{len(rows)}",
8
    "worklist_wall_s": round(elapsed, 2),
9
    "structures_per_hour": round(throughput, 2),
10
    "structures_per_gpu_hour": round(throughput / replicas, 2),
11
    "example": {
12
        "record_id": row["__record_id"],
13
        "output_path": row["output_path"],
14
        "ptm": scores["ptm"],
15
        "mean_plddt": round(sum(scores["plddt"]) / len(scores["plddt"]), 2),
16
        "pae_shape": [len(scores["pae"]), len(scores["pae"][0])],
17
    },
18
}
19
print(json.dumps(summary, indent=2))

The summary reports:

• completed_structures — the number of successful outputs compared with the number of submitted requests.
• worklist_wall_s — the elapsed time for materializing the complete worklist.
• structures_per_hour — worklist throughput across all visible GPUs.
• structures_per_gpu_hour — aggregate throughput divided by the number of replicas.
• example — the output path and confidence metrics from the first result, without printing the complete pLDDT and PAE arrays.

Refer to Outputs for the complete output-row schema.

Run the Example

From the repository root, run the maintained example:

$
python examples/quickstart/boltz2_ray.py

If you copied the example into a standalone file, run that file instead:

$
python quickstart_ray.py

Example Output

The output path starts under the directory where you launch the script. Timing and scores can vary across GPUs, systems, and releases.

One NVIDIA H100 80 GB HBM3

This captured run used one 700 W H100:

1
{
2
  "visible_gpus": 1,
3
  "ray_gpus": 1.0,
4
  "completed_structures": "1/1",
5
  "worklist_wall_s": 24.95,
6
  "structures_per_hour": 144.31,
7
  "structures_per_gpu_hour": 144.31,
8
  "example": {
9
    "record_id": "1UBQ-1",
10
    "output_path": "/bioir/output/ray-1gpu/1UBQ-1.cif",
11
    "ptm": 0.9149,
12
    "mean_plddt": 0.93,
13
    "pae_shape": [
14
      76,
15
      76
16
    ]
17
  }
18
}

Eight NVIDIA H200 NVL GPUs

This captured run used eight 600 W H200 NVL GPUs:

1
{
2
  "visible_gpus": 8,
3
  "ray_gpus": 8.0,
4
  "completed_structures": "8/8",
5
  "worklist_wall_s": 54.25,
6
  "structures_per_hour": 530.91,
7
  "structures_per_gpu_hour": 66.36,
8
  "example": {
9
    "record_id": "1UBQ-1",
10
    "output_path": "/bioir/output/ray-8gpu/1UBQ-1.cif",
11
    "ptm": 0.9149,
12
    "mean_plddt": 0.93,
13
    "pae_shape": [
14
      76,
15
      76
16
    ]
17
  }
18
}

Treat these values as successful-run examples, not benchmarks. Each run uses only one short request per GPU and includes pipeline and model setup in the worklist time. The stable result is that completed_structures matches the request count and every reported output path contains a non-empty CIF file.

Next Steps

• Configure stage resources and larger worklists in Ray Multi-GPU Replicas.
• Review controlled performance measurements in Benchmarks.
• Check supported accelerators in the Support Matrix.