refactor + chore: shuffle end of pipeline steps to enable dataproc for local users (#1009)

* Support gcs dirs in rsync

* ws

* Add create dataproc cluster task

* add dataproc

* ruff

* requirements

* still struggling

* Gencode refactor to remove gcs

* bump reqs

* Run dataproc job

* lib

* running

* merge requirements

* Flip'em

* Better exception handling

* Cleaner approach if less generalizable

* write a test

* Fix tests

* lint

* Add test for success

* refactor to use a base class... better for adding support for multiple jobs

* cleanup

* ruff

* Fix missing mock

* Fix flapping test

* first commit

* Finish test and cleanup

* Allow any order

* First commit

* ruff

* ruff

* finish off

* A few minor tweaks

Author: bpblanken

v03_pipeline/lib/model/feature_flag.py

@@ -11,6 +11,7 @@
 INCLUDE_PIPELINE_VERSION_IN_PREFIX = (
     os.environ.get('INCLUDE_PIPELINE_VERSION_IN_PREFIX') == '1'
 )
+RUN_PIPELINE_ON_DATAPROC = os.environ.get('RUN_PIPELINE_ON_DATAPROC') == '1'
 SHOULD_TRIGGER_HAIL_BACKEND_RELOAD = (
     os.environ.get('SHOULD_TRIGGER_HAIL_BACKEND_RELOAD') == '1'
 )
@@ -23,4 +24,5 @@ class FeatureFlag:
     EXPECT_TDR_METRICS: bool = EXPECT_TDR_METRICS
     EXPECT_WES_FILTERS: bool = EXPECT_WES_FILTERS
     INCLUDE_PIPELINE_VERSION_IN_PREFIX: bool = INCLUDE_PIPELINE_VERSION_IN_PREFIX
+    RUN_PIPELINE_ON_DATAPROC: bool = RUN_PIPELINE_ON_DATAPROC
     SHOULD_TRIGGER_HAIL_BACKEND_RELOAD: bool = SHOULD_TRIGGER_HAIL_BACKEND_RELOAD

v03_pipeline/lib/tasks/__init__.py

@@ -11,6 +11,7 @@
 from v03_pipeline.lib.tasks.reference_data.update_variant_annotations_table_with_updated_reference_dataset import (
     UpdateVariantAnnotationsTableWithUpdatedReferenceDataset,
 )
+from v03_pipeline.lib.tasks.run_pipeline import RunPipelineTask
 from v03_pipeline.lib.tasks.update_lookup_table import (
     UpdateLookupTableTask,
 )
@@ -46,6 +47,7 @@
     'DeleteProjectTablesTask',
     'MigrateAllLookupTablesTask',
     'MigrateAllVariantAnnotationsTablesTask',
+    'RunPipelineTask',
     'UpdateProjectTableTask',
     'UpdateProjectTablesWithDeletedFamiliesTask',
     'UpdateLookupTableTask',

v03_pipeline/lib/tasks/dataproc/base_run_job_on_dataproc.py

@@ -33,12 +33,12 @@ def __init__(self, *args, **kwargs):
         )
 
     @property
-    def task_name(self):
-        return self.get_task_family().split('.')[-1]
+    def task(self):
+        raise NotImplementedError
 
     @property
     def job_id(self):
-        return f'{self.task_name}-{self.run_id}'
+        return f'{self.task.task_family}-{self.run_id}'
 
     def requires(self) -> [luigi.Task]:
         return [self.clone(CreateDataprocClusterTask)]
@@ -58,7 +58,7 @@ def complete(self) -> bool:
         except google.api_core.exceptions.NotFound:
             return False
         if job.status.state == ERROR_STATE:
-            msg = f'Job {self.task_name}-{self.run_id} entered ERROR state'
+            msg = f'Job {self.task.task_family}-{self.run_id} entered ERROR state'
             logger.error(msg)
             logger.error(job.status.details)
         return job.status.state == DONE_STATE
@@ -81,7 +81,7 @@ def run(self):
                     'pyspark_job': {
                         'main_python_file_uri': f'{SEQR_PIPELINE_RUNNER_BUILD}/bin/run_task.py',
                         'args': [
-                            self.task_name,
+                            self.task.task_family,
                             '--local-scheduler',
                             *to_kebab_str_args(self),
                         ],

v03_pipeline/lib/tasks/dataproc/misc_test.py

@@ -3,8 +3,8 @@
 
 from v03_pipeline.lib.model import DatasetType, ReferenceGenome, SampleType
 from v03_pipeline.lib.tasks.dataproc.misc import to_kebab_str_args
-from v03_pipeline.lib.tasks.dataproc.write_success_file_on_dataproc import (
-    WriteSuccessFileOnDataprocTask,
+from v03_pipeline.lib.tasks.dataproc.rsync_to_seqr_app_dirs import (
+    RsyncToSeqrAppDirsTask,
 )
 
 
@@ -13,7 +13,7 @@
 )
 class MiscTest(unittest.TestCase):
     def test_to_kebab_str_args(self, _: Mock):
-        t = WriteSuccessFileOnDataprocTask(
+        t = RsyncToSeqrAppDirsTask(
             reference_genome=ReferenceGenome.GRCh38,
             dataset_type=DatasetType.SNV_INDEL,
             sample_type=SampleType.WGS,

v03_pipeline/lib/tasks/dataproc/rsync_to_seqr_app_dirs.py

@@ -9,6 +9,9 @@
 from v03_pipeline.lib.tasks.base.base_loading_run_params import (
     BaseLoadingRunParams,
 )
+from v03_pipeline.lib.tasks.dataproc.run_pipeline_on_dataproc import (
+    RunPipelineOnDataprocTask,
+)
 
 
 def hail_search_value(value: str) -> str:
@@ -38,6 +41,9 @@ def output(self) -> None:
     def complete(self) -> bool:
         return self.done
 
+    def requires(self) -> luigi.Task:
+        return self.clone(RunPipelineOnDataprocTask)
+
     def run(self) -> None:
         if not (
             Env.SEQR_APP_HAIL_SEARCH_DATA_DIR and Env.SEQR_APP_REFERENCE_DATASETS_DIR

v03_pipeline/lib/tasks/dataproc/rsync_to_seqr_app_dirs_test.py

@@ -13,14 +13,20 @@
 from v03_pipeline.lib.tasks.dataproc.rsync_to_seqr_app_dirs import (
     RsyncToSeqrAppDirsTask,
 )
+from v03_pipeline.lib.test.mock_complete_task import MockCompleteTask
 
 
 class RsyncToSeqrAppDirsTaskTest(unittest.TestCase):
+    @patch(
+        'v03_pipeline.lib.tasks.dataproc.rsync_to_seqr_app_dirs.RunPipelineOnDataprocTask',
+    )
     @patch('v03_pipeline.lib.tasks.dataproc.rsync_to_seqr_app_dirs.subprocess')
     def test_rsync_to_seqr_app_dirs_no_sync(
         self,
         mock_subprocess: Mock,
+        mock_run_pipeline_task: Mock,
     ) -> None:
+        mock_run_pipeline_task.return_value = MockCompleteTask()
         worker = luigi.worker.Worker()
         task = RsyncToSeqrAppDirsTask(
             reference_genome=ReferenceGenome.GRCh38,
@@ -37,6 +43,9 @@ def test_rsync_to_seqr_app_dirs_no_sync(
         self.assertTrue(task.complete())
         mock_subprocess.call.assert_not_called()
 
+    @patch(
+        'v03_pipeline.lib.tasks.dataproc.rsync_to_seqr_app_dirs.RunPipelineOnDataprocTask',
+    )
     @patch('v03_pipeline.lib.tasks.dataproc.rsync_to_seqr_app_dirs.subprocess')
     @patch.object(Env, 'HAIL_SEARCH_DATA_DIR', 'gs://test-hail-search-dir')
     @patch.object(Env, 'REFERENCE_DATASETS_DIR', 'gs://test-reference-data-dir')
@@ -58,7 +67,9 @@ def test_rsync_to_seqr_app_dirs_no_sync(
     def test_rsync_to_seqr_app_dirs_sync(
         self,
         mock_subprocess: Mock,
+        mock_run_pipeline_task: Mock,
     ) -> None:
+        mock_run_pipeline_task.return_value = MockCompleteTask()
         worker = luigi.worker.Worker()
         task = RsyncToSeqrAppDirsTask(
             reference_genome=ReferenceGenome.GRCh38,

v03_pipeline/lib/tasks/dataproc/run_pipeline_on_dataproc.py

@@ -0,0 +1,16 @@
+import luigi
+
+from v03_pipeline.lib.tasks.base.base_loading_run_params import (
+    BaseLoadingRunParams,
+)
+from v03_pipeline.lib.tasks.dataproc.base_run_job_on_dataproc import (
+    BaseRunJobOnDataprocTask,
+)
+from v03_pipeline.lib.tasks.run_pipeline import RunPipelineTask
+
+
+@luigi.util.inherits(BaseLoadingRunParams)
+class RunPipelineOnDataprocTask(BaseRunJobOnDataprocTask):
+    @property
+    def task(self) -> luigi.Task:
+        return RunPipelineTask

v03_pipeline/lib/tasks/dataproc/write_success_file_on_dataproc_test.py renamed to v03_pipeline/lib/tasks/dataproc/run_pipeline_on_dataproc_test.py

@@ -6,8 +6,8 @@
 import luigi
 
 from v03_pipeline.lib.model import DatasetType, ReferenceGenome, SampleType
-from v03_pipeline.lib.tasks.dataproc.write_success_file_on_dataproc import (
-    WriteSuccessFileOnDataprocTask,
+from v03_pipeline.lib.tasks.dataproc.run_pipeline_on_dataproc import (
+    RunPipelineOnDataprocTask,
 )
 from v03_pipeline.lib.test.mock_complete_task import MockCompleteTask
 
@@ -38,7 +38,7 @@ def test_job_already_exists_failed(
             google.api_core.exceptions.AlreadyExists('job exists')
         )
         worker = luigi.worker.Worker()
-        task = WriteSuccessFileOnDataprocTask(
+        task = RunPipelineOnDataprocTask(
             reference_genome=ReferenceGenome.GRCh38,
             dataset_type=DatasetType.SNV_INDEL,
             sample_type=SampleType.WGS,
@@ -54,7 +54,7 @@ def test_job_already_exists_failed(
         mock_logger.error.assert_has_calls(
             [
                 call(
-                    'Job WriteSuccessFileOnDataprocTask-manual__2024-04-03 entered ERROR state',
+                    'Job RunPipelineTask-manual__2024-04-03 entered ERROR state',
                 ),
             ],
         )
@@ -70,7 +70,7 @@ def test_job_already_exists_success(
             status=SimpleNamespace(state='DONE'),
         )
         worker = luigi.worker.Worker()
-        task = WriteSuccessFileOnDataprocTask(
+        task = RunPipelineOnDataprocTask(
             reference_genome=ReferenceGenome.GRCh38,
             dataset_type=DatasetType.SNV_INDEL,
             sample_type=SampleType.WGS,
@@ -102,7 +102,7 @@ def test_job_failed(
             'FailedPrecondition: 400 Job failed with message',
         )
         worker = luigi.worker.Worker()
-        task = WriteSuccessFileOnDataprocTask(
+        task = RunPipelineOnDataprocTask(
             reference_genome=ReferenceGenome.GRCh38,
             dataset_type=DatasetType.SNV_INDEL,
             sample_type=SampleType.WGS,
@@ -118,7 +118,7 @@ def test_job_failed(
         mock_logger.info.assert_has_calls(
             [
                 call(
-                    'Waiting for job completion WriteSuccessFileOnDataprocTask-manual__2024-04-05',
+                    'Waiting for job completion RunPipelineTask-manual__2024-04-05',
                 ),
             ],
         )
@@ -141,7 +141,7 @@ def test_job_success(
         operation = mock_client.submit_job_as_operation.return_value
         operation.done.side_effect = [False, True]
         worker = luigi.worker.Worker()
-        task = WriteSuccessFileOnDataprocTask(
+        task = RunPipelineOnDataprocTask(
             reference_genome=ReferenceGenome.GRCh38,
             dataset_type=DatasetType.SNV_INDEL,
             sample_type=SampleType.WGS,

v03_pipeline/lib/tasks/dataproc/write_success_file_on_dataproc.py

@@ -0,0 +1,32 @@
+import luigi
+import luigi.util
+
+from v03_pipeline.lib.tasks.base.base_loading_run_params import (
+    BaseLoadingRunParams,
+)
+from v03_pipeline.lib.tasks.update_variant_annotations_table_with_new_samples import (
+    UpdateVariantAnnotationsTableWithNewSamplesTask,
+)
+from v03_pipeline.lib.tasks.write_metadata_for_run import WriteMetadataForRunTask
+from v03_pipeline.lib.tasks.write_project_family_tables import (
+    WriteProjectFamilyTablesTask,
+)
+
+
+@luigi.util.inherits(BaseLoadingRunParams)
+class RunPipelineTask(luigi.WrapperTask):
+    def requires(self):
+        requirements = [
+            self.clone(WriteMetadataForRunTask),
+            self.clone(UpdateVariantAnnotationsTableWithNewSamplesTask),
+        ]
+        return [
+            *requirements,
+            *[
+                self.clone(
+                    WriteProjectFamilyTablesTask,
+                    project_i=i,
+                )
+                for i in range(len(self.project_guids))
+            ],
+        ]