diff --git a/Singularity.def b/Singularity.def index eb26b80..8ed1d4f 100644 --- a/Singularity.def +++ b/Singularity.def @@ -1,5 +1,5 @@ Bootstrap: docker -From: nvidia/cuda:11.0-cudnn8-runtime-ubuntu18.04 +From: nvidia/cuda:11.1.1-cudnn8-runtime-ubuntu18.04 Stage: spython-base %files @@ -25,7 +25,7 @@ Stage: spython-base apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \ build-essential \ cmake \ -cuda-command-line-tools-11-0 \ +cuda-command-line-tools-11-1 \ git \ hmmer \ kalign \ @@ -55,29 +55,37 @@ PATH="/opt/conda/bin:$PATH" conda update -qy conda \ && conda install -y -c conda-forge \ openmm=7.5.1 \ -cudatoolkit==11.0.3 \ +cudatoolkit==11.1.1 \ pdbfixer \ pip \ python=3.7 +### /bin/cp -r . /app/alphafold + wget -q -P /app/alphafold/alphafold/common/ \ https://git.scicore.unibas.ch/schwede/openstructure/-/raw/7102c63615b64735c4941278d92b554ec94415f8/modules/mol/alg/src/stereo_chemical_props.txt # Install pip packages. +# N.B. The URL specifies the list of jaxlib releases. The URL in alphafold release 2.2.2 is incorrect +# as it specifies the non-CUDA releases. pip3 install --upgrade pip \ && pip3 install -r /app/alphafold/requirements.txt \ -&& pip3 install --upgrade jax jaxlib==0.1.69+cuda110 -f \ -https://storage.googleapis.com/jax-releases/jax_releases.html +&& pip3 install --upgrade jax==0.2.14 jaxlib==0.1.69+cuda111 -f \ +https://storage.googleapis.com/jax-releases/jax_cuda_releases.html # Apply OpenMM patch. cd /opt/conda/lib/python3.7/site-packages patch -p0 < /app/alphafold/docker/openmm.patch +# Add SETUID bit to the ldconfig binary so that non-root users can run it. +chmod u+s /sbin/ldconfig.real + %environment export PATH="/opt/conda/bin:$PATH" %runscript cd /app/alphafold +ldconfig exec python /app/alphafold/run_alphafold.py "$@" -%startscript -cd /app/alphafold -exec python /app/alphafold/run_alphafold.py "$@" +# %startscript +# cd /app/alphafold +# exec python /app/alphafold/run_alphafold.py "$@" diff --git a/example_slurm_job.sh b/example_slurm_job.sh index ca13e6f..09dc79b 100644 --- a/example_slurm_job.sh +++ b/example_slurm_job.sh @@ -1,32 +1,45 @@ #!/bin/bash -#SBATCH --partition=gpu -#SBATCH --time=2:00:00 -#SBATCH --gres=gpu:1 +#SBATCH -p gpu +#SBATCH --time=18:00:00 +#SBATCH --gpus=4 #SBATCH --cpus-per-gpu=12 -#SBATCH --mem-per-gpu=30G +#SBATCH --mem=140G -module load alphafold +module load alphafold/2.2.2 +module load python/gcc/3.10 -# the alphafold modulefile should define: -# * ALPHAFOLD_DIR -- install location AlphaFold -# * ALPHAFOLD_DATADIR -- the DOWLOAD_DIR set in scripts/download_all_data.sh +### Check values of some environment variables +echo SLURM_JOB_GPUS=$SLURM_JOB_GPUS +echo ALPHAFOLD_DIR=$ALPHAFOLD_DIR +echo ALPHAFOLD_DATADIR=$ALPHAFOLD_DATADIR -# Run AlphaFold; default is to use GPUs, i.e. "--use_gpu" -python3 ${ALPHAFOLD_DIR}/singularity/run_singularity.py \ - --fasta_paths=T1050.fasta \ - --max_template_date=2020-05-14 \ - --preset=reduced_dbs +### +### README This runs AlphaFold 2.2.2 on the T1050.fasta file +### -# AlphaFold should use all GPU devices available to the job. +# AlphaFold should use all GPU devices available to the job by default. # To explicitly specify use of GPUs, and the GPU devices to use, add # --use_gpu --gpu_devices=${SLURM_JOB_GPUS} - +# # To run the CASP14 evaluation, use: -# --preset=casp14 - -# To benchmark, running multiple JAX model evaluations: +# --model_preset=monomer_casp14 +# +# To benchmark, running multiple JAX model evaluations (NB this +# significantly increases run time): # --benchmark -# Copy all output from AlphaFold back to directory where "sbatch" command was issued -cp -R $TMPDIR $SLURM_SUBMIT_DIR +# Run AlphaFold; default is to use GPUs, i.e. "--use_gpu" can be omitted. +python3 ${ALPHAFOLD_DIR}/singularity/run_singularity.py \ + --use_gpu --gpu_devices=${SLURM_JOB_GPUS} \ + --data_dir=${ALPHAFOLD_DATADIR} \ + --fasta_paths=T1050.fasta \ + --max_template_date=2020-05-14 \ + --model_preset=monomer_casp14 \ + --benchmark + +echo INFO: AlphaFold returned $? + +### Copy Alphafold output back to directory where "sbatch" command was issued. +mkdir $SLURM_SUBMIT_DIR/Output-$SLURM_JOB_ID +cp -R $TMPDIR $SLURM_SUBMIT_DIR/Output-$SLURM_JOB_ID diff --git a/run_singularity.py b/run_singularity.py index 7f787ac..bf1f1b3 100644 --- a/run_singularity.py +++ b/run_singularity.py @@ -15,100 +15,97 @@ """Singularity launch script for Alphafold Singularity image.""" import os +import pathlib import signal -from pathlib import Path from typing import Tuple -import tempfile from absl import app from absl import flags from absl import logging -from spython.main import Client +import tempfile +from spython.main import Client #### USER CONFIGURATION #### -# Set to target of scripts/download_all_databases.sh -DOWNLOAD_DIR = 'SET ME' - -# AlphaFold Singularity image. -singularity_image = Client.load('alphafold.sif') +# Path to AlphaFold Singularity image. This relies on +# the environment variable ALPHAFOLD_DIR which is the +# directory where AlphaFold is installed. +singularity_image = Client.load(os.path.join(os.environ['ALPHAFOLD_DIR'], 'alphafold.sif')) # Path to a directory that will store the results. if 'TMPDIR' in os.environ: - output_dir = os.environ['TMPDIR'] + output_dir = os.environ['TMPDIR'] else: - output_dir = tempfile.mkdtemp(dir='/tmp', prefix='alphafold-') - -# Names of models to use. -model_names = [ - 'model_1', - 'model_2', - 'model_3', - 'model_4', - 'model_5', -] - -# You can individually override the following paths if you have placed the -# data in locations other than the DOWNLOAD_DIR. - -# Path to directory of supporting data, contains 'params' dir. -data_dir = DOWNLOAD_DIR - -# Path to the Uniref90 database for use by JackHMMER. -uniref90_database_path = os.path.join( - DOWNLOAD_DIR, 'uniref90', 'uniref90.fasta') - -# Path to the MGnify database for use by JackHMMER. -mgnify_database_path = os.path.join( - DOWNLOAD_DIR, 'mgnify', 'mgy_clusters_2018_12.fa') - -# Path to the BFD database for use by HHblits. -bfd_database_path = os.path.join( - DOWNLOAD_DIR, 'bfd', 'bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt') - -# Path to the Small BFD database for use by JackHMMER. -small_bfd_database_path = os.path.join( - DOWNLOAD_DIR, 'small_bfd', 'bfd-first_non_consensus_sequences.fasta') - -# Path to the Uniclust30 database for use by HHblits. -uniclust30_database_path = os.path.join( - DOWNLOAD_DIR, 'uniclust30', 'uniclust30_2018_08', 'uniclust30_2018_08') - -# Path to the PDB70 database for use by HHsearch. -pdb70_database_path = os.path.join(DOWNLOAD_DIR, 'pdb70', 'pdb70') - -# Path to a directory with template mmCIF structures, each named .cif') -template_mmcif_dir = os.path.join(DOWNLOAD_DIR, 'pdb_mmcif', 'mmcif_files') - -# Path to a file mapping obsolete PDB IDs to their replacements. -obsolete_pdbs_path = os.path.join(DOWNLOAD_DIR, 'pdb_mmcif', 'obsolete.dat') - -#### END OF USER CONFIGURATION #### - - -flags.DEFINE_bool('use_gpu', True, 'Enable NVIDIA runtime to run with GPUs.') -flags.DEFINE_string('gpu_devices', 'all', 'Comma separated list of devices to ' - 'pass to NVIDIA_VISIBLE_DEVICES.') -flags.DEFINE_list('fasta_paths', None, 'Paths to FASTA files, each containing ' - 'one sequence. Paths should be separated by commas. ' - 'All FASTA paths must have a unique basename as the ' - 'basename is used to name the output directories for ' - 'each prediction.') -flags.DEFINE_string('max_template_date', None, 'Maximum template release date ' - 'to consider (ISO-8601 format - i.e. YYYY-MM-DD). ' - 'Important if folding historical test sets.') -flags.DEFINE_enum('preset', 'full_dbs', - ['reduced_dbs', 'full_dbs', 'casp14'], - 'Choose preset model configuration - no ensembling and ' - 'smaller genetic database config (reduced_dbs), no ' - 'ensembling and full genetic database config (full_dbs) or ' - 'full genetic database config and 8 model ensemblings ' - '(casp14).') -flags.DEFINE_boolean('benchmark', False, 'Run multiple JAX model evaluations ' - 'to obtain a timing that excludes the compilation time, ' - 'which should be more indicative of the time required for ' - 'inferencing many proteins.') + output_dir = tempfile.mkdtemp(dir='/tmp', prefix='alphafold-') + +#### END USER CONFIGURATION #### + + +flags.DEFINE_bool( + 'use_gpu', True, 'Enable NVIDIA runtime to run with GPUs.') +flags.DEFINE_boolean( + 'run_relax', True, + 'Whether to run the final relaxation step on the predicted models. Turning ' + 'relax off might result in predictions with distracting stereochemical ' + 'violations but might help in case you are having issues with the ' + 'relaxation stage.') +flags.DEFINE_bool( + 'enable_gpu_relax', True, 'Run relax on GPU if GPU is enabled.') +flags.DEFINE_string( + 'gpu_devices', 'all', + 'Comma separated list of devices to pass to NVIDIA_VISIBLE_DEVICES.') +flags.DEFINE_list( + 'fasta_paths', None, 'Paths to FASTA files, each containing a prediction ' + 'target that will be folded one after another. If a FASTA file contains ' + 'multiple sequences, then it will be folded as a multimer. Paths should be ' + 'separated by commas. All FASTA paths must have a unique basename as the ' + 'basename is used to name the output directories for each prediction.') +flags.DEFINE_string( + 'output_dir', '/tmp/alphafold', + 'Path to a directory that will store the results.') +flags.DEFINE_string( + 'data_dir', None, + 'Path to directory with supporting data: AlphaFold parameters and genetic ' + 'and template databases. Set to the target of download_all_databases.sh.') +flags.DEFINE_string( + 'docker_image_name', 'alphafold', 'Name of the AlphaFold Docker image.') +flags.DEFINE_string( + 'max_template_date', None, + 'Maximum template release date to consider (ISO-8601 format: YYYY-MM-DD). ' + 'Important if folding historical test sets.') +flags.DEFINE_enum( + 'db_preset', 'full_dbs', ['full_dbs', 'reduced_dbs'], + 'Choose preset MSA database configuration - smaller genetic database ' + 'config (reduced_dbs) or full genetic database config (full_dbs)') +flags.DEFINE_enum( + 'model_preset', 'monomer', + ['monomer', 'monomer_casp14', 'monomer_ptm', 'multimer'], + 'Choose preset model configuration - the monomer model, the monomer model ' + 'with extra ensembling, monomer model with pTM head, or multimer model') +flags.DEFINE_integer('num_multimer_predictions_per_model', 5, 'How many ' + 'predictions (each with a different random seed) will be ' + 'generated per model. E.g. if this is 2 and there are 5 ' + 'models then there will be 10 predictions per input. ' + 'Note: this FLAG only applies if model_preset=multimer') +flags.DEFINE_boolean( + 'benchmark', False, + 'Run multiple JAX model evaluations to obtain a timing that excludes the ' + 'compilation time, which should be more indicative of the time required ' + 'for inferencing many proteins.') +flags.DEFINE_boolean( + 'use_precomputed_msas', False, + 'Whether to read MSAs that have been written to disk instead of running ' + 'the MSA tools. The MSA files are looked up in the output directory, so it ' + 'must stay the same between multiple runs that are to reuse the MSAs. ' + 'WARNING: This will not check if the sequence, database or configuration ' + 'have changed.') +flags.DEFINE_string( + 'docker_user', f'{os.geteuid()}:{os.getegid()}', + 'UID:GID with which to run the Docker container. The output directories ' + 'will be owned by this user:group. By default, this is the current user. ' + 'Valid options are: uid or uid:gid, non-numeric values are not recognised ' + 'by Docker unless that user has been created within the container.') FLAGS = flags.FLAGS @@ -135,6 +132,55 @@ def main(argv): if len(argv) > 1: raise app.UsageError('Too many command-line arguments.') + # You can individually override the following paths if you have placed the + # data in locations other than the FLAGS.data_dir. + + # Path to the Uniref90 database for use by JackHMMER. + uniref90_database_path = os.path.join( + FLAGS.data_dir, 'uniref90', 'uniref90.fasta') + + # Path to the Uniprot database for use by JackHMMER. + uniprot_database_path = os.path.join( + FLAGS.data_dir, 'uniprot', 'uniprot.fasta') + + # Path to the MGnify database for use by JackHMMER. + mgnify_database_path = os.path.join( + FLAGS.data_dir, 'mgnify', 'mgy_clusters_2018_12.fa') + + # Path to the BFD database for use by HHblits. + bfd_database_path = os.path.join( + FLAGS.data_dir, 'bfd', + 'bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt') + + # Path to the Small BFD database for use by JackHMMER. + small_bfd_database_path = os.path.join( + FLAGS.data_dir, 'small_bfd', 'bfd-first_non_consensus_sequences.fasta') + + # Path to the Uniclust30 database for use by HHblits. + uniclust30_database_path = os.path.join( + FLAGS.data_dir, 'uniclust30', 'uniclust30_2018_08', 'uniclust30_2018_08') + + # Path to the PDB70 database for use by HHsearch. + pdb70_database_path = os.path.join(FLAGS.data_dir, 'pdb70', 'pdb70') + + # Path to the PDB seqres database for use by hmmsearch. + pdb_seqres_database_path = os.path.join( + FLAGS.data_dir, 'pdb_seqres', 'pdb_seqres.txt') + + # Path to a directory with template mmCIF structures, each named .cif. + template_mmcif_dir = os.path.join(FLAGS.data_dir, 'pdb_mmcif', 'mmcif_files') + + # Path to a file mapping obsolete PDB IDs to their replacements. + obsolete_pdbs_path = os.path.join(FLAGS.data_dir, 'pdb_mmcif', 'obsolete.dat') + + alphafold_path = pathlib.Path(__file__).parent.parent + data_dir_path = pathlib.Path(FLAGS.data_dir) + if alphafold_path == data_dir_path or alphafold_path in data_dir_path.parents: + raise app.UsageError( + f'The download directory {FLAGS.data_dir} should not be a subdirectory ' + f'in the AlphaFold repository directory. If it is, the Docker build is ' + f'slow since the large databases are copied during the image creation.') + binds = [] command_args = [] @@ -149,20 +195,25 @@ def main(argv): database_paths = [ ('uniref90_database_path', uniref90_database_path), ('mgnify_database_path', mgnify_database_path), - ('pdb70_database_path', pdb70_database_path), - ('data_dir', data_dir), + ('data_dir', FLAGS.data_dir), ('template_mmcif_dir', template_mmcif_dir), ('obsolete_pdbs_path', obsolete_pdbs_path), ] - if FLAGS.preset == 'reduced_dbs': + if FLAGS.model_preset == 'multimer': + database_paths.append(('uniprot_database_path', uniprot_database_path)) + database_paths.append(('pdb_seqres_database_path', + pdb_seqres_database_path)) + else: + database_paths.append(('pdb70_database_path', pdb70_database_path)) + + if FLAGS.db_preset == 'reduced_dbs': database_paths.append(('small_bfd_database_path', small_bfd_database_path)) else: database_paths.extend([ - ('uniclust30_database_path', uniclust30_database_path), - ('bfd_database_path', bfd_database_path), + ('uniclust30_database_path', uniclust30_database_path), + ('bfd_database_path', bfd_database_path), ]) - for name, path in database_paths: if path: bind, target_path = _create_bind(name, path) @@ -172,13 +223,19 @@ def main(argv): output_target_path = os.path.join(_ROOT_MOUNT_DIRECTORY, 'output') binds.append(f'{output_dir}:{output_target_path}') + use_gpu_relax = FLAGS.enable_gpu_relax and FLAGS.use_gpu + command_args.extend([ - f'--output_dir={output_target_path}', - f'--model_names={",".join(model_names)}', - f'--max_template_date={FLAGS.max_template_date}', - f'--preset={FLAGS.preset}', - f'--benchmark={FLAGS.benchmark}', - '--logtostderr', + f'--output_dir={output_target_path}', + f'--max_template_date={FLAGS.max_template_date}', + f'--db_preset={FLAGS.db_preset}', + f'--model_preset={FLAGS.model_preset}', + f'--benchmark={FLAGS.benchmark}', + f'--use_precomputed_msas={FLAGS.use_precomputed_msas}', + f'--num_multimer_predictions_per_model={FLAGS.num_multimer_predictions_per_model}', + f'--run_relax={FLAGS.run_relax}', + f'--use_gpu_relax={use_gpu_relax}', + '--logtostderr', ]) options = [ @@ -188,31 +245,21 @@ def main(argv): '--env', 'OPENMM_CPU_THREADS=12' ] - if FLAGS.use_gpu: - if FLAGS.gpu_devices: - options.extend(['--env', f'NVIDIA_VISIBLE_DEVICES={FLAGS.gpu_devices}']) - elif os.environ['SLURM_JOB_GPUS']: - options.extend(['--env', f'NVIDIA_VISIBLE_DEVICES={os.environ["SLURM_JOB_GPUS"]}']) - - - # result is a dict with keys "message" (value = all output as a single string), + # Run the container. + # Result is a dict with keys "message" (value = all output as a single string), # and "return_code" (value = integer return code) result = Client.run( - singularity_image, - command_args, - nv=True if FLAGS.use_gpu else None, - return_result=True, - options=options - ) - - for line in result["message"].strip().split('\n'): - print(f'{line}') - - return result['return_code'] + singularity_image, + command_args, + nv=True if FLAGS.use_gpu else None, + return_result=True, + options=options + ) if __name__ == '__main__': flags.mark_flags_as_required([ + 'data_dir', 'fasta_paths', 'max_template_date', ])