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,59 @@ 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}"
+            )
+
+            StepHeartbeatWorker(step_id=step_run.id).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:
+                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
+            except:  # noqa: E722
+                output_utils.remove_artifact_dirs(
+                    artifact_uris=list(output_artifact_uris.values())
                 )
-        except:  # noqa: E722
-            output_utils.remove_artifact_dirs(
-                artifact_uris=list(output_artifact_uris.values())
-            )
-            raise
+                raise
 
-        duration = time.time() - start_time
-        logger.info(
-            f"Step `{self._invocation_id}` has finished in "
-            f"`{string_utils.get_human_readable_time(duration)}`."
-        )
+            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,52 +17,30 @@
 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__)
 
 
 class StepHeartBeatTerminationException(Exception):
     """Custom exception class for heartbeat termination."""
-
     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
@@ -79,7 +57,7 @@ def interval(self) -> int:
         Returns:
             The heartbeat polling interval value.
         """
-        return self.options.interval
+        return self.STEP_HEARTBEAT_INTERVAL_SECONDS
 
     @property
     def name(self) -> str:
@@ -88,7 +66,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 +75,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 +90,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 +99,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 +111,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 +128,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 {

src/zenml/utils/exception_utils.py

@@ -18,7 +18,9 @@
 import re
 import textwrap
 import traceback
-from typing import TYPE_CHECKING, Optional
+from contextlib import ContextDecorator
+from types import TracebackType
+from typing import TYPE_CHECKING, Optional, Type
 
 from zenml.constants import MEDIUMTEXT_MAX_LENGTH
 from zenml.logger import get_logger
@@ -91,3 +93,84 @@ def collect_exception_information(
         traceback=tb_bytes.decode(errors="ignore"),
         step_code_line=line_number,
     )
+
+
+class ContextReraise(ContextDecorator):
+    """Utility class. Capture & reraise exceptions within a context."""
+    def __init__(
+        self,
+        source_exceptions: list[Type[BaseException]],
+        target_exception: Type[BaseException],
+        message: str | None = None,
+        propagate_traceback: bool = True,
+    ) -> None:
+        """ContextReraise constructor.
+
+        Args:
+            source_exceptions: A list of exception types to capture.
+            target_exception: The exception to re-raise.
+            message: The target exception message. If None, will be inferred from source exception arguments.
+            propagate_traceback: Whether to propagate exception traceback when re-raising.
+        """
+        self._source_exceptions = source_exceptions
+        self._target_exception = target_exception
+        self._propagate_traceback = propagate_traceback
+        self._message = message
+
+    def __enter__(self) -> "ContextReraise":
+        """Context manager enter magic method.
+
+        Returns:
+            The context manager.
+        """
+        return self
+
+    @staticmethod
+    def _get_exc_message(exc: BaseException) -> str:
+        """Helper function that attempts to extract a formatted exception message.
+
+        Args:
+            exc: The exception to extract the message from.
+
+        Returns:
+            A formatted exception message (or a fallback stringified version).
+        """
+        if len(exc.args) > 1 and isinstance(exc.args[0], str):
+            try:
+                return exc.args[0] % exc.args[1:]
+            except Exception:
+                pass
+        return str(exc)
+
+    def __exit__(
+        self,
+        exc_type: Type[BaseException] | None,
+        exc_value: BaseException | None,
+        trace: TracebackType | None,
+    ) -> bool:
+        """Context manager exit magic method.
+
+        Args:
+            exc_type: The exception type.
+            exc_value: The exception.
+            trace: The exception trace.
+
+        Returns:
+            Return False if context exits normally or re-raises exception.
+        """
+        if exc_type is None:
+            return False
+
+        if any(isinstance(exc_value, exc) for exc in self._source_exceptions):
+            if self._message:
+                exc_ = self._target_exception(self._message)
+            else:
+                exc_ = self._target_exception(self._get_exc_message(exc_value))
+
+            if self._propagate_traceback:
+                exc_ = exc_.with_traceback(trace)
+
+            raise exc_
+        else:
+            raise exc_value
+