Build AlphaFold2 cluster on Azure CycleCloud

Since released from July last year, AlphaFold2 protein folding algorithm is often used by more researchers and companies to drive more innovations for molecular analysis, drug discovery & etc. To build an AlphaFold2 computing cluster rapidly on the cloud will be the necessary step to leverage agility of cloud computing without CAPEX ahead.

Azure HPC stack has complete portfolio suitable for running AlphaFold2 in large scale, including GPU, storage and orchestrator service. This blog brings detailed steps of building AlphaFold2 HPC cluster on Azure to fasten your process.

Architecture

 Build Steps

1. Prerequisites
1. Check GPU quota and Azure NetApp Files(ANF) quota. SKU of NCsv3_T4 will be used in this building and NV_A10_v5 SKU (in preview) will also be suitable in the next.
2. Create a storage account with unique name (eg. saAlphaFold2 ) for CycleCloud using.
3. Prepare a SSH key pair in Azure portal.
4. Determine your working region with consideration of ANF service availability (eg. Southeast Asia).
5. Create a resource group in selected region (eg. rgAlphaFold).
6. Set Azure Cloud Shell ready for use.
2. Build CycleCloud environment following ARM template method. Set the VNet name as "vnetprotein". Use the "saAlphaFold2" as the related storage account. After all the resources are built, you can find the CycleCloud UI portal address in console "Home->Virtual Machines->cyclecloud->Overview->DNS name". Go through the first login process using your username and password.
3. Config ANF storage. Follow the steps to set up an ANF volume. Consider the dataset size of AlphaFold2, suggest to set the capacity pool and volume size as 4TB at least. Set the volume name as "volprotein" and create a dedicate subnet with CIDR “10.0.2.0/24” in Visual Network "vnetprotein". In "Protocol" settings, set file path also as "volprotein" and select "NFSv4.1". After volume is ready, remember the "Mount path" info like "10.0.2.4:/volprotein".
4. Prepare the VM image.
   1. Boot a VM using the "CentOS-based 7.9 HPC - x64 Gen2" marketplace image and change the OS disk size as 128GB.
   2. Connect the VM by SSH and install AlphaFold2 components using below commands.sudo yum install epel-release python3 -y sudo yum install aria2 -y sudo yum remove moby-cli.x86_64 moby-containerd.x86_64 moby-engine.x86_64 moby-runc.x86_64 -y sudo yum-config-manager --add-repo=https://download.docker.com/linux/centos/docker-ce.repo sudo yum install -y https://download.docker.com/linux/centos/7/x86_64/stable/Packages/containerd.io-1.4.13-3.1.el7.x86_64.rpm sudo yum install docker-ce -y sudo systemctl --now enable docker distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \ && curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.repo | sudo tee /etc/yum.repos.d/nvidia-container-toolkit.repo sudo yum clean expire-cache sudo yum install -y nvidia-docker2 sudo systemctl restart docker sudo usermod -aGdocker $USER newgrp docker sudo su cd /opt git clone https://github.com/deepmind/alphafold.git cd alphafold/ sed -i '/SHELL ["/bin/bash", "-c"]/a\RUN gpg --keyserver keyserver.ubuntu.com --recv A4B469963BF863CC && gpg --export --armor A4B469963BF863CC | apt-key add -' docker/Dockerfile docker build -f docker/Dockerfile -t alphafold . pip3 install -r docker/requirements.txt​Check the "docker images" to confirm the "alphafold:latest" is ready in the list.
   3. Build the custom image. Keep in the same SSH window and execute as below and go ahead with a 'y' confirmation. sudo waagent -deprovision+user​Back to Cloud Shell. Execute these commands to produce the custom image. export myVM=vmImgAlpha export myImage=imgAlphaFold2 export myResourceGroup=rgAlphaFold az vm deallocate --resource-group $myResourceGroup --name $myVM az vm generalize --resource-group $myResourceGroup --name $myVM az image create --resource-group $myResourceGroup --name $myImage --source $myVM --hyper-v-generation V2​After accomplished, find the image's "Resource ID" in console "Home->Images->Properties" page and remember it for further usage, which the form is as "/subscriptions/xxxx-xxxx-x…/resourceGroups/…/providers/Microsoft.Compute/images/imgAlphaFold2".
5. Create HPC cluster for Alphafold2. 
   1. Create a new cluster in CycleCloud and select "Slurm" as the scheduler type. Set parameter as below with other as is. Save the configuration then.
      • "Require setting" page - HPC VM Type: Standard_NC8as_T4_v3, Max HPC Cores: 24, Subnet ID: vnetprotein-compute.
      • "Network Attached Storage" page - Add NFS Mount: clicked, NFS IP: 10.0.2.4, NFS Mount point: /volprotein, NFS Export Path: /volprotein.
      • "Advanced Settings" page - Scheduler & HPC OS both with “Custom image” option clicked and stuff with custom image resource ID string in step 4.
   2. Start the cluster and wait several minutes to wait cluster in ready.
   3. Login scheduler. Below steps aim to prepare dataset. Total size of the Alphafold2 dataset is ~2.2TB. Suggest to execute each download sentence in download_all_data.sh if you want to save some time, such as download_pdb70.sh, download_uniref90.sh & etc.. Dataset preparation may need several hours as expected.mkdir /volprotein/AlphaFold2 mkdir /volprotein/AlphaFold2/input mkdir /volprotein/AlphaFold2/result sudo chmod +w /volprotein/AlphaFold2 /opt/alphafold/scripts/download_all_data.sh /volprotein/AlphaFold2/​
6. Run samples
   1. A sample Slurm job script is as below. Save it as run.sh.#!/bin/bash #SBATCH -o job%j.out #SBATCH --job-name=AlphaFold #SBATCH --nodes=1 #SBATCH --cpus-per-task=4 #SBATCH --gres=gpu:1 INPUT_FILE=$1 WORKDIR=/opt/alphafold INPUTDIR=/volprotein/AlphaFold2/input OUTPUTDIR=/volprotein/AlphaFold2/result DATABASEDIR=/volprotein/AlphaFold2/ sudo python3 $WORKDIR/docker/run_docker.py --fasta_paths=$INPUTDIR/$INPUT_FILE --output_dir=$OUTPUTDIR --max_template_date=2020-05-14 --data_dir=$DATABASEDIR --db_preset=reduced_dbs​
   2. Now we can submit the AlphaFold2 computing jobs! Submit this job with a test sample (*.fa or *.fasta) in /volprotein/AlphaFold2/input. At the first running, cluster need several minutes waiting compute nodes get ready. Parallel jobs can be submitted and will be running on different compute node according Slurm's allocation. Then we can use "squeue" to check the Slurm queue status. Meanwhile, there are resource monitoring graphic in CycleCloud UI to grasp the performance status of this AlphaFold2 cluster. After certain job is done, check the info in .out file and the pdb result file in /volprotein/AlphaFold2/result.sbatch run.sh input.fa sbatch run.sh P05067.fasta​
   3. Tear down. When no need to use this cluster, directly delete the resource group "rgAlphaFold" will tear down the related resources in it.

Reference links

deepmind/alphafold: Open source code for AlphaFold. (github.com)

Azure CycleCloud Documentation - Azure CycleCloud | Microsoft Docs

Azure NetApp Files documentation | Microsoft Docs