Add Blackwell GPU support using NVIDIA PyTorch container

Build_instructions_blackwell.md

@@ -0,0 +1,32 @@
+# Build Instructions for Dockerfile.blackwell
+
+## Basic build command:
+
+```bash
+docker build -f Dockerfile.blackwell -t openfold-3-blackwell:latest .
+```
+
+This will create a Docker image named `openfold-3-blackwell` with the `latest` tag.
+
+
+## test Pytorch and CUDA
+
+```bash
+docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 openfold-3-blackwell:latest   python -c "import torch; print('CUDA:', torch.version.cuda); print('PyTorch:', torch.__version__)"  
+```
+
+Should print something like:
+CUDA: 12.8
+PyTorch: 2.7.0a0+ecf3bae40a.nv25.02
+
+
+## test run_openfold inference example
+
+docker run --gpus all -it --ipc=host --ulimit memlock=-1 \
+    -v $(pwd):/output \
+    -w /output openfold-3-blackwell:latest \
+    run_openfold predict \
+    --query_json=/opt/openfold-3/examples/example_inference_inputs/query_ubiquitin.json \
+    --num_diffusion_samples=1 \
+    --num_model_seeds=1 \
+    --use_templates=false 

Dockerfile.blackwell

@@ -0,0 +1,81 @@
+# Simple OpenFold3 Dockerfile using NVIDIA PyTorch container
+FROM nvcr.io/nvidia/pytorch:25.02-py3
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    wget \
+    git \
+    hmmer \
+    kalign \
+    libxrender1 \
+    libxext6 \
+    libsm6 \
+    libxft2 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Clone OpenFold3 source and modify environment file
+WORKDIR /opt
+RUN git clone https://github.com/aqlaboratory/openfold-3.git && \
+    cd openfold-3 && \
+    cp -p environments/production.yml environments/production.yml.backup && \
+    grep -v "pytorch::pytorch" environments/production.yml > environments/production.yml.tmp && \
+    mv environments/production.yml.tmp environments/production.yml
+
+WORKDIR /opt/openfold-3
+
+# Install Python dependencies
+RUN pip install --no-cache-dir \
+    biopython \
+    numpy \
+    pandas \
+    PyYAML \
+    requests \
+    scipy \
+    tqdm \
+    typing-extensions \
+    wandb \
+    modelcif \
+    ml-collections \
+    rdkit \
+    boto3 \
+    lmdb \
+    ijson \
+    deepspeed \
+    pdbeccdutils \
+    pytorch-lightning \
+    awscli \
+    memory_profiler \
+    func_timeout \
+    biotite==1.2.0 \
+    "nvidia-cutlass<4" \
+    "cuda-python<12.9.1"
+
+# Install CUTLASS for DeepSpeed Evoformer attention kernel
+WORKDIR /opt
+RUN git clone https://github.com/NVIDIA/cutlass --branch v3.6.0 --depth 1
+
+# Install OpenFold3 package itself (provides run_openfold command)
+WORKDIR /opt/openfold-3
+RUN python3 setup.py install
+
+# Set environment variables including CUDA architecture for Blackwell
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    KMP_AFFINITY=none \
+    CUTLASS_PATH=/opt/cutlass \
+    TORCH_CUDA_ARCH_LIST="12.0"
+
+# Pre-compile DeepSpeed operations for Blackwell GPUs to avoid runtime compilation
+# Create necessary cache directories
+RUN python3 -c "import os; os.makedirs('/root/.triton/autotune', exist_ok=True)"
+
+# Create a Python sitecustomize.py to set TORCH_CUDA_ARCH_LIST before any imports
+# This ensures the variable is set before PyTorch's cpp_extension checks it
+RUN mkdir -p /usr/local/lib/python3.12/site-packages && \
+    echo 'import os' > /usr/local/lib/python3.12/site-packages/sitecustomize.py && \
+    echo 'os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")' >> /usr/local/lib/python3.12/site-packages/sitecustomize.py && \
+    echo 'os.environ.setdefault("CUTLASS_PATH", "/opt/cutlass")' >> /usr/local/lib/python3.12/site-packages/sitecustomize.py && \
+    echo 'os.environ.setdefault("KMP_AFFINITY", "none")' >> /usr/local/lib/python3.12/site-packages/sitecustomize.py
+
+# Default command
+CMD ["/bin/bash"]