bugfix cluster and job state

Author: bpblanken

v03_pipeline/lib/tasks/dataproc/base_run_job_on_dataproc.py

@@ -1,6 +1,7 @@
 import time
 
 import google.api_core.exceptions
+import google.cloud.dataproc_v1.types.jobs
 import luigi
 from google.cloud import dataproc_v1 as dataproc
 
@@ -14,8 +15,6 @@
 )
 from v03_pipeline.lib.tasks.dataproc.misc import get_cluster_name, to_kebab_str_args
 
-DONE_STATE = 'DONE'
-ERROR_STATE = 'ERROR'
 SEQR_PIPELINE_RUNNER_BUILD = f'gs://seqr-pipeline-runner-builds/{Env.DEPLOYMENT_TYPE}/{Env.PIPELINE_RUNNER_APP_VERSION}'
 TIMEOUT_S = 172800  # 2 days
 
@@ -57,11 +56,17 @@ def complete(self) -> bool:
             )
         except google.api_core.exceptions.NotFound:
             return False
-        if job.status.state == ERROR_STATE:
-            msg = f'Job {self.task.task_family}-{self.run_id} entered ERROR state'
+        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,
+        }:
+            msg = f'Job {self.task.task_family}-{self.run_id} entered {job.status.state!s} state'
             logger.error(msg)
             logger.error(job.status.details)
-        return job.status.state == DONE_STATE
+        return (
+            job.status.state == google.cloud.dataproc_v1.types.jobs.JobStatus.State.DONE
+        )
 
     def run(self):
         operation = self.client.submit_job_as_operation(

v03_pipeline/lib/tasks/dataproc/create_dataproc_cluster.py

@@ -1,6 +1,7 @@
 import time
 
 import google.api_core.exceptions
+import google.cloud.dataproc_v1.types.clusters
 import hail as hl
 import luigi
 from google.cloud import dataproc_v1 as dataproc
@@ -15,11 +16,9 @@
 from v03_pipeline.lib.tasks.dataproc.misc import get_cluster_name
 
 DEBIAN_IMAGE = '2.1.33-debian11'
-ERROR_STATE = 'ERROR'
 HAIL_VERSION = hl.version().split('-')[0]
 INSTANCE_TYPE = 'n1-highmem-8'
 PKGS = '|'.join(pip_freeze.freeze())
-RUNNING_STATE = 'RUNNING'
 TIMEOUT_S = 900
 
 logger = get_logger(__name__)
@@ -160,11 +159,20 @@ def complete(self) -> bool:
             )
         except google.api_core.exceptions.NotFound:
             return False
-        if cluster.status.state == ERROR_STATE:
-            msg = f'Cluster {cluster.cluster_name} entered ERROR state'
+        if cluster.status.state in {
+            google.cloud.dataproc_v1.types.clusters.ClusterStatus.State.UNKNOWN,
+            google.cloud.dataproc_v1.types.clusters.ClusterStatus.State.ERROR,
+            google.cloud.dataproc_v1.types.clusters.ClusterStatus.State.ERROR_DUE_TO_UPDATE,
+        }:
+            msg = (
+                f'Cluster {cluster.cluster_name} entered {cluster.status.state!s} state'
+            )
             logger.error(msg)
         # This will return False when the cluster is "CREATING"
-        return cluster.status.state == RUNNING_STATE
+        return (
+            cluster.status.state
+            == google.cloud.dataproc_v1.types.clusters.ClusterStatus.State.RUNNING
+        )
 
     def run(self):
         if not Env.GCLOUD_PROJECT or not Env.GCLOUD_REGION or not Env.GCLOUD_ZONE:

v03_pipeline/lib/tasks/dataproc/create_dataproc_cluster_test.py

@@ -3,6 +3,7 @@
 from unittest.mock import Mock, call, patch
 
 import google.api_core.exceptions
+import google.cloud.dataproc_v1.types.clusters
 import luigi
 
 from v03_pipeline.lib.misc.gcp import SOCIAL_AUTH_GOOGLE_OAUTH2_SCOPE
@@ -48,7 +49,9 @@ def test_spinup_cluster_already_exists_failed(
     ) -> None:
         mock_client = mock_cluster_controller.return_value
         mock_client.get_cluster.return_value = SimpleNamespace(
-            status=SimpleNamespace(state='ERROR'),
+            status=SimpleNamespace(
+                state=google.cloud.dataproc_v1.types.clusters.ClusterStatus.State.ERROR
+            ),
             cluster_name='abc',
         )
         mock_client.create_cluster.side_effect = (
@@ -71,7 +74,9 @@ def test_spinup_cluster_already_exists_success(
     ) -> None:
         mock_client = mock_cluster_controller.return_value
         mock_client.get_cluster.return_value = SimpleNamespace(
-            status=SimpleNamespace(state='RUNNING'),
+            status=SimpleNamespace(
+                state=google.cloud.dataproc_v1.types.clusters.ClusterStatus.State.RUNNING
+            ),
         )
         mock_client.create_cluster.side_effect = (
             google.api_core.exceptions.AlreadyExists('cluster exists')

v03_pipeline/lib/tasks/dataproc/run_pipeline_on_dataproc_test.py

@@ -3,6 +3,7 @@
 from unittest.mock import Mock, call, patch
 
 import google.api_core.exceptions
+import google.cloud.dataproc_v1.types.jobs
 import luigi
 
 from v03_pipeline.lib.model import DatasetType, ReferenceGenome, SampleType
@@ -30,7 +31,7 @@ def test_job_already_exists_failed(
         mock_client = mock_job_controller_client.return_value
         mock_client.get_job.return_value = SimpleNamespace(
             status=SimpleNamespace(
-                state='ERROR',
+                state=google.cloud.dataproc_v1.types.jobs.JobStatus.State.ERROR,
                 details='Google Cloud Dataproc Agent reports job failure. If logs are available, they can be found at...',
             ),
         )
@@ -54,7 +55,7 @@ def test_job_already_exists_failed(
         mock_logger.error.assert_has_calls(
             [
                 call(
-                    'Job RunPipelineTask-manual__2024-04-03 entered ERROR state',
+                    'Job RunPipelineTask-manual__2024-04-03 entered State.ERROR state',
                 ),
             ],
         )
@@ -67,7 +68,7 @@ def test_job_already_exists_success(
         mock_create_dataproc_cluster.return_value = MockCompleteTask()
         mock_client = mock_job_controller_client.return_value
         mock_client.get_job.return_value = SimpleNamespace(
-            status=SimpleNamespace(state='DONE'),
+            status=SimpleNamespace(state=google.cloud.dataproc_v1.types.jobs.JobStatus.State.DONE),
         )
         worker = luigi.worker.Worker()
         task = RunPipelineOnDataprocTask(
@@ -135,7 +136,7 @@ def test_job_success(
                 'job not found',
             ),
             SimpleNamespace(
-                status=SimpleNamespace(state='DONE'),
+                status=SimpleNamespace(state=google.cloud.dataproc_v1.types.jobs.JobStatus.State.DONE),
             ),
         ]
         operation = mock_client.submit_job_as_operation.return_value