Batch Search (Experimental)#

Quick Start#

Create local directories for the NIM cache and batch output:

export LOCAL_NIM_CACHE=~/.cache/nim
mkdir -p "$LOCAL_NIM_CACHE"
sudo chmod 0777 -R "$LOCAL_NIM_CACHE"

export MSA_OUTPUT_DIR=~/msa-outputs
mkdir -p "$MSA_OUTPUT_DIR"
# Make the output mount writable by the container (UID may differ from the host user)
sudo chmod 0777 -R "$MSA_OUTPUT_DIR"

Start the NIM with experimental features and file output enabled:

docker run -it --rm \
  --runtime=nvidia \
  --gpus all \
  -e NGC_API_KEY \
  -e NIM_EXPERIMENTAL=1 \
  -e NIM_MSA_OUTPUT_DIRECTORY=/msa-outputs \
  -v "$LOCAL_NIM_CACHE:/opt/nim/.cache" \
  -v "$MSA_OUTPUT_DIR:/msa-outputs" \
  -p 8000:8000 \
  nvcr.io/nim/colabfold/msa-search:2

Confirm the service is ready:

curl http://localhost:8000/v1/health/ready

Submit a batch request:

curl http://localhost:8000/biology/colabfold/msa-search/predict \
  -H "Content-Type: application/json" \
  -d '{
    "sequence": ["SGSMKTAISLPDETFDRVSRRASELGMSRSEFFTKAAQR", "MTTQAPTFTQPLQSVVV"],
    "e_value": 0.0001,
    "search_type": "colabfold"
  }'

Benchmarking Batch Performance#

The MSA Search NIM includes a benchmarking script that can measure batch throughput. To run a scaling benchmark against a running NIM instance, do the following:

Ensure the NIM is running as described in the previous section.
Execute the benchmark:

docker run -it --net host --entrypoint "" \
    nvcr.io/nim/colabfold/msa-search:2 \
    /opt/nim/benchmark.py --benchmark-type scaling

The benchmark measures throughput (sequences per second) and speedup across increasing batch sizes. Use it to determine the optimal batch size for your specific CPU/RAM/GPU combination. Throughput typically increases with batch size but may degrade if the batch exceeds available RAM.

Custom batch sizes can be specified with --batch-sizes:

docker run -it --net host --entrypoint "" \
    nvcr.io/nim/colabfold/msa-search:2 \
    /opt/nim/benchmark.py --benchmark-type scaling --batch-sizes 2,50,150,300

Refer to Sample Benchmarking Scripts for additional benchmark options.

When the NIM is started with NIM_EXPERIMENTAL=1, all three search endpoints accept batch inputs. For monomer/template endpoints, pass sequence as a list of strings. For paired/multimer, pass sequences as a list of complexes (each complex is a list of chain sequences). The response wraps individual results in a results array with per-item status.

Requirements: Set NIM_EXPERIMENTAL=1 environment variable when starting the NIM.

Monomer Batch#

Endpoint: /biology/colabfold/msa-search/predict

Send sequence as a list of strings:

{
  "sequence": ["MARL...", "MTTQ...", "SGSMK..."],
  "e_value": 0.0001,
  "search_type": "colabfold"
}

Response:

{
  "results": [
    {"index": 0, "status": "success", "alignments": {}, "metrics": {}},
    {"index": 1, "status": "success", "alignments": {}, "metrics": {}},
    {"index": 2, "status": "failed", "error": "..."}
  ]
}

Template Batch#

Endpoint: /biology/colabfold/msa-search/structure-templates/predict

Send sequence as a list of strings:

{
  "sequence": ["MARL...", "MTTQ..."],
  "structural_template_databases": ["pdb70_220313"]
}

Response:

{
  "results": [
    {"index": 0, "status": "success", "alignments": {}, "search_hits": {}, "structures": {}},
    {"index": 1, "status": "failed", "error": "..."}
  ]
}

Paired/Multimer Batch#

Endpoint: /biology/colabfold/msa-search/paired/predict

Send sequences as a list of lists (each inner list is one complex):

{
  "sequences": [
    ["CHAIN_A_SEQ", "CHAIN_B_SEQ"],
    ["CHAIN_C_SEQ", "CHAIN_D_SEQ"]
  ]
}

Each result item includes alignments_by_chain.

Batch Limits#

Maximum batch size: NIM_MSA_API_MAX_BATCH_SIZE (default: 300)
All items in a batch share the same parameters (including databases and e-value thresholds)
Failed items do not affect other items in the batch

File Output Mode (Optional)#

When NIM_EXPERIMENTAL=1 and NIM_MSA_OUTPUT_DIRECTORY are set at deployment time, NIM writes alignment files to disk and alignment content is returned as a file reference instead of inline content. The top-level batch response shape stays the same ({"results": [...]}); only the nested alignment field changes. That field is polymorphic (string for inline content or {"ref": "file://..."}) and lives at results[i].alignments[<database>][<format>] for monomer and template batch responses (or under results[i].alignments_by_chain[<chain>][<database>][<format>] for paired/multimer batch responses).

File references are returned for batch requests only. Non-batch requests, such as a monomer or template search with sequence as a string or a paired search with sequences as a list of chain strings, always return inline alignment content, even when NIM_MSA_OUTPUT_DIRECTORY is set.

NIM_MSA_OUTPUT_DIRECTORY=/mnt/msa-outputs

Output layout: <NIM_MSA_OUTPUT_DIRECTORY>/<batch-id>/seq_<NNN>_<db>.<fmt>

Example monomer batch response without NIM_MSA_OUTPUT_DIRECTORY (default). The value at results[0].alignments["uniref30_2302"]["a3m"] is an inline AlignmentFileRecord:

{
  "results": [
    {
      "index": 0,
      "status": "success",
      "alignments": {
        "uniref30_2302": {
          "a3m": {
            "alignment": ">query\nMVPSAGQLALFALGIV...\n>hit1\n...",
            "format": "a3m"
          }
        }
      },
      "metrics": {}
    }
  ]
}

Example monomer batch response with NIM_MSA_OUTPUT_DIRECTORY=/mnt/outputs. The same nested path now carries a file reference instead of inline content:

{
  "results": [
    {
      "index": 0,
      "status": "success",
      "alignments": {
        "uniref30_2302": {
          "a3m": {
            "alignment": {"ref": "file:///mnt/outputs/018f6b2c-abc123/seq_000_uniref30_2302.a3m"},
            "format": "a3m"
          }
        }
      },
      "metrics": {}
    }
  ]
}

The alignment field is polymorphic: either a string (inline content) or an object with a ref key containing a file:// URI. Clients should check the type of the alignment value at results[i].alignments[<database>][<format>].alignment (or the equivalent alignments_by_chain path for paired/multimer) and resolve the reference if it is an object.

Atomicity: Files are written to a temporary directory first, then renamed into place using POSIX rename(). On filesystems that support atomic rename (local disk, NFS), readers will not observe partial content. Verify that your storage backend supports POSIX rename semantics if atomicity is required.

Important

The NIM is designed to only write to this directory. Lifecycle management — cleanup, rotation, archival to object storage (such as S3 and GCS) — is the responsibility of the external orchestrator or control plane. This design allows mounting any POSIX-compatible filesystem (local disk, NFS, FUSE-based S3 gateways) without NIM needing awareness of the storage backend.