Feature:4151 Heartbeat activation logic (#4216)

* Heartbeat activation logic
* Remove usused compile_spec method
* Fix heartbeat status conditions

Author: Json-Andriopoulos

src/zenml/config/compiler.py

@@ -28,16 +28,17 @@
 
 from zenml import __version__
 from zenml.config.base_settings import BaseSettings, ConfigurationLevel
-from zenml.config.pipeline_configurations import PipelineConfiguration
 from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
 from zenml.config.pipeline_spec import OutputSpec, PipelineSpec
 from zenml.config.settings_resolver import SettingsResolver
 from zenml.config.step_configurations import (
     InputSpec,
     Step,
+    StepConfiguration,
     StepConfigurationUpdate,
     StepSpec,
 )
+from zenml.enums import StepRuntime
 from zenml.environment import get_run_environment_dict
 from zenml.exceptions import StackValidationError
 from zenml.models import PipelineSnapshotBase
@@ -138,7 +139,7 @@ def compile(
                 invocation=invocation,
                 stack=stack,
                 step_config=(run_configuration.steps or {}).get(invocation_id),
-                pipeline_configuration=pipeline.configuration,
+                pipeline=pipeline,
                 skip_input_validation=skip_input_validation,
             )
             for invocation_id, invocation in self._get_sorted_invocations(
@@ -177,38 +178,6 @@ def compile(
 
         return snapshot
 
-    def compile_spec(self, pipeline: "Pipeline") -> PipelineSpec:
-        """Compiles a ZenML pipeline to a pipeline spec.
-
-        This method can be used when a pipeline spec is needed but the full
-        snapshot including stack information is not required.
-
-        Args:
-            pipeline: The pipeline to compile.
-
-        Returns:
-            The compiled pipeline spec.
-        """
-        logger.debug(
-            "Compiling pipeline spec for pipeline `%s`.", pipeline.name
-        )
-        # Copy the pipeline before we connect the steps, so we don't mess with
-        # the pipeline object/step objects in any way
-        pipeline = copy.deepcopy(pipeline)
-
-        invocations = [
-            self._get_step_spec(invocation=invocation)
-            for _, invocation in self._get_sorted_invocations(
-                pipeline=pipeline
-            )
-        ]
-
-        pipeline_spec = self._compute_pipeline_spec(
-            pipeline=pipeline, step_specs=invocations
-        )
-        logger.debug("Compiled pipeline spec: %s", pipeline_spec)
-        return pipeline_spec
-
     def _apply_run_configuration(
         self, pipeline: "Pipeline", config: PipelineRunConfiguration
     ) -> None:
@@ -487,11 +456,13 @@ def _filter_and_validate_settings(
     def _get_step_spec(
         self,
         invocation: "StepInvocation",
+        enable_heartbeat: bool,
     ) -> StepSpec:
         """Gets the spec for a step invocation.
 
         Args:
             invocation: The invocation for which to get the spec.
+            enable_heartbeat: Whether to enable the heartbeat.
 
         Returns:
             The step spec.
@@ -508,14 +479,40 @@ def _get_step_spec(
             upstream_steps=sorted(invocation.upstream_steps),
             inputs=inputs,
             invocation_id=invocation.id,
+            enable_heartbeat=enable_heartbeat,
         )
 
+    @staticmethod
+    def _get_heartbeat_flag(
+        pipeline: "Pipeline", stack: "Stack", step_config: "StepConfiguration"
+    ) -> bool:
+        if stack.orchestrator.flavor == "local":
+            return False
+        elif not pipeline.is_dynamic:
+            # containerized static pipeline
+            return True
+        else:
+            # dynamic pipelines
+            from zenml.execution.pipeline.dynamic.runner import (
+                get_step_runtime,
+            )
+
+            step_runtime = get_step_runtime(
+                step_config=step_config,
+                pipeline_docker_settings=pipeline.configuration.docker_settings,
+            )
+            if step_runtime == StepRuntime.ISOLATED:
+                # dynamic pipelines & isolated execution
+                return True
+            # dynamic pipelines & inline execution
+            return False
+
     def _compile_step_invocation(
         self,
         invocation: "StepInvocation",
         stack: "Stack",
         step_config: Optional["StepConfigurationUpdate"],
-        pipeline_configuration: "PipelineConfiguration",
+        pipeline: "Pipeline",
         skip_input_validation: bool = False,
     ) -> Step:
         """Compiles a ZenML step.
@@ -524,7 +521,7 @@ def _compile_step_invocation(
             invocation: The step invocation to compile.
             stack: The stack on which the pipeline will be run.
             step_config: Run configuration for the step.
-            pipeline_configuration: Configuration for the pipeline.
+            pipeline: Configuration for the pipeline.
             skip_input_validation: If True, will skip the input validation.
 
         Returns:
@@ -548,7 +545,6 @@ def _compile_step_invocation(
             convert_component_shortcut_settings_keys(
                 step.configuration.settings, stack=stack
             )
-            step_spec = self._get_step_spec(invocation=invocation)
             step_secrets = secret_utils.resolve_and_verify_secrets(
                 step.configuration.secrets
             )
@@ -572,9 +568,18 @@ def _compile_step_invocation(
             )
         full_step_config = (
             step_configuration_overrides.apply_pipeline_configuration(
-                pipeline_configuration=pipeline_configuration
+                pipeline_configuration=pipeline.configuration
             )
         )
+
+        step_spec = self._get_step_spec(
+            invocation=invocation,
+            enable_heartbeat=self._get_heartbeat_flag(
+                pipeline=pipeline,
+                stack=stack,
+                step_config=full_step_config,
+            ),
+        )
         return Step(
             spec=step_spec,
             config=full_step_config,

src/zenml/config/step_configurations.py

@@ -416,6 +416,7 @@ class StepSpec(FrozenBaseModel):
     upstream_steps: List[str]
     inputs: Dict[str, InputSpec] = {}
     invocation_id: str
+    enable_heartbeat: bool = False
 
     @model_validator(mode="before")
     @classmethod

src/zenml/execution/pipeline/dynamic/runner.py

@@ -34,7 +34,6 @@
 from zenml.artifacts.in_memory_cache import InMemoryArtifactCache
 from zenml.client import Client
 from zenml.config.compiler import Compiler
-from zenml.config.step_configurations import Step
 from zenml.enums import ExecutionMode, StepRuntime
 from zenml.execution.pipeline.dynamic.outputs import (
     ArtifactFuture,
@@ -68,7 +67,7 @@
 
 if TYPE_CHECKING:
     from zenml.config import DockerSettings
-    from zenml.config.step_configurations import Step
+    from zenml.config.step_configurations import Step, StepConfiguration
     from zenml.steps import BaseStep
 
 
@@ -384,11 +383,12 @@ def _await_and_validate_input(input: Any) -> Any:
         model_artifacts_or_metadata={},
         client_lazy_loaders={},
     )
+
     return Compiler()._compile_step_invocation(
         invocation=pipeline.invocations[invocation_id],
         stack=Client().active_stack,
         step_config=None,
-        pipeline_configuration=pipeline.configuration,
+        pipeline=pipeline,
     )
 
 
@@ -440,7 +440,10 @@ def _should_retry_locally(
     if step.config.step_operator:
         return True
 
-    runtime = get_step_runtime(step, pipeline_docker_settings)
+    runtime = get_step_runtime(
+        step_config=step.config,
+        pipeline_docker_settings=pipeline_docker_settings,
+    )
     if runtime == StepRuntime.INLINE or step.config.step_operator:
         return True
     else:
@@ -451,29 +454,30 @@ def _should_retry_locally(
 
 
 def get_step_runtime(
-    step: "Step", pipeline_docker_settings: "DockerSettings"
+    step_config: "StepConfiguration",
+    pipeline_docker_settings: "DockerSettings",
 ) -> StepRuntime:
     """Determine if a step should be run in process.
 
     Args:
-        step: The step.
+        step_config: The step configuration.
         pipeline_docker_settings: The Docker settings of the parent pipeline.
 
     Returns:
         The runtime for the step.
     """
-    if step.config.step_operator:
+    if step_config.step_operator:
         return StepRuntime.ISOLATED
 
     if not Client().active_stack.orchestrator.can_run_isolated_steps:
         return StepRuntime.INLINE
 
-    runtime = step.config.runtime
+    runtime = step_config.runtime
 
     if runtime is None:
-        if not step.config.resource_settings.empty:
+        if not step_config.resource_settings.empty:
             runtime = StepRuntime.ISOLATED
-        elif step.config.docker_settings != pipeline_docker_settings:
+        elif step_config.docker_settings != pipeline_docker_settings:
             runtime = StepRuntime.ISOLATED
         else:
             runtime = StepRuntime.INLINE

src/zenml/orchestrators/step_launcher.py

@@ -43,7 +43,6 @@
 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.time_utils import utc_now
 
@@ -108,23 +107,20 @@ def __init__(
         snapshot: PipelineSnapshotResponse,
         step: Step,
         orchestrator_run_id: str,
-        heartbeat_enabled: bool = False,
     ):
         """Initializes the launcher.
 
         Args:
             snapshot: The pipeline snapshot.
             step: The step to launch.
             orchestrator_run_id: The orchestrator pipeline run id.
-            heartbeat_enabled: Flag - if set will start heartbeat thread worker
 
         Raises:
             RuntimeError: If the snapshot has no associated stack.
         """
         self._snapshot = snapshot
         self._step = step
         self._orchestrator_run_id = orchestrator_run_id
-        self._heartbeat_enabled = heartbeat_enabled
 
         if not snapshot.stack:
             raise RuntimeError(
@@ -433,9 +429,6 @@ def _run_step(
             step_run: The model of the current step run.
             force_write_logs: The context for the step logs.
 
-        Raises:
-            StepHeartBeatTerminationException: if step heartbeat is enabled and the step is remotely stopped.
-            KeyboardInterrupt: Will capture, evaluate and reraise keyboard interrupts.
         """
         from zenml.deployers.server import runtime
 
@@ -460,19 +453,6 @@ def _run_step(
 
         start_time = time.time()
 
-        # 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.
-
-        heartbeat_worker = StepHeartbeatWorker(step_id=step_run.id)
-
-        if self._heartbeat_enabled:
-            logger.info(
-                "Initiating heartbeat for step: %s", self._invocation_id
-            )
-            heartbeat_worker.start()
-
         try:
             if self._step.config.step_operator:
                 step_operator_name = None
@@ -497,7 +477,7 @@ def _run_step(
                 )
 
                 step_runtime = get_step_runtime(
-                    step=self._step,
+                    step_config=self._step.config,
                     pipeline_docker_settings=self._snapshot.pipeline_configuration.docker_settings,
                 )
 
@@ -524,22 +504,11 @@ def _run_step(
                     self._run_step_with_dynamic_orchestrator(
                         step_run_info=step_run_info
                     )
-        except KeyboardInterrupt:
-            if heartbeat_worker.is_terminated:
-                msg = f"Step {self._invocation_id} has been remotely stopped - terminating"
-                logger.info(msg)
-                output_utils.remove_artifact_dirs(
-                    artifact_uris=list(output_artifact_uris.values())
-                )
-                raise StepHeartBeatTerminationException(msg)
-            raise
         except:  # noqa: E722
             output_utils.remove_artifact_dirs(
                 artifact_uris=list(output_artifact_uris.values())
             )
             raise
-        finally:
-            heartbeat_worker.stop()
 
         duration = time.time() - start_time
         logger.info(