Feature:4135 Orchestrator generic graceful stopping

- Early stopping on unhealthy heartbeat
- Use pipeline status for heartbeat
- Some extra utils and improvements

Author: Json-Andriopoulos

src/zenml/client.py

@@ -4947,7 +4947,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,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/models/v2/core/step_run.py

@@ -823,3 +823,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

@@ -678,27 +678,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,
@@ -720,13 +709,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/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__)
 
 
@@ -161,8 +166,43 @@ def _heartbeat(self) -> None:
 
         response = store.update_step_heartbeat(step_run_id=self.step_id)
 
-        if response.status in {
+        if response.pipeline_run_status in {
             ExecutionStatus.STOPPED,
             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

src/zenml/utils/pagination_utils.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Pagination utilities."""
 
-from typing import Any, Callable, List, TypeVar
+from typing import Any, Callable, Generator, List, TypeVar
 
 from zenml.models import BaseIdentifiedResponse, Page
 
@@ -40,3 +40,28 @@ def depaginate(
         items += list(page.items)
 
     return items
+
+
+def depaginate_stream(
+    list_method: Callable[..., Page[AnyResponse]], **kwargs: Any
+) -> Generator[AnyResponse, None, None]:
+    """Depaginate the results from a client or store method that returns pages.
+
+    Args:
+        list_method: The list method to depaginate.
+        **kwargs: Arguments for the list method.
+
+    Yields:
+        A generator of the corresponding Response Models.
+    """
+    page = list_method(**kwargs)
+
+    for item in page.items:
+        yield item
+
+    while page.index < page.total_pages:
+        kwargs["page"] = page.index + 1
+        page = list_method(**kwargs)
+
+        for item in page.items:
+            yield item

src/zenml/zen_server/routers/steps_endpoints.py

@@ -226,44 +226,41 @@ def update_heartbeat(
     """
     step = zen_store().get_run_step(step_run_id, hydrate=False)
 
-    # Avoid using status.is_finished as it invalidates useful statuses for heartbeat
-    # such as STOPPED.
-    if step.status.is_failed or step.status.is_successful:
+    if step.is_finished:
         raise HTTPException(
             status_code=422,
             detail=f"Step {step.id} is finished - can not update heartbeat.",
         )
 
-    def validate_token_access(
-        ctx: AuthContext, step_: StepRunResponse
-    ) -> None:
+    pipeline_run = zen_store().get_run(step.pipeline_run_id, hydrate=False)
+
+    def validate_token_access(ctx: AuthContext) -> None:
         token_run_id = ctx.access_token.pipeline_run_id  # type: ignore[union-attr]
         token_schedule_id = ctx.access_token.schedule_id  # type: ignore[union-attr]
 
         if token_run_id:
-            if step_.pipeline_run_id != token_run_id:
+            if step.pipeline_run_id != token_run_id:
                 raise AuthorizationException(
-                    f"Authentication token provided is invalid for step: {step_.id}"
+                    f"Authentication token provided is invalid for step: {step.id}"
                 )
         elif token_schedule_id:
-            pipeline_run = zen_store().get_run(
-                step_.pipeline_run_id, hydrate=False
-            )
-
             if not (
                 pipeline_run.schedule
                 and pipeline_run.schedule.id == token_schedule_id
             ):
                 raise AuthorizationException(
-                    f"Authentication token provided is invalid for step: {step_.id}"
+                    f"Authentication token provided is invalid for step: {step.id}"
                 )
         else:
             # un-scoped token. Soon to-be-deprecated, we will ignore validation temporarily.
             pass
 
-    validate_token_access(ctx=auth_context, step_=step)
+    validate_token_access(ctx=auth_context)
+
+    hb = zen_store().update_step_heartbeat(step_run_id=step_run_id)
+    hb.pipeline_run_status = pipeline_run.status
 
-    return zen_store().update_step_heartbeat(step_run_id=step_run_id)
+    return hb
 
 
 @router.get(

tests/integration/functional/steps/test_heartbeat.py

@@ -88,6 +88,7 @@ def test_heartbeat_rest_functionality():
 
     assert hb_response.status == ExecutionStatus.RUNNING
     assert hb_response.latest_heartbeat is not None
+    assert hb_response.pipeline_run_status == ExecutionStatus.RUNNING
 
     assert (
         client.zen_store.get_run_step(step_run_id=step_run.id).latest_heartbeat