Merge branch 'main' of github.com:broadinstitute/seqr-loading-pipelines into dev

bpblanken · bpblanken · commit 61eda3bd2d7a · 2024-09-30T17:07:55.000-04:00
diff --git a/v03_pipeline/bin/download_vep_reference_data.bash b/v03_pipeline/bin/download_vep_reference_data.bash
@@ -49,13 +49,15 @@ if [ -f $VEP_REFERENCE_DATASETS_DIR/$REFERENCE_GENOME/_SUCCESS ]; then
 fi
 
 mkdir -p $VEP_REFERENCE_DATASETS_DIR/$REFERENCE_GENOME;
+rm -rf $VEP_REFERENCE_DATASETS_DIR/$REFERENCE_GENOME/*;
+
 for vep_reference_data_file in ${VEP_REFERENCE_DATA_FILES[@]}; do
     if  [[ $vep_reference_data_file == *.tar.gz ]]; then
         echo "Downloading and extracting" $vep_reference_data_file;
         gsutil cat $vep_reference_data_file | tar -xzf - -C $VEP_REFERENCE_DATASETS_DIR/$REFERENCE_GENOME/ &
     else 
         echo "Downloading" $vep_reference_data_file;
-        gsutil cp $vep_reference_data_file $VEP_DATA/$REFERENCE_GENOME/ &
+        gsutil cp $vep_reference_data_file $VEP_REFERENCE_DATASETS_DIR/$REFERENCE_GENOME/ &
     fi
 done;
 wait
diff --git a/v03_pipeline/lib/misc/project_tables.py b/v03_pipeline/lib/misc/project_tables.py
@@ -0,0 +1,31 @@
+import os
+
+import hail as hl
+import hailtop.fs as hfs
+
+from v03_pipeline.lib.model import DatasetType, ReferenceGenome, SampleType
+from v03_pipeline.lib.paths import project_table_path
+
+
+def get_valid_project_tables(
+    reference_genome: ReferenceGenome,
+    dataset_type: DatasetType,
+    project_guid: str,
+) -> set[tuple[hl.Table, SampleType]]:
+    project_tables = set()
+    for sample_type in SampleType:
+        project_ht_path = project_table_path(
+            reference_genome,
+            dataset_type,
+            sample_type,
+            project_guid,
+        )
+        if hfs.exists(project_ht_path) and hfs.exists(
+            os.path.join(project_ht_path, '_SUCCESS'),
+        ):
+            project_ht = hl.read_table(project_ht_path)
+            project_tables.add((project_ht, sample_type))
+    if len(project_tables) == 0:
+        msg = f'No project tables found for {project_guid}'
+        raise RuntimeError(msg)
+    return project_tables
diff --git a/v03_pipeline/lib/model/environment.py b/v03_pipeline/lib/model/environment.py
@@ -11,10 +11,13 @@
     'gs://hail-common/references/grch38_to_grch37.over.chain.gz',
 )
 HAIL_TMP_DIR = os.environ.get('HAIL_TMP_DIR', '/tmp')  # noqa: S108
-HAIL_SEARCH_DATA_DIR = os.environ.get('HAIL_SEARCH_DATA_DIR', '/seqr/hail-search-data')
+HAIL_SEARCH_DATA_DIR = os.environ.get(
+    'HAIL_SEARCH_DATA_DIR',
+    '/seqr/seqr-hail-search-data',
+)
 LOADING_DATASETS_DIR = os.environ.get('LOADING_DATASETS_DIR', '/seqr/seqr-loading-temp')
-PRIVATE_REFERENCE_DATASETS_DIR_DIR = os.environ.get(
-    'PRIVATE_REFERENCE_DATASETS_DIR_DIR',
+PRIVATE_REFERENCE_DATASETS_DIR = os.environ.get(
+    'PRIVATE_REFERENCE_DATASETS_DIR',
     '/seqr/seqr-reference-data-private',
 )
 REFERENCE_DATASETS_DIR = os.environ.get(
@@ -50,7 +53,7 @@ class Env:
     GRCH37_TO_GRCH38_LIFTOVER_REF_PATH: str = GRCH37_TO_GRCH38_LIFTOVER_REF_PATH
     GRCH38_TO_GRCH37_LIFTOVER_REF_PATH: str = GRCH38_TO_GRCH37_LIFTOVER_REF_PATH
     LOADING_DATASETS_DIR: str = LOADING_DATASETS_DIR
-    PRIVATE_REFERENCE_DATASETS_DIR_DIR: str = PRIVATE_REFERENCE_DATASETS_DIR_DIR
+    PRIVATE_REFERENCE_DATASETS_DIR: str = PRIVATE_REFERENCE_DATASETS_DIR
     PROJECT_ID: str | None = PROJECT_ID
     REFERENCE_DATASETS_DIR: str = REFERENCE_DATASETS_DIR
     SHOULD_REGISTER_ALLELES: bool = SHOULD_REGISTER_ALLELES
diff --git a/v03_pipeline/lib/paths.py b/v03_pipeline/lib/paths.py
@@ -32,7 +32,7 @@ def _v03_reference_data_prefix(
     reference_genome: ReferenceGenome,
 ) -> str:
     root = (
-        Env.PRIVATE_REFERENCE_DATASETS_DIR_DIR
+        Env.PRIVATE_REFERENCE_DATASETS_DIR
         if access_control == AccessControl.PRIVATE
         else Env.REFERENCE_DATASETS_DIR
     )
diff --git a/v03_pipeline/lib/paths_test.py b/v03_pipeline/lib/paths_test.py
@@ -48,7 +48,7 @@ def test_family_table_path(self) -> None:
                 SampleType.WES,
                 'franklin',
             ),
-            '/seqr/hail-search-data/v3.1/GRCh37/SNV_INDEL/families/WES/franklin.ht',
+            '/seqr/seqr-hail-search-data/v3.1/GRCh37/SNV_INDEL/families/WES/franklin.ht',
         )
         with patch('v03_pipeline.lib.paths.Env') as mock_env:
             mock_env.HAIL_SEARCH_DATA_DIR = 'gs://seqr-datasets/'
@@ -90,7 +90,7 @@ def test_project_table_path(self) -> None:
                 SampleType.WES,
                 'R0652_pipeline_test',
             ),
-            '/seqr/hail-search-data/v3.1/GRCh38/MITO/projects/WES/R0652_pipeline_test.ht',
+            '/seqr/seqr-hail-search-data/v3.1/GRCh38/MITO/projects/WES/R0652_pipeline_test.ht',
         )
 
     def test_valid_reference_dataset_collection_path(self) -> None:
@@ -119,7 +119,7 @@ def test_lookup_table_path(self) -> None:
                 ReferenceGenome.GRCh37,
                 DatasetType.SV,
             ),
-            '/seqr/hail-search-data/v3.1/GRCh37/SV/lookup.ht',
+            '/seqr/seqr-hail-search-data/v3.1/GRCh37/SV/lookup.ht',
         )
 
     def test_sex_check_table_path(self) -> None:
@@ -159,7 +159,7 @@ def test_metadata_for_run_path(self) -> None:
                 DatasetType.SNV_INDEL,
                 'manual__2023-06-26T18:30:09.349671+00:00',
             ),
-            '/seqr/hail-search-data/v3.1/GRCh38/SNV_INDEL/runs/manual__2023-06-26T18:30:09.349671+00:00/metadata.json',
+            '/seqr/seqr-hail-search-data/v3.1/GRCh38/SNV_INDEL/runs/manual__2023-06-26T18:30:09.349671+00:00/metadata.json',
         )
 
     def test_variant_annotations_table_path(self) -> None:
@@ -168,7 +168,7 @@ def test_variant_annotations_table_path(self) -> None:
                 ReferenceGenome.GRCh38,
                 DatasetType.GCNV,
             ),
-            '/seqr/hail-search-data/v3.1/GRCh38/GCNV/annotations.ht',
+            '/seqr/seqr-hail-search-data/v3.1/GRCh38/GCNV/annotations.ht',
         )
 
     def test_remapped_and_subsetted_callset_path(self) -> None:
@@ -218,7 +218,7 @@ def test_new_variants_table_path(self) -> None:
                 DatasetType.SNV_INDEL,
                 'manual__2023-06-26T18:30:09.349671+00:00',
             ),
-            '/seqr/hail-search-data/v3.1/GRCh38/SNV_INDEL/runs/manual__2023-06-26T18:30:09.349671+00:00/new_variants.ht',
+            '/seqr/seqr-hail-search-data/v3.1/GRCh38/SNV_INDEL/runs/manual__2023-06-26T18:30:09.349671+00:00/new_variants.ht',
         )
 
     def test_project_remap_path(self) -> None:
diff --git a/v03_pipeline/lib/tasks/__init__.py b/v03_pipeline/lib/tasks/__init__.py
@@ -3,7 +3,7 @@
 from v03_pipeline.lib.tasks.delete_project_family_tables import (
     DeleteProjectFamilyTablesTask,
 )
-from v03_pipeline.lib.tasks.delete_project_table import DeleteProjectTableTask
+from v03_pipeline.lib.tasks.delete_project_tables import DeleteProjectTablesTask
 from v03_pipeline.lib.tasks.migrate_all_lookup_tables import MigrateAllLookupTablesTask
 from v03_pipeline.lib.tasks.migrate_all_variant_annotations_tables import (
     MigrateAllVariantAnnotationsTablesTask,
@@ -21,8 +21,8 @@
     UpdateLookupTableWithDeletedProjectTask,
 )
 from v03_pipeline.lib.tasks.update_project_table import UpdateProjectTableTask
-from v03_pipeline.lib.tasks.update_project_table_with_deleted_families import (
-    UpdateProjectTableWithDeletedFamiliesTask,
+from v03_pipeline.lib.tasks.update_project_tables_with_deleted_families import (
+    UpdateProjectTablesWithDeletedFamiliesTask,
 )
 from v03_pipeline.lib.tasks.update_variant_annotations_table_with_deleted_families import (
     UpdateVariantAnnotationsTableWithDeletedFamiliesTask,
@@ -42,11 +42,11 @@
     'DeleteFamilyTableTask',
     'DeleteFamilyTablesTask',
     'DeleteProjectFamilyTablesTask',
-    'DeleteProjectTableTask',
+    'DeleteProjectTablesTask',
     'MigrateAllLookupTablesTask',
     'MigrateAllVariantAnnotationsTablesTask',
     'UpdateProjectTableTask',
-    'UpdateProjectTableWithDeletedFamiliesTask',
+    'UpdateProjectTablesWithDeletedFamiliesTask',
     'UpdateLookupTableTask',
     'UpdateLookupTableWithDeletedProjectTask',
     'UpdateLookupTableWithDeletedFamiliesTask',
diff --git a/v03_pipeline/lib/tasks/delete_project_family_tables.py b/v03_pipeline/lib/tasks/delete_project_family_tables.py
@@ -1,15 +1,12 @@
 import hail as hl
-import hailtop.fs as hfs
 import luigi
 import luigi.util
 
-from v03_pipeline.lib.model import SampleType
-from v03_pipeline.lib.paths import project_table_path
+from v03_pipeline.lib.misc.project_tables import get_valid_project_tables
 from v03_pipeline.lib.tasks.base.base_loading_pipeline_params import (
     BaseLoadingPipelineParams,
 )
 from v03_pipeline.lib.tasks.delete_family_table import DeleteFamilyTableTask
-from v03_pipeline.lib.tasks.files import HailTableTask
 
 
 @luigi.util.inherits(BaseLoadingPipelineParams)
@@ -27,23 +24,11 @@ def complete(self) -> bool:
         )
 
     def run(self):
-        project_tables = set()
-        for sample_type in SampleType:
-            project_ht_path = project_table_path(
-                self.reference_genome,
-                self.dataset_type,
-                sample_type,
-                self.project_guid,
-            )
-            if hfs.exists(project_ht_path):
-                project_table_task: luigi.Target = yield HailTableTask(project_ht_path)
-                project_ht = hl.read_table(project_table_task.path)
-                project_tables.add((project_ht, sample_type))
-
-        if len(project_tables) == 0:
-            msg = f'No project tables found for {self.project_guid}'
-            raise RuntimeError(msg)
-
+        project_tables = get_valid_project_tables(
+            self.reference_genome,
+            self.dataset_type,
+            self.project_guid,
+        )
         for project_ht, sample_type in project_tables:
             family_guids = hl.eval(project_ht.globals.family_guids)
             for family_guid in family_guids:
diff --git a/v03_pipeline/lib/tasks/delete_project_tables.py b/v03_pipeline/lib/tasks/delete_project_tables.py
@@ -0,0 +1,34 @@
+import luigi
+
+from v03_pipeline.lib.model import SampleType
+from v03_pipeline.lib.tasks.base.base_loading_pipeline_params import (
+    BaseLoadingPipelineParams,
+)
+from v03_pipeline.lib.tasks.delete_project_table import DeleteProjectTableTask
+
+
+@luigi.util.inherits(BaseLoadingPipelineParams)
+class DeleteProjectTablesTask(luigi.Task):
+    project_guid = luigi.Parameter()
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.dynamic_delete_project_table_tasks = set()
+
+    def complete(self) -> bool:
+        return len(self.dynamic_delete_project_table_tasks) >= 1 and all(
+            delete_project_table_task.complete()
+            for delete_project_table_task in self.dynamic_delete_project_table_tasks
+        )
+
+    def run(self):
+        for sample_type in SampleType:
+            self.dynamic_delete_project_table_tasks.add(
+                DeleteProjectTableTask(
+                    reference_genome=self.reference_genome,
+                    dataset_type=self.dataset_type,
+                    sample_type=sample_type,
+                    project_guid=self.project_guid,
+                ),
+            )
+        yield self.dynamic_delete_project_table_tasks
diff --git a/v03_pipeline/lib/tasks/update_lookup_table.py b/v03_pipeline/lib/tasks/update_lookup_table.py
@@ -2,7 +2,7 @@
 import luigi
 import luigi.util
 
-from v03_pipeline.lib.misc.io import remap_pedigree_hash
+from v03_pipeline.lib.misc.io import checkpoint, remap_pedigree_hash
 from v03_pipeline.lib.misc.lookup import (
     compute_callset_lookup_ht,
     join_lookup_hts,
@@ -111,4 +111,5 @@ def update_table(self, ht: hl.Table) -> hl.Table:
                 ),
                 migrations=ht.migrations,
             )
+            ht, _ = checkpoint(ht)
         return ht
diff --git a/v03_pipeline/lib/tasks/update_project_tables_with_deleted_families.py b/v03_pipeline/lib/tasks/update_project_tables_with_deleted_families.py
@@ -0,0 +1,44 @@
+import luigi
+import luigi.util
+
+from v03_pipeline.lib.misc.project_tables import get_valid_project_tables
+from v03_pipeline.lib.tasks.base.base_loading_pipeline_params import (
+    BaseLoadingPipelineParams,
+)
+from v03_pipeline.lib.tasks.update_project_table_with_deleted_families import (
+    UpdateProjectTableWithDeletedFamiliesTask,
+)
+
+
+@luigi.util.inherits(BaseLoadingPipelineParams)
+class UpdateProjectTablesWithDeletedFamiliesTask(luigi.Task):
+    project_guid = luigi.Parameter()
+    family_guids = luigi.ListParameter()
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.dynamic_update_project_table_tasks = set()
+
+    def complete(self) -> bool:
+        return len(self.dynamic_update_project_table_tasks) >= 1 and all(
+            update_project_table_task.complete()
+            for update_project_table_task in self.dynamic_update_project_table_tasks
+        )
+
+    def run(self):
+        project_tables = get_valid_project_tables(
+            self.reference_genome,
+            self.dataset_type,
+            self.project_guid,
+        )
+        for _, sample_type in project_tables:
+            self.dynamic_update_project_table_tasks.add(
+                UpdateProjectTableWithDeletedFamiliesTask(
+                    reference_genome=self.reference_genome,
+                    dataset_type=self.dataset_type,
+                    sample_type=sample_type,
+                    project_guid=self.project_guid,
+                    family_guids=self.family_guids,
+                ),
+            )
+        yield self.dynamic_update_project_table_tasks
diff --git a/v03_pipeline/var/luigi_config/luigi-defaults.conf b/v03_pipeline/var/luigi_config/luigi-defaults.conf
@@ -1,2 +1,9 @@
 [core]
-default_scheduler_url=http://pipeline-runner-ui:8082/
+default_scheduler_url=http://pipeline-runner-ui:8082/
+
+[scheduler]
+record_task_history = True
+state_path = /seqr/luigi-state/luigi-state.pickle
+
+[task_history]
+db_connection = sqlite:////seqr/luigi-state/luigi-task-hist.db

Original file line number	Diff line number	Diff line change
`@@ -32,7 +32,7 @@ def _v03_reference_data_prefix(`
`32`	`32`	`reference_genome: ReferenceGenome,`
`33`	`33`	`) -> str:`
`34`	`34`	`root = (`
`35`		`- Env.PRIVATE_REFERENCE_DATASETS_DIR_DIR`
	`35`	`+ Env.PRIVATE_REFERENCE_DATASETS_DIR`
`36`	`36`	`if access_control == AccessControl.PRIVATE`
`37`	`37`	`else Env.REFERENCE_DATASETS_DIR`
`38`	`38`	`)`