Feature:4135 Handle heartbeat interrupts properly

Author: Json-Andriopoulos

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

@@ -43,6 +43,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.time_utils import utc_now
 
@@ -361,13 +362,22 @@ def _bypass() -> None:
                 except RunStoppedException as e:
                     raise e
                 except BaseException as e:  # noqa: E722
-                    logger.error(
-                        "Failed to run step `%s`: %s",
-                        self._invocation_id,
-                        e,
-                    )
-                    publish_utils.publish_failed_step_run(step_run.id)
-                    raise
+                    step_run = Client().get_run_step(step_run_id=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
             else:
                 logger.info(
                     f"Using cached version of step `{self._invocation_id}`."

src/zenml/orchestrators/step_runner.py

@@ -37,7 +37,7 @@
     ENV_ZENML_STEP_OPERATOR,
     handle_bool_env_var,
 )
-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
@@ -251,46 +251,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()