Step context changes

Author: schustmi

src/zenml/orchestrators/step_runner.py

@@ -184,8 +184,6 @@ def run(
 
             self._stack.prepare_step_run(info=step_run_info)
 
-            # Initialize the step context singleton
-            StepContext._clear()
             step_context = StepContext(
                 pipeline_run=pipeline_run,
                 step_run=step_run,
@@ -196,75 +194,78 @@ def run(
                 },
             )
 
-            # Parse the inputs for the entrypoint function.
-            function_params = self._parse_inputs(
-                args=spec.args,
-                annotations=spec.annotations,
-                input_artifacts=input_artifacts,
-            )
+            with step_context:
+                function_params = self._parse_inputs(
+                    args=spec.args,
+                    annotations=spec.annotations,
+                    input_artifacts=input_artifacts,
+                )
 
-            # Get all step environment variables. For most orchestrators, the
-            # non-secret environment variables have been set before by the
-            # orchestrator. But for some orchestrators, this is not possible and
-            # we therefore make sure to set them here so they're at least
-            # available for the user code.
-            step_environment = env_utils.get_step_environment(
-                step_config=step_run.config, stack=self._stack
-            )
-            secret_environment = env_utils.get_step_secret_environment(
-                step_config=step_run.config, stack=self._stack
-            )
-            step_environment.update(secret_environment)
-
-            step_failed = False
-            try:
-                if (
-                    # TODO: do we need to disable this for dynamic pipelines?
-                    pipeline_run.snapshot
-                    and self._stack.orchestrator.run_init_cleanup_at_step_level
-                ):
-                    self._stack.orchestrator.run_init_hook(
-                        snapshot=pipeline_run.snapshot
-                    )
+                # Get all step environment variables. For most orchestrators, the
+                # non-secret environment variables have been set before by the
+                # orchestrator. But for some orchestrators, this is not possible and
+                # we therefore make sure to set them here so they're at least
+                # available for the user code.
+                step_environment = env_utils.get_step_environment(
+                    step_config=step_run.config, stack=self._stack
+                )
+                secret_environment = env_utils.get_step_secret_environment(
+                    step_config=step_run.config, stack=self._stack
+                )
+                step_environment.update(secret_environment)
 
-                with env_utils.temporary_environment(step_environment):
-                    return_values = step_instance.call_entrypoint(
-                        **function_params
-                    )
-            except BaseException as step_exception:  # noqa: E722
-                step_failed = True
+                step_failed = False
+                try:
+                    if (
+                        # TODO: do we need to disable this for dynamic pipelines?
+                        pipeline_run.snapshot
+                        and self._stack.orchestrator.run_init_cleanup_at_step_level
+                    ):
+                        self._stack.orchestrator.run_init_hook(
+                            snapshot=pipeline_run.snapshot
+                        )
 
-                exception_info = exception_utils.collect_exception_information(
-                    step_exception, step_instance
-                )
+                    with env_utils.temporary_environment(step_environment):
+                        return_values = step_instance.call_entrypoint(
+                            **function_params
+                        )
+                except BaseException as step_exception:  # noqa: E722
+                    step_failed = True
 
-                if ENV_ZENML_STEP_OPERATOR in os.environ:
-                    # We're running in a step operator environment, so we can't
-                    # depend on the step launcher to publish the exception info
-                    Client().zen_store.update_run_step(
-                        step_run_id=step_run_info.step_run_id,
-                        step_run_update=StepRunUpdate(
-                            exception_info=exception_info,
-                        ),
+                    exception_info = (
+                        exception_utils.collect_exception_information(
+                            step_exception, step_instance
+                        )
                     )
-                else:
-                    # This will be published by the step launcher
-                    step_exception_info.set(exception_info)
 
-                if not step_run.is_retriable:
-                    if (
-                        failure_hook_source
-                        := self.configuration.failure_hook_source
-                    ):
-                        logger.info("Detected failure hook. Running...")
-                        with env_utils.temporary_environment(step_environment):
-                            load_and_run_hook(
-                                failure_hook_source,
-                                step_exception=step_exception,
-                            )
-                raise
-            finally:
-                try:
+                    if ENV_ZENML_STEP_OPERATOR in os.environ:
+                        # We're running in a step operator environment, so we can't
+                        # depend on the step launcher to publish the exception info
+                        Client().zen_store.update_run_step(
+                            step_run_id=step_run_info.step_run_id,
+                            step_run_update=StepRunUpdate(
+                                exception_info=exception_info,
+                            ),
+                        )
+                    else:
+                        # This will be published by the step launcher
+                        step_exception_info.set(exception_info)
+
+                    if not step_run.is_retriable:
+                        if (
+                            failure_hook_source
+                            := self.configuration.failure_hook_source
+                        ):
+                            logger.info("Detected failure hook. Running...")
+                            with env_utils.temporary_environment(
+                                step_environment
+                            ):
+                                load_and_run_hook(
+                                    failure_hook_source,
+                                    step_exception=step_exception,
+                                )
+                    raise
+                finally:
                     step_run_metadata = self._stack.get_step_run_metadata(
                         info=step_run_info,
                     )
@@ -338,12 +339,6 @@ def run(
                             snapshot=pipeline_run.snapshot
                         )
 
-                finally:
-                    step_context._cleanup_registry.execute_callbacks(
-                        raise_on_exception=False
-                    )
-                    StepContext._clear()  # Remove the step context singleton
-
             # Update the status and output artifacts of the step run.
             output_artifact_ids = {
                 output_name: [

src/zenml/steps/step_context.py

@@ -13,21 +13,24 @@
 #  permissions and limitations under the License.
 """Step context class."""
 
+import contextvars
 from typing import (
     TYPE_CHECKING,
     Any,
     Dict,
     List,
     Mapping,
     Optional,
+    Self,
     Sequence,
     Type,
 )
 
 from zenml.exceptions import StepContextError
 from zenml.logger import get_logger
+from zenml.utils import context_utils
 from zenml.utils.callback_registry import CallbackRegistry
-from zenml.utils.singleton import SingletonMetaClass, ThreadLocalSingleton
+from zenml.utils.singleton import SingletonMetaClass
 
 if TYPE_CHECKING:
     from zenml.artifacts.artifact_config import ArtifactConfig
@@ -54,8 +57,9 @@ def get_step_context() -> "StepContext":
     Raises:
         RuntimeError: If no step is currently running.
     """
-    if StepContext._exists():
-        return StepContext()  # type: ignore
+    if ctx := StepContext.get():
+        return ctx
+
     raise RuntimeError(
         "The step context is only available inside a step function."
     )
@@ -110,13 +114,9 @@ def initialize(self, state: Optional[Any]) -> None:
         self.initialized = True
 
 
-# TODO: use base context class
-class StepContext(metaclass=ThreadLocalSingleton):
+class StepContext(context_utils.BaseContext):
     """Provides additional context inside a step function.
 
-    This singleton class is used to access information about the current run,
-    step run, or its outputs inside a step function.
-
     Usage example:
 
     ```python
@@ -139,6 +139,8 @@ def my_trainer_step() -> Any:
     ```
     """
 
+    __context_var__ = contextvars.ContextVar("step_context")
+
     def __init__(
         self,
         pipeline_run: "PipelineRunResponse",
@@ -165,6 +167,8 @@ def __init__(
         """
         from zenml.client import Client
 
+        super().__init__()
+
         try:
             pipeline_run = Client().get_pipeline_run(pipeline_run.id)
         except KeyError:
@@ -462,6 +466,30 @@ def remove_output_tags(
             return
         output.tags = [tag for tag in output.tags if tag not in tags]
 
+    def __enter__(self) -> Self:
+        """Enter the step context.
+
+        Raises:
+            RuntimeError: If the step context has already been entered.
+
+        Returns:
+            The step context object.
+        """
+        if self._token is not None:
+            raise RuntimeError(
+                "Running a step from within another step is not allowed."
+            )
+        return super().__enter__()
+
+    def __exit__(self, *_: Any) -> None:
+        """Exit the step context.
+
+        Raises:
+            RuntimeError: If the step context has not been entered.
+        """
+        self._cleanup_registry.execute_callbacks(raise_on_exception=False)
+        super().__exit__(*_)
+
 
 class StepContextOutput:
     """Represents a step output in the step context."""

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

@@ -26,7 +26,6 @@ def upgrade() -> None:
         )
 
     with op.batch_alter_table("step_configuration", schema=None) as batch_op:
-        # TODO: missing check constraint
         batch_op.add_column(
             sa.Column(
                 "step_run_id", sqlmodel.sql.sqltypes.GUID(), nullable=True