===== Boltz2 =====

  * [[https://build.nvidia.com/mit/boltz2|Boltz2]]
  * [[https://boltz.bio/boltz2|Introducing Boltz-2]]

=== Boltz2 Apptainer File Image ===

Boltz2 Apptainer File Image:

<code>
/hpc/share/containers/apptainer/boltz2/1.4.0/boltz2-1.4.0.sif
</code>

=== Boltz2 python script ===

Download the ''Boltz2'' script file ''boltz2.py'' and save it:

<code python boltz2.py>
import requests
import json
from typing import Dict, Any

import os
port = os.getenv('NIM_HTTP_API_PORT', '8000')
print(port)

SEQUENCE = "MTEYKLVVVGACGVGKSALTIQLIQNHFVDEYDPTIEDSYRKQVVID"

def query_boltz2_nim(
    input_data: Dict[str, Any],
    base_url: str = f"http://localhost:{port}"
) -> Dict[str, Any]:
    """
    Query the Boltz2 NIM with input data.

    Args:
        input_data: Dictionary containing the prediction request data
        base_url: Base URL of the NIM service (default: http://localhost:8000)

    Returns:
        Dictionary containing the prediction response
    """
    # Construct the full URL
    url = f"{base_url}/biology/mit/boltz2/predict"

    # Set headers
    headers = {
        "Content-Type": "application/json"
    }

    try:
        # Make the POST request
        response = requests.post(url, json=input_data, headers=headers)

        # Check if request was successful
        response.raise_for_status()

        # Return the JSON response
        return response.json()

    except requests.exceptions.RequestException as e:
        print(f"Error querying NIM: {e}")
        if hasattr(e.response, 'text'):
            print(f"Response text: {e.response.text}")
        raise

# Example usage
if __name__ == "__main__":
    # Example input data - modify this according to your BoltzPredictionRequest structure
    example_input = {"polymers":[
        {
            "id": "A",
            "molecule_type": "protein",
            "sequence": SEQUENCE
        }
    ]}

    try:
        # Query the NIM
        result = query_boltz2_nim(example_input)

        # Print the result
        print("Prediction result:")
        print(json.dumps(result, indent=2))

    except Exception as e:
        print(f"Failed to get prediction: {e}")
</code>

=== Boltz2 GPU job ===

Script ''slurm-boltz2-gpu-a100_40g.sh'' to run ''boltz2'' on 1 node with 1 A100 (40 GB) GPU (8 tasks per node):

<code bash slurm-boltz2-gpu-a100_40g.sh>
#!/bin/bash --login
#SBATCH --job-name=boltz2
#SBATCH --output=af_output/%x.d%j/%x.o%j
#SBATCH --error=af_output/%x.d%j/%x.e%j
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=8
#SBATCH --time=0-02:00:00
#SBATCH --mem=40G
#SBATCH --partition=gpu
#SBATCH --qos=gpu
#SBATCH --gres=gpu:a100_40g:1
##SBATCH --account=<account>

shopt -q login_shell || exit 1
test -n "$SLURM_NODELIST" || exit 1
test $SLURM_NNODES -eq 1 || exit 1

module load apptainer
module load boltz2

export NGC_API_KEY="INSERT API KEY HERE"
export NIM_HTTP_API_PORT=$(hpc-find-free-tcp4-port 2>/dev/null || echo 8000)
export TMPDIR=$HOME/nim-cache
export APPTAINERENV_NGC_API_KEY=$NGC_API_KEY
export APPTAINERENV_LOCAL_NIM_CACHE=/nim-cache
export APPTAINERENV_NIM_CACHE=/nim-cache
export APPTAINERENV_NIM_WORKSPACE=/nim-cache/workspace
export APPTAINERENV_NIM_HTTP_API_PORT=$NIM_HTTP_API_PORT

mkdir -p $HOME/nim-cache/workspace

apptainer instance run \
  --nv \
  --bind $HOME/nim-cache:/nim-cache \
  --bind $HOME/nim-cache:/opt/nim/.cache \
  --bind $HOME/nim-cache/workspace:/opt/nim/workspace \
  $BOLTZ2_CONTAINER $SLURM_JOB_ID

apptainer instance list

# Waiting for the boltz2 server to start
until curl -sSf http://localhost:$NIM_HTTP_API_PORT/health/ready; do
    echo "Server not ready, waiting 10 seconds..."
    sleep 10
done
echo "Boltz2 server ready!"

# Testing boltz2
module load miniconda3
source "$CONDA_PREFIX/etc/profile.d/conda.sh"
conda activate boltz2-env
python $HOME/boltz2.py
</code>