add task to reload hail backend globals (#902)

* make run id shared

* add task to reload hail backend globals

* ruff

* simpler

* fix run_id"

* A fix

* another test

* A few more

* Few more

* Unnecessary

* ruf

* run id

* string

* unused

* missed one

* Fix it correctly

* missed one

* Last one!

* add dependency and base params

* add new param...

* comments

---------

Co-authored-by: Benjamin Blankenmeister <bblanken@broadinstitute.org>

Author: jklugherz

v03_pipeline/bin/pipeline_worker.py

@@ -7,6 +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.paths import (
     loading_pipeline_queue_path,
     project_pedigree_path,
@@ -16,6 +17,7 @@
     UpdateVariantAnnotationsTableWithNewSamplesTask,
     WriteProjectFamilyTablesTask,
 )
+from v03_pipeline.lib.tasks.trigger_hail_backend_reload import TriggerHailBackendReload
 
 logger = get_logger(__name__)
 
@@ -48,14 +50,15 @@ def main():
             task_kwargs = {
                 k: v for k, v in lpr.model_dump().items() if k != 'projects_to_run'
             }
+            run_id = datetime.datetime.now(datetime.timezone.utc).strftime(
+                '%Y%m%d-%H%M%S',
+            )
             tasks = [
                 UpdateVariantAnnotationsTableWithNewSamplesTask(
                     project_guids=lpr.projects_to_run,
                     project_remap_paths=project_remap_paths,
                     project_pedigree_paths=project_pedigree_paths,
-                    run_id=datetime.datetime.now(datetime.timezone.utc).strftime(
-                        '%Y%m%d-%H%M%S',
-                    ),
+                    run_id=run_id,
                     **task_kwargs,
                 ),
                 *[
@@ -68,6 +71,16 @@ def main():
                     for i in range(len(lpr.projects_to_run))
                 ],
             ]
+            if Env.SHOULD_TRIGGER_HAIL_BACKEND_RELOAD:
+                tasks.append(
+                    TriggerHailBackendReload(
+                        project_guids=lpr.projects_to_run,
+                        project_remap_paths=project_remap_paths,
+                        project_pedigree_paths=project_pedigree_paths,
+                        run_id=run_id,
+                        **task_kwargs,
+                    ),
+                )
             luigi.build(tasks)
         except Exception:
             logger.exception('Unhandled Exception')

v03_pipeline/lib/model/environment.py

@@ -28,6 +28,8 @@
     'VEP_REFERENCE_DATASETS_DIR',
     '/seqr/vep-reference-data',
 )
+HAIL_BACKEND_SERVICE_HOSTNAME = os.environ.get('HAIL_BACKEND_SERVICE_HOSTNAME')
+HAIL_BACKEND_SERVICE_PORT = int(os.environ.get('HAIL_BACKEND_SERVICE_PORT', '5000'))
 
 # Allele registry secrets :/
 ALLELE_REGISTRY_SECRET_NAME = os.environ.get('ALLELE_REGISTRY_SECRET_NAME', None)
@@ -40,6 +42,9 @@
 CHECK_SEX_AND_RELATEDNESS = os.environ.get('CHECK_SEX_AND_RELATEDNESS') == '1'
 EXPECT_WES_FILTERS = os.environ.get('EXPECT_WES_FILTERS') == '1'
 SHOULD_REGISTER_ALLELES = os.environ.get('SHOULD_REGISTER_ALLELES') == '1'
+SHOULD_TRIGGER_HAIL_BACKEND_RELOAD = (
+    os.environ.get('SHOULD_TRIGGER_HAIL_BACKEND_RELOAD') == '1'
+)
 
 
 @dataclass
@@ -48,6 +53,8 @@ class Env:
     ALLELE_REGISTRY_SECRET_NAME: str | None = ALLELE_REGISTRY_SECRET_NAME
     CHECK_SEX_AND_RELATEDNESS: bool = CHECK_SEX_AND_RELATEDNESS
     EXPECT_WES_FILTERS: bool = EXPECT_WES_FILTERS
+    HAIL_BACKEND_SERVICE_HOSTNAME: str | None = HAIL_BACKEND_SERVICE_HOSTNAME
+    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
     GRCH37_TO_GRCH38_LIFTOVER_REF_PATH: str = GRCH37_TO_GRCH38_LIFTOVER_REF_PATH
@@ -57,4 +64,5 @@ class Env:
     PROJECT_ID: str | None = PROJECT_ID
     REFERENCE_DATASETS_DIR: str = REFERENCE_DATASETS_DIR
     SHOULD_REGISTER_ALLELES: bool = SHOULD_REGISTER_ALLELES
+    SHOULD_TRIGGER_HAIL_BACKEND_RELOAD: bool = SHOULD_TRIGGER_HAIL_BACKEND_RELOAD
     VEP_REFERENCE_DATASETS_DIR: str = VEP_REFERENCE_DATASETS_DIR

v03_pipeline/lib/tasks/base/base_project_info_params.py

@@ -0,0 +1,11 @@
+import luigi
+import luigi.util
+
+from v03_pipeline.lib.tasks.base.base_loading_run_params import BaseLoadingRunParams
+
+
+@luigi.util.inherits(BaseLoadingRunParams)
+class BaseLoadingRunWithProjectInfoParams(luigi.Task):
+    project_guids = luigi.ListParameter()
+    project_remap_paths = luigi.ListParameter()
+    project_pedigree_paths = luigi.ListParameter()

v03_pipeline/lib/tasks/trigger_hail_backend_reload.py

@@ -0,0 +1,31 @@
+import luigi
+import luigi.util
+import requests
+
+from v03_pipeline.lib.logger import get_logger
+from v03_pipeline.lib.model import Env
+from v03_pipeline.lib.tasks import UpdateVariantAnnotationsTableWithNewSamplesTask
+from v03_pipeline.lib.tasks.base.base_project_info_params import (
+    BaseLoadingRunWithProjectInfoParams,
+)
+
+logger = get_logger(__name__)
+
+
+@luigi.util.inherits(BaseLoadingRunWithProjectInfoParams)
+class TriggerHailBackendReload(luigi.Task):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.done = False
+
+    def requires(self):
+        return self.clone(UpdateVariantAnnotationsTableWithNewSamplesTask)
+
+    def run(self):
+        url = f'{Env.HAIL_BACKEND_SERVICE_HOSTNAME}:{Env.HAIL_BACKEND_SERVICE_PORT}/reload_globals'
+        res = requests.post(url, headers={'From': 'pipeline-runner'}, timeout=300)
+        res.raise_for_status()
+        self.done = True
+
+    def complete(self):
+        return self.done

v03_pipeline/lib/tasks/trigger_hail_backend_reload_test.py

@@ -0,0 +1,90 @@
+from unittest.mock import Mock, patch
+
+import hail as hl
+import luigi.worker
+import requests
+
+from v03_pipeline.lib.misc.io import remap_pedigree_hash
+from v03_pipeline.lib.model import DatasetType, ReferenceGenome, SampleType
+from v03_pipeline.lib.paths import variant_annotations_table_path
+from v03_pipeline.lib.tasks.trigger_hail_backend_reload import TriggerHailBackendReload
+from v03_pipeline.lib.test.mocked_dataroot_testcase import MockedDatarootTestCase
+
+TEST_VCF = 'v03_pipeline/var/test/callsets/1kg_30variants.vcf'
+TEST_REMAP = 'v03_pipeline/var/test/remaps/test_remap_1.tsv'
+TEST_PEDIGREE_3 = 'v03_pipeline/var/test/pedigrees/test_pedigree_3.tsv'
+
+
+class TriggerHailBackendReloadTestCase(MockedDatarootTestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        annotations_ht = hl.Table.parallelize(
+            [],
+            hl.tstruct(
+                locus=hl.tlocus(ReferenceGenome.GRCh38.value),
+                alleles=hl.tarray(hl.tstr),
+            ),
+            key=('locus', 'alleles'),
+            globals=hl.Struct(
+                paths=hl.Struct(),
+                versions=hl.Struct(),
+                enums=hl.Struct(),
+                updates={
+                    hl.Struct(
+                        callset=TEST_VCF,
+                        project_guid='R0113_test_project',
+                        remap_pedigree_hash=hl.eval(
+                            remap_pedigree_hash(TEST_REMAP, TEST_PEDIGREE_3),
+                        ),
+                    ),
+                },
+                migrations=hl.empty_array(hl.tstr),
+            ),
+        )
+        annotations_ht.write(
+            variant_annotations_table_path(
+                ReferenceGenome.GRCh38,
+                DatasetType.SNV_INDEL,
+            ),
+        )
+
+    @patch.object(requests, 'post')
+    def test_success(self, mock_post: Mock):
+        mock_resp = requests.models.Response()
+        mock_resp.status_code = 200
+        mock_post.return_value = mock_resp
+
+        worker = luigi.worker.Worker()
+        task = TriggerHailBackendReload(
+            reference_genome=ReferenceGenome.GRCh38,
+            dataset_type=DatasetType.SNV_INDEL,
+            sample_type=SampleType.WES,
+            callset_path=TEST_VCF,
+            project_guids=['R0113_test_project'],
+            run_id='manual__2024-09-20',
+            project_remap_paths=[TEST_REMAP],
+            project_pedigree_paths=[TEST_PEDIGREE_3],
+        )
+        worker.add(task)
+        worker.run()
+        self.assertTrue(task.complete())
+
+    @patch.object(requests, 'post')
+    def test_failure(self, mock_post: Mock):
+        mock_resp = requests.models.Response()
+        mock_resp.status_code = 500
+        mock_post.return_value = mock_resp
+
+        worker = luigi.worker.Worker()
+        task = TriggerHailBackendReload(
+            reference_genome=ReferenceGenome.GRCh38,
+            dataset_type=DatasetType.SNV_INDEL,
+            sample_type=SampleType.WES,
+            callset_path=TEST_VCF,
+            project_guids=['R0113_test_project'],
+            run_id='manual__2024-09-20',
+            project_remap_paths=[TEST_REMAP],
+            project_pedigree_paths=[TEST_PEDIGREE_3],
+        )
+        worker.add(task)
+        self.assertFalse(task.complete())

v03_pipeline/lib/tasks/update_lookup_table.py

@@ -10,7 +10,9 @@
 )
 from v03_pipeline.lib.model.constants import PROJECTS_EXCLUDED_FROM_LOOKUP
 from v03_pipeline.lib.paths import remapped_and_subsetted_callset_path
-from v03_pipeline.lib.tasks.base.base_loading_run_params import BaseLoadingRunParams
+from v03_pipeline.lib.tasks.base.base_project_info_params import (
+    BaseLoadingRunWithProjectInfoParams,
+)
 from v03_pipeline.lib.tasks.base.base_update_lookup_table import (
     BaseUpdateLookupTableTask,
 )
@@ -19,12 +21,8 @@
 )
 
 
-@luigi.util.inherits(BaseLoadingRunParams)
+@luigi.util.inherits(BaseLoadingRunWithProjectInfoParams)
 class UpdateLookupTableTask(BaseUpdateLookupTableTask):
-    project_guids = luigi.ListParameter()
-    project_remap_paths = luigi.ListParameter()
-    project_pedigree_paths = luigi.ListParameter()
-
     def complete(self) -> bool:
         return super().complete() and hl.eval(
             hl.bind(

v03_pipeline/lib/tasks/update_variant_annotations_table_with_new_samples.py

@@ -9,21 +9,19 @@
     lookup_table_path,
     new_variants_table_path,
 )
-from v03_pipeline.lib.tasks.base.base_loading_run_params import BaseLoadingRunParams
+from v03_pipeline.lib.tasks.base.base_project_info_params import (
+    BaseLoadingRunWithProjectInfoParams,
+)
 from v03_pipeline.lib.tasks.base.base_update_variant_annotations_table import (
     BaseUpdateVariantAnnotationsTableTask,
 )
 from v03_pipeline.lib.tasks.write_new_variants_table import WriteNewVariantsTableTask
 
 
-@luigi.util.inherits(BaseLoadingRunParams)
+@luigi.util.inherits(BaseLoadingRunWithProjectInfoParams)
 class UpdateVariantAnnotationsTableWithNewSamplesTask(
     BaseUpdateVariantAnnotationsTableTask,
 ):
-    project_guids = luigi.ListParameter()
-    project_remap_paths = luigi.ListParameter()
-    project_pedigree_paths = luigi.ListParameter()
-
     def requires(self) -> list[luigi.Task]:
         return [
             *super().requires(),

v03_pipeline/lib/tasks/write_metadata_for_run.py

@@ -5,19 +5,17 @@
 import luigi.util
 
 from v03_pipeline.lib.paths import metadata_for_run_path
-from v03_pipeline.lib.tasks.base.base_loading_run_params import BaseLoadingRunParams
+from v03_pipeline.lib.tasks.base.base_project_info_params import (
+    BaseLoadingRunWithProjectInfoParams,
+)
 from v03_pipeline.lib.tasks.files import GCSorLocalTarget
 from v03_pipeline.lib.tasks.write_remapped_and_subsetted_callset import (
     WriteRemappedAndSubsettedCallsetTask,
 )
 
 
-@luigi.util.inherits(BaseLoadingRunParams)
+@luigi.util.inherits(BaseLoadingRunWithProjectInfoParams)
 class WriteMetadataForRunTask(luigi.Task):
-    project_guids = luigi.ListParameter()
-    project_remap_paths = luigi.ListParameter()
-    project_pedigree_paths = luigi.ListParameter()
-
     def output(self) -> luigi.Target:
         return GCSorLocalTarget(
             metadata_for_run_path(

v03_pipeline/lib/tasks/write_new_variants_table.py

@@ -24,7 +24,9 @@
     load_gencode_ensembl_to_refseq_id,
     load_gencode_gene_symbol_to_gene_id,
 )
-from v03_pipeline.lib.tasks.base.base_loading_run_params import BaseLoadingRunParams
+from v03_pipeline.lib.tasks.base.base_project_info_params import (
+    BaseLoadingRunWithProjectInfoParams,
+)
 from v03_pipeline.lib.tasks.base.base_write import BaseWriteTask
 from v03_pipeline.lib.tasks.files import GCSorLocalTarget
 from v03_pipeline.lib.tasks.reference_data.update_variant_annotations_table_with_updated_reference_dataset import (
@@ -43,12 +45,8 @@
 GENCODE_FOR_VEP_RELEASE = 44
 
 
-@luigi.util.inherits(BaseLoadingRunParams)
+@luigi.util.inherits(BaseLoadingRunWithProjectInfoParams)
 class WriteNewVariantsTableTask(BaseWriteTask):
-    project_guids = luigi.ListParameter()
-    project_remap_paths = luigi.ListParameter()
-    project_pedigree_paths = luigi.ListParameter()
-
     @property
     def annotation_dependencies(self) -> dict[str, hl.Table]:
         deps = get_rdc_annotation_dependencies(self.dataset_type, self.reference_genome)