Boltz2 Boltz2

Boltz2 Apptainer File Image:

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

Download the Boltz2 scipt file boltz2.py and save i:

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
 
ip="$(hostname -s).hpc.unipr.it"
read lowerport upperport < /proc/sys/net/ipv4/ip_local_port_range
while :; do
    port=$(shuf -i $lowerport-$upperport -n 1)
    ss -lpna | grep -q ":$port " || break
done
#port=$(((RANDOM % (65535-49152)) + 49152))
 
export NIM_HTTP_API_PORT=$port
export APPTAINERENV_NIM_HTTP_API_PORT=$NIM_HTTP_API_PORT
echo "ip: $ip port: $port"
 
#  --network-args "portmap=$NIM_HTTP_API_PORT:8000" \
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
 
# Attendo che il server boltz2 sia partito
until curl -sSf http://localhost:$NIM_HTTP_API_PORT/health/ready; do
    echo "Server non pronto, attendo 10s..."
    sleep 10
done
echo "Boltz2 server pronto!"
 
# Testing boltz2
module load miniconda3
source "$CONDA_PREFIX/etc/profile.d/conda.sh"
conda activate boltz2-env
python $HOME/boltz2.py