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/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,58 @@ 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,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,19 @@ 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 = 30
+
+    def __init__(self, step_id: UUID):
         """Heartbeat worker constructor.
 
         Args:
             options: Parameter group - polling interval, step id, etc.
         """
-        self.options = options
+
+        self._step_id = step_id
 
         self._thread: threading.Thread | None = None
         self._running: bool = False
@@ -79,7 +59,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 +68,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,7 +77,7 @@ 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
 

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,41 @@ def collect_exception_information(
         traceback=tb_bytes.decode(errors="ignore"),
         step_code_line=line_number,
     )
+
+
+class ContextReraise(ContextDecorator):
+    def __init__(
+        self,
+        source_exceptions: list[Type[BaseException]],
+        target_exception: Type[BaseException],
+        message: str,
+        propagate_traceback: bool = True,
+    ) -> None:
+        self._source_exceptions = source_exceptions
+        self._target_exception = target_exception
+        self._message = message
+        self._propagate_traceback = propagate_traceback
+
+    @property
+    def message(self) -> str:
+        return self._message
+
+    def __enter__(self) -> "ContextReraise":
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Type[BaseException] | None,
+        exc_value: BaseException | None,
+        trace: TracebackType | None,
+    ) -> bool:
+        if exc_type is None:
+            return False
+        if any(isinstance(exc_value, exc) for exc in self._source_exceptions):
+            if self._propagate_traceback:
+                raise self._target_exception(self._message).with_traceback(
+                    trace
+                )
+            else:
+                raise self._target_exception(self._message)
+        return False

src/zenml/zen_stores/migrations/versions/a5a17015b681_add_heartbeat_column_for_step_runs.py

@@ -11,7 +11,7 @@
 
 # revision identifiers, used by Alembic.
 revision = "a5a17015b681"
-down_revision = "0.90.0"
+down_revision = "0.91.0"
 branch_labels = None
 depends_on = None
 

src/zenml/zen_stores/schemas/step_run_schemas.py

@@ -407,6 +407,7 @@ def to_model(
             is_retriable=self.is_retriable,
             start_time=self.start_time,
             end_time=self.end_time,
+            latest_heartbeat=self.latest_heartbeat,
             created=self.created,
             updated=self.updated,
             model_version_id=self.model_version_id,

tests/unit/utils/test_exception_utils.py

@@ -5,7 +5,10 @@
 
 import pytest
 
-from zenml.utils.exception_utils import collect_exception_information
+from zenml.utils.exception_utils import (
+    ContextReraise,
+    collect_exception_information,
+)
 
 
 def test_regex_pattern_no_syntax_warning():
@@ -101,3 +104,44 @@ def test_regex_pattern_matches_windows_paths_and_special_chars():
         r'  File "C:\Other\path\file.py", line 123, in some_function'
     )
     assert line_pattern_win.search(non_match_line) is None
+
+
+def test_context_reraise():
+    # test source errors are captured and re-raised
+
+    class CustomError(Exception):
+        pass
+
+    with pytest.raises(CustomError):
+        with ContextReraise(
+            source_exceptions=[ValueError, TypeError],
+            target_exception=CustomError,
+            message="Oh no",
+        ):
+            raise ValueError("VALUE ERROR")
+
+    # test other errors propagate normally
+
+    with pytest.raises(ZeroDivisionError):
+        with ContextReraise(
+            source_exceptions=[ValueError, TypeError],
+            target_exception=RuntimeError,
+            message="Oh no",
+        ):
+            _ = 1 / 0
+
+    # test inheritance works
+
+    class CustomValueError(ValueError):
+        pass
+
+    class CustomTypeError(TypeError):
+        pass
+
+    with pytest.raises(CustomTypeError):
+        with ContextReraise(
+            source_exceptions=[ValueError, TypeError],
+            target_exception=CustomTypeError,
+            message="Oh no",
+        ):
+            raise CustomValueError("VALUE ERROR")