Add MMseqs Colabfold Search GPU Support

CHANGELOG.md

@@ -120,6 +120,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - [[#503](https://github.com/nf-core/proteinfold/issues/503)] - Add checkIfExists validation to user-provided database paths across all prepare DB subworkflows.
 - [[#507](https://github.com/nf-core/proteinfold/issues/507)] - Implement missing full tests and check that the others work before release 2.0.0.
 - [[PR #509](https://github.com/nf-core/proteinfold/pulls/509)] - Setup gpu environment for AWS full tests.
+- [[PR #497](https://github.com/nf-core/proteinfold/pull/497)] - Add MMseqs Colabfold Search GPU Support.
 - [[PR #538](https://github.com/nf-core/proteinfold/pulls/538)] - Fix alphafold2_random_seed type.
 
 ### Parameters

conf/dbs.config

@@ -83,7 +83,9 @@ params {
 
     // Colabfold paths
     colabfold_envdb_path             = "${params.colabfold_db}/colabfold_envdb/*"
+    colabfold_envdb_path_padded      = "${params.colabfold_db}/colabfold_envdb_padded/*"
     colabfold_uniref30_path          = "${params.colabfold_db}/colabfold_uniref30/*"
+    colabfold_uniref30_path_padded   = "${params.colabfold_db}/colabfold_uniref30_padded/*"
     // Are all these params options needed?
     colabfold_alphafold2_params_tags = [
         "alphafold2_multimer_v1" : "alphafold_params_colab_2021-10-27",

conf/modules_colabfold.config

@@ -33,6 +33,9 @@ process {
     withName: '.*PREPARE_COLABFOLD_DBS:ARIA2_UNIREF30:UNTAR' {
         ext.prefix = 'colabfold_uniref30'
     }
+    withName: '.*PREPARE_COLABFOLD_DBS:MMSEQS_CREATEINDEX_UNIPROT30_PADDED' {
+        ext.prefix = 'colabfold_uniref30_gpu'
+    }
 }
 
 process {

docs/gpu-dbs.md

@@ -0,0 +1,237 @@
+# Using padded databases in proteinfold
+
+Proteinfold can make use of GPU MSA search for faster searching. However, this requires creating padded databases for the GPU hardware you wish to use, and setting the appropriate flags.
+
+## Requirements
+
+- mmseqs-gpu
+- uniref30 database
+- colabfold envdb database
+- NVIDIA Ampere GPU or newer
+- CUDA 12.4 or newer
+
+## Obtaining MMseqs-GPU
+
+MMseqs has two x86 builds available for Linux. This requires the GPU version of MMseqs. It can be obtained via this command:
+
+```bash
+wget https://github.com/soedinglab/MMseqs2/releases/download/18-8cc5c/mmseqs-linux-gpu.tar.gz
+tar xvf mmseqs-linux-gpu.tar.gz
+```
+
+## Database structure
+
+Proteinfold can make use of the `--db` flag to load in all required databases. In order to load the padded databases, the database should be structured as such. Below is a truncated version of the database tree. It is important to note that the padded database files have the same prefix as the CPU files.
+
+```
+.
+├── colabfold_envdb
+│   ├── colabfold_envdb_202108_sample_h.tsv
+│   ├── colabfold_envdb_202108_sample_seq.tsv
+│   └── colabfold_envdb_202108_sample.tsv
+├── colabfold_envdb_padded
+│   ├── colabfold_envdb_202108_db_seq_h.index
+│   ├── colabfold_envdb_202108_db_seq.index
+│   └── colabfold_envdb_202108_db_seq.lookup
+├── colabfold_uniref30
+│   ├── uniref30_2302_db_seq_h.dbtype
+│   ├── uniref30_2302_db_seq_h.index
+│   ├── uniref30_2302_db_seq.index
+│   └── uniref30_2302_db_taxonomy
+└── colabfold_uniref30_padded
+    ├── uniref30_2302_db_seq_h.index
+    ├── uniref30_2302_db_seq.index
+    └── uniref30_2302_db_taxonomy
+```
+
+## Downloading UniRef30 Database
+
+Firstly, you must obtain the UniRef database. The database file is approx. 55GB. You may be able to get faster downloads by using `aria2c` with the `-x 8` option.
+
+```bash
+wget https://opendata.mmseqs.org/colabfold/uniref30_2302.db.tar.gz
+tar xvf uniref30_2302.db.tar.gz
+```
+
+## Downloading the Colabfold envdb
+
+Next, you will need to download the Colabfold envdb. This database is approx. 120GB.
+
+```bash
+wget https://opendata.mmseqs.org/colabfold/colabfold_envdb_202108.db.tar.gz
+tar xvf colabfold_envdb_202108.db.tar.gz
+
+```
+
+## CPU Database structure
+
+By now, your directory structure should look something like this
+
+```
+.
+├── colabfold_envdb
+│   ├── colabfold_envdb_202108_sample_h.tsv
+│   ├── colabfold_envdb_202108_sample_seq.tsv
+│   └── colabfold_envdb_202108_sample.tsv
+└── colabfold_uniref30
+    ├── uniref30_2302_db_seq_h.dbtype
+    ├── uniref30_2302_db_seq_h.index
+    ├── uniref30_2302_db_seq.index
+    └── uniref30_2302_db_taxonomy
+```
+## Create padded database
+
+Next, we need to create the padded databases. For this, it is recommended to duplicate the databases.
+
+```bash
+mkdir colabfold_uniref30_padded
+mmseqs makepaddedseqdb ./colabfold_uniref30/uniref30_2302_db_seq ./colabfold_uniref30_padded/uniref30_2302_db_seq
+mmseqs makepaddedseqdb ./colabfold_uniref30/uniref30_2302_db ./colabfold_uniref30_padded/uniref30_2302_db
+mkdir colabfold_envdb_padded
+mmseqs makepaddedseqdb ./colabfold_envdb/colabfold_envdb_202108_db ./colabfold_envdb_padded/colabfold_envdb_202108_db
+mmseqs makepaddedseqdb ./colabfold_envdb/colabfold_envdb_202108_db_seq ./colabfold_envdb_padded/colabfold_envdb_202108_db_seq
+cp ./colabfold_envdb/colabfold_envdb_202108_db_aln.* ./colabfold_envdb_padded/
+```
+
+You should now have a directory structure that looks something similar to this
+
+```
+.
+├── colabfold_envdb
+│   ├── colabfold_envdb_202108_db.0
+│   ├── colabfold_envdb_202108_db.1
+│   ├── colabfold_envdb_202108_db.10
+│   ├── colabfold_envdb_202108_db.11
+│   ├── colabfold_envdb_202108_db.12
+│   ├── colabfold_envdb_202108_db.13
+│   ├── colabfold_envdb_202108_db.14
+│   ├── colabfold_envdb_202108_db.15
+│   ├── colabfold_envdb_202108_db.2
+│   ├── colabfold_envdb_202108_db.3
+│   ├── colabfold_envdb_202108_db.4
+│   ├── colabfold_envdb_202108_db.5
+│   ├── colabfold_envdb_202108_db.6
+│   ├── colabfold_envdb_202108_db.7
+│   ├── colabfold_envdb_202108_db.8
+│   ├── colabfold_envdb_202108_db.9
+│   ├── colabfold_envdb_202108_db_aln.0
+│   ├── colabfold_envdb_202108_db_aln.1
+│   ├── colabfold_envdb_202108_db_aln.10
+│   ├── colabfold_envdb_202108_db_aln.11
+│   ├── colabfold_envdb_202108_db_aln.12
+│   ├── colabfold_envdb_202108_db_aln.13
+│   ├── colabfold_envdb_202108_db_aln.14
+│   ├── colabfold_envdb_202108_db_aln.15
+│   ├── colabfold_envdb_202108_db_aln.2
+│   ├── colabfold_envdb_202108_db_aln.3
+│   ├── colabfold_envdb_202108_db_aln.4
+│   ├── colabfold_envdb_202108_db_aln.5
+│   ├── colabfold_envdb_202108_db_aln.6
+│   ├── colabfold_envdb_202108_db_aln.7
+│   ├── colabfold_envdb_202108_db_aln.8
+│   ├── colabfold_envdb_202108_db_aln.9
+│   ├── colabfold_envdb_202108_db_aln.dbtype
+│   ├── colabfold_envdb_202108_db_aln.index
+│   ├── colabfold_envdb_202108_db.dbtype
+│   ├── colabfold_envdb_202108_db_h
+│   ├── colabfold_envdb_202108_db_h.dbtype
+│   ├── colabfold_envdb_202108_db_h.index
+│   ├── colabfold_envdb_202108_db.idx
+│   ├── colabfold_envdb_202108_db.idx.dbtype
+│   ├── colabfold_envdb_202108_db.idx.index
+│   ├── colabfold_envdb_202108_db.index
+│   ├── colabfold_envdb_202108_db_seq.0
+│   ├── colabfold_envdb_202108_db_seq.1
+│   ├── colabfold_envdb_202108_db_seq.10
+│   ├── colabfold_envdb_202108_db_seq.11
+│   ├── colabfold_envdb_202108_db_seq.12
+│   ├── colabfold_envdb_202108_db_seq.13
+│   ├── colabfold_envdb_202108_db_seq.14
+│   ├── colabfold_envdb_202108_db_seq.15
+│   ├── colabfold_envdb_202108_db_seq.2
+│   ├── colabfold_envdb_202108_db_seq.3
+│   ├── colabfold_envdb_202108_db_seq.4
+│   ├── colabfold_envdb_202108_db_seq.5
+│   ├── colabfold_envdb_202108_db_seq.6
+│   ├── colabfold_envdb_202108_db_seq.7
+│   ├── colabfold_envdb_202108_db_seq.8
+│   ├── colabfold_envdb_202108_db_seq.9
+│   ├── colabfold_envdb_202108_db_seq.dbtype
+│   ├── colabfold_envdb_202108_db_seq_h
+│   ├── colabfold_envdb_202108_db_seq_h.dbtype
+│   ├── colabfold_envdb_202108_db_seq_h.index
+│   ├── colabfold_envdb_202108_db_seq.index
+│   ├── colabfold_envdb_202108_sample_aln.tsv
+│   ├── colabfold_envdb_202108_sample_h.tsv
+│   ├── colabfold_envdb_202108_sample_seq.tsv
+│   └── colabfold_envdb_202108_sample.tsv
+├── colabfold_envdb_padded
+│   ├── colabfold_envdb_202108_db
+│   ├── colabfold_envdb_202108_db.dbtype
+│   ├── colabfold_envdb_202108_db_h
+│   ├── colabfold_envdb_202108_db_h.dbtype
+│   ├── colabfold_envdb_202108_db_h.index
+│   ├── colabfold_envdb_202108_db.index
+│   ├── colabfold_envdb_202108_db.lookup
+│   ├── colabfold_envdb_202108_db_seq
+│   ├── colabfold_envdb_202108_db_seq.dbtype
+│   ├── colabfold_envdb_202108_db_seq_h
+│   ├── colabfold_envdb_202108_db_seq_h.dbtype
+│   ├── colabfold_envdb_202108_db_seq_h.index
+│   ├── colabfold_envdb_202108_db_seq.index
+│   └── colabfold_envdb_202108_db_seq.lookup
+├── colabfold_uniref30
+│   ├── uniref30_2302_db
+│   ├── uniref30_2302_db_aln
+│   ├── uniref30_2302_db_aln.dbtype
+│   ├── uniref30_2302_db_aln.index
+│   ├── uniref30_2302_db.dbtype
+│   ├── uniref30_2302_db.GPU_READY
+│   ├── uniref30_2302_db_h
+│   ├── uniref30_2302_db_h.dbtype
+│   ├── uniref30_2302_db_h.index
+│   ├── uniref30_2302_db.idx
+│   ├── uniref30_2302_db.idx.dbtype
+│   ├── uniref30_2302_db.idx.index
+│   ├── uniref30_2302_db.index
+│   ├── uniref30_2302_db.lookup
+│   ├── uniref30_2302_db_mapping
+│   ├── uniref30_2302_db_seq
+│   ├── uniref30_2302_db_seq.dbtype
+│   ├── uniref30_2302_db_seq_h
+│   ├── uniref30_2302_db_seq_h.dbtype
+│   ├── uniref30_2302_db_seq_h.index
+│   ├── uniref30_2302_db_seq.index
+│   └── uniref30_2302_db_taxonomy
+└── colabfold_uniref30_padded
+    ├── uniref30_2302_db
+    ├── uniref30_2302_db.dbtype
+    ├── uniref30_2302_db_h
+    ├── uniref30_2302_db_h.dbtype
+    ├── uniref30_2302_db_h.index
+    ├── uniref30_2302_db.index
+    ├── uniref30_2302_db.lookup
+    ├── uniref30_2302_db_seq
+    ├── uniref30_2302_db_seq.dbtype
+    ├── uniref30_2302_db_seq_h
+    ├── uniref30_2302_db_seq_h.dbtype
+    ├── uniref30_2302_db_seq_h.index
+    ├── uniref30_2302_db_seq.index
+    └── uniref30_2302_db_seq.lookup
+```
+
+## Running colabfold
+
+You will need to set the `--colabfold_enable_gpu_search true` flag. Below is an example command you can use to run with GPU search enabled:
+
+```bash
+nextflow run ./main.nf \
+    --input "samplesheet.csv" \
+    --outdir "output" \
+    --mode "colabfold" \
+    --use_gpu \
+    --db /path/to/db/root \
+    --use_msa_server false \
+    --colabfold_enable_gpu_search true \
+    --colabfold_model_preset alphafold2_ptm
+```

docs/usage.md

@@ -183,7 +183,7 @@ To provide the predownloaded AlphaFold3 databases and parameters you can specify
 
 </details>
 
-Colabfold mode can be used with local database search using the following command:
+Colabfold mode can be used with local database search. GPU MSA search is enabled by default. You may disable GPU MSA search with `--colabfold_enable_gpu_search false`:
 
 ```bash
 nextflow run nf-core/proteinfold \

main.nf

@@ -247,7 +247,10 @@ workflow NFCORE_PROTEINFOLD {
             params.use_msa_server,
             params.colabfold_alphafold2_params_path,
             params.colabfold_envdb_path,
+            params.colabfold_envdb_path_padded,
             params.colabfold_uniref30_path,
+            params.colabfold_uniref30_path_padded,
+            params.colabfold_enable_gpu_search,
             params.colabfold_alphafold2_params_link,
             params.colabfold_db_link,
             params.colabfold_uniref30_link,
@@ -258,16 +261,27 @@ workflow NFCORE_PROTEINFOLD {
         //
         // WORKFLOW: Run nf-core/colabfold workflow
         //
-        COLABFOLD (
-            ch_samplesheet,
-            ch_versions,
-            params.colabfold_model_preset,
-            PREPARE_COLABFOLD_DBS.out.params,
-            PREPARE_COLABFOLD_DBS.out.colabfold_db,
-            PREPARE_COLABFOLD_DBS.out.uniref30,
-            params.colabfold_num_recycles
-        )
-
+        if (params.colabfold_enable_gpu_search) {
+            COLABFOLD (
+                ch_samplesheet,
+                ch_versions,
+                params.colabfold_model_preset,
+                PREPARE_COLABFOLD_DBS.out.params,
+                PREPARE_COLABFOLD_DBS.out.colabfold_db_padded,
+                PREPARE_COLABFOLD_DBS.out.uniref30_padded,
+                params.colabfold_num_recycles
+            )
+        } else {
+            COLABFOLD (
+                ch_samplesheet,
+                ch_versions,
+                params.colabfold_model_preset,
+                PREPARE_COLABFOLD_DBS.out.params,
+                PREPARE_COLABFOLD_DBS.out.colabfold_db,
+                PREPARE_COLABFOLD_DBS.out.uniref30,
+                params.colabfold_num_recycles
+            )
+        }
         ch_multiqc          = ch_multiqc.mix(COLABFOLD.out.multiqc_report)
         ch_versions         = ch_versions.mix(COLABFOLD.out.versions)
         ch_report_input     = ch_report_input

modules/local/mmseqs_colabfoldsearch/Dockerfile

@@ -22,10 +22,10 @@ RUN apt-get update && \
     pip install --no-cache-dir --break-system-packages \
         "colabfold @ git+https://github.com/sokrypton/ColabFold.git@e8ebd9a" && \
     \
-    wget -q https://github.com/soedinglab/MMseqs2/releases/download/${MMSEQS2_VERSION}/mmseqs-linux-sse41.tar.gz && \
-    tar xzf mmseqs-linux-sse41.tar.gz && \
+    wget -q https://github.com/soedinglab/MMseqs2/releases/download/${MMSEQS2_VERSION}/mmseqs-linux-gpu.tar.gz && \
+    tar xzf mmseqs-linux-gpu.tar.gz && \
     cp mmseqs/bin/* /usr/local/bin/ && \
-    rm -rf mmseqs mmseqs-linux-sse41.tar.gz && \
+    rm -rf mmseqs mmseqs-linux-gpu.tar.gz && \
     \
     apt-get remove -y \
         python3-dev \

modules/local/mmseqs_colabfoldsearch/main.nf

@@ -3,12 +3,13 @@ process MMSEQS_COLABFOLDSEARCH {
     label 'process_high_memory'
     label 'process_high'
 
-    container "nf-core/proteinfold_mmseqs_colabfoldsearch:2.0.0"
+    container "docker.io/nbtmsh/mmseqs_colabfoldsearch:latest"
 
     input:
     tuple val(meta), path(fasta)
     path ('db/*')
     path ('db/*')
+    val colabfold_enable_gpu_search
 
     output:
     tuple val(meta), path("**.a3m"), emit: a3m
@@ -25,8 +26,13 @@ process MMSEQS_COLABFOLDSEARCH {
     def args = task.ext.args ?: ''
 
     """
+    GPU_ARG=""
+    if [ "${colabfold_enable_gpu_search}" == "1" ]; then
+        GPU_ARG="--gpu 1"
+    fi
     colabfold_search \\
         $args \\
+        \${GPU_ARG} \\
         --threads $task.cpus ${fasta} \\
         ./db \\
         --af3-json \\

nextflow.config

@@ -110,14 +110,17 @@ params {
     colabfold_db_load_mode      = 0
     colabfold_use_templates     = false
     colabfold_create_index      = false
+    colabfold_enable_gpu_search = false
 
     // Colabfold links
     colabfold_db_link       = null
     colabfold_uniref30_link = null
 
     // Colabfold paths
     colabfold_envdb_path       = null
+    colabfold_envdb_path_padded = null
     colabfold_uniref30_path = null
+    colabfold_uniref30_path_padded = null
 
     // Esmfold parameters
     esmfold_db           = null