• Boltz2
• Introducing Boltz-2

Boltz2 Apptainer File Image:

/hpc/share/containers/apptainer/boltz2/1.4.0/boltz2-1.4.0.sif

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

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}")

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

slurm-alphafold-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 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
 
mkdir -p $HOME/nim-cache/workspace
 
export NIM_HTTP_API_PORT=$(hpc-find-free-tcp4-port)
export APPTAINERENV_NIM_HTTP_API_PORT=$NIM_HTTP_API_PORT
 
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