Feature:4135 Orchestrator generic graceful stopping

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/orchestrators/base_orchestrator.py

@@ -52,6 +52,7 @@
 from zenml.stack import Flavor, Stack, StackComponent, StackComponentConfig
 from zenml.steps.step_context import RunContext, get_or_create_run_context
 from zenml.utils.env_utils import temporary_environment
+from zenml.utils.pagination_utils import depaginate_stream
 from zenml.utils.pydantic_utils import before_validator_handler
 
 if TYPE_CHECKING:
@@ -678,19 +679,8 @@ 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(
@@ -720,13 +710,84 @@ 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.
+        """
+
+        if graceful:
+            self._stop_run_gracefully(pipeline_run=run)
+        else:
+            self._stop_run_forcefully(pipeline_run=run)
+
+    @staticmethod
+    def _stop_run_gracefully(pipeline_run: "PipelineRunResponse") -> None:
+        """Graceful pipeline shutdown.
+
+        Iterates running steps and sets their status to STOPPING.
+        This is turn will cause heartbeat workers to interrupt execution of the step containers.
+
+        Args:
+            pipeline_run: A pipeline run response to stop (already updated to STOPPING status).
 
         Raises:
-            NotImplementedError: If any orchestrator inheriting from the base
-                class does not implement this logic.
+            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
+
+        client = Client()
+
+        steps_stopped: int = 0
+        steps_failed: int = 0
+        steps_skipped: int = 0
+
+        for step in depaginate_stream(
+            client.list_run_steps,
+            pipeline_run_id=pipeline_run.id,
+            status=ExecutionStatus.RUNNING,
+            size=50,
+        ):
+            if not step.spec.enable_heartbeat:
+                steps_skipped += 1
+            try:
+                client.zen_store.update_run_step(
+                    step_run_id=step.id,
+                    step_run_update=StepRunUpdate(
+                        status=ExecutionStatus.STOPPING
+                    ),
+                )
+                steps_stopped += 1
+            except Exception as exc:
+                logger.debug(
+                    "Could update step % status to %: %s",
+                    step.id,
+                    ExecutionStatus.STOPPING,
+                    str(exc),
+                )
+                steps_failed += 1
+
+        operation_summary = (
+            f"{steps_stopped}, {steps_skipped=}, {steps_failed=}"
+        )
+
+        logger.info("Graceful stopping statistics: %s", operation_summary)
+
+        if steps_failed or steps_skipped:
+            if not steps_stopped:
+                # If nothing was stopped successfully, raise an error
+                raise RuntimeError(
+                    f"Failed to stop pipeline run: {operation_summary}"
+                )
+            else:
+                # If some things were stopped but others failed, raise an error
+                raise RuntimeError(
+                    f"Partial stop operation completed with errors: {operation_summary}"
+                )
+
+    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/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.
+
+    Returns:
+        A list 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