Starting to implement GPU support for polishing

- add test cases
- add gpu parameters
- add functionality to test for gpu availability

config/base.config

@@ -3,9 +3,13 @@
 
 // PROCESS RESOURCES
 process {
+
 	withName: "POLISH_CLUSTER" {
 		memory = { 2.GB * task.attempt }
 		cpus = 1
+		if(params.use_gpu){
+			container = "ontresearch/medaka:latest"
+		}
 	}
 
 	errorStrategy =  'retry'

lib/processes/polish_cluster.nf

@@ -2,6 +2,7 @@ cpus=2
 process POLISH_CLUSTER {
     cpus "${cpus}"
     tag "${sample}"
+    label "gpu_possible"
 
     input:
         tuple val( sample ), val( target ), path( smolecule_clusters_fastq )
@@ -13,9 +14,6 @@ process POLISH_CLUSTER {
     script:
     """
         medaka smolecule \
-          --threads $cpus \
-          --length 50 \
-          --depth 2 \
           --model ${params.medaka_model} \
           --method spoa . \
           --qualities \

lib/workflows/umi-pipeline.nf

@@ -4,6 +4,24 @@ requiredParams = [
     'input', 'reference', 'reference_fai', 'bed', 'output'
 ]
 
+if(params.use_gpu){
+    try {
+        // Execute the nvidia-smi command
+        def process = "nvidia-smi".execute()
+        def output = new StringBuffer()
+        def error = new StringBuffer()
+
+        // Capture the output and error streams
+        process.consumeProcessOutput(output, error)
+        process.waitFor()
+
+        println process.exitValue() == 0 && output.toString().contains("NVIDIA-SMI")
+    } catch (Exception e) {
+        println false
+    }
+
+}
+
 for (param in requiredParams) {
     if (params[param] == null) {
       exit 1, "Parameter ${param} is required."

main.nf

@@ -47,6 +47,7 @@ if(params.help){
               --write_reports                 Write stats of cluster and cluster filtering [default: true]
               --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]
               --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"]

nextflow.config

@@ -51,6 +51,8 @@ params {
     min_overlap                           = 0.95
     include_secondary_reads               = false
     balance_strands                       = true
+    use_gpu                               = false
+    gpus                                  = "all"
     medaka_model                          = "r1041_e82_400bps_hac_g615"
     fwd_umi                               = "TTTVVVVTTVVVVTTVVVVTTVVVVTTT"
     rev_umi                               = "AAABBBBAABBBBAABBBBAABBBBAAA"
@@ -67,7 +69,7 @@ params {
 // Load base.config by default for all pipelines
 includeConfig "config/base.config"
 
-process.container                         = 'quay.io/genepi/umi-pipeline-nf:v0.2.1'
+process.container = 'quay.io/genepi/umi-pipeline-nf:v0.2.1'
 
 profiles {
 
@@ -94,3 +96,20 @@ profiles {
     }
 
 }
+
+/*
+if(params.use_gpu){
+    println("use GPU")
+    process {
+        container                         = 'quay.io/genepi/umi-pipeline-nf:v0.2.1'
+        withLabel: gpu_possible {
+            container = 'ontresearch/medaka:latest'
+            // Set environment variables
+            // docker.envWhitelist = "TF_FORCE_GPU_ALLOW_GROWTH = true"
+            docker.runOptions = "--rm --gpus ${params.gpus}"
+            docker.runOptions = "--rm --gpus all"
+            println("All set")
+        }
+    }
+}
+*/

tests/lib/processes/polish_cluster.nf.test

@@ -16,7 +16,7 @@ nextflow_process {
                         ["barcode06", "raw", "$baseDir/tests/input/cluster/barcode06/detected_umis.fastq"],
                         ["barcode07", "raw", "$baseDir/tests/input/cluster/barcode07/detected_umis.fastq"]
                     )
-                    input[1] = Channel.of("raw", "raw", "raw", "raw", "raw")
+                    input[1] = "raw"
                     """
                 }
             }
@@ -28,7 +28,7 @@ nextflow_process {
                 process {
                     """
                     input[0] = CLUSTER.out.cluster_fastas
-                    input[1] = Channel.of("raw", "raw", "raw", "raw", "raw")
+                    input[1] = "raw"
                     input[2] = file("$baseDir/bin/parse_clusters.py")
                     """
                 }

tests/lib/processes/polish_cluster_minimal.nf.test

@@ -0,0 +1,27 @@
+nextflow_process {
+
+    name "Test Process POLISH_CLUSTER_minimal"
+    script "lib/processes/polish_cluster.nf"
+    process "POLISH_CLUSTER"
+
+    test("Should use gpu") {
+        when {
+            params{
+                use_gpu = true
+            }
+            process {
+                """
+                input[0] = Channel.of(["barcode03", "raw", "$baseDir/tests/input/polish_cluster/barcode03/smolecule0sub12.fastq"])
+                input[1] = "Consensus"
+                """
+            }
+        }
+
+        then {
+            assert process.success
+            assert snapshot(process.out).match()
+        }
+
+    }
+
+}

tests/lib/processes/polish_cluster_minimal.nf.test.snap

@@ -0,0 +1,27 @@
+{
+    "Should use gpu": {
+        "content": [
+            {
+                "0": [
+                    [
+                        "barcode03",
+                        "raw",
+                        "smolecule0sub12_consensus.fastq:md5,458ecdaf2e49929a67c52389a01b5f7c"
+                    ]
+                ],
+                "consensus_fastq": [
+                    [
+                        "barcode03",
+                        "raw",
+                        "smolecule0sub12_consensus.fastq:md5,458ecdaf2e49929a67c52389a01b5f7c"
+                    ]
+                ]
+            }
+        ],
+        "meta": {
+            "nf-test": "0.9.0-rc2",
+            "nextflow": "24.05.0"
+        },
+        "timestamp": "2024-06-21T11:14:11.515754"
+    }
+}

tests/main.nf.test

@@ -27,4 +27,30 @@ nextflow_pipeline {
         }
 
     }
+
+    test("Should execute umi-pipeline-nf with gpu") {
+
+        when {
+
+            params {
+                input 			= "$baseDir/tests/input/pipeline/fastq_pass/"
+                output 			= "test_umi-pipeline-nf"
+                reference 		= "$baseDir/tests/input/pipeline/ref/lpa-ref2645.fasta"
+                reference_fai 	= "$baseDir/tests/input/pipeline/ref/lpa-ref2645.fasta.fai"
+                bed 			= "$baseDir/tests/input/pipeline/ref/lpa-ref2645.bed"
+
+                min_reads_per_cluster 	= 10
+                max_reads_per_cluster 	= 10
+                min_reads_per_barcode = 0
+                call_variants 		= true
+                variant_caller 		= "freebayes"
+                use_gpu             = true
+            }
+        }
+
+        then {
+            assert workflow.success
+        }
+
+    }
 }

tests/nextflow.config

@@ -3,10 +3,9 @@
     Nextflow config file for running tests
 ========================================================================================
 */
-
-process {
+process{
     withName: '.*' {
         cpus = 1
         memory = 2.GB
     }
-}
+}