fixup! Improvements and bug fixes

- Updates migration down revision refs
- context-reraise exception
- changes in the step-heartbeat logic
- fix null heartbeat in list/get endpoints

Author: Json-Andriopoulos

src/zenml/constants.py

@@ -440,7 +440,7 @@ def handle_int_env_var(var: str, default: int = 0) -> int:
 STATUS = "/status"
 STEP_CONFIGURATION = "/step-configuration"
 STEPS = "/steps"
-HEARTBEAT = "heartbeat"
+HEARTBEAT = "/heartbeat"
 STOP = "/stop"
 TAGS = "/tags"
 TAG_RESOURCES = "/tag_resources"

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

@@ -813,9 +813,9 @@ def get_custom_filters(
 # ------------------ Heartbeat Model ---------------
 
 
-class StepHeartbeatResponse(BaseModel):
+class StepHeartbeatResponse(BaseModel, use_enum_values=True):
     """Light-weight model for Step Heartbeat responses."""
 
     id: UUID
-    status: str
+    status: ExecutionStatus
     latest_heartbeat: datetime

src/zenml/orchestrators/step_launcher.py

@@ -43,7 +43,9 @@
 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, StepHeartbeatWorker
 from zenml.utils import env_utils, exception_utils, string_utils
+from zenml.utils.exception_utils import ContextReraise
 from zenml.utils.time_utils import utc_now
 
 if TYPE_CHECKING:
@@ -167,7 +169,6 @@ def signal_handler(signum: int, frame: Any) -> None:
 
             try:
                 client = Client()
-                pipeline_run = None
 
                 if self._step_run:
                     pipeline_run = client.get_pipeline_run(
@@ -443,35 +444,66 @@ def _run_step(
         )
 
         start_time = time.time()
-        try:
-            if self._step.config.step_operator:
-                step_operator_name = None
-                if isinstance(self._step.config.step_operator, str):
-                    step_operator_name = self._step.config.step_operator
-
-                self._run_step_with_step_operator(
-                    step_operator_name=step_operator_name,
-                    step_run_info=step_run_info,
+
+        # To have a cross-platform compatible handling of main thread termination
+        # we use Python's interrupt_main instead of termination signals (not Windows supported).
+        # Since interrupt_main raises KeyboardInterrupt we want in this context to capture it
+        # and handle it as a custom exception.
+
+        with ContextReraise(
+            source_exceptions=[KeyboardInterrupt],
+            target_exception=StepHeartBeatTerminationException,
+            message=f"Step {self._invocation_id} has been remotely stopped - terminating",
+            propagate_traceback=False,
+        ) as ctx_reraise:
+            logger.info(
+                f"Initiating heartbeat for step: {self._invocation_id}"
+            )
+
+            heartbeat_worker = StepHeartbeatWorker(step_id=step_run.id)
+            heartbeat_worker.start()
+
+            try:
+                if self._step.config.step_operator:
+                    step_operator_name = None
+                    if isinstance(self._step.config.step_operator, str):
+                        step_operator_name = self._step.config.step_operator
+
+                    self._run_step_with_step_operator(
+                        step_operator_name=step_operator_name,
+                        step_run_info=step_run_info,
+                    )
+                else:
+                    self._run_step_without_step_operator(
+                        pipeline_run=pipeline_run,
+                        step_run=step_run,
+                        step_run_info=step_run_info,
+                        input_artifacts=step_run.regular_inputs,
+                        output_artifact_uris=output_artifact_uris,
+                    )
+            except StepHeartBeatTerminationException as exc:
+                logger.info(ctx_reraise.message)
+                output_utils.remove_artifact_dirs(
+                    artifact_uris=list(output_artifact_uris.values())
                 )
-            else:
-                self._run_step_without_step_operator(
-                    pipeline_run=pipeline_run,
-                    step_run=step_run,
-                    step_run_info=step_run_info,
-                    input_artifacts=step_run.regular_inputs,
-                    output_artifact_uris=output_artifact_uris,
+                raise (
+                    exc
+                    if heartbeat_worker.is_terminated
+                    else KeyboardInterrupt
                 )
-        except:  # noqa: E722
-            output_utils.remove_artifact_dirs(
-                artifact_uris=list(output_artifact_uris.values())
-            )
-            raise
+            except:  # noqa: E722
+                output_utils.remove_artifact_dirs(
+                    artifact_uris=list(output_artifact_uris.values())
+                )
+                raise
 
-        duration = time.time() - start_time
-        logger.info(
-            f"Step `{self._invocation_id}` has finished in "
-            f"`{string_utils.get_human_readable_time(duration)}`."
-        )
+            heartbeat_worker.stop()
+
+            duration = time.time() - start_time
+            logger.info(
+                f"Step `{self._invocation_id}` has finished in "
+                f"`{string_utils.get_human_readable_time(duration)}`."
+            )
 
     def _run_step_with_step_operator(
         self,

src/zenml/steps/__init__.py

@@ -28,6 +28,7 @@
 
 from zenml.steps.base_step import BaseStep
 from zenml.config.resource_settings import ResourceSettings
+from zenml.steps.heartbeat import StepHeartbeatWorker, StepHeartBeatTerminationException
 from zenml.steps.step_context import StepContext, get_step_context
 from zenml.steps.step_decorator import step
 
@@ -36,5 +37,7 @@
     "ResourceSettings",
     "StepContext",
     "step",
-    "get_step_context"
+    "get_step_context",
+    "StepHeartbeatWorker",
+    "StepHeartBeatTerminationException",
 ]

src/zenml/steps/heartbeat.py

@@ -17,11 +17,8 @@
 import logging
 import threading
 import time
-from typing import Annotated
 from uuid import UUID
 
-from pydantic import BaseModel, conint, model_validator
-
 from zenml.enums import ExecutionStatus
 
 logger = logging.getLogger(__name__)
@@ -33,36 +30,18 @@ class StepHeartBeatTerminationException(Exception):
     pass
 
 
-class StepHeartBeatOptions(BaseModel):
-    """Options group for step heartbeat execution."""
-
-    step_id: UUID
-    interval: Annotated[int, conint(ge=10, le=60)]
-    name: str | None = None
-
-    @model_validator(mode="after")
-    def set_default_name(self) -> "StepHeartBeatOptions":
-        """Model validator - set name value if missing.
-
-        Returns:
-            The validated step heartbeat options.
-        """
-        if not self.name:
-            self.name = f"HeartBeatWorker-{self.step_id}"
-
-        return self
-
-
-class HeartbeatWorker:
+class StepHeartbeatWorker:
     """Worker class implementing heartbeat polling and remote termination."""
 
-    def __init__(self, options: StepHeartBeatOptions):
+    STEP_HEARTBEAT_INTERVAL_SECONDS = 60
+
+    def __init__(self, step_id: UUID):
         """Heartbeat worker constructor.
 
         Args:
-            options: Parameter group - polling interval, step id, etc.
+            step_id: The step id heartbeat is running for.
         """
-        self.options = options
+        self._step_id = step_id
 
         self._thread: threading.Thread | None = None
         self._running: bool = False
@@ -72,14 +51,23 @@ def __init__(self, options: StepHeartBeatOptions):
 
     # properties
 
+    @property
+    def is_terminated(self) -> bool:
+        """Property function for termination signal.
+
+        Returns:
+            True if the worker has been terminated.
+        """
+        return self._terminated
+
     @property
     def interval(self) -> int:
         """Property function for heartbeat interval.
 
         Returns:
             The heartbeat polling interval value.
         """
-        return self.options.interval
+        return self.STEP_HEARTBEAT_INTERVAL_SECONDS
 
     @property
     def name(self) -> str:
@@ -88,7 +76,7 @@ def name(self) -> str:
         Returns:
             The name of the heartbeat worker.
         """
-        return str(self.options.name)
+        return f"HeartBeatWorker-{self.step_id}"
 
     @property
     def step_id(self) -> UUID:
@@ -97,14 +85,13 @@ def step_id(self) -> UUID:
         Returns:
             The id of the step heartbeat is running for.
         """
-        return self.options.step_id
+        return self._step_id
 
     # public functions
 
     def start(self) -> None:
         """Start the heartbeat worker on a background thread."""
         if self._thread and self._thread.is_alive():
-            logger.info("%s already running; start() is a no-op", self.name)
             return
 
         self._running = True
@@ -113,7 +100,7 @@ def start(self) -> None:
             target=self._run, name=self.name, daemon=True
         )
         self._thread.start()
-        logger.info(
+        logger.debug(
             "Daemon thread %s started (interval=%s)", self.name, self.interval
         )
 
@@ -122,7 +109,7 @@ def stop(self) -> None:
         if not self._running:
             return
         self._running = False
-        logger.info("%s stop requested", self.name)
+        logger.debug("%s stop requested", self.name)
 
     def is_alive(self) -> bool:
         """Liveness of the heartbeat worker thread.
@@ -134,7 +121,7 @@ def is_alive(self) -> bool:
         return bool(t and t.is_alive())
 
     def _run(self) -> None:
-        logger.info("%s run() loop entered", self.name)
+        logger.debug("%s run() loop entered", self.name)
         try:
             while self._running:
                 try:
@@ -151,22 +138,21 @@ def _run(self) -> None:
                         _thread.interrupt_main()  # raises KeyboardInterrupt in main thread
                     # Ensure we stop our own loop as well.
                     self._running = False
-                except Exception:
+                except Exception as exc:
                     # Log-and-continue policy for all other errors.
-                    logger.exception(
-                        "%s heartbeat() failed; continuing", self.name
+                    logger.debug(
+                        "%s heartbeat() failed with %s", self.name, str(exc)
                     )
                 # Sleep after each attempt (even after errors, unless stopped).
                 if self._running:
                     time.sleep(self.interval)
         finally:
-            logger.info("%s run() loop exiting", self.name)
+            logger.debug("%s run() loop exiting", self.name)
 
     def _heartbeat(self) -> None:
         from zenml.config.global_config import GlobalConfiguration
 
         store = GlobalConfiguration().zen_store
-
         response = store.update_step_heartbeat(step_run_id=self.step_id)
 
         if response.status in {