Add MaxRetries parameter to CloudRunJob block. (#235)

* 🔨 add MaxRetries parameter to CloudRunJob block * 📝 update CHANGELOG.md file * ♻️ refactor max_retries field --------- Co-authored-by: nate nowack <[email protected]>
PrefectHQ · Feb 27, 2024 · 0898eeb · 0898eeb
1 parent ea32a84
commit 0898eeb
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/prefect_gcp/cloud_run.py b/prefect_gcp/cloud_run.py
@@ -299,6 +299,16 @@ class CloudRunJob(Infrastructure):
             "before raising an exception."
         ),
     )
+    max_retries: Optional[int] = Field(
+        default=3,
+        ge=0,
+        le=10,
+        title="Max Retries",
+        description=(
+            "The maximum retries setting specifies the number of times a task is "
+            "allowed to restart in case of failure before being failed permanently."
+        ),
+    )
     # For private use
     _job_name: str = None
     _execution: Optional[Execution] = None
@@ -679,6 +689,7 @@ def _jobs_body(self) -> dict:
                             "spec": {
                                 "containers": containers,
                                 "timeoutSeconds": timeout_seconds,
+                                "maxRetries": self.max_retries,
                             }  # TaskSpec
                         }
                     },

diff --git a/tests/test_cloud_run.py b/tests/test_cloud_run.py
@@ -520,6 +520,15 @@ def test_timeout_added_correctly(self, cloud_run_job):
             "timeoutSeconds"
         ] == str(timeout)
 
+    def test_max_retries_added_correctly(self, cloud_run_job):
+        max_retries = 2
+        cloud_run_job.max_retries = max_retries
+        result = cloud_run_job._jobs_body()
+        assert (
+            result["spec"]["template"]["spec"]["template"]["spec"]["maxRetries"]
+            == max_retries
+        )
+
     def test_vpc_connector_name_added_correctly(self, cloud_run_job):
         cloud_run_job.vpc_connector_name = "vpc_name"
         result = cloud_run_job._jobs_body()