Merge pull request #628 from broadinstitute/dp-bbnorm

Integrate bbnorm normalization and optimize resource allocation

Author: dpark01

pipes/WDL/tasks/tasks_assembly.wdl

@@ -15,7 +15,8 @@ task assemble {
       String   sample_name = basename(basename(reads_unmapped_bam, ".bam"), ".taxfilt")
 
       Int?     machine_mem_gb
-      String   docker = "quay.io/broadinstitute/viral-assemble:2.5.18.0"
+      Int?     cpu
+      String   docker = "quay.io/broadinstitute/viral-assemble:2.5.21.0"
     }
     parameter_meta{
       reads_unmapped_bam: {
@@ -101,9 +102,9 @@ task assemble {
 
     runtime {
         docker: docker
-        memory: select_first([machine_mem_gb, 63]) + " GB"
-        cpu: 4
-        disks:  "local-disk " + disk_size + " HDD"
+        memory: select_first([machine_mem_gb, 32]) + " GB"
+        cpu:    select_first([cpu, 8])
+        disks:  "local-disk " + disk_size + " SSD"
         disk: disk_size + " GB" # TES
         dx_instance_type: "mem1_ssd1_v2_x8"
         maxRetries: 2
@@ -124,7 +125,7 @@ task select_references {
     Int?          skani_c
     Int?          skani_n
 
-    String        docker = "quay.io/broadinstitute/viral-assemble:2.5.18.0"
+    String        docker = "quay.io/broadinstitute/viral-assemble:2.5.21.0"
     Int           machine_mem_gb = 4
     Int           cpu = 2
     Int           disk_size = 100
@@ -193,7 +194,7 @@ task select_references {
     docker: docker
     memory: machine_mem_gb + " GB"
     cpu:    cpu
-    disks:  "local-disk " + disk_size + " HDD"
+    disks:  "local-disk " + disk_size + " SSD"
     disk: disk_size + " GB" # TESs
     dx_instance_type: "mem1_ssd1_v2_x2"
     preemptible: 2
@@ -204,7 +205,7 @@ task select_references {
 task scaffold {
     input {
       File         contigs_fasta
-      File         reads_bam
+      File?        reads_bam
       Array[File]+ reference_genome_fasta
 
       String       aligner="muscle"
@@ -224,16 +225,22 @@ task scaffold {
       Float?       scaffold_min_pct_contig_aligned
 
       Int?         machine_mem_gb
-      String       docker="quay.io/broadinstitute/viral-assemble:2.5.18.0"
+      String       docker="quay.io/broadinstitute/viral-assemble:2.5.21.0"
 
       # do this in multiple steps in case the input doesn't actually have "assembly1-x" in the name
       String       sample_name = basename(basename(contigs_fasta, ".fasta"), ".assembly1-spades")
     }
+
+    # Determine whether to run Gap2Seq based on reads_bam size
+    # Gap2Seq can take 100+ min for large BAMs (>1GB), providing diminishing returns
+    Float reads_bam_size_gb = if defined(reads_bam) then size(select_first([reads_bam]), "GB") else 0.0
+    Boolean run_gap2seq = defined(reads_bam) && reads_bam_size_gb < 1.0
+
     parameter_meta {
       reads_bam: {
-        description: "Reads in BAM format.",
+        description: "Reads in BAM format. If provided, Gap2Seq will attempt to fill gaps using reads. Skipping this for large BAMs (>1GB) can save significant runtime.",
         patterns: ["*.bam"],
-        category: "required"
+        category: "optional"
       }
 
       contigs_fasta: {
@@ -367,13 +374,19 @@ task scaffold {
         fi
         grep '^>' "~{sample_name}".scaffolding_chosen_ref.fasta | cut -c 2- | cut -f 1 -d ' ' > "~{sample_name}".scaffolding_chosen_refs.txt
 
-        assembly.py gapfill_gap2seq \
-          "~{sample_name}".intermediate_scaffold.fasta \
-          "~{reads_bam}" \
-          "~{sample_name}".intermediate_gapfill.fasta \
-          --memLimitGb $mem_in_gb \
-          --maskErrors \
-          --loglevel=DEBUG
+        # Run Gap2Seq only if reads_bam is provided and smaller than 1GB
+        if ~{true='true' false='false' run_gap2seq}; then
+            assembly.py gapfill_gap2seq \
+              "~{sample_name}".intermediate_scaffold.fasta \
+              "~{reads_bam}" \
+              "~{sample_name}".intermediate_gapfill.fasta \
+              --memLimitGb $mem_in_gb \
+              --maskErrors \
+              --loglevel=DEBUG
+        else
+            echo "Skipping Gap2Seq: reads_bam not provided or >= 1GB (~{reads_bam_size_gb} GB)" >&2
+            cp "~{sample_name}".intermediate_scaffold.fasta "~{sample_name}".intermediate_gapfill.fasta
+        fi
 
         set +e +o pipefail
         grep -v '^>' "~{sample_name}".intermediate_gapfill.fasta | tr -d '\n' | wc -c | tee assembly_preimpute_length
@@ -435,9 +448,9 @@ task scaffold {
 
     runtime {
         docker: docker
-        memory: select_first([machine_mem_gb, 63]) + " GB"
+        memory: select_first([machine_mem_gb, 20]) + " GB"
         cpu: 4
-        disks:  "local-disk " + disk_size + " HDD"
+        disks:  "local-disk " + disk_size + " SSD"
         disk: disk_size + " GB" # TES
         dx_instance_type: "mem1_ssd1_v2_x8"
         maxRetries: 2
@@ -457,7 +470,7 @@ task skani_triangle {
     Int     compression_factor = 10
     Int     min_aligned_frac = 15
 
-    String  docker = "quay.io/broadinstitute/viral-assemble:2.5.18.0"
+    String  docker = "quay.io/broadinstitute/viral-assemble:2.5.21.0"
     Int     machine_mem_gb = 8
     Int     cpu = 4
     Int     disk_size = 100
@@ -696,7 +709,7 @@ task align_reads {
 
     Int?     cpu
     Int?     machine_mem_gb
-    String   docker = "quay.io/broadinstitute/viral-core:2.5.20"
+    String   docker = "quay.io/broadinstitute/viral-core:2.5.21"
 
     String   sample_name = basename(basename(basename(reads_unmapped_bam, ".bam"), ".taxfilt"), ".clean")
   }
@@ -707,8 +720,8 @@ task align_reads {
   # Linear scaling: 8 + (input_GB / 15) * 56, capped at 64, rounded to nearest multiple of 4
   Float        cpu_unclamped = 8.0 + (size(reads_unmapped_bam, "GB") / 15.0) * 56.0
   Int          cpu_actual = select_first([cpu, floor(((if cpu_unclamped > 64.0 then 64.0 else cpu_unclamped) + 2.0) / 4.0) * 4])
-  # Memory scales with CPU at 2x ratio (default), or use override
-  Int          machine_mem_gb_actual = select_first([machine_mem_gb, cpu_actual * 2])
+  # Memory scales with CPU at 3x ratio (default), or use override
+  Int          machine_mem_gb_actual = select_first([machine_mem_gb, cpu_actual * 3])
 
   parameter_meta {
     reference_fasta: {
@@ -811,7 +824,7 @@ task align_reads {
     docker: docker
     memory: machine_mem_gb_actual + " GB"
     cpu: cpu_actual
-    disks:  "local-disk " + disk_size + " SSD"
+    disks:  "local-disk " + disk_size + " LOCAL"
     disk: disk_size + " GB" # TES
     dx_instance_type: "mem1_ssd1_v2_x8"
     preemptible: 1
@@ -834,8 +847,8 @@ task refine_assembly_with_aligned_reads {
       Float    major_cutoff = 0.5
       Int      min_coverage = 3
 
-      Int      machine_mem_gb = 15
-      String   docker = "quay.io/broadinstitute/viral-assemble:2.5.18.0"
+      Int      machine_mem_gb = 8
+      String   docker = "quay.io/broadinstitute/viral-assemble:2.5.21.0"
     }
 
     Int disk_size = 375
@@ -972,7 +985,7 @@ task run_discordance {
       String out_basename = "run"
       Int    min_coverage = 4
 
-      String docker = "quay.io/broadinstitute/viral-core:2.5.20"
+      String docker = "quay.io/broadinstitute/viral-core:2.5.21"
     }
     parameter_meta {
       reads_aligned_bam: {
@@ -1221,7 +1234,7 @@ task wgsim {
         Int?   random_seed
 
         Int    machine_mem_gb = 7
-        String docker = "quay.io/broadinstitute/viral-assemble:2.5.18.0"
+        String docker = "quay.io/broadinstitute/viral-assemble:2.5.21.0"
     }
 
     parameter_meta {

pipes/WDL/tasks/tasks_demux.wdl

@@ -6,7 +6,7 @@ task merge_tarballs {
     String       out_filename
 
     Int?         machine_mem_gb
-    String       docker = "quay.io/broadinstitute/viral-core:2.5.20"
+    String       docker = "quay.io/broadinstitute/viral-core:2.5.21"
   }
 
   Int disk_size = 2625
@@ -181,7 +181,7 @@ task illumina_demux {
     # --- options for VM shape ----------------------
     Int?    machine_mem_gb
     Int     disk_size = 2625
-    String  docker    = "quay.io/broadinstitute/viral-core:2.5.20"
+    String  docker    = "quay.io/broadinstitute/viral-core:2.5.21"
   }
 
   parameter_meta {
@@ -823,7 +823,7 @@ task get_illumina_run_metadata {
     String? sequencing_center
 
     Int?   machine_mem_gb
-    String docker = "quay.io/broadinstitute/viral-core:2.5.20"
+    String docker = "quay.io/broadinstitute/viral-core:2.5.21"
   }
 
   parameter_meta {
@@ -927,7 +927,7 @@ task demux_fastqs {
     Int?    machine_mem_gb
     Int     max_cpu = 32       # Maximum CPU cap for autoscaling (use 16 for 2-barcode, 64 for 3-barcode)
     Int     disk_size = 750
-    String  docker = "quay.io/broadinstitute/viral-core:2.5.20"
+    String  docker = "quay.io/broadinstitute/viral-core:2.5.21"
   }
 
   # Calculate total input size for autoscaling
@@ -1056,7 +1056,7 @@ task merge_demux_metrics {
   input {
     Array[File]+ metrics_files
     String       output_filename = "merged_demux_metrics.txt"
-    String       docker = "quay.io/broadinstitute/viral-core:2.5.20"
+    String       docker = "quay.io/broadinstitute/viral-core:2.5.21"
   }
 
   parameter_meta {

pipes/WDL/tasks/tasks_interhost.wdl

@@ -351,7 +351,7 @@ task index_ref {
     File?  novocraft_license
 
     Int?   machine_mem_gb
-    String docker = "quay.io/broadinstitute/viral-core:2.5.20"
+    String docker = "quay.io/broadinstitute/viral-core:2.5.21"
   }
 
   Int disk_size = 100

pipes/WDL/tasks/tasks_intrahost.wdl

@@ -178,7 +178,7 @@ task lofreq {
     docker: docker
     cpu:    2
     memory: "3 GB"
-    disks:  "local-disk " + disk_size + " HDD"
+    disks:  "local-disk " + disk_size + " SSD"
     disk: disk_size + " GB" # TES
     dx_instance_type: "mem1_ssd1_v2_x2"
     maxRetries: 2

pipes/WDL/tasks/tasks_megablast.wdl

@@ -15,7 +15,7 @@ task trim_rmdup_subsamp {
         Int cpu            = 16
         Int disk_size_gb   = 100 
 
-        String docker      = "quay.io/broadinstitute/viral-assemble:2.5.18.0"
+        String docker      = "quay.io/broadinstitute/viral-core:2.5.21"
     }
 
     parameter_meta {
@@ -36,17 +36,14 @@ task trim_rmdup_subsamp {
     command <<<
         set -ex o pipefail
         assembly.py --version | tee VERSION
-        #BAM ->FASTQ-> OutBam? https://github.com/broadinstitute/viral-assemble:2.5.18.0
         assembly.py trim_rmdup_subsamp \
-        "~{inBam}" \
-        "~{clipDb}" \
-        "$(pwd)/outbam.bam" \
-        ~{'--n_reads=' + n_reads}
+            "~{inBam}" \
+            "~{clipDb}" \
+            outbam.bam \
+            ~{'--n_reads=' + n_reads}
 
-
-        #samtools [OutBam -> FASTA]
         #-f 4 (f = include only) (4 = unmapped reads) https://broadinstitute.github.io/picard/explain-flags.html
-        samtools fasta "$(pwd)/outbam.bam" > "~{bam_basename}.fasta"
+        samtools fasta outbam.bam > "~{bam_basename}.fasta"
     >>>
 
     output {
@@ -58,7 +55,6 @@ task trim_rmdup_subsamp {
         memory: machine_mem_gb + "GB"
         cpu:    cpu
         disks:  "local-disk " + disk_size_gb + " LOCAL"
-
         dx_instance_type: "n2-highmem-4"
     }
 }

pipes/WDL/tasks/tasks_metagenomics.wdl

@@ -11,7 +11,7 @@ task krakenuniq {
     File        krona_taxonomy_db_tgz  # taxonomy.tab
 
     Int         machine_mem_gb = 320
-    String      docker = "quay.io/broadinstitute/viral-classify:2.2.4.0" #skip-global-version-pin
+    String      docker = "quay.io/broadinstitute/viral-classify:2.1.33.0" #skip-global-version-pin
   }
 
   Int disk_size = 750
@@ -148,7 +148,7 @@ task build_krakenuniq_db {
     Int?     zstd_compression_level
 
     Int      machine_mem_gb = 240
-    String   docker = "quay.io/broadinstitute/viral-classify:2.2.4.0" #skip-global-version-pin
+    String   docker = "quay.io/broadinstitute/viral-classify:2.1.33.0" #skip-global-version-pin
   }
 
   Int disk_size = 750
@@ -243,7 +243,12 @@ task kraken2 {
   }
 
   String out_basename = basename(basename(reads_bam, '.bam'), '.fasta')
-  Int disk_size = 750
+
+  # Disk autoscaling: BAM->FASTQ expansion is ~7-8x, plus kraken2 reads output (~1x input),
+  # plus kraken2 database (1x localized tarball + 2x decompressed = 3x), plus overhead for krona and temp files.
+  # Minimum 375GB to accommodate typical database sizes.
+  Int disk_size_auto = ceil((8 * size(reads_bam, "GB") + 3 * size(kraken2_db_tgz, "GB") + 50) / 375.0) * 375
+  Int disk_size = if disk_size_auto < 375 then 375 else disk_size_auto
 
   command <<<
     set -ex -o pipefail
@@ -329,7 +334,7 @@ task kraken2 {
     memory: machine_mem_gb + " GB"
     cpu: 16
     cpuPlatform: "Intel Ice Lake"
-    disks:  "local-disk " + disk_size + " HDD"
+    disks:  "local-disk " + disk_size + " LOCAL"
     disk: disk_size + " GB" # TESs
     dx_instance_type: "mem3_ssd1_v2_x8"
     preemptible: 2
@@ -378,7 +383,7 @@ task report_primary_kraken_taxa {
     docker: docker
     memory: machine_mem_gb + " GB"
     cpu: 1
-    disks:  "local-disk " + disk_size + " LOCAL"
+    disks:  "local-disk " + disk_size + " HDD"
     disk: disk_size + " GB" # TESs
     dx_instance_type: "mem1_ssd1_v2_x2"
     preemptible: 2

pipes/WDL/tasks/tasks_ncbi.wdl

@@ -183,7 +183,7 @@ task sequencing_platform_from_bam {
   input {
     File    bam
 
-    String  docker = "quay.io/broadinstitute/viral-core:2.5.20"
+    String  docker = "quay.io/broadinstitute/viral-core:2.5.21"
   }
 
   command <<<
@@ -293,7 +293,7 @@ task structured_comments {
 
     File?  filter_to_ids
 
-    String docker = "quay.io/broadinstitute/viral-core:2.5.20"
+    String docker = "quay.io/broadinstitute/viral-core:2.5.21"
   }
   String out_base = basename(assembly_stats_tsv, '.txt')
   command <<<
@@ -346,7 +346,7 @@ task structured_comments_from_aligned_bam {
     String  out_basename = basename(aligned_bam, '.bam')
     Boolean is_genome_assembly = true
     Boolean sanitize_ids = true
-    String  docker = "quay.io/broadinstitute/viral-core:2.5.20"
+    String  docker = "quay.io/broadinstitute/viral-core:2.5.21"
   }
   # see https://www.ncbi.nlm.nih.gov/genbank/structuredcomment/
   command <<<
@@ -465,7 +465,7 @@ task rename_fasta_header {
 
     String out_basename = basename(genome_fasta, ".fasta")
 
-    String docker = "quay.io/broadinstitute/viral-core:2.5.20"
+    String docker = "quay.io/broadinstitute/viral-core:2.5.21"
   }
   command <<<
     set -e
@@ -630,7 +630,7 @@ task sra_meta_prep {
     Boolean     paired
 
     String      out_name = "sra_metadata.tsv"
-    String      docker="quay.io/broadinstitute/viral-core:2.5.20"
+    String      docker="quay.io/broadinstitute/viral-core:2.5.21"
   }
   Int disk_size = 100
   parameter_meta {

pipes/WDL/tasks/tasks_nextstrain.wdl

@@ -332,7 +332,7 @@ task derived_cols {
         String?       lab_highlight_loc
         Array[File]   table_map = []
 
-        String        docker = "quay.io/broadinstitute/viral-core:2.5.20"
+        String        docker = "quay.io/broadinstitute/viral-core:2.5.21"
         Int           disk_size = 50
     }
     parameter_meta {
@@ -900,7 +900,7 @@ task filter_sequences_to_list {
 
         String       out_fname = sub(sub(basename(sequences, ".zst"), ".vcf", ".filtered.vcf"), ".fasta$", ".filtered.fasta")
         # Prior docker image: "nextstrain/base:build-20240318T173028Z"
-        String       docker = "quay.io/broadinstitute/viral-core:2.5.20"
+        String       docker = "quay.io/broadinstitute/viral-core:2.5.21"
         Int          disk_size = 750
     }
     parameter_meta {