Add hint for parallelism in registry files.

Author: tomwhite

eggo/dag.py

@@ -309,45 +309,10 @@ class ADAMPartitionTask(Task):
 
     adam_command = Parameter()
     allowed_file_formats = Parameter()
+    source_edition = Parameter()
+    edition = Parameter()
     partition_strategy_file = Parameter()
-    source_edition = 'basic'
-    edition = 'locuspart'
-
-    def requires(self):
-        return ADAMBasicTask(adam_command=self.adam_command,
-                             allowed_file_formats=self.allowed_file_formats)
-
-    def run(self):
-        adam_cmd = ('{hadoop_home}/bin/hadoop jar {adam_partitioning_jar}'
-                    ' CrunchPartitionTool -D mapreduce.job.reduces={parallelism}'
-                    ' {partition_strategy_file} {source} {target}').format(
-            hadoop_home=os.environ['HADOOP_HOME'],
-            adam_partitioning_jar=os.environ['ADAM_PARTITIONING_JAR'],
-            parallelism=1,
-            partition_strategy_file=self.partition_strategy_file,
-            source=target_s3n_url(ToastConfig().config['name'],
-                                  edition=self.source_edition),
-            target=target_s3n_url(ToastConfig().config['name'],
-                                  edition=self.edition))
-        p = Popen(adam_cmd, shell=True)
-        p.wait()
-
-        if p.returncode == 0:
-            create_SUCCESS_file(target_s3_url(ToastConfig().config['name'],
-                                              edition=self.edition))
-
-    def output(self):
-        return S3FlagTarget(target_s3_url(ToastConfig().config['name'],
-                                          edition=self.edition))
-
-
-class ADAMFlattenPartitionTask(Task):
-
-    adam_command = Parameter()
-    allowed_file_formats = Parameter()
-    partition_strategy_file = Parameter()
-    source_edition = 'flat'
-    edition = 'flat_locuspart'
+    parallelism = Parameter()
 
     def requires(self):
         return ADAMBasicTask(adam_command=self.adam_command,
@@ -359,7 +324,7 @@ def run(self):
                     ' {partition_strategy_file} {source} {target}').format(
             hadoop_home=os.environ['HADOOP_HOME'],
             adam_partitioning_jar=os.environ['ADAM_PARTITIONING_JAR'],
-            parallelism=1,
+            parallelism=self.parallelism,
             partition_strategy_file=self.partition_strategy_file,
             source=target_s3n_url(ToastConfig().config['name'],
                                   edition=self.source_edition),
@@ -376,20 +341,27 @@ def output(self):
         return S3FlagTarget(target_s3_url(ToastConfig().config['name'],
                                           edition=self.edition))
 
-
 class VCF2ADAMTask(Task):
 
     def requires(self):
+        conf = ToastConfig().config
+        parallelism = conf['numPartitionsHint'] if 'numPartitionsHint' in conf else 1
         basic = ADAMBasicTask(adam_command='vcf2adam',
                               allowed_file_formats=['vcf'])
         flat = ADAMFlattenTask(adam_command='vcf2adam',
                                allowed_file_formats=['vcf'])
         locuspart = ADAMPartitionTask(adam_command='vcf2adam',
                                       allowed_file_formats=['vcf'],
-                                      partition_strategy_file='genotypes-partition-strategy')
-        flat_locuspart = ADAMFlattenPartitionTask(adam_command='vcf2adam',
-                                                  allowed_file_formats=['vcf'],
-                                                  partition_strategy_file='flat-genotypes-partition-strategy')
+                                      source_edition='basic',
+                                      edition='locuspart',
+                                      partition_strategy_file='genotypes-partition-strategy',
+                                      parallelism=parallelism)
+        flat_locuspart = ADAMPartitionTask(adam_command='vcf2adam',
+                                           allowed_file_formats=['vcf'],
+                                           source_edition='flat',
+                                           edition='flat_locuspart',
+                                           partition_strategy_file='flat-genotypes-partition-strategy',
+                                           parallelism=parallelism)
         dependencies = [basic]
         for edition in ToastConfig().config['editions']:
             if edition == 'basic':
@@ -412,16 +384,24 @@ def output(self):
 class BAM2ADAMTask(Task):
 
     def requires(self):
+        conf = ToastConfig().config
+        parallelism = conf['numPartitionsHint'] if 'numPartitionsHint' in conf else 1
         basic = ADAMBasicTask(adam_command='transform',
                               allowed_file_formats=['sam', 'bam'])
         flat = ADAMFlattenTask(adam_command='transform',
                                allowed_file_formats=['sam', 'bam'])
         locuspart = ADAMPartitionTask(adam_command='transform',
                                       allowed_file_formats=['sam', 'bam'],
-                                      partition_strategy_file='alignments-partition-strategy')
-        flat_locuspart = ADAMFlattenPartitionTask(adam_command='transform',
-                                                  allowed_file_formats=['sam', 'bam'],
-                                                  partition_strategy_file='flat-alignments-partition-strategy')
+                                      source_edition='basic',
+                                      edition='locuspart',
+                                      partition_strategy_file='alignments-partition-strategy',
+                                      parallelism=parallelism)
+        flat_locuspart = ADAMPartitionTask(adam_command='transform',
+                                           allowed_file_formats=['sam', 'bam'],
+                                           source_edition='flat',
+                                           edition='flat_locuspart',
+                                           partition_strategy_file='flat-alignments-partition-strategy',
+                                           parallelism=parallelism)
         dependencies = [basic]
         for edition in ToastConfig().config['editions']:
             if edition == 'basic':

test/registry/test-1kg-genotypes-subset.json

@@ -3,6 +3,7 @@
     "title": "Test 1000 Genomes Project VCF data",
     "dag": "VCF2ADAMTask",
     "editions": ["basic", "flat", "locuspart", "flat_locuspart"],
+    "numPartitionsHint": 36,
     "sources": [
       {"format": "vcf", "compression": true, "url": "ftp://ftp-trace.ncbi.nih.gov/1000genomes/ftp/release/20110521/ALL.chr22.phase1_release_v3.20101123.snps_indels_svs.genotypes.vcf.gz"}
     ]

test/registry/test-alignments.json

@@ -3,6 +3,7 @@
     "title": "Test SAM data",
     "dag": "BAM2ADAMTask",
     "editions": ["basic", "flat", "locuspart", "flat_locuspart"],
+    "numPartitionsHint": 1,
     "sources": [
         {"format": "sam", "compression": false, "url": "https://github.com/bigdatagenomics/eggo/raw/master/test/resources/small.sam"}
     ]

test/registry/test-genotypes.json

@@ -4,6 +4,7 @@
     "target": "test/genotypes",
     "dag": "VCF2ADAMTask",
     "editions": ["basic", "flat", "locuspart", "flat_locuspart"],
+    "numPartitionsHint": 1,
     "sources": [
         {"format": "vcf", "compression": true, "url": "https://github.com/bigdatagenomics/eggo/raw/master/test/resources/chr22.small.vcf.gz"}
     ]