Reworked the init/cleanup hooks and the run context

Author: stefannica

src/zenml/deployers/server/service.py

@@ -43,6 +43,7 @@
     PipelineRunTriggerInfo,
     PipelineSnapshotResponse,
 )
+from zenml.orchestrators.base_orchestrator import BaseOrchestrator
 from zenml.orchestrators.local.local_orchestrator import (
     LocalOrchestrator,
     LocalOrchestratorConfig,
@@ -56,10 +57,12 @@
 
 
 class SharedLocalOrchestrator(LocalOrchestrator):
-    """Local orchestrator that uses a separate run id for each request.
+    """Local orchestrator tweaked for deployments.
 
-    This is a slight modification of the LocalOrchestrator to allow for
-    request-scoped orchestrator run ids by storing them in contextvars.
+    This is a slight modification of the LocalOrchestrator:
+    - uses request-scoped orchestrator run ids by storing them in contextvars
+    - bypasses the init/cleanup hook execution because they are run globally by
+    the deployment service
     """
 
     # Use contextvars for thread-safe, request-scoped state
@@ -79,6 +82,28 @@ def get_orchestrator_run_id(self) -> str:
             self._shared_orchestrator_run_id.set(run_id)
         return run_id
 
+    @classmethod
+    def run_init_hook(cls, snapshot: "PipelineSnapshotResponse") -> None:
+        """Runs the init hook.
+
+        Args:
+            snapshot: The snapshot to run the init hook for.
+        """
+        # Bypass the init hook execution because it is run globally by
+        # the deployment service
+        pass
+
+    @classmethod
+    def run_cleanup_hook(cls, snapshot: "PipelineSnapshotResponse") -> None:
+        """Runs the cleanup hook.
+
+        Args:
+            snapshot: The snapshot to run the cleanup hook for.
+        """
+        # Bypass the cleanup hook execution because it is run globally by
+        # the deployment service
+        pass
+
 
 class PipelineDeploymentService:
     """Pipeline deployment service."""
@@ -97,7 +122,6 @@ def __init__(self, deployment_id: Union[str, UUID]) -> None:
             deployment_id = UUID(deployment_id)
 
         self._client = Client()
-        self.pipeline_state: Optional[Any] = None
 
         # Execution tracking
         self.service_start_time = time.time()
@@ -137,9 +161,7 @@ def initialize(self) -> None:
         """
         try:
             # Execute init hook
-            self._execute_init_hook()
-
-            self._orchestrator.set_shared_run_state(self.pipeline_state)
+            BaseOrchestrator.run_init_hook(self.snapshot)
 
             # Log success
             self._log_initialization_success()
@@ -150,28 +172,8 @@ def initialize(self) -> None:
             raise
 
     def cleanup(self) -> None:
-        """Execute cleanup hook if present.
-
-        Raises:
-            Exception: If the cleanup hook cannot be executed.
-        """
-        cleanup_hook_source = (
-            self.snapshot
-            and self.snapshot.pipeline_configuration.cleanup_hook_source
-        )
-
-        if not cleanup_hook_source:
-            return
-
-        logger.info("Executing pipeline's cleanup hook...")
-        try:
-            with env_utils.temporary_environment(
-                self.snapshot.pipeline_configuration.environment
-            ):
-                load_and_run_hook(cleanup_hook_source)
-        except Exception as e:
-            logger.exception(f"Failed to execute cleanup hook: {e}")
-            raise
+        """Execute cleanup hook if present."""
+        BaseOrchestrator.run_cleanup_hook(self.snapshot)
 
     def execute_pipeline(
         self,

src/zenml/exceptions.py

@@ -248,3 +248,7 @@ def __init__(
 
 class HookValidationException(ZenMLBaseException):
     """Exception raised when hook validation fails."""
+
+
+class HookExecutionException(ZenMLBaseException):
+    """Exception raised when hook execution fails."""

src/zenml/hooks/hook_validators.py

@@ -430,7 +430,11 @@ def _parse_hook_inputs(
         resolved_type = resolve_type_annotation(arg_type) if arg_type else None
 
         # Handle BaseException parameters - inject step_exception
-        if resolved_type and issubclass(resolved_type, BaseException):
+        if (
+            resolved_type
+            and isinstance(resolved_type, type)
+            and issubclass(resolved_type, BaseException)
+        ):
             function_params[arg] = step_exception
             continue
 

src/zenml/orchestrators/base_orchestrator.py

@@ -37,10 +37,12 @@
 )
 from zenml.enums import ExecutionMode, ExecutionStatus, StackComponentType
 from zenml.exceptions import (
+    HookExecutionException,
     IllegalOperationError,
     RunMonitoringError,
     RunStoppedException,
 )
+from zenml.hooks.hook_validators import load_and_run_hook
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
 from zenml.orchestrators.publish_utils import (
@@ -51,7 +53,8 @@
 from zenml.orchestrators.step_launcher import StepLauncher
 from zenml.orchestrators.utils import get_config_environment_vars
 from zenml.stack import Flavor, Stack, StackComponent, StackComponentConfig
-from zenml.steps.step_context import RunContext
+from zenml.steps.step_context import RunContext, get_or_create_run_context
+from zenml.utils.env_utils import temporary_environment
 from zenml.utils.pydantic_utils import before_validator_handler
 
 if TYPE_CHECKING:
@@ -393,13 +396,11 @@ def run(
     def run_step(
         self,
         step: "Step",
-        run_context: Optional[RunContext] = None,
     ) -> None:
         """Runs the given step.
 
         Args:
             step: The step to run.
-            run_context: A shared run context.
 
         Raises:
             RunStoppedException: If the run was stopped.
@@ -413,7 +414,6 @@ def _launch_step() -> None:
                 snapshot=self._active_snapshot,
                 step=step,
                 orchestrator_run_id=self.get_orchestrator_run_id(),
-                run_context=run_context,
             )
             launcher.launch()
 
@@ -498,6 +498,96 @@ def supported_execution_modes(self) -> List[ExecutionMode]:
         """
         return [ExecutionMode.CONTINUE_ON_FAILURE]
 
+    @property
+    def run_init_cleanup_at_step_level(self) -> bool:
+        """Whether the orchestrator runs the init and cleanup hooks at step level.
+
+        For orchestrators that run their steps in isolated step environments,
+        the run context cannot be shared between steps. In this case, the init
+        and cleanup hooks need to be run at step level for each individual step.
+
+        For orchestrators that run their steps in a shared environment with a
+        shared memory (e.g. the local orchestrator), the init and cleanup hooks
+        can be run at run level and this property should be overridden to return
+        True.
+
+        Returns:
+            Whether the orchestrator runs the init and cleanup hooks at step
+            level.
+        """
+        return True
+
+    @classmethod
+    def run_init_hook(cls, snapshot: "PipelineSnapshotResponse") -> None:
+        """Runs the init hook.
+
+        Args:
+            snapshot: The snapshot to run the init hook for.
+
+        Raises:
+            HookExecutionException: If the init hook fails.
+        """
+        # The lifetime of the run context starts when the init hook is executed
+        # and ends when the cleanup hook is executed
+        run_context = get_or_create_run_context()
+        init_hook_source = snapshot.pipeline_configuration.init_hook_source
+        init_hook_kwargs = snapshot.pipeline_configuration.init_hook_kwargs
+
+        # We only run the init hook once, if the (thread-local) run context
+        # associated with the current run has not been initialized yet. This
+        # allows us to run the init hook only once per run per execution
+        # environment (process, container, etc.).
+        if not run_context.initialized:
+            if not init_hook_source:
+                run_context.initialize(None)
+                return
+
+            logger.info("Executing the pipeline's init hook...")
+            try:
+                with temporary_environment(
+                    snapshot.pipeline_configuration.environment
+                ):
+                    run_state = load_and_run_hook(
+                        init_hook_source,
+                        hook_parameters=init_hook_kwargs,
+                        raise_on_error=True,
+                    )
+            except Exception as e:
+                raise HookExecutionException(
+                    f"Failed to execute init hook for pipeline "
+                    f"{snapshot.pipeline_configuration.name}"
+                ) from e
+
+            run_context.initialize(run_state)
+
+    @classmethod
+    def run_cleanup_hook(cls, snapshot: "PipelineSnapshotResponse") -> None:
+        """Runs the cleanup hook.
+
+        Args:
+            snapshot: The snapshot to run the cleanup hook for.
+        """
+        # The lifetime of the run context starts when the init hook is executed
+        # and ends when the cleanup hook is executed
+        if not RunContext._exists():
+            return
+
+        if (
+            cleanup_hook_source
+            := snapshot.pipeline_configuration.cleanup_hook_source
+        ):
+            logger.info("Executing the pipeline's cleanup hook...")
+            with temporary_environment(
+                snapshot.pipeline_configuration.environment
+            ):
+                load_and_run_hook(
+                    cleanup_hook_source,
+                    raise_on_error=False,
+                )
+
+        # Destroy the run context, so it's created anew for the next run
+        RunContext._clear()
+
     def _validate_execution_mode(
         self, snapshot: "PipelineSnapshotResponse"
     ) -> None:

src/zenml/orchestrators/local/local_orchestrator.py

@@ -14,11 +14,10 @@
 """Implementation of the ZenML local orchestrator."""
 
 import time
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type
+from typing import TYPE_CHECKING, Dict, List, Optional, Type
 from uuid import uuid4
 
 from zenml.enums import ExecutionMode
-from zenml.hooks.hook_validators import load_and_run_hook
 from zenml.logger import get_logger
 from zenml.orchestrators import (
     BaseOrchestrator,
@@ -27,7 +26,6 @@
     SubmissionResult,
 )
 from zenml.stack import Stack
-from zenml.steps.step_context import RunContext
 from zenml.utils import string_utils
 from zenml.utils.env_utils import temporary_environment
 
@@ -45,15 +43,25 @@ class LocalOrchestrator(BaseOrchestrator):
     """
 
     _orchestrator_run_id: Optional[str] = None
-    _run_context: Optional[RunContext] = None
 
-    def set_shared_run_state(self, state: Optional[Any]) -> None:
-        """Sets the state to be shared between all steps of all runs executed by this orchestrator.
+    @property
+    def run_init_cleanup_at_step_level(self) -> bool:
+        """Whether the orchestrator runs the init and cleanup hooks at step level.
 
-        Args:
-            state: the state to be shared
+        For orchestrators that run their steps in isolated step environments,
+        the run context cannot be shared between steps. In this case, the init
+        and cleanup hooks need to be run at step level for each individual step.
+
+        For orchestrators that run their steps in a shared environment with a
+        shared memory (e.g. the local orchestrator), the init and cleanup hooks
+        can be run at run level and this property should be overridden to return
+        True.
+
+        Returns:
+            Whether the orchestrator runs the init and cleanup hooks at step
+            level.
         """
-        self._run_context = RunContext(state=state)
+        return False
 
     def submit_pipeline(
         self,
@@ -100,25 +108,10 @@ def submit_pipeline(
         execution_mode = snapshot.pipeline_configuration.execution_mode
 
         failed_steps: List[str] = []
+        step_exception: Optional[Exception] = None
         skipped_steps: List[str] = []
 
-        # If the run context is not set globally, we initialize it by running
-        # the init hook
-        if self._run_context:
-            run_context = self._run_context
-        else:
-            state = None
-            if (
-                init_hook_source
-                := snapshot.pipeline_configuration.init_hook_source
-            ):
-                logger.info("Executing the pipeline's init hook...")
-                state = load_and_run_hook(
-                    init_hook_source,
-                    hook_parameters=snapshot.pipeline_configuration.init_hook_kwargs,
-                    raise_on_error=True,
-                )
-            run_context = RunContext(state=state)
+        self.run_init_hook(snapshot=snapshot)
 
         # Run each step
         for step_name, step in snapshot.step_configurations.items():
@@ -170,32 +163,21 @@ def submit_pipeline(
             step_environment = step_environments[step_name]
             try:
                 with temporary_environment(step_environment):
-                    self.run_step(step=step, run_context=run_context)
-            except Exception:
+                    self.run_step(step=step)
+            except Exception as e:
                 logger.exception("Failed to execute step %s.", step_name)
                 failed_steps.append(step_name)
                 logger.exception("Step %s failed.", step_name)
 
                 if execution_mode == ExecutionMode.FAIL_FAST:
-                    raise
-
-            finally:
-                try:
-                    # If the run context is not set globally, we also run the
-                    # cleanup hook
-                    if not self._run_context:
-                        if (
-                            cleanup_hook_source
-                            := snapshot.pipeline_configuration.cleanup_hook_source
-                        ):
-                            logger.info(
-                                "Executing the pipeline's cleanup hook..."
-                            )
-                            load_and_run_hook(
-                                cleanup_hook_source,
-                            )
-                except Exception:
-                    logger.exception("Failed to execute cleanup hook.")
+                    step_exception = e
+                    break
+
+        self.run_cleanup_hook(snapshot=snapshot)
+
+        if execution_mode == ExecutionMode.FAIL_FAST and failed_steps:
+            assert step_exception is not None
+            raise step_exception
 
         if failed_steps:
             raise RuntimeError(