Threadsafe static pipeline compilation

Author: schustmi

src/zenml/pipelines/compilation_context.py

@@ -0,0 +1,51 @@
+import contextvars
+from typing import TYPE_CHECKING, Self
+
+from zenml.utils import context_utils
+
+if TYPE_CHECKING:
+    from zenml.pipelines.pipeline_definition import Pipeline
+
+
+class PipelineCompilationContext(context_utils.BaseContext):
+    """Pipeline compilation context."""
+
+    __context_var__ = contextvars.ContextVar("pipeline_compilation_context")
+
+    def __init__(
+        self,
+        pipeline: "Pipeline",
+    ) -> None:
+        """Initialize the pipeline compilation context.
+
+        Args:
+            pipeline: The pipeline that is being compiled.
+        """
+        super().__init__()
+        self._pipeline = pipeline
+
+    @property
+    def pipeline(self) -> "Pipeline":
+        """The pipeline that is being compiled.
+
+        Returns:
+            The pipeline that is being compiled.
+        """
+        return self._pipeline
+
+    def __enter__(self) -> Self:
+        """Enter the pipeline compilation context.
+
+        Raises:
+            RuntimeError: If the pipeline compilation context has already been
+                entered.
+
+        Returns:
+            The pipeline compilation context object.
+        """
+        if self._token is not None:
+            raise RuntimeError(
+                "Compiling a pipeline while another pipeline is being compiled "
+                "is not allowed."
+            )
+        return super().__enter__()

src/zenml/pipelines/dynamic/context.py renamed to src/zenml/pipelines/dynamic/run_context.py

@@ -1,54 +1,30 @@
+#  Copyright (c) ZenML GmbH 2025. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Dynamic pipeline run context."""
+
 import contextvars
-from typing import TYPE_CHECKING, Any, ClassVar, Optional, Self, cast
+from typing import TYPE_CHECKING, Self
+
+from zenml.utils import context_utils
 
 if TYPE_CHECKING:
     from zenml.models import PipelineRunResponse, PipelineSnapshotResponse
     from zenml.pipelines.dynamic.pipeline_definition import DynamicPipeline
     from zenml.pipelines.dynamic.runner import DynamicPipelineRunner
 
 
-class BaseContext:
-    """Base context class."""
-
-    __context_var__: ClassVar[contextvars.ContextVar[Self]]
-
-    def __init__(self) -> None:
-        """Initialize the context."""
-        self._token: Optional[contextvars.Token[Any]] = None
-
-    @classmethod
-    def get(cls: type[Self]) -> Optional[Self]:
-        """Get the active context for the current thread.
-
-        Returns:
-            The active context for the current thread.
-        """
-        return cast(Optional[Self], cls.__context_var__.get(None))
-
-    def __enter__(self) -> Self:
-        """Enter the context.
-
-        Returns:
-            The context object.
-        """
-        self._token = self.__context_var__.set(self)
-        return self
-
-    def __exit__(self, *_: Any) -> None:
-        """Exit the context.
-
-        Raises:
-            RuntimeError: If the context has not been entered.
-        """
-        if not self._token:
-            raise RuntimeError(
-                f"Can't exit {self.__class__.__name__} because it has not been "
-                "entered."
-            )
-        self.__context_var__.reset(self._token)
-
-
-class DynamicPipelineRunContext(BaseContext):
+class DynamicPipelineRunContext(context_utils.BaseContext):
     """Dynamic pipeline run context."""
 
     __context_var__ = contextvars.ContextVar("dynamic_pipeline_run_context")
@@ -76,18 +52,38 @@ def __init__(
 
     @property
     def pipeline(self) -> "DynamicPipeline":
+        """The pipeline that is being executed.
+
+        Returns:
+            The pipeline that is being executed.
+        """
         return self._pipeline
 
     @property
     def run(self) -> "PipelineRunResponse":
+        """The pipeline run.
+
+        Returns:
+            The pipeline run.
+        """
         return self._run
 
     @property
     def snapshot(self) -> "PipelineSnapshotResponse":
+        """The snapshot of the pipeline.
+
+        Returns:
+            The snapshot of the pipeline.
+        """
         return self._snapshot
 
     @property
     def runner(self) -> "DynamicPipelineRunner":
+        """The runner executing the pipeline.
+
+        Returns:
+            The runner executing the pipeline.
+        """
         return self._runner
 
     def __enter__(self) -> Self:
@@ -105,7 +101,3 @@ def __enter__(self) -> Self:
                 "Calling a pipeline within a dynamic pipeline is not allowed."
             )
         return super().__enter__()
-
-
-def executing_dynamic_pipeline() -> bool:
-    return DynamicPipelineRunContext.get() is not None

src/zenml/pipelines/dynamic/runner.py

@@ -1,3 +1,18 @@
+#  Copyright (c) ZenML GmbH 2025. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Dynamic pipeline runner."""
+
 import contextvars
 import inspect
 import time
@@ -34,8 +49,8 @@
     publish_successful_pipeline_run,
 )
 from zenml.orchestrators.step_launcher import StepLauncher
-from zenml.pipelines.dynamic.context import DynamicPipelineRunContext
 from zenml.pipelines.dynamic.pipeline_definition import DynamicPipeline
+from zenml.pipelines.dynamic.run_context import DynamicPipelineRunContext
 from zenml.pipelines.run_utils import create_placeholder_run
 from zenml.stack import Stack
 from zenml.steps.entrypoint_function_utils import StepArtifact

src/zenml/pipelines/pipeline_context.py

@@ -33,9 +33,11 @@ def get_pipeline_context() -> "PipelineContext":
         RuntimeError: If no active pipeline is found.
         RuntimeError: If inside a running step.
     """
-    from zenml.pipelines.pipeline_definition import Pipeline
+    from zenml.pipelines.compilation_context import PipelineCompilationContext
 
-    if Pipeline.ACTIVE_PIPELINE is None:
+    context = PipelineCompilationContext.get()
+
+    if context is None:
         try:
             from zenml.steps.step_context import get_step_context
 
@@ -49,7 +51,7 @@ def get_pipeline_context() -> "PipelineContext":
             )
 
     return PipelineContext(
-        pipeline_configuration=Pipeline.ACTIVE_PIPELINE.configuration
+        pipeline_configuration=context.pipeline.configuration
     )
 
 

src/zenml/pipelines/pipeline_definition.py

@@ -22,7 +22,6 @@
     TYPE_CHECKING,
     Any,
     Callable,
-    ClassVar,
     Dict,
     Iterator,
     List,
@@ -81,6 +80,7 @@
     ScheduleRequest,
 )
 from zenml.pipelines import build_utils
+from zenml.pipelines.compilation_context import PipelineCompilationContext
 from zenml.pipelines.run_utils import (
     create_placeholder_run,
     should_prevent_pipeline_execution,
@@ -131,11 +131,6 @@
 class Pipeline:
     """ZenML pipeline class."""
 
-    # The active pipeline is the pipeline to which step invocations will be
-    # added when a step is called. It is set using a context manager when a
-    # pipeline is called (see Pipeline.__call__ for more context)
-    ACTIVE_PIPELINE: ClassVar[Optional["Pipeline"]] = None
-
     def __init__(
         self,
         name: str,
@@ -627,10 +622,7 @@ def pipeline_(param_name: str):
                 if k not in kwargs:
                     kwargs[k] = v_config
 
-        with self:
-            # Enter the context manager, so we become the active pipeline. This
-            # means that all steps that get called while the entrypoint function
-            # is executed will be added as invocation to this pipeline instance.
+        with PipelineCompilationContext(pipeline=self):
             self._call_entrypoint(*args, **kwargs)
 
     def register(self) -> "PipelineResponse":
@@ -1361,7 +1353,15 @@ def add_step_invocation(
         Returns:
             The step invocation ID.
         """
-        if not self.is_dynamic and Pipeline.ACTIVE_PIPELINE != self:
+        from zenml.pipelines.dynamic.run_context import (
+            DynamicPipelineRunContext,
+        )
+
+        context = (
+            PipelineCompilationContext.get() or DynamicPipelineRunContext.get()
+        )
+
+        if not context or context.pipeline != self:
             raise RuntimeError(
                 "A step invocation can only be added to an active pipeline."
             )
@@ -1429,32 +1429,6 @@ def _compute_invocation_id(
 
         raise RuntimeError("Unable to find step ID")
 
-    def __enter__(self) -> Self:
-        """Activate the pipeline context.
-
-        Raises:
-            RuntimeError: If a different pipeline is already active.
-
-        Returns:
-            The pipeline instance.
-        """
-        if Pipeline.ACTIVE_PIPELINE:
-            raise RuntimeError(
-                "Unable to enter pipeline context. A different pipeline "
-                f"{Pipeline.ACTIVE_PIPELINE.name} is already active."
-            )
-
-        Pipeline.ACTIVE_PIPELINE = self
-        return self
-
-    def __exit__(self, *args: Any) -> None:
-        """Deactivates the pipeline context.
-
-        Args:
-            *args: The arguments passed to the context exit handler.
-        """
-        Pipeline.ACTIVE_PIPELINE = None
-
     def _parse_config_file(
         self, config_path: Optional[str], matcher: List[str]
     ) -> Dict[str, Any]:
@@ -1590,7 +1564,7 @@ def __call__(
             `entrypoint` method. Otherwise, returns the pipeline run or `None`
             if running with a schedule.
         """
-        if Pipeline.ACTIVE_PIPELINE:
+        if PipelineCompilationContext.is_active():
             # Calling a pipeline inside a pipeline, we return the potential
             # outputs of the entrypoint function