Feature flag reference data update (#692)

* Feature flag reference data update

* ruff

* Add flag to unit tests

* Update update_variant_annotations_table_with_new_samples.py

Author: bpblanken

.github/workflows/unit-tests.yml

@@ -34,6 +34,7 @@ jobs:
         run: ruff . --output-format github
       - name: Unit Tests
         run: |
+          export REFERENCE_DATA_AUTO_UPDATE=1
           export ACCESS_PRIVATE_REFERENCE_DATASETS=1
           export PYSPARK_SUBMIT_ARGS='--driver-memory 8G pyspark-shell'
           nosetests --with-coverage --cover-package v03_pipeline/lib v03_pipeline/lib

v03_pipeline/lib/model/environment.py

@@ -6,6 +6,7 @@
     os.environ.get('ACCESS_PRIVATE_REFERENCE_DATASETS') == '1'
 )
 CHECK_SEX_AND_RELATEDNESS = os.environ.get('CHECK_SEX_AND_RELATEDNESS') == '1'
+REFERENCE_DATA_AUTO_UPDATE = os.environ.get('REFERENCE_DATA_AUTO_UPDATE') == '1'
 HAIL_TMPDIR = os.environ.get('HAIL_TMPDIR', '/tmp')  # noqa: S108
 HAIL_SEARCH_DATA = os.environ.get('HAIL_SEARCH_DATA', '/hail-search-data')
 LOADING_DATASETS = os.environ.get('LOADING_DATASETS', '/seqr-loading-temp')
@@ -23,6 +24,7 @@
 class Env:
     ACCESS_PRIVATE_REFERENCE_DATASETS: bool = ACCESS_PRIVATE_REFERENCE_DATASETS
     CHECK_SEX_AND_RELATEDNESS: bool = CHECK_SEX_AND_RELATEDNESS
+    REFERENCE_DATA_AUTO_UPDATE: bool = REFERENCE_DATA_AUTO_UPDATE
     HAIL_TMPDIR: str = HAIL_TMPDIR
     HAIL_SEARCH_DATA: str = HAIL_SEARCH_DATA
     LOADING_DATASETS: str = LOADING_DATASETS

v03_pipeline/lib/tasks/base/base_variant_annotations_table.py

@@ -1,13 +1,13 @@
 import hail as hl
 import luigi
 
-from v03_pipeline.lib.model import ReferenceDatasetCollection
+from v03_pipeline.lib.model import Env, ReferenceDatasetCollection
 from v03_pipeline.lib.paths import (
     valid_reference_dataset_collection_path,
     variant_annotations_table_path,
 )
 from v03_pipeline.lib.tasks.base.base_update_task import BaseUpdateTask
-from v03_pipeline.lib.tasks.files import GCSorLocalTarget
+from v03_pipeline.lib.tasks.files import GCSorLocalTarget, HailTableTask
 from v03_pipeline.lib.tasks.reference_data.updated_reference_dataset_collection import (
     UpdatedReferenceDatasetCollectionTask,
 )
@@ -40,11 +40,21 @@ def output(self) -> luigi.Target:
 
     def requires(self) -> list[luigi.Task]:
         return [
-            UpdatedReferenceDatasetCollectionTask(
-                self.reference_genome,
-                self.dataset_type,
-                self.sample_type,
-                rdc,
+            (
+                UpdatedReferenceDatasetCollectionTask(
+                    self.reference_genome,
+                    self.dataset_type,
+                    self.sample_type,
+                    rdc,
+                )
+                if Env.REFERENCE_DATA_AUTO_UPDATE
+                else HailTableTask(
+                    valid_reference_dataset_collection_path(
+                        self.reference_genome,
+                        self.dataset_type,
+                        rdc,
+                    ),
+                )
             )
             for rdc in ReferenceDatasetCollection.for_reference_genome_dataset_type(
                 self.reference_genome,

v03_pipeline/lib/tasks/update_variant_annotations_table_with_new_samples.py

@@ -8,7 +8,7 @@
 from v03_pipeline.lib.annotations.fields import get_fields
 from v03_pipeline.lib.misc.math import constrain
 from v03_pipeline.lib.misc.util import callset_project_pairs
-from v03_pipeline.lib.model import ReferenceDatasetCollection
+from v03_pipeline.lib.model import Env, ReferenceDatasetCollection
 from v03_pipeline.lib.paths import (
     remapped_and_subsetted_callset_path,
     sample_lookup_table_path,
@@ -76,13 +76,16 @@ def other_annotation_dependencies(self) -> dict[str, hl.Table]:
         return annotation_dependencies
 
     def requires(self) -> list[luigi.Task]:
-        upstream_table_tasks: list[luigi.Task] = [
-            UpdateVariantAnnotationsTableWithUpdatedReferenceDataset(
-                self.reference_genome,
-                self.dataset_type,
-                self.sample_type,
-            ),
-        ]
+        if Env.REFERENCE_DATA_AUTO_UPDATE:
+            upstream_table_tasks: list[luigi.Task] = [
+                UpdateVariantAnnotationsTableWithUpdatedReferenceDataset(
+                    self.reference_genome,
+                    self.dataset_type,
+                    self.sample_type,
+                ),
+            ]
+        else:
+            upstream_table_tasks: list[luigi.Task] = []
         if self.dataset_type.has_sample_lookup_table:
             # NB: the sample lookup table task has remapped and subsetted callset tasks as dependencies.
             upstream_table_tasks.extend(