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

Author: bpblanken

v03_pipeline/lib/tasks/update_new_variants_with_caids.py

@@ -0,0 +1,62 @@
+import hail as hl
+import luigi
+import luigi.util
+
+from v03_pipeline.lib.misc.allele_registry import register_alleles_in_chunks
+from v03_pipeline.lib.misc.io import checkpoint
+from v03_pipeline.lib.model import (
+    Env,
+)
+from v03_pipeline.lib.paths import (
+    new_variants_table_path,
+)
+from v03_pipeline.lib.tasks.base.base_loading_run_params import (
+    BaseLoadingRunParams,
+)
+from v03_pipeline.lib.tasks.base.base_update import BaseUpdateTask
+from v03_pipeline.lib.tasks.files import GCSorLocalTarget
+from v03_pipeline.lib.tasks.write_new_variants_table import WriteNewVariantsTableTask
+
+
+@luigi.util.inherits(BaseLoadingRunParams)
+class UpdateNewVariantsWithCAIDsTask(BaseUpdateTask):
+    def output(self) -> luigi.Target:
+        return GCSorLocalTarget(
+            new_variants_table_path(
+                self.reference_genome,
+                self.dataset_type,
+                self.run_id,
+            ),
+        )
+
+    def requires(self) -> list[luigi.Task]:
+        return [
+            self.clone(WriteNewVariantsTableTask),
+        ]
+
+    def complete(self) -> bool:
+        return super().complete() and hasattr(hl.read_table(self.output().path), 'CAID')
+
+    def update_table(self, ht: hl.Table) -> hl.Table:
+        # Register the new variant alleles to the Clingen Allele Registry
+        # and annotate new_variants table with CAID.
+        if not (
+            Env.CLINGEN_ALLELE_REGISTRY_LOGIN and Env.CLINGEN_ALLELE_REGISTRY_PASSWORD
+        ):
+            return ht.annotate(CAID=hl.missing(hl.tstr))
+        ar_ht = hl.Table.parallelize(
+            [],
+            hl.tstruct(
+                locus=hl.tlocus(self.reference_genome.value),
+                alleles=hl.tarray(hl.tstr),
+                CAID=hl.tstr,
+            ),
+            key=('locus', 'alleles'),
+        )
+        for ar_ht_chunk in register_alleles_in_chunks(
+            ht,
+            self.reference_genome,
+        ):
+            ar_ht = ar_ht.union(ar_ht_chunk)
+            ar_ht, _ = checkpoint(ar_ht)
+        return ht.join(ar_ht, 'left')

v03_pipeline/lib/tasks/update_variant_annotations_table_with_new_samples.py

@@ -15,6 +15,9 @@
 from v03_pipeline.lib.tasks.base.base_update_variant_annotations_table import (
     BaseUpdateVariantAnnotationsTableTask,
 )
+from v03_pipeline.lib.tasks.update_new_variants_with_caids import (
+    UpdateNewVariantsWithCAIDsTask,
+)
 from v03_pipeline.lib.tasks.write_new_variants_table import WriteNewVariantsTableTask
 
 
@@ -25,7 +28,9 @@ class UpdateVariantAnnotationsTableWithNewSamplesTask(
     def requires(self) -> list[luigi.Task]:
         return [
             *super().requires(),
-            self.clone(WriteNewVariantsTableTask),
+            self.clone(UpdateNewVariantsWithCAIDsTask)
+            if self.dataset_type.should_send_to_allele_registry
+            else self.clone(WriteNewVariantsTableTask),
         ]
 
     def complete(self) -> bool:

v03_pipeline/lib/tasks/update_variant_annotations_table_with_new_samples_test.py

@@ -140,7 +140,10 @@ def test_missing_interval_reference_dataset(
             self.assertFalse(uvatwns_task.complete())
 
     @responses.activate
-    @patch('v03_pipeline.lib.tasks.write_new_variants_table.register_alleles_in_chunks')
+    @patch(
+        'v03_pipeline.lib.tasks.update_new_variants_with_caids.register_alleles_in_chunks',
+    )
+    @patch('v03_pipeline.lib.tasks.update_new_variants_with_caids.Env')
     @patch('v03_pipeline.lib.tasks.write_new_variants_table.Env')
     @patch(
         'v03_pipeline.lib.tasks.write_new_variants_table.UpdateVariantAnnotationsTableWithUpdatedReferenceDataset',
@@ -160,7 +163,8 @@ def test_multiple_update_vat(
         mock_vep: Mock,
         mock_standard_contigs: Mock,
         mock_update_vat_with_rd_task: Mock,
-        mock_env: Mock,
+        mock_env_new_variants: Mock,
+        mock_env_caids: Mock,
         mock_register_alleles: Mock,
     ) -> None:
         mock_update_vat_with_rd_task.return_value = (
@@ -174,9 +178,11 @@ def test_multiple_update_vat(
             {'ENST00000327044': 'NM_015658.4'},
         )
         # make register_alleles return CAIDs for 4 of 30 variants
-        mock_env.GRCH38_TO_GRCH37_LIFTOVER_REF_PATH = GRCH38_TO_GRCH37_LIFTOVER_REF_PATH
-        mock_env.CLINGEN_ALLELE_REGISTRY_LOGIN = 'login'
-        mock_env.CLINGEN_ALLELE_REGISTRY_PASSWORD = 'password1'  # noqa: S105
+        mock_env_new_variants.GRCH38_TO_GRCH37_LIFTOVER_REF_PATH = (
+            GRCH38_TO_GRCH37_LIFTOVER_REF_PATH
+        )
+        mock_env_caids.CLINGEN_ALLELE_REGISTRY_LOGIN = 'login'
+        mock_env_caids.CLINGEN_ALLELE_REGISTRY_PASSWORD = 'password1'  # noqa: S105
         mock_register_alleles.side_effect = [
             iter(
                 [
@@ -543,7 +549,9 @@ def test_multiple_update_vat(
             )
 
     @responses.activate
-    @patch('v03_pipeline.lib.tasks.write_new_variants_table.register_alleles_in_chunks')
+    @patch(
+        'v03_pipeline.lib.tasks.update_new_variants_with_caids.register_alleles_in_chunks',
+    )
     @patch(
         'v03_pipeline.lib.tasks.write_new_variants_table.UpdateVariantAnnotationsTableWithUpdatedReferenceDataset',
     )
@@ -698,7 +706,9 @@ def test_update_vat_grch37(
             )
 
     @responses.activate
-    @patch('v03_pipeline.lib.tasks.write_new_variants_table.register_alleles_in_chunks')
+    @patch(
+        'v03_pipeline.lib.tasks.update_new_variants_with_caids.register_alleles_in_chunks',
+    )
     @patch(
         'v03_pipeline.lib.tasks.write_new_variants_table.UpdateVariantAnnotationsTableWithUpdatedReferenceDataset',
     )
@@ -771,7 +781,9 @@ def test_update_vat_without_accessing_private_datasets(
             )
 
     @responses.activate
-    @patch('v03_pipeline.lib.tasks.write_new_variants_table.register_alleles_in_chunks')
+    @patch(
+        'v03_pipeline.lib.tasks.update_new_variants_with_caids.register_alleles_in_chunks',
+    )
     @patch(
         'v03_pipeline.lib.tasks.write_new_variants_table.UpdateVariantAnnotationsTableWithUpdatedReferenceDataset',
     )

v03_pipeline/lib/tasks/write_new_variants_table.py

@@ -5,9 +5,8 @@
 import luigi.util
 
 from v03_pipeline.lib.annotations.fields import get_fields
-from v03_pipeline.lib.misc.allele_registry import register_alleles_in_chunks
 from v03_pipeline.lib.misc.callsets import get_callset_ht
-from v03_pipeline.lib.misc.io import checkpoint, remap_pedigree_hash
+from v03_pipeline.lib.misc.io import remap_pedigree_hash
 from v03_pipeline.lib.misc.math import constrain
 from v03_pipeline.lib.model import (
     Env,
@@ -191,33 +190,6 @@ def create_table(self) -> hl.Table:
                 },
             )
             new_variants_ht = new_variants_ht.join(reference_dataset_ht, 'left')
-
-        # Register the new variant alleles to the Clingen Allele Registry
-        # and annotate new_variants table with CAID.
-        if (
-            Env.CLINGEN_ALLELE_REGISTRY_LOGIN
-            and Env.CLINGEN_ALLELE_REGISTRY_PASSWORD
-            and self.dataset_type.should_send_to_allele_registry
-        ):
-            ar_ht = hl.Table.parallelize(
-                [],
-                hl.tstruct(
-                    locus=hl.tlocus(self.reference_genome.value),
-                    alleles=hl.tarray(hl.tstr),
-                    CAID=hl.tstr,
-                ),
-                key=('locus', 'alleles'),
-            )
-            for ar_ht_chunk in register_alleles_in_chunks(
-                new_variants_ht,
-                self.reference_genome,
-            ):
-                ar_ht = ar_ht.union(ar_ht_chunk)
-                ar_ht, _ = checkpoint(ar_ht)
-            new_variants_ht = new_variants_ht.join(ar_ht, 'left')
-        elif self.dataset_type.should_send_to_allele_registry:
-            new_variants_ht = new_variants_ht.annotate(CAID=hl.missing(hl.tstr))
-
         return new_variants_ht.select_globals(
             updates={
                 hl.Struct(