refactor: Move feature flags to FeatureFlag enum. (#1002)

* refactor: Move feature flags out of environment to their own dataclass

* lint: ruff

* ruff

Author: bpblanken

v03_pipeline/bin/pipeline_worker.py

@@ -7,7 +7,7 @@
 
 from v03_pipeline.api.model import LoadingPipelineRequest
 from v03_pipeline.lib.logger import get_logger
-from v03_pipeline.lib.model import Env
+from v03_pipeline.lib.model import FeatureFlag
 from v03_pipeline.lib.paths import (
     loading_pipeline_queue_path,
     project_pedigree_path,
@@ -54,7 +54,7 @@ def main():
                 'run_id': run_id,
                 **{k: v for k, v in lpr.model_dump().items() if k != 'projects_to_run'},
             }
-            if Env.SHOULD_TRIGGER_HAIL_BACKEND_RELOAD:
+            if FeatureFlag.SHOULD_TRIGGER_HAIL_BACKEND_RELOAD:
                 tasks = [
                     TriggerHailBackendReload(**loading_run_task_params),
                 ]

v03_pipeline/lib/model/__init__.py

@@ -7,11 +7,13 @@
     Sex,
 )
 from v03_pipeline.lib.model.environment import Env
+from v03_pipeline.lib.model.feature_flag import FeatureFlag
 
 __all__ = [
     'AccessControl',
     'DatasetType',
     'Env',
+    'FeatureFlag',
     'Sex',
     'PipelineVersion',
     'ReferenceGenome',

v03_pipeline/lib/model/environment.py

@@ -51,30 +51,12 @@
 GCLOUD_REGION = os.environ.get('GCLOUD_REGION')
 PIPELINE_RUNNER_APP_VERSION = os.environ.get('PIPELINE_RUNNER_APP_VERSION', 'latest')
 
-# Feature Flags
-ACCESS_PRIVATE_REFERENCE_DATASETS = (
-    os.environ.get('ACCESS_PRIVATE_REFERENCE_DATASETS') == '1'
-)
-CHECK_SEX_AND_RELATEDNESS = os.environ.get('CHECK_SEX_AND_RELATEDNESS') == '1'
-EXPECT_TDR_METRICS = os.environ.get('EXPECT_TDR_METRICS') == '1'
-EXPECT_WES_FILTERS = os.environ.get('EXPECT_WES_FILTERS') == '1'
-INCLUDE_PIPELINE_VERSION_IN_PREFIX = (
-    os.environ.get('INCLUDE_PIPELINE_VERSION_IN_PREFIX') == '1'
-)
-SHOULD_TRIGGER_HAIL_BACKEND_RELOAD = (
-    os.environ.get('SHOULD_TRIGGER_HAIL_BACKEND_RELOAD') == '1'
-)
-
 
 @dataclass
 class Env:
-    ACCESS_PRIVATE_REFERENCE_DATASETS: bool = ACCESS_PRIVATE_REFERENCE_DATASETS
-    CHECK_SEX_AND_RELATEDNESS: bool = CHECK_SEX_AND_RELATEDNESS
     CLINGEN_ALLELE_REGISTRY_LOGIN: str | None = CLINGEN_ALLELE_REGISTRY_LOGIN
     CLINGEN_ALLELE_REGISTRY_PASSWORD: str | None = CLINGEN_ALLELE_REGISTRY_PASSWORD
     DEPLOYMENT_TYPE: Literal['dev', 'prod'] = DEPLOYMENT_TYPE
-    EXPECT_TDR_METRICS: bool = EXPECT_TDR_METRICS
-    EXPECT_WES_FILTERS: bool = EXPECT_WES_FILTERS
     GCLOUD_DATAPROC_SECONDARY_WORKERS: str = GCLOUD_DATAPROC_SECONDARY_WORKERS
     GCLOUD_PROJECT: str | None = GCLOUD_PROJECT
     GCLOUD_ZONE: str | None = GCLOUD_ZONE
@@ -85,10 +67,8 @@ class Env:
     HAIL_BACKEND_SERVICE_PORT: int = HAIL_BACKEND_SERVICE_PORT
     HAIL_TMP_DIR: str = HAIL_TMP_DIR
     HAIL_SEARCH_DATA_DIR: str = HAIL_SEARCH_DATA_DIR
-    INCLUDE_PIPELINE_VERSION_IN_PREFIX: bool = INCLUDE_PIPELINE_VERSION_IN_PREFIX
     LOADING_DATASETS_DIR: str = LOADING_DATASETS_DIR
     PIPELINE_RUNNER_APP_VERSION: str = PIPELINE_RUNNER_APP_VERSION
     PRIVATE_REFERENCE_DATASETS_DIR: str = PRIVATE_REFERENCE_DATASETS_DIR
     REFERENCE_DATASETS_DIR: str = REFERENCE_DATASETS_DIR
-    SHOULD_TRIGGER_HAIL_BACKEND_RELOAD: bool = SHOULD_TRIGGER_HAIL_BACKEND_RELOAD
     VEP_REFERENCE_DATASETS_DIR: str = VEP_REFERENCE_DATASETS_DIR

v03_pipeline/lib/model/feature_flag.py

@@ -0,0 +1,26 @@
+import os
+from dataclasses import dataclass
+
+# Feature Flags
+ACCESS_PRIVATE_REFERENCE_DATASETS = (
+    os.environ.get('ACCESS_PRIVATE_REFERENCE_DATASETS') == '1'
+)
+CHECK_SEX_AND_RELATEDNESS = os.environ.get('CHECK_SEX_AND_RELATEDNESS') == '1'
+EXPECT_TDR_METRICS = os.environ.get('EXPECT_TDR_METRICS') == '1'
+EXPECT_WES_FILTERS = os.environ.get('EXPECT_WES_FILTERS') == '1'
+INCLUDE_PIPELINE_VERSION_IN_PREFIX = (
+    os.environ.get('INCLUDE_PIPELINE_VERSION_IN_PREFIX') == '1'
+)
+SHOULD_TRIGGER_HAIL_BACKEND_RELOAD = (
+    os.environ.get('SHOULD_TRIGGER_HAIL_BACKEND_RELOAD') == '1'
+)
+
+
+@dataclass
+class FeatureFlag:
+    ACCESS_PRIVATE_REFERENCE_DATASETS: bool = ACCESS_PRIVATE_REFERENCE_DATASETS
+    CHECK_SEX_AND_RELATEDNESS: bool = CHECK_SEX_AND_RELATEDNESS
+    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
+    SHOULD_TRIGGER_HAIL_BACKEND_RELOAD: bool = SHOULD_TRIGGER_HAIL_BACKEND_RELOAD

v03_pipeline/lib/paths.py

@@ -6,6 +6,7 @@
     AccessControl,
     DatasetType,
     Env,
+    FeatureFlag,
     PipelineVersion,
     ReferenceGenome,
     SampleType,
@@ -21,7 +22,7 @@ def _pipeline_prefix(
     reference_genome: ReferenceGenome,
     dataset_type: DatasetType,
 ) -> str:
-    if Env.INCLUDE_PIPELINE_VERSION_IN_PREFIX:
+    if FeatureFlag.INCLUDE_PIPELINE_VERSION_IN_PREFIX:
         return os.path.join(
             root,
             PipelineVersion.V3_1.value,
@@ -45,7 +46,7 @@ def _v03_reference_data_prefix(
         if access_control == AccessControl.PRIVATE
         else Env.REFERENCE_DATASETS_DIR
     )
-    if Env.INCLUDE_PIPELINE_VERSION_IN_PREFIX:
+    if FeatureFlag.INCLUDE_PIPELINE_VERSION_IN_PREFIX:
         return os.path.join(
             root,
             PipelineVersion.V03.value,
@@ -68,7 +69,7 @@ def _v03_reference_dataset_prefix(
         if access_control == AccessControl.PRIVATE
         else Env.REFERENCE_DATASETS_DIR
     )
-    if Env.INCLUDE_PIPELINE_VERSION_IN_PREFIX:
+    if FeatureFlag.INCLUDE_PIPELINE_VERSION_IN_PREFIX:
         return os.path.join(
             root,
             PipelineVersion.V3_1.value,
@@ -287,7 +288,7 @@ def valid_filters_path(
     callset_path: str,
 ) -> str | None:
     if (
-        not Env.EXPECT_WES_FILTERS
+        not FeatureFlag.EXPECT_WES_FILTERS
         or not dataset_type.expect_filters(sample_type)
         or 'part_one_outputs' not in callset_path
     ):

v03_pipeline/lib/paths_test.py

@@ -36,7 +36,9 @@ def test_family_table_path(self) -> None:
             ),
             '/var/seqr/seqr-hail-search-data/v3.1/GRCh37/SNV_INDEL/families/WES/franklin.ht',
         )
-        with patch('v03_pipeline.lib.paths.Env') as mock_env:
+        with patch('v03_pipeline.lib.paths.Env') as mock_env, patch(
+            'v03_pipeline.lib.paths.FeatureFlag',
+        ) as mock_ff:
             mock_env.HAIL_SEARCH_DATA_DIR = 'gs://seqr-datasets/'
             self.assertEqual(
                 family_table_path(
@@ -47,7 +49,7 @@ def test_family_table_path(self) -> None:
                 ),
                 'gs://seqr-datasets/v3.1/GRCh37/SNV_INDEL/families/WES/franklin.ht',
             )
-            mock_env.INCLUDE_PIPELINE_VERSION_IN_PREFIX = False
+            mock_ff.INCLUDE_PIPELINE_VERSION_IN_PREFIX = False
             self.assertEqual(
                 family_table_path(
                     ReferenceGenome.GRCh37,
@@ -67,8 +69,8 @@ def test_valid_filters_path(self) -> None:
             ),
             None,
         )
-        with patch('v03_pipeline.lib.paths.Env') as mock_env:
-            mock_env.EXPECT_WES_FILTERS = True
+        with patch('v03_pipeline.lib.paths.FeatureFlag') as mock_ff:
+            mock_ff.EXPECT_WES_FILTERS = True
             self.assertEqual(
                 valid_filters_path(
                     DatasetType.SNV_INDEL,

v03_pipeline/lib/reference_datasets/reference_dataset.py

@@ -10,7 +10,12 @@
     validate_allele_type,
     validate_no_duplicate_variants,
 )
-from v03_pipeline.lib.model import AccessControl, DatasetType, Env, ReferenceGenome
+from v03_pipeline.lib.model import (
+    AccessControl,
+    DatasetType,
+    FeatureFlag,
+    ReferenceGenome,
+)
 from v03_pipeline.lib.reference_datasets import clinvar, dbnsfp
 from v03_pipeline.lib.reference_datasets.misc import (
     compress_floats,
@@ -41,7 +46,7 @@ def for_reference_genome_dataset_type(
             for dataset, config in CONFIG.items()
             if dataset_type in config.get(reference_genome, {}).get(DATASET_TYPES, [])
         ]
-        if not Env.ACCESS_PRIVATE_REFERENCE_DATASETS:
+        if not FeatureFlag.ACCESS_PRIVATE_REFERENCE_DATASETS:
             return {
                 dataset
                 for dataset in reference_datasets

v03_pipeline/lib/tasks/dataproc/create_dataproc_cluster.py

@@ -7,7 +7,7 @@
 
 from v03_pipeline.lib.logger import get_logger
 from v03_pipeline.lib.misc.gcp import get_service_account_credentials
-from v03_pipeline.lib.model import Env, ReferenceGenome
+from v03_pipeline.lib.model import Env, FeatureFlag, ReferenceGenome
 from v03_pipeline.lib.tasks.base.base_loading_pipeline_params import (
     BaseLoadingPipelineParams,
 )
@@ -88,18 +88,21 @@ def get_cluster_config(reference_genome: ReferenceGenome, run_id: str):
                     'spark:spark.executorEnv.HAIL_WORKER_OFF_HEAP_MEMORY_PER_CORE_MB': '6323',
                     'spark:spark.speculation': 'true',
                     'spark-env:ACCESS_PRIVATE_REFERENCE_DATASETS': '1'
-                    if Env.ACCESS_PRIVATE_REFERENCE_DATASETS
+                    if FeatureFlag.ACCESS_PRIVATE_REFERENCE_DATASETS
                     else '0',
                     'spark-env:CHECK_SEX_AND_RELATEDNESS': '1'
-                    if Env.CHECK_SEX_AND_RELATEDNESS
+                    if FeatureFlag.CHECK_SEX_AND_RELATEDNESS
+                    else '0',
+                    'spark-env:EXPECT_TDR_METRICS': '1'
+                    if FeatureFlag.EXPECT_TDR_METRICS
                     else '0',
                     'spark-env:EXPECT_WES_FILTERS': '1'
-                    if Env.EXPECT_WES_FILTERS
+                    if FeatureFlag.EXPECT_WES_FILTERS
                     else '0',
                     'spark-env:HAIL_SEARCH_DATA_DIR': Env.HAIL_SEARCH_DATA_DIR,
                     'spark-env:HAIL_TMP_DIR': Env.HAIL_TMP_DIR,
                     'spark-env:INCLUDE_PIPELINE_VERSION_IN_PREFIX': '1'
-                    if Env.INCLUDE_PIPELINE_VERSION_IN_PREFIX
+                    if FeatureFlag.INCLUDE_PIPELINE_VERSION_IN_PREFIX
                     else '0',
                     'spark-env:LOADING_DATASETS_DIR': Env.LOADING_DATASETS_DIR,
                     'spark-env:PRIVATE_REFERENCE_DATASETS_DIR': Env.PRIVATE_REFERENCE_DATASETS_DIR,

v03_pipeline/lib/tasks/update_variant_annotations_table_with_new_samples_test.py

@@ -712,7 +712,7 @@ def test_update_vat_grch37(
     @patch(
         'v03_pipeline.lib.tasks.write_new_variants_table.UpdateVariantAnnotationsTableWithUpdatedReferenceDataset',
     )
-    @patch('v03_pipeline.lib.reference_datasets.reference_dataset.Env')
+    @patch('v03_pipeline.lib.reference_datasets.reference_dataset.FeatureFlag')
     @patch('v03_pipeline.lib.vep.hl.vep')
     @patch(
         'v03_pipeline.lib.tasks.write_new_variants_table.load_gencode_ensembl_to_refseq_id',
@@ -721,7 +721,7 @@ def test_update_vat_without_accessing_private_datasets(
         self,
         mock_load_gencode_ensembl_to_refseq_id: Mock,
         mock_vep: Mock,
-        mock_rd_env: Mock,
+        mock_rd_ff: Mock,
         mock_update_vat_with_rd_task: Mock,
         mock_register_alleles: Mock,
     ) -> None:
@@ -740,7 +740,7 @@ def test_update_vat_without_accessing_private_datasets(
                 ReferenceDataset.hgmd,
             ),
         )
-        mock_rd_env.ACCESS_PRIVATE_REFERENCE_DATASETS = False
+        mock_rd_ff.ACCESS_PRIVATE_REFERENCE_DATASETS = False
         mock_vep.side_effect = lambda ht, **_: ht.annotate(vep=MOCK_38_VEP_DATA)
         mock_register_alleles.side_effect = None
 

v03_pipeline/lib/tasks/validate_callset.py

@@ -10,7 +10,7 @@
     validate_no_duplicate_variants,
     validate_sample_type,
 )
-from v03_pipeline.lib.model.environment import Env
+from v03_pipeline.lib.model.feature_flag import FeatureFlag
 from v03_pipeline.lib.paths import (
     imported_callset_path,
     sex_check_table_path,
@@ -41,7 +41,7 @@ def get_validation_dependencies(self) -> dict[str, hl.Table]:
             ),
         )
         if (
-            Env.CHECK_SEX_AND_RELATEDNESS
+            FeatureFlag.CHECK_SEX_AND_RELATEDNESS
             and self.dataset_type.check_sex_and_relatedness
             and not self.skip_check_sex_and_relatedness
         ):
@@ -86,7 +86,7 @@ def requires(self) -> list[luigi.Task]:
                 ),
             ]
         if (
-            Env.CHECK_SEX_AND_RELATEDNESS
+            FeatureFlag.CHECK_SEX_AND_RELATEDNESS
             and self.dataset_type.check_sex_and_relatedness
             and not self.skip_check_sex_and_relatedness
         ):