Implement GPU support for medaka with docker

genepi · Sep 13, 2024 · bb5bf98 · bb5bf98
1 parent c13cf3f
commit bb5bf98
Show file tree

Hide file tree

Showing 13 changed files with 1,283 additions and 33 deletions.
diff --git a/README.md b/README.md
@@ -33,6 +33,8 @@ Umi-pipeline-nf orignates from a snakemake-based analysis pipeline ([pipeline-um
 * **Three commonly used variant callers** ([freebayes](https://github.com/freebayes/freebayes), [lofreq](http://csb5.github.io/lofreq/) or [mutserve](https://mitoverse.readthedocs.io/mutserve/mutserve/)) are supported by the pipeline.
 * The raw reads can be optionally **subsampled**.
 * The raw reads can be **filtered by read length and quality**.
+* **GPU acceleration for cluster polishing by Medaka** is available when using the `docker` profile. The GPU driver, [nvidia-toolkit](https://developer.nvidia.com/cuda-toolkit), and [nvidia-container-toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) must be installed. Tested with an RTX 4080 SUPER GPU (16 GB). Note: GPU acceleration is not compatible with cluster profiles.
+
 
 > See the [usage documentation](docs/usage.md) for all of the available parameters of the pipeline.
 

diff --git a/config/base.config b/config/base.config
@@ -8,7 +8,9 @@ process {
 		memory = { 2.GB * task.attempt }
 		cpus = 1
 		if(params.use_gpu){
+			maxForks = params.gpu_memory
 			container = "ontresearch/medaka:latest"
+			containerOptions = 	"--rm --gpus all -e TF_FORCE_GPU_ALLOW_GROWTH=true"
 		}
 	}
 

diff --git a/docs/usage.md b/docs/usage.md
@@ -29,6 +29,8 @@ This document describes the parameter options used by the pipeline.
     * [`--min_overlap`](#--min_overlap-arg)
     * [`--include_secondary_reads`](#--include_secdondary_reads-arg)
     * [`--balance_strands`](#--balance_strands-arg)
+    * [`--use_gpu`](#--use_gpu-arg)
+    * [`--gpu_memory`](#--gpu_memory-arg)
     * [`--medaka_model`](#--medaka_model-arg)
     * [`--fwd_umi`](#--fwd_umi-arg)
     * [`--rev_umi`](#--rev_umi-arg)
@@ -156,6 +158,14 @@ Specify if secondary mappings should be included in the analysis. [default: fals
 ### `--balance_strands <ARG>`
 Specify if the number of forward and reverse reads per cluster should be equalized. [default: true]
 
+### `--use_gpu`
+Specify if GPU acceleration for Medaka should be enabled. [default: false]  
+Note: GPU acceleration is only available with the `docker` profile and has been tested with an RTX 4080 SUPER GPU (16 GB).
+
+### `--gpu_driver`
+Specify the path to the GPU driver.  
+The NVIDIA toolkit (https://developer.nvidia.com/cuda-toolkit) and NVIDIA container toolkit (https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) must be installed for GPU usage.
+
 ### `--medaka_model <ARG>`
 Specify the medaka model that is used for cluster polishing. [default: "r1041_e82_400bps_hac_g615"]  
 Note: The models are specific for Chemistry, basecalling algorithm and sequencing speed. 

diff --git a/main.nf b/main.nf
@@ -48,6 +48,7 @@ if(params.help){
               --min_overlap                   Min overlap with target region [default: 0.90]
               --balance_strands               Balance forward and reverse raw reads in clusters [default: true]
               --use_gpu                       Use GPU acceleration for cluster polishing with medaka [default: false]
+              --gpu_memory                    Available GPU memory in GB [default: 16]
               --medaka_model                  Medaka model used to compute consensus reads [default: "r1041_e82_400bps_hac_g615"]
               --fwd_umi                       Forward UMI (Ftail...UMI...primer) [default: "TTTVVVVTTVVVVTTVVVVTTVVVVTTT"]
               --rev_umi                       Reverse UMI (Rtail...UMI...primer) [default: "AAABBBBAABBBBAABBBBAABBBBAAA"]

diff --git a/nextflow.config b/nextflow.config
@@ -52,6 +52,7 @@ params {
     include_secondary_reads               = false
     balance_strands                       = true
     use_gpu                               = false
+    gpu_memory                            = 16
     gpus                                  = "all"
     medaka_model                          = "r1041_e82_400bps_hac_g615"
     fwd_umi                               = "TTTVVVVTTVVVVTTVVVVTTVVVVTTT"

diff --git a/tests/input/pipeline/ref/lpa-ref2645.bed b/tests/input/pipeline/ref/lpa-ref2645.bed
diff --git a/tests/input/pipeline/ref/lpa-ref2645.fasta b/tests/input/pipeline/ref/lpa-ref2645.fasta
diff --git a/tests/lib/processes/freebayes.nf.test.snap b/tests/lib/processes/freebayes.nf.test.snap
@@ -3,38 +3,38 @@
         "content": [
             {
                 "0": [
-                    "final.vcf:md5,076ec150e130dd090f605e2118bd6bc1",
-                    "final.vcf:md5,af49c1b8586ffc70b93798b59ca9a336",
-                    "final.vcf:md5,c58124f72006a34ca175b2dad280ea4a"
+                    "final.vcf:md5,1baff5303e86173c69757556e49c3797",
+                    "final.vcf:md5,27860e0e8e90ce3ba861c61149bca350",
+                    "final.vcf:md5,8fd6d0cf7573385cf8fda83270b9c6a6"
                 ],
                 "variants": [
-                    "final.vcf:md5,076ec150e130dd090f605e2118bd6bc1",
-                    "final.vcf:md5,af49c1b8586ffc70b93798b59ca9a336",
-                    "final.vcf:md5,c58124f72006a34ca175b2dad280ea4a"
+                    "final.vcf:md5,1baff5303e86173c69757556e49c3797",
+                    "final.vcf:md5,27860e0e8e90ce3ba861c61149bca350",
+                    "final.vcf:md5,8fd6d0cf7573385cf8fda83270b9c6a6"
                 ]
             }
         ],
         "meta": {
-            "nf-test": "0.9.0-rc2",
-            "nextflow": "24.05.0"
+            "nf-test": "0.9.0",
+            "nextflow": "22.10.3"
         },
-        "timestamp": "2024-06-20T09:55:07.96638"
+        "timestamp": "2024-09-13T15:43:06.822646"
     },
     "Should run without failures": {
         "content": [
             {
                 "0": [
-                    "final.vcf:md5,076ec150e130dd090f605e2118bd6bc1"
+                    "final.vcf:md5,27860e0e8e90ce3ba861c61149bca350"
                 ],
                 "variants": [
-                    "final.vcf:md5,076ec150e130dd090f605e2118bd6bc1"
+                    "final.vcf:md5,27860e0e8e90ce3ba861c61149bca350"
                 ]
             }
         ],
         "meta": {
-            "nf-test": "0.9.0-rc2",
-            "nextflow": "24.05.0"
+            "nf-test": "0.9.0",
+            "nextflow": "22.10.3"
         },
-        "timestamp": "2024-06-20T09:54:51.891187"
+        "timestamp": "2024-09-13T15:43:02.164783"
     }
 }
diff --git a/tests/lib/processes/lofreq.nf.test.snap b/tests/lib/processes/lofreq.nf.test.snap
@@ -3,38 +3,38 @@
         "content": [
             {
                 "0": [
-                    "final.vcf:md5,8761e5c55f8e722dd33472968f33ecd6",
-                    "final.vcf:md5,b48b1996596787daa8a83dc973120006",
-                    "final.vcf:md5,de4ad8ac00fddc259406af289d16961b"
+                    "final.vcf:md5,22906ceda834065e1d62878ed521f18c",
+                    "final.vcf:md5,9d0d55c36e3f6f72b65371b3ea534260",
+                    "final.vcf:md5,9da201124f97ba4ffcea42720219ce88"
                 ],
                 "variants": [
-                    "final.vcf:md5,8761e5c55f8e722dd33472968f33ecd6",
-                    "final.vcf:md5,b48b1996596787daa8a83dc973120006",
-                    "final.vcf:md5,de4ad8ac00fddc259406af289d16961b"
+                    "final.vcf:md5,22906ceda834065e1d62878ed521f18c",
+                    "final.vcf:md5,9d0d55c36e3f6f72b65371b3ea534260",
+                    "final.vcf:md5,9da201124f97ba4ffcea42720219ce88"
                 ]
             }
         ],
         "meta": {
-            "nf-test": "0.9.0-rc2",
-            "nextflow": "24.05.0"
+            "nf-test": "0.9.0",
+            "nextflow": "22.10.3"
         },
-        "timestamp": "2024-06-20T09:55:30.290847"
+        "timestamp": "2024-09-13T15:43:14.912982"
     },
     "Should run without failures": {
         "content": [
             {
                 "0": [
-                    "final.vcf:md5,b48b1996596787daa8a83dc973120006"
+                    "final.vcf:md5,9d0d55c36e3f6f72b65371b3ea534260"
                 ],
                 "variants": [
-                    "final.vcf:md5,b48b1996596787daa8a83dc973120006"
+                    "final.vcf:md5,9d0d55c36e3f6f72b65371b3ea534260"
                 ]
             }
         ],
         "meta": {
-            "nf-test": "0.9.0-rc2",
-            "nextflow": "24.05.0"
+            "nf-test": "0.9.0",
+            "nextflow": "22.10.3"
         },
-        "timestamp": "2024-06-20T09:55:18.706674"
+        "timestamp": "2024-09-13T15:43:10.490675"
     }
 }
diff --git a/tests/lib/processes/polish_cluster.nf.test b/tests/lib/processes/polish_cluster.nf.test
@@ -53,4 +53,23 @@ nextflow_process {
 
     }
 
+    test("Should use GPU") {
+        when {            
+            params{
+                use_gpu = true
+            }
+            process {
+                """
+                input[0] = REFORMAT_FILTER_CLUSTER.out.smolecule_cluster_fastqs.transpose()
+                input[1] = "Consensus"
+                """
+            }
+        }
+
+        then {
+            assert process.success
+            assert snapshot(process.out).match()
+        }
+    }
+
 }