openfold2

Generate API Key

Pull and run openfold/openfold2 using Docker (this will download the full model and run it in your local environment)

Copy Code
$ docker login nvcr.io
Username: $oauthtoken
Password: <PASTE_API_KEY_HERE>

Start NIM

1. Export the NGC_API_KEY environment variable.

Copy Code
export NGC_API_KEY=<your personal NGC key>

2. NIM container automatically downloads models. For OpenFold2, download should take about 5min. To save time and bandwidth it is recommended to provide local cache directory. This way NIM will be able to reuse already downloaded models. Execute the following command to setup cache directory.

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

3. Run the NIM container with the following command. Note: You might need to adjust your GPU's indices.

Copy Code
docker run -it \
    --runtime=nvidia \
    --gpus='"device=0"' \
    -p 8000:8000 \
    -e NGC_API_KEY \
    -v "$LOCAL_NIM_CACHE":/opt/nim/.cache \
    nvcr.io/nim/openfold/openfold2:1.0.0

This command will start the NIM container and expose port 8000 for the user to interact with the NIM.

4. Open a new terminal, leaving the terminal open with the just launched service. In the new terminal, wait until the health check end point returns {"status":"ready"} before proceeding. This may take a couple of minutes. You can use the following command to query the health check.

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

Structure Prediction Example

The following examples predict 3D atomic structures from input protein sequence, optionally include multiple-sequence-alignments.

Python client example

The following is an example of how you can use a Python client to input a protein sequence and two multiple sequence alignments. Two model parameters sets are included. The response includes a prediction for each provided parameters set, ordered by confidence. The relaxation step is not applied.

1. Save following Python example to a file named nim_client.py.

Copy Code
#!/usr/bin/env python3
import requests
import os
import json
from pathlib import Path

output_file="output.json"
url = "http://localhost:8000/biology/openfold/openfold2/predict-structure-from-msa-and-template"
selected_models = [1, 2]
sequence = (
    "GGSKENEISHHAKEIERLQKEIERHKQSIKKLKQSEQSNPPPNP"
    "EGTRQARRNRRRRWRERQRQKENEISHHAKEIERLQKEIERHKQSIKKLKQSEC"
)
uniref90_alignment_in_a3m_trunc10=\
""">BQXYMDHSRWGGVPIWVK
GGSKENEISHHAKEIERLQKEIERHKQSIKKLKQSEQSNPPPNPEGTRQARRNRRRRWRERQRQKENEISHHAKEIERLQKEIERHKQSIKKLKQSEC
>UniRef90_A0A221IUG4
--------------------------QTVKLVKRLYQSNPPPNPEGTRQARRNRRRRWRERQRQ----------------------------------
>UniRef90_A7KWE0
---------------------------TVRLIKQLYQSNPPPNPEGTRQARRNRRRRWRERQRQ----------------------------------
>UniRef90_A0A2I6UE91
---------------------------TVKLIKEIYQSNPPPNPEGTRQARRNRRRRWRERQRQ----------------------------------
>UniRef90_D6NY33
---------------------------AVRLIKQIYQSNPPPNPEGTRQARRNRRRRWRERQRQ----------------------------------
>UniRef90_A0A221IUJ5
--------------------------QTVKLIKRLYQSNPPPNPEGTRQARRNRRRRWREKQRQ----------------------------------
>UniRef90_D6NYR3
---------------------------TVRLVKQLYQSNPPPNPEGTRQARRNRRRRWRERQRQ----------------------------------
>UniRef90_I6Y2C4
---------------------------TVRLIKRIYQSNPPPNPEGTRQARRNRRRRWRERQRQIQN-------------------------------
>UniRef90_A0A161CVP3
--------------------------QTIRLIKLLYQSNPPPNPEGTRQARRNRRRRWRERQRQ----------------------------------
>UniRef90_Q6EFX9
--------------------------QTVRLIKLLYQSNPPPNPEGTRQARRNRRRRWRERQRQ----------------------------------
>UniRef90_A0A2I6UAR5
--------------------------ETVKIIKYLYQSNPPPNPEGTRQARRNRRRRWRERQRQ----------------------------------
"""
small_bfd_alignment_in_a3m = \
""">BQXYMDHSRWGGVPIWVK
GGSKENEISHHAKEIERLQKEIERHKQSIKKLKQSEQSNPPPNPEGTRQARRNRRRRWRERQRQKENEISHHAKEIERLQKEIERHKQSIKKLKQSEC
>A0A076V4A1_9HIV1
------------------------------------QSNPPPNHEGTRQARRNRRRRWRERQRQ----------------------------------
"""
data = {
    "sequence": sequence,
    "alignments": {
        "uniref90": {
            "a3m": {
                "alignment": uniref90_alignment_in_a3m_trunc10, 
                "format": "a3m",
            }
        },
        "small_bfd": {
            "a3m": {
                "alignment": small_bfd_alignment_in_a3m, 
                "format": "a3m",
            }
        },
    },
    "selected_models": selected_models,
    "relax_prediction": False,
}
r = requests.post(url=url, json=data)
print(r, "Saving to output.json:\n", r.text[:200], "...")
Path(output_file).write_text(r.text)

2. Execute the example.

Copy Code
chmod +x nim_client.py

./nim_client.py

3. The example saves results to the output.json file in json format. You can quickly view the file using the following command.

Copy Code
less output.json

Shell client example

The following is an example of how you can use a shell client to input a protein sequence. Two model parameters sets are included. The response includes a prediction for each provided parameters set, ordered by confidence. The relaxation step is not applied.

1. Save the following Shell example to a file named nim_client.sh.

Copy Code
#!/usr/bin/env bash
set -e

URL='http://localhost:8000/biology/openfold/openfold2/predict-structure-from-msa-and-template'

SEQUENCE="GGSKENEISHHAKEIERLQKEIERHKQSIKKLKQSEQSNPPPNPEGTRQ\
ARRNRRRRWRERQRQKENEISHHAKEIERLQKEIERHKQSIKKLKQSEC"

# Prepare the JSON data
DATA=$(cat <<EOF
{
    "sequence": "${SEQUENCE}",
    "selected_models": [1,2]
}
EOF
)

echo "Generated JSON data:"
echo "DATA=${DATA}"

# Make the request
echo "Submitting request...the response should arrive in less than a 1 minute"
POLL_SECONDS=300
curl -s -X POST "${URL}" \
    -H "content-type: application/json" \
    -H "NVCF-POLL-SECONDS: ${POLL_SECONDS}" \
    -d "${DATA}"

2. Execute the example. The example displays the results to the terminal in JSON format.

Copy Code
chmod +x nim_client.sh

./nim_client.sh

For more details on getting started with this NIM, visit the NVIDIA NIM Docs.