Feature:4135 Orchestrator generic graceful stopping (#4247)

* Feature:4135 Orchestrator generic graceful stopping

- Early stopping on unhealthy heartbeat
- Token invalidation
- Graceful step termination on pipeline stop event

Author: Json-Andriopoulos

docs/book/how-to/steps-pipelines/advanced_features.md

@@ -147,6 +147,40 @@ If steps 2, 3, and 4 execute in parallel and step 2 fails:
 All three execution modes are currently only supported by the `local`, `local_docker`, and `kubernetes` orchestrator flavors. For any other orchestrator flavor, the default (and only available) behavior is `CONTINUE_ON_FAILURE`. If you would like to see any of the other orchestrators extended to support the other execution modes, reach out to us in [Slack](https://zenml.io/slack-invite). 
 {% endhint %}
 
+### Step Heartbeat
+
+Step heartbeat is a background mechanism that runs alongside step executions and performs two core functions:
+
+ - Periodically pings the ZenML server to refresh the step's heartbeat value.
+ - Retrieves the current pipeline and step status, and terminates the step if the pipeline has entered a stopping state.
+
+This enables ZenML to:
+
+ - Track the liveness of a step execution and assess its health based on incoming heartbeats.
+ - Gracefully interrupt running steps when a pipeline is being stopped.
+
+*Scope and current behavior*
+
+ - Heartbeats are enabled only for steps executed in isolated environments. This excludes:
+    - `Inline` steps in `dynamic` pipelines.
+    - Steps run via the `local` orchestrator.
+- A step that becomes unhealthy automatically triggers a graceful shutdown (currently supported for the `kubernetes` orchestrator).
+- When using `CONTINUE_ON_FAILURE` execution mode, heartbeat status is also used to decide whether execution tokens should be invalidated.
+
+*Configuration*
+
+You can configure how long a step may go without sending a heartbeat before it is considered unhealthy using the `heartbeat_healthy_threshold` step parameter.
+The default value currently applied is the system's maximum allowed value (30 minutes).
+
+```python
+from zenml import step
+
+@step(heartbeat_healthy_threshold=30)
+def my_step():
+    ...
+
+```
+
 ## Data & Output Management
 
 ## Type annotations

src/zenml/client.py

@@ -4959,7 +4959,7 @@ def list_run_steps(
             cache_expired: Whether the cache expiration time of the step run
                 has passed.
             code_hash: The code hash of the step run to filter by.
-            status: The name of the run to filter by.
+            status: The status of the step run.
             run_metadata: Filter by run metadata.
             exclude_retried: Whether to exclude retried step runs.
             hydrate: Flag deciding whether to hydrate the output model(s)

src/zenml/config/step_configurations.py

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

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

@@ -612,6 +612,24 @@ 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:
+                return is_heartbeat_unhealthy(
+                    step_run_id=sr_.items[0].id,
+                    status=sr_.items[0].status,
+                    start_time=sr_.items[0].start_time,
+                    heartbeat_threshold=sr_.items[0].heartbeat_threshold,
+                    latest_heartbeat=sr_.items[0].latest_heartbeat,
+                )
+
+            return False
+
         def check_job_status(node: Node) -> NodeStatus:
             """Check the status of a job.
 
@@ -652,6 +670,13 @@ 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,
+                )
+                stop_step(node=node)
+                return NodeStatus.FAILED
             else:
                 return NodeStatus.RUNNING
 

src/zenml/models/v2/core/step_run.py

@@ -617,6 +617,17 @@ def latest_heartbeat(self) -> Optional[datetime]:
         """
         return self.get_body().latest_heartbeat
 
+    @property
+    def heartbeat_threshold(self) -> Optional[int]:
+        """The `heartbeat_threshold` property.
+
+        Returns:
+            the value of the property.
+        """
+        if self.get_metadata().spec.enable_heartbeat:
+            return self.get_metadata().config.heartbeat_healthy_threshold
+        return None
+
     @property
     def snapshot_id(self) -> UUID:
         """The `snapshot_id` property.
@@ -843,3 +854,4 @@ class StepHeartbeatResponse(BaseModel, use_enum_values=True):
     id: UUID
     status: ExecutionStatus
     latest_heartbeat: datetime
+    pipeline_run_status: ExecutionStatus | None = None

src/zenml/orchestrators/base_orchestrator.py

@@ -682,27 +682,16 @@ def stop_run(
                 If False, forces immediate termination. Default is False.
 
         Raises:
-            NotImplementedError: If any orchestrator inheriting from the base
-                class does not implement this logic.
             IllegalOperationError: If the run has no orchestrator run id yet.
         """
-        # Check if the orchestrator supports cancellation
-        if (
-            getattr(self._stop_run, "__func__", None)
-            is BaseOrchestrator._stop_run
-        ):
-            raise NotImplementedError(
-                f"The '{self.__class__.__name__}' orchestrator does not "
-                "support stopping pipeline runs."
-            )
-
         if not run.orchestrator_run_id:
             raise IllegalOperationError(
                 "Cannot stop a pipeline run that has no orchestrator run id "
                 "yet."
             )
 
         # Update pipeline status to STOPPING before calling concrete implementation
+        # Initiates graceful termination.
         publish_pipeline_run_status_update(
             pipeline_run_id=run.id,
             status=ExecutionStatus.STOPPING,
@@ -724,13 +713,24 @@ def _stop_run(
             run: A pipeline run response to stop (already updated to STOPPING status).
             graceful: If True, allows for graceful shutdown where possible.
                 If False, forces immediate termination. Default is True.
-
-        Raises:
-            NotImplementedError: If any orchestrator inheriting from the base
-                class does not implement this logic.
         """
+        if graceful:
+            # This should work out of the box for HeartBeat step termination.
+            # Orchestrators should extend the functionality to cover other scenarios.
+            self._stop_run_gracefully(pipeline_run=run)
+        else:
+            self._stop_run_forcefully(pipeline_run=run)
+
+    def _stop_run_gracefully(
+        self, pipeline_run: "PipelineRunResponse"
+    ) -> None:
+        pass
+
+    def _stop_run_forcefully(
+        self, pipeline_run: "PipelineRunResponse"
+    ) -> None:
         raise NotImplementedError(
-            "The stop run functionality is not implemented for the "
+            "The forceful stop run functionality is not implemented for the "
             f"'{self.__class__.__name__}' orchestrator."
         )
 

src/zenml/orchestrators/publish_utils.py

@@ -118,6 +118,23 @@ def publish_failed_step_run(step_run_id: "UUID") -> "StepRunResponse":
     )
 
 
+def publish_stopped_step_run(step_run_id: "UUID") -> "StepRunResponse":
+    """Publishes a stopped step run.
+
+    Args:
+        step_run_id: The ID of the step run to update.
+
+    Returns:
+        The updated step run.
+    """
+    return publish_step_run_status_update(
+        step_run_id=step_run_id,
+        status=ExecutionStatus.STOPPED,
+        end_time=utc_now(),
+        exception_info=step_exception_info.get(),
+    )
+
+
 def publish_successful_pipeline_run(
     pipeline_run_id: "UUID",
 ) -> "PipelineRunResponse":

src/zenml/orchestrators/step_launcher.py

@@ -39,6 +39,7 @@
 from zenml.orchestrators import utils as orchestrator_utils
 from zenml.orchestrators.step_runner import StepRunner
 from zenml.stack import Stack
+from zenml.steps import StepHeartBeatTerminationException
 from zenml.utils import env_utils, exception_utils, string_utils
 from zenml.utils.logging_utils import (
     is_step_logging_enabled,
@@ -325,19 +326,32 @@ def launch(self) -> StepRunResponse:
                     except RunStoppedException as e:
                         raise e
                     except BaseException as e:  # noqa: E722
-                        logger.error(
-                            "Failed to run step `%s`: %s",
-                            self._invocation_id,
-                            e,
+                        step_run = Client().get_run_step(
+                            step_run_id=step_run.id
                         )
-                        publish_utils.publish_failed_step_run(step_run.id)
+
+                        if (
+                            isinstance(e, StepHeartBeatTerminationException)
+                            or step_run.status == ExecutionStatus.STOPPING
+                        ):
+                            # Handle as a non-failure as exception is a propagation of graceful termination.
+                            publish_utils.publish_stopped_step_run(step_run.id)
+
+                        else:
+                            logger.error(
+                                "Failed to run step `%s`: %s",
+                                self._invocation_id,
+                                e,
+                            )
+                            publish_utils.publish_failed_step_run(step_run.id)
                         raise
 
                 duration = time.time() - start_time
                 logger.info(
                     f"Step `{self._invocation_id}` has finished in "
                     f"`{string_utils.get_human_readable_time(duration)}`."
                 )
+
             else:
                 logger.info(
                     f"Using cached version of step `{self._invocation_id}`."

src/zenml/orchestrators/step_runner.py

@@ -35,7 +35,7 @@
 from zenml.constants import (
     ENV_ZENML_STEP_OPERATOR,
 )
-from zenml.enums import ArtifactSaveType
+from zenml.enums import ArtifactSaveType, ExecutionStatus
 from zenml.exceptions import StepInterfaceError
 from zenml.hooks.hook_validators import load_and_run_hook
 from zenml.logger import get_logger
@@ -236,46 +236,55 @@ def run(
                         )
                 except BaseException as step_exception:  # noqa: E722
                     step_failed = True
-
-                    exception_info = (
-                        exception_utils.collect_exception_information(
-                            step_exception, step_instance
-                        )
-                    )
-
-                    if ENV_ZENML_STEP_OPERATOR in os.environ:
-                        # We're running in a step operator environment, so we can't
-                        # depend on the step launcher to publish the exception info
+                    if (
+                        isinstance(step_exception, KeyboardInterrupt)
+                        and heartbeat_worker.is_terminated
+                    ):
                         Client().zen_store.update_run_step(
                             step_run_id=step_run_info.step_run_id,
                             step_run_update=StepRunUpdate(
-                                exception_info=exception_info,
+                                status=ExecutionStatus.STOPPING,
                             ),
                         )
-                    else:
-                        # This will be published by the step launcher
-                        step_exception_info.set(exception_info)
 
-                    if not step_run.is_retriable:
-                        if (
-                            failure_hook_source
-                            := self.configuration.failure_hook_source
-                        ):
-                            logger.info("Detected failure hook. Running...")
-                            with env_utils.temporary_environment(
-                                step_environment
-                            ):
-                                load_and_run_hook(
-                                    failure_hook_source,
-                                    step_exception=step_exception,
-                                )
-                    if (
-                        isinstance(step_exception, KeyboardInterrupt)
-                        and heartbeat_worker.is_terminated
-                    ):
                         raise StepHeartBeatTerminationException(
                             "Remotely stopped step - terminating execution."
                         )
+                    else:
+                        exception_info = (
+                            exception_utils.collect_exception_information(
+                                step_exception, step_instance
+                            )
+                        )
+
+                        if ENV_ZENML_STEP_OPERATOR in os.environ:
+                            # We're running in a step operator environment, so we can't
+                            # depend on the step launcher to publish the exception info
+                            Client().zen_store.update_run_step(
+                                step_run_id=step_run_info.step_run_id,
+                                step_run_update=StepRunUpdate(
+                                    exception_info=exception_info,
+                                ),
+                            )
+                        else:
+                            # This will be published by the step launcher
+                            step_exception_info.set(exception_info)
+
+                        if not step_run.is_retriable:
+                            if (
+                                failure_hook_source
+                                := self.configuration.failure_hook_source
+                            ):
+                                logger.info(
+                                    "Detected failure hook. Running..."
+                                )
+                                with env_utils.temporary_environment(
+                                    step_environment
+                                ):
+                                    load_and_run_hook(
+                                        failure_hook_source,
+                                        step_exception=step_exception,
+                                    )
                     raise step_exception
                 finally:
                     heartbeat_worker.stop()