Feature:4015 Early stopping on unhealthy heartbeat

Author: Json-Andriopoulos

src/zenml/config/step_configurations.py

@@ -221,6 +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,
+        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.",
+        ge=1,
+    )
 
     outputs: Mapping[str, PartialArtifactConfiguration] = {}
 

src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py

@@ -601,6 +601,20 @@ def stop_step(node: Node) -> None:
                     )
                     break
 
+        def is_node_heartbeat_unhealthy(node: Node) -> bool:
+            from zenml.steps.heartbeat import is_heartbeat_unhealthy
+
+            sr_ = client.list_run_steps(
+                name=node.id, pipeline_run_id=pipeline_run.id
+            )
+
+            if sr_.items:
+                sr_ = sr_[0]
+
+                return is_heartbeat_unhealthy(step_run=sr_)
+
+            return False
+
         def check_job_status(node: Node) -> NodeStatus:
             """Check the status of a job.
 
@@ -641,6 +655,12 @@ def check_job_status(node: Node) -> NodeStatus:
                     error_message,
                 )
                 return NodeStatus.FAILED
+            elif is_node_heartbeat_unhealthy(node):
+                logger.error(
+                    "Heartbeat for step `%s` indicates unhealthy status.",
+                    step_name,
+                )
+                return NodeStatus.FAILED
             else:
                 return NodeStatus.RUNNING
 

src/zenml/orchestrators/base_orchestrator.py

@@ -681,7 +681,6 @@ def stop_run(
         Raises:
             IllegalOperationError: If the run has no orchestrator run id yet.
         """
-
         if not run.orchestrator_run_id:
             raise IllegalOperationError(
                 "Cannot stop a pipeline run that has no orchestrator run id "
@@ -711,7 +710,6 @@ def _stop_run(
             graceful: If True, allows for graceful shutdown where possible.
                 If False, forces immediate termination. Default is True.
         """
-
         if graceful:
             self._stop_run_gracefully(pipeline_run=run)
         else:
@@ -730,7 +728,6 @@ def _stop_run_gracefully(pipeline_run: "PipelineRunResponse") -> None:
         Raises:
             RuntimeError: If steps fail to be set to STOPPING or steps don't have heartbeat enabled.
         """
-
         from zenml.client import Client
         from zenml.models import StepRunUpdate
 

src/zenml/steps/heartbeat.py

@@ -17,10 +17,15 @@
 import logging
 import threading
 import time
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING
 from uuid import UUID
 
 from zenml.enums import ExecutionStatus
 
+if TYPE_CHECKING:
+    from zenml.models import StepRunResponse
+
 logger = logging.getLogger(__name__)
 
 
@@ -166,3 +171,38 @@ def _heartbeat(self) -> None:
             ExecutionStatus.STOPPING,
         }:
             self._terminated = True
+
+
+def is_heartbeat_unhealthy(step_run: "StepRunResponse") -> bool:
+    """Utility function - Checks if step heartbeats indicate un-healthy execution.
+
+    Args:
+        step_run: Information regarding a step run.
+
+    Returns:
+        True if the step heartbeat is unhealthy, False otherwise.
+    """
+    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:
+        heartbeat_diff = step_run.end_time - (
+            step_run.latest_heartbeat or step_run.start_time
+        )
+    else:
+        heartbeat_diff = datetime.now(tz=timezone.utc) - (
+            step_run.latest_heartbeat or step_run.start_time
+        )
+
+    logger.info("%s heartbeat diff=%s", step_run.name, heartbeat_diff)
+
+    if (
+        heartbeat_diff.total_seconds()
+        > step_run.config.heartbeat_healthy_threshold * 60
+    ):
+        return True
+
+    return False