E2e configurable heartbeat_healthy_threshold

Author: Json-Andriopoulos

src/zenml/config/step_configurations.py

@@ -221,12 +221,13 @@ class StepConfigurationUpdate(FrozenBaseModel):
         "run inline unless a step operator or docker/resource settings "
         "are configured. This is only applicable for dynamic pipelines.",
     )
-    heartbeat_healthy_threshold: int | None = Field(
-        default=None,
+    heartbeat_healthy_threshold: int = Field(
+        default=30,
         description="The amount of time (in minutes) that a running step "
-        "has not received heartbeat and is considered healthy. Set null value"
-        "disable healthiness checks via heartbeat.",
+        "has not received heartbeat and is considered healthy. By default, "
+        "set to the maximum value (30 minutes).",
         ge=1,
+        le=30,
     )
 
     outputs: Mapping[str, PartialArtifactConfiguration] = {}

src/zenml/steps/base_step.py

@@ -133,6 +133,7 @@ def __init__(
         substitutions: Optional[Dict[str, str]] = None,
         cache_policy: Optional[CachePolicyOrString] = None,
         runtime: Optional[StepRuntime] = None,
+        heartbeat_healthy_threshold: Optional[int] = None,
     ) -> None:
         """Initializes a step.
 
@@ -170,6 +171,9 @@ def __init__(
                 run inline unless a step operator or docker/resource settings
                 are configured. This is only applicable for dynamic
                 pipelines.
+            heartbeat_healthy_threshold: The amount of time (in minutes) that a
+                running step has not received heartbeat and is considered healthy.
+                By default, set to the maximum value (30 minutes).",
         """
         from zenml.config.step_configurations import PartialStepConfiguration
 
@@ -237,6 +241,7 @@ def __init__(
             substitutions=substitutions,
             cache_policy=cache_policy,
             runtime=runtime,
+            heartbeat_healthy_threshold=heartbeat_healthy_threshold,
         )
 
         notebook_utils.try_to_save_notebook_cell_code(self.source_object)
@@ -876,6 +881,7 @@ def configure(
         cache_policy: Optional[CachePolicyOrString] = None,
         runtime: Optional[StepRuntime] = None,
         merge: bool = True,
+        heartbeat_healthy_threshold: Optional[int] = None,
     ) -> T:
         """Configures the step.
 
@@ -925,6 +931,9 @@ def configure(
                 configurations. If `False` the given configurations will
                 overwrite all existing ones. See the general description of this
                 method for an example.
+            heartbeat_healthy_threshold: The amount of time (in minutes) that a
+                running step has not received heartbeat and is considered healthy.
+                By default, set to the maximum value (30 minutes).",
 
         Returns:
             The step instance that this method was called on.
@@ -999,6 +1008,7 @@ def _convert_to_tuple(value: Any) -> Tuple[Source, ...]:
                 "substitutions": substitutions,
                 "cache_policy": cache_policy,
                 "runtime": runtime,
+                "heartbeat_healthy_threshold": heartbeat_healthy_threshold,
             }
         )
         config = StepConfigurationUpdate(**values)

src/zenml/steps/heartbeat.py

@@ -185,9 +185,6 @@ def is_heartbeat_unhealthy(step_run: "StepRunResponse") -> bool:
     if not step_run.spec.enable_heartbeat:
         return False
 
-    if not step_run.config.heartbeat_healthy_threshold:
-        return False
-
     if step_run.status.is_finished:
         return False
 

src/zenml/steps/step_decorator.py

@@ -107,6 +107,7 @@ def step(
     substitutions: Optional[Dict[str, str]] = None,
     cache_policy: Optional["CachePolicyOrString"] = None,
     runtime: Optional[StepRuntime] = None,
+    heartbeat_healthy_threshold: Optional[int] = None,
 ) -> Union["BaseStep", Callable[["F"], "BaseStep"]]:
     """Decorator to create a ZenML step.
 
@@ -146,6 +147,9 @@ def step(
             run inline unless a step operator or docker/resource settings
             are configured. This is only applicable for dynamic
             pipelines.
+        heartbeat_healthy_threshold: The amount of time (in minutes) that a
+            running step has not received heartbeat and is considered healthy.
+            By default, set to the maximum value (30 minutes).",
 
     Returns:
         The step instance.
@@ -184,6 +188,7 @@ def inner_decorator(func: "F") -> "BaseStep":
             substitutions=substitutions,
             cache_policy=cache_policy,
             runtime=runtime,
+            heartbeat_healthy_threshold=heartbeat_healthy_threshold,
         )
 
         return step_instance

tests/unit/steps/test_heartbeat_worker.py

@@ -188,15 +188,6 @@ def now(cls, tz=None):
         step_run.body.latest_heartbeat = None
         assert not heartbeat.is_heartbeat_unhealthy(step_run)
 
-        # if no threshold set = healthy (default response)
-        step_run.body.start_time = non_passable_diff_time
-        step_run.body.latest_heartbeat = None
-        assert heartbeat.is_heartbeat_unhealthy(step_run)
-        step_run.metadata.config = StepConfiguration(
-            name="test", heartbeat_healthy_threshold=None
-        )
-        assert not heartbeat.is_heartbeat_unhealthy(step_run)
-
         # if step heartbeat not enabled = healthy (default response)
         step_run.metadata.config = StepConfiguration(
             name="test", heartbeat_healthy_threshold=5