Allow runtime configuration

Author: schustmi

src/zenml/config/step_configurations.py

@@ -44,6 +44,7 @@
 from zenml.config.frozen_base_model import FrozenBaseModel
 from zenml.config.retry_config import StepRetryConfig
 from zenml.config.source import Source, SourceWithValidator
+from zenml.enums import StepRuntime
 from zenml.logger import get_logger
 from zenml.model.lazy_load import ModelVersionDataLazyLoader
 from zenml.model.model import Model
@@ -214,12 +215,11 @@ class StepConfigurationUpdate(FrozenBaseModel):
         default=None,
         description="The cache policy for the step.",
     )
-    in_process: Optional[bool] = Field(
+    runtime: Optional[StepRuntime] = Field(
         default=None,
-        description="Whether to run the step in process. This is only "
-        "applicable for dynamic pipelines. If not set, the step will by "
-        "default run in-process unless it requires a different Docker image "
-        "or has special resource requirements.",
+        description="The step runtime. If not configured, the step will "
+        "run inline unless a step operator or docker/resource settings "
+        "are configured. This is only applicable for dynamic pipelines.",
     )
 
     outputs: Mapping[str, PartialArtifactConfiguration] = {}

src/zenml/enums.py

@@ -545,3 +545,10 @@ class PipelineRunTriggeredByType(StrEnum):
 
     STEP_RUN = "step_run"
     DEPLOYMENT = "deployment"
+
+
+class StepRuntime(StrEnum):
+    """All possible runtime modes for a step."""
+
+    INLINE = "inline"
+    ISOLATED = "isolated"

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

@@ -35,7 +35,7 @@
 from zenml.client import Client
 from zenml.config.compiler import Compiler
 from zenml.config.step_configurations import Step
-from zenml.enums import ExecutionMode
+from zenml.enums import ExecutionMode, StepRuntime
 from zenml.execution.pipeline.dynamic.outputs import (
     ArtifactFuture,
     OutputArtifact,
@@ -438,43 +438,45 @@ def _should_retry_locally(
     if step.config.step_operator:
         return True
 
-    if should_run_in_process(step, pipeline_docker_settings):
+    runtime = get_step_runtime(step, pipeline_docker_settings)
+    if runtime == StepRuntime.INLINE or step.config.step_operator:
         return True
     else:
-        # Running out of process with the orchestrator
+        # Running in isolated mode with the orchestrator
         return (
             not Client().active_stack.orchestrator.config.handles_step_retries
         )
 
 
-def should_run_in_process(
+def get_step_runtime(
     step: "Step", pipeline_docker_settings: "DockerSettings"
-) -> bool:
+) -> StepRuntime:
     """Determine if a step should be run in process.
 
     Args:
         step: The step.
         pipeline_docker_settings: The Docker settings of the parent pipeline.
 
     Returns:
-        Whether the step should be run in process.
+        The runtime for the step.
     """
     if step.config.step_operator:
-        return False
+        return StepRuntime.ISOLATED
 
     if not Client().active_stack.orchestrator.can_launch_dynamic_steps:
-        return True
+        return StepRuntime.INLINE
 
-    if step.config.in_process is False:
-        return False
-    elif step.config.in_process is None:
-        if not step.config.resource_settings.empty:
-            return False
+    runtime = step.config.runtime
 
-        if step.config.docker_settings != pipeline_docker_settings:
-            return False
+    if runtime is None:
+        if not step.config.resource_settings.empty:
+            runtime = StepRuntime.ISOLATED
+        elif step.config.docker_settings != pipeline_docker_settings:
+            runtime = StepRuntime.ISOLATED
+        else:
+            runtime = StepRuntime.INLINE
 
-    return True
+    return runtime
 
 
 def get_config_template(

src/zenml/orchestrators/step_launcher.py

@@ -26,7 +26,7 @@
     ENV_ZENML_STEP_OPERATOR,
     handle_bool_env_var,
 )
-from zenml.enums import ExecutionMode, ExecutionStatus
+from zenml.enums import ExecutionMode, ExecutionStatus, StepRuntime
 from zenml.environment import get_run_environment_dict
 from zenml.exceptions import RunInterruptedException, RunStoppedException
 from zenml.logger import get_logger
@@ -472,19 +472,22 @@ def _run_step(
                 )
             else:
                 from zenml.execution.pipeline.dynamic.runner import (
-                    should_run_in_process,
+                    get_step_runtime,
                 )
 
-                if should_run_in_process(
-                    self._step,
-                    self._snapshot.pipeline_configuration.docker_settings,
-                ):
-                    if self._step.config.in_process is False:
-                        # The step was configured to run out of process, but
-                        # the orchestrator doesn't support it.
+                step_runtime = get_step_runtime(
+                    step=self._step,
+                    pipeline_docker_settings=self._snapshot.pipeline_configuration.docker_settings,
+                )
+
+                if step_runtime == StepRuntime.INLINE:
+                    if self._step.config.runtime == StepRuntime.ISOLATED:
+                        # The step was configured to run in an isolated runtime,
+                        # but the orchestrator doesn't support it.
                         logger.warning(
-                            "The %s does not support running dynamic out of "
-                            "process steps. Running step `%s` locally instead.",
+                            "The %s does not support running dynamic steps "
+                            "in isolated runtimes. Running step `%s` in inline "
+                            "runtime instead.",
                             self._stack.orchestrator.__class__.__name__,
                             self._invocation_id,
                         )

src/zenml/steps/base_step.py

@@ -47,6 +47,7 @@
     ENV_ZENML_RUN_SINGLE_STEPS_WITHOUT_STACK,
     handle_bool_env_var,
 )
+from zenml.enums import StepRuntime
 from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
@@ -131,7 +132,7 @@ def __init__(
         retry: Optional[StepRetryConfig] = None,
         substitutions: Optional[Dict[str, str]] = None,
         cache_policy: Optional[CachePolicyOrString] = None,
-        in_process: Optional[bool] = None,
+        runtime: Optional[StepRuntime] = None,
     ) -> None:
         """Initializes a step.
 
@@ -165,8 +166,10 @@ def __init__(
             retry: Configuration for retrying the step in case of failure.
             substitutions: Extra placeholders to use in the name template.
             cache_policy: Cache policy for this step.
-            in_process: Whether to run the step in process. This is only
-                applicable for dynamic pipelines.
+            runtime: The step runtime. If not configured, the step will
+                run inline unless a step operator or docker/resource settings
+                are configured. This is only applicable for dynamic
+                pipelines.
         """
         from zenml.config.step_configurations import PartialStepConfiguration
 
@@ -233,7 +236,7 @@ def __init__(
             retry=retry,
             substitutions=substitutions,
             cache_policy=cache_policy,
-            in_process=in_process,
+            runtime=runtime,
         )
 
         notebook_utils.try_to_save_notebook_cell_code(self.source_object)
@@ -717,7 +720,7 @@ def configure(
         retry: Optional[StepRetryConfig] = None,
         substitutions: Optional[Dict[str, str]] = None,
         cache_policy: Optional[CachePolicyOrString] = None,
-        in_process: Optional[bool] = None,
+        runtime: Optional[StepRuntime] = None,
         merge: bool = True,
     ) -> T:
         """Configures the step.
@@ -761,8 +764,8 @@ def configure(
             retry: Configuration for retrying the step in case of failure.
             substitutions: Extra placeholders to use in the name template.
             cache_policy: Cache policy for this step.
-            in_process: Whether to run the step in process. This is only
-                applicable for dynamic pipelines.
+            runtime: The step runtime. This is only applicable for dynamic
+                pipelines.
             merge: If `True`, will merge the given dictionary configurations
                 like `parameters` and `settings` with existing
                 configurations. If `False` the given configurations will
@@ -841,7 +844,7 @@ def _convert_to_tuple(value: Any) -> Tuple[Source, ...]:
                 "retry": retry,
                 "substitutions": substitutions,
                 "cache_policy": cache_policy,
-                "in_process": in_process,
+                "runtime": runtime,
             }
         )
         config = StepConfigurationUpdate(**values)
@@ -870,7 +873,7 @@ def with_options(
         retry: Optional[StepRetryConfig] = None,
         substitutions: Optional[Dict[str, str]] = None,
         cache_policy: Optional[CachePolicyOrString] = None,
-        in_process: Optional[bool] = None,
+        runtime: Optional[StepRuntime] = None,
         merge: bool = True,
     ) -> "BaseStep":
         """Copies the step and applies the given configurations.
@@ -904,8 +907,8 @@ def with_options(
             retry: Configuration for retrying the step in case of failure.
             substitutions: Extra placeholders for the step name.
             cache_policy: Cache policy for this step.
-            in_process: Whether to run the step in process. This is only
-                applicable for dynamic pipelines.
+            runtime: The step runtime. This is only applicable for dynamic
+                pipelines.
             merge: If `True`, will merge the given dictionary configurations
                 like `parameters` and `settings` with existing
                 configurations. If `False` the given configurations will
@@ -935,7 +938,7 @@ def with_options(
             retry=retry,
             substitutions=substitutions,
             cache_policy=cache_policy,
-            in_process=in_process,
+            runtime=runtime,
             merge=merge,
         )
         return step_copy

src/zenml/steps/step_decorator.py

@@ -29,6 +29,7 @@
 )
 from uuid import UUID
 
+from zenml.enums import StepRuntime
 from zenml.logger import get_logger
 
 if TYPE_CHECKING:
@@ -80,7 +81,7 @@ def step(
     retry: Optional["StepRetryConfig"] = None,
     substitutions: Optional[Dict[str, str]] = None,
     cache_policy: Optional["CachePolicyOrString"] = None,
-    in_process: Optional[bool] = None,
+    runtime: Optional[StepRuntime] = None,
 ) -> Callable[["F"], "BaseStep"]: ...
 
 
@@ -105,7 +106,7 @@ def step(
     retry: Optional["StepRetryConfig"] = None,
     substitutions: Optional[Dict[str, str]] = None,
     cache_policy: Optional["CachePolicyOrString"] = None,
-    in_process: Optional[bool] = None,
+    runtime: Optional[StepRuntime] = None,
 ) -> Union["BaseStep", Callable[["F"], "BaseStep"]]:
     """Decorator to create a ZenML step.
 
@@ -141,8 +142,10 @@ def step(
         retry: configuration of step retry in case of step failure.
         substitutions: Extra placeholders for the step name.
         cache_policy: Cache policy for this step.
-        in_process: Whether to run the step in process. This is only
-            applicable for dynamic pipelines.
+        runtime: The step runtime. If not configured, the step will
+            run inline unless a step operator or docker/resource settings
+            are configured. This is only applicable for dynamic
+            pipelines.
 
     Returns:
         The step instance.
@@ -180,7 +183,7 @@ def inner_decorator(func: "F") -> "BaseStep":
             retry=retry,
             substitutions=substitutions,
             cache_policy=cache_policy,
-            in_process=in_process,
+            runtime=runtime,
         )
 
         return step_instance