feat: restart resilience (#1188)

* retry resilience

* more restart resilience

* this works!

* refactor

* job

* fix tests

* ruff

Author: bpblanken

v03_pipeline/api/app_test.py

@@ -98,6 +98,7 @@ async def test_loading_pipeline_enqueue(self):
                     'skip_check_sex_and_relatedness': False,
                     'skip_validation': False,
                     'skip_expect_tdr_metrics': False,
+                    'attempt_id': 0,
                 },
             },
         )

v03_pipeline/api/model.py

@@ -1,12 +1,13 @@
 import hailtop.fs as hfs
-from pydantic import AliasChoices, BaseModel, Field, field_validator
+from pydantic import AliasChoices, BaseModel, Field, conint, field_validator
 
 from v03_pipeline.lib.core import DatasetType, ReferenceGenome, SampleType
 
+MAX_LOADING_PIPELINE_ATTEMPTS = 3
 VALID_FILE_TYPES = ['vcf', 'vcf.gz', 'vcf.bgz', 'mt']
 
 
-class PipelineRunnerRequest(BaseModel, frozen=True):
+class PipelineRunnerRequest(BaseModel):
     request_type: str
 
     def __init_subclass__(cls, **kwargs):
@@ -15,6 +16,7 @@ def __init_subclass__(cls, **kwargs):
 
 
 class LoadingPipelineRequest(PipelineRunnerRequest):
+    attempt_id: conint(ge=0, le=MAX_LOADING_PIPELINE_ATTEMPTS - 1) = 0
     callset_path: str
     project_guids: list[str] = Field(
         min_length=1,
@@ -28,6 +30,12 @@ class LoadingPipelineRequest(PipelineRunnerRequest):
     skip_check_sex_and_relatedness: bool = False
     skip_expect_tdr_metrics: bool = False
 
+    def incr_attempt(self):
+        if self.attempt_id == (MAX_LOADING_PIPELINE_ATTEMPTS - 1):
+            return False
+        self.attempt_id += 1
+        return True
+
     @field_validator('callset_path')
     @classmethod
     def check_valid_callset_path(cls, callset_path: str) -> str:

v03_pipeline/api/model_test.py

@@ -20,6 +20,7 @@ def test_valid_loading_pipeline_requests(self) -> None:
         self.assertEqual(lpr.reference_genome, ReferenceGenome.GRCh38)
         self.assertEqual(lpr.project_guids, ['project_a'])
         self.assertEqual(lpr.request_type, 'LoadingPipelineRequest')
+        self.assertEqual(lpr.attempt_id, 0)
 
         # Test wildcard VCF
         raw_request['callset_path'] = CALLSET_PATH.replace(
@@ -35,12 +36,13 @@ def test_invalid_loading_pipeline_requests(self) -> None:
             'sample_type': 'BLENDED',
             'reference_genome': ReferenceGenome.GRCh38.value,
             'dataset_type': DatasetType.SNV_INDEL.value,
+            'attempt_id': 5,
         }
         with self.assertRaises(ValueError) as cm:
             LoadingPipelineRequest.model_validate(raw_request)
         self.assertTrue(
             str(cm.exception).startswith(
-                '3 validation errors for LoadingPipelineRequest',
+                '4 validation errors for LoadingPipelineRequest',
             ),
         )
 

v03_pipeline/api/request_handlers.py

@@ -12,11 +12,10 @@
     PipelineRunnerRequest,
     RebuildGtStatsRequest,
 )
-from v03_pipeline.lib.core import DatasetType, FeatureFlag, ReferenceGenome
+from v03_pipeline.lib.core import DatasetType, ReferenceGenome
 from v03_pipeline.lib.logger import get_logger
 from v03_pipeline.lib.misc.clickhouse import (
     delete_family_guids,
-    load_complete_run,
     rebuild_gt_stats,
 )
 from v03_pipeline.lib.misc.retry import retry
@@ -75,39 +74,22 @@ def run_loading_pipeline(
     local_scheduler: bool,
     *_: Any,
 ):
-    for attempt_id in range(3):
-        luigi_task_result = luigi.build(
-            [
-                WriteSuccessFileTask(
-                    run_id=run_id,
-                    attempt_id=attempt_id,
-                    **lpr.model_dump(exclude='request_type'),
-                ),
-            ],
-            detailed_summary=True,
-            local_scheduler=local_scheduler,
-        )
-        if luigi_task_result.status in {
-            luigi.execution_summary.LuigiStatusCode.SUCCESS,
-            luigi.execution_summary.LuigiStatusCode.SUCCESS_WITH_RETRY,
-        }:
-            break
-    else:
-        raise RuntimeError(luigi_task_result.status.value[1])
-    if FeatureFlag.CLICKHOUSE_LOADER_DISABLED:
-        project_guids, family_guids = fetch_run_metadata(
-            lpr.reference_genome,
-            lpr.dataset_type,
-            run_id,
-        )
-        load_complete_run(
-            lpr.reference_genome,
-            lpr.dataset_type,
-            run_id,
-            project_guids,
-            family_guids,
-        )
-        write_success_file(lpr.reference_genome, lpr.dataset_type, run_id)
+    luigi_task_result = luigi.build(
+        [
+            WriteSuccessFileTask(
+                run_id=run_id,
+                **lpr.model_dump(exclude='request_type'),
+            ),
+        ],
+        detailed_summary=True,
+        local_scheduler=local_scheduler,
+    )
+    if luigi_task_result.status in {
+        luigi.execution_summary.LuigiStatusCode.SUCCESS,
+        luigi.execution_summary.LuigiStatusCode.SUCCESS_WITH_RETRY,
+    }:
+        return
+    raise RuntimeError(luigi_task_result.status.value[1])
 
 
 def run_delete_families(dpr: DeleteFamiliesRequest, run_id: str, *_: Any):

v03_pipeline/bin/pipeline_worker.py

@@ -22,6 +22,7 @@
 from v03_pipeline.lib.paths import (
     loading_pipeline_deadletter_queue_dir,
     loading_pipeline_deadletter_queue_path,
+    loading_pipeline_queue_path,
 )
 
 logger = get_logger(__name__)
@@ -65,13 +66,19 @@ def process_queue(local_scheduler=False):
             return
         prr, run_id = parse_latest_queue_path(latest_queue_path)
         REQUEST_HANDLER_MAP[type(prr)](prr, run_id, local_scheduler)
+        os.remove(latest_queue_path)
         safe_post_to_slack_success(
             run_id,
             prr,
         )
     except Exception as e:
         logger.exception('Unhandled Exception')
-        if run_id is not None:
+        if run_id is None:
+            return
+        if hasattr(prr, 'attempt_id') and prr.incr_attempt():
+            with open(loading_pipeline_queue_path(run_id), 'w') as f:
+                f.write(prr.model_dump_json())
+        else:
             safe_post_to_slack_failure(
                 run_id,
                 prr,
@@ -80,16 +87,13 @@ def process_queue(local_scheduler=False):
             os.makedirs(loading_pipeline_deadletter_queue_dir(), exist_ok=True)
             with open(loading_pipeline_deadletter_queue_path(run_id), 'w') as f:
                 f.write(prr.model_dump_json())
-    finally:
-        if latest_queue_path is not None and os.path.exists(latest_queue_path):
-            os.remove(latest_queue_path)
-        logger.info('Looking for more work')
-        time.sleep(1)
 
 
 def main():
     while True:
         process_queue()
+        logger.info('Looking for more work')
+        time.sleep(1)
 
 
 if __name__ == '__main__':

v03_pipeline/bin/pipeline_worker_test.py

@@ -58,7 +58,7 @@ def test_process_queue(
             json.dump(raw_request, f)
         process_queue(local_scheduler=True)
         mock_safe_post_to_slack.assert_called_once_with(
-            ':white_check_mark: Pipeline Runner Request Success! :white_check_mark:\nRun ID: 20250916-200704-123456\n```{\n    "callset_path": "v03_pipeline/var/test/callsets/1kg_30variants.vcf",\n    "dataset_type": "SNV_INDEL",\n    "project_guids": [\n        "project_a"\n    ],\n    "reference_genome": "GRCh38",\n    "request_type": "LoadingPipelineRequest",\n    "sample_type": "WGS",\n    "skip_check_sex_and_relatedness": false,\n    "skip_expect_tdr_metrics": false,\n    "skip_validation": false\n}```',
+            ':white_check_mark: Pipeline Runner Request Success! :white_check_mark:\nRun ID: 20250916-200704-123456\n```{\n    "attempt_id": 0,\n    "callset_path": "v03_pipeline/var/test/callsets/1kg_30variants.vcf",\n    "dataset_type": "SNV_INDEL",\n    "project_guids": [\n        "project_a"\n    ],\n    "reference_genome": "GRCh38",\n    "request_type": "LoadingPipelineRequest",\n    "sample_type": "WGS",\n    "skip_check_sex_and_relatedness": false,\n    "skip_expect_tdr_metrics": false,\n    "skip_validation": false\n}```',
         )
 
     @patch('v03_pipeline.lib.misc.slack._safe_post_to_slack')
@@ -92,13 +92,17 @@ def test_process_failure(
         ) as f:
             json.dump(raw_request, f)
         process_queue(local_scheduler=True)
+        process_queue(local_scheduler=True)
+        process_queue(local_scheduler=True)
         mock_safe_post_to_slack.assert_called_once_with(
-            ':failed: Pipeline Runner Request Failed :failed:\nRun ID: 20250918-200704-123456\n```{\n    "callset_path": "v03_pipeline/var/test/callsets/1kg_30variants.vcf",\n    "dataset_type": "SNV_INDEL",\n    "project_guids": [\n        "project_a"\n    ],\n    "reference_genome": "GRCh38",\n    "request_type": "LoadingPipelineRequest",\n    "sample_type": "WGS",\n    "skip_check_sex_and_relatedness": false,\n    "skip_expect_tdr_metrics": false,\n    "skip_validation": false\n}```\nReason: there were failed tasks',
+            ':failed: Pipeline Runner Request Failed :failed:\nRun ID: 20250918-200704-123456\n```{\n    "attempt_id": 2,\n    "callset_path": "v03_pipeline/var/test/callsets/1kg_30variants.vcf",\n    "dataset_type": "SNV_INDEL",\n    "project_guids": [\n        "project_a"\n    ],\n    "reference_genome": "GRCh38",\n    "request_type": "LoadingPipelineRequest",\n    "sample_type": "WGS",\n    "skip_check_sex_and_relatedness": false,\n    "skip_expect_tdr_metrics": false,\n    "skip_validation": false\n}```\nReason: there were failed tasks',
         )
         with open(
             os.path.join(
                 loading_pipeline_deadletter_queue_dir(),
                 'request_20250918-200704-123456.json',
             ),
         ) as f:
-            self.assertEqual(json.load(f)['request_type'], 'LoadingPipelineRequest')
+            r = json.load(f)
+            self.assertEqual(r['request_type'], 'LoadingPipelineRequest')
+            self.assertEqual(r['attempt_id'], 2)

v03_pipeline/lib/core/environment.py

@@ -3,7 +3,7 @@
 from typing import Literal
 
 # NB: using os.environ.get inside the dataclass defaults gives a lint error.
-LOCAL_DISK_MOUNT_PATH = os.environ.get('LOCAL_DISK_MOUNT_PATH', '/var/seqr')
+LOCAL_DISK_MOUNT_DIR = os.environ.get('LOCAL_DISK_MOUNT_DIR', '/var/seqr')
 HAIL_TMP_DIR = os.environ.get('HAIL_TMP_DIR', '/tmp')  # noqa: S108
 PIPELINE_DATA_DIR = os.environ.get(
     'PIPELINE_DATA_DIR',
@@ -78,7 +78,7 @@ class Env:
     GCLOUD_ZONE: str | None = GCLOUD_ZONE
     GCLOUD_REGION: str | None = GCLOUD_REGION
     HAIL_TMP_DIR: str = HAIL_TMP_DIR
-    LOCAL_DISK_MOUNT_PATH: str = LOCAL_DISK_MOUNT_PATH
+    LOCAL_DISK_MOUNT_DIR: str = LOCAL_DISK_MOUNT_DIR
     PIPELINE_DATA_DIR: str = PIPELINE_DATA_DIR
     LOADING_DATASETS_DIR: str = LOADING_DATASETS_DIR
     LOADING_QUEUE_LIMIT: int = LOADING_QUEUE_LIMIT

v03_pipeline/lib/paths.py

@@ -406,15 +406,15 @@ def loading_pipeline_queue_dir() -> str:
     Returns the directory where loading pipeline requests are queued.
     """
     return os.path.join(
-        Env.LOCAL_DISK_MOUNT_PATH,
+        Env.LOCAL_DISK_MOUNT_DIR,
         'loading_pipeline_queue',
     )
 
 
 # https://en.wikipedia.org/wiki/Dead_letter_queue
 def loading_pipeline_deadletter_queue_dir() -> str:
     return os.path.join(
-        Env.LOCAL_DISK_MOUNT_PATH,
+        Env.LOCAL_DISK_MOUNT_DIR,
         'loading_pipeline_deadletter_queue',
     )
 

v03_pipeline/lib/tasks/dataproc/base_run_job_on_dataproc.py

@@ -15,6 +15,11 @@
 )
 from v03_pipeline.lib.tasks.dataproc.misc import get_cluster_name, to_kebab_str_args
 
+FAILURE_STATUSES = {
+    google.cloud.dataproc_v1.types.jobs.JobStatus.State.CANCELLED,
+    google.cloud.dataproc_v1.types.jobs.JobStatus.State.ERROR,
+    google.cloud.dataproc_v1.types.jobs.JobStatus.State.ATTEMPT_FAILURE,
+}
 SEQR_PIPELINE_RUNNER_BUILD = f'gs://seqr-pipeline-runner-builds/{Env.DEPLOYMENT_TYPE}/{Env.PIPELINE_RUNNER_APP_VERSION}'
 TIMEOUT_S = 172800  # 2 days
 
@@ -44,7 +49,9 @@ def job_id(self):
     def requires(self) -> [luigi.Task]:
         return [self.clone(CreateDataprocClusterTask)]
 
-    def complete(self) -> bool:
+    def safely_get_job(
+        self,
+    ):
         try:
             job = self.client.get_job(
                 request={
@@ -54,12 +61,15 @@ def complete(self) -> bool:
                 },
             )
         except google.api_core.exceptions.NotFound:
+            return None
+        else:
+            return job
+
+    def complete(self) -> bool:
+        job = self.safely_get_job()
+        if not job:
             return False
-        if job.status.state in {
-            google.cloud.dataproc_v1.types.jobs.JobStatus.State.CANCELLED,
-            google.cloud.dataproc_v1.types.jobs.JobStatus.State.ERROR,
-            google.cloud.dataproc_v1.types.jobs.JobStatus.State.ATTEMPT_FAILURE,
-        }:
+        if job.status.state in FAILURE_STATUSES:
             msg = f'Job {self.job_id} entered {job.status.state.name} state'
             logger.error(msg)
             logger.error(job.status.details)
@@ -68,43 +78,52 @@ def complete(self) -> bool:
         )
 
     def run(self):
-        operation = self.client.submit_job_as_operation(
-            request={
-                'project_id': Env.GCLOUD_PROJECT,
-                'region': Env.GCLOUD_REGION,
-                'job': {
-                    'reference': {
-                        'job_id': self.job_id,
-                    },
-                    'placement': {
-                        'cluster_name': get_cluster_name(
-                            self.reference_genome,
-                            self.run_id,
-                        ),
-                    },
-                    'pyspark_job': {
-                        'main_python_file_uri': f'{SEQR_PIPELINE_RUNNER_BUILD}/bin/run_task.py',
-                        'args': [
-                            self.task.task_family,
-                            '--local-scheduler',
-                            *to_kebab_str_args(self),
-                        ],
-                        'python_file_uris': [
-                            f'{SEQR_PIPELINE_RUNNER_BUILD}/pyscripts.zip',
-                        ],
+        job = self.safely_get_job()
+        if not job:
+            self.client.submit_job_as_operation(
+                request={
+                    'project_id': Env.GCLOUD_PROJECT,
+                    'region': Env.GCLOUD_REGION,
+                    'job': {
+                        'reference': {
+                            'job_id': self.job_id,
+                        },
+                        'placement': {
+                            'cluster_name': get_cluster_name(
+                                self.reference_genome,
+                                self.run_id,
+                            ),
+                        },
+                        'pyspark_job': {
+                            'main_python_file_uri': f'{SEQR_PIPELINE_RUNNER_BUILD}/bin/run_task.py',
+                            'args': [
+                                self.task.task_family,
+                                '--local-scheduler',
+                                *to_kebab_str_args(self),
+                            ],
+                            'python_file_uris': [
+                                f'{SEQR_PIPELINE_RUNNER_BUILD}/pyscripts.zip',
+                            ],
+                        },
                     },
                 },
-            },
-        )
+            )
         wait_s = 0
         while wait_s < TIMEOUT_S:
-            if operation.done():
-                operation.result()  # Will throw on failure!
-                msg = f'Finished {self.job_id}'
+            job = self.safely_get_job()
+            if (
+                job.status.state
+                == google.cloud.dataproc_v1.types.jobs.JobStatus.State.DONE
+            ):
+                msg = f'Job {self.job_id} is complete'
                 logger.info(msg)
                 break
+            if job.status.state in FAILURE_STATUSES:
+                msg = f'Job {self.job_id} entered {job.status.state.name} state'
+                logger.error(msg)
+                raise RuntimeError(msg)
             logger.info(
-                f'Waiting for job completion {self.job_id}',
+                f'Waiting for Job completion {self.job_id}',
             )
             time.sleep(3)
             wait_s += 3