Misc fixes

Author: schustmi

src/zenml/entrypoints/step_entrypoint_configuration.py

@@ -161,6 +161,20 @@ def load_snapshot(self) -> "PipelineSnapshotResponse":
             snapshot_id=snapshot_id, step_configuration_filter=[step_name]
         )
 
+    def get_step_configuration(
+        self, snapshot: "PipelineSnapshotResponse", step_name: str
+    ) -> "Step":
+        """Gets the step configuration.
+
+        Args:
+            snapshot: The snapshot.
+            step_name: The name of the step.
+
+        Returns:
+            The step configuration.
+        """
+        return snapshot.step_configurations[step_name]
+
     def run(self) -> None:
         """Prepares the environment and runs the configured step."""
         snapshot = self.load_snapshot()
@@ -188,7 +202,9 @@ def run(self) -> None:
 
         pipeline_name = snapshot.pipeline_configuration.name
 
-        step = snapshot.step_configurations[step_name]
+        step = self.get_step_configuration(
+            snapshot=snapshot, step_name=step_name
+        )
         self._run_step(step, snapshot=snapshot)
 
         self.post_run(

src/zenml/step_operators/step_operator_entrypoint_configuration.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Abstract base class for entrypoint configurations that run a single step."""
 
-from typing import TYPE_CHECKING, Any, Dict, List
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 from uuid import UUID
 
 from zenml.client import Client
@@ -64,6 +64,60 @@ def get_entrypoint_arguments(
             kwargs[STEP_RUN_ID_OPTION],
         ]
 
+    def _should_download_code(
+        self,
+        snapshot: "PipelineSnapshotResponse",
+        step_name: Optional[str] = None,
+    ) -> bool:
+        """Checks whether code should be downloaded.
+
+        Args:
+            snapshot: The snapshot to check.
+            step_name: Name of the step to be run. This will be used to
+                determine whether code download is necessary. If not given,
+                the DockerSettings of the pipeline will be used to make that
+                decision instead.
+
+        Returns:
+            Whether code should be downloaded.
+        """
+        step_run_id = UUID(self.entrypoint_args[STEP_RUN_ID_OPTION])
+        step_run = Client().zen_store.get_run_step(step_run_id)
+
+        docker_settings = step_run.config.docker_settings
+
+        if (
+            snapshot.code_reference
+            and docker_settings.allow_download_from_code_repository
+        ):
+            return True
+
+        if (
+            snapshot.code_path
+            and docker_settings.allow_download_from_artifact_store
+        ):
+            return True
+
+        return False
+
+    def get_step_configuration(
+        self, snapshot: "PipelineSnapshotResponse", step_name: str
+    ) -> "Step":
+        """Gets the step configuration.
+
+        Args:
+            snapshot: The snapshot.
+            step_name: The name of the step.
+
+        Returns:
+            The step configuration.
+        """
+        step_run_id = UUID(self.entrypoint_args[STEP_RUN_ID_OPTION])
+        step_run = Client().zen_store.get_run_step(step_run_id)
+        from zenml.config.step_configurations import Step
+
+        return Step(spec=step_run.spec, config=step_run.config)
+
     def _run_step(
         self,
         step: "Step",

src/zenml/steps/base_step.py

@@ -125,6 +125,7 @@ def __init__(
         retry: Optional[StepRetryConfig] = None,
         substitutions: Optional[Dict[str, str]] = None,
         cache_policy: Optional[CachePolicyOrString] = None,
+        in_process: Optional[bool] = None,
     ) -> None:
         """Initializes a step.
 
@@ -158,6 +159,8 @@ 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.
         """
         from zenml.config.step_configurations import PartialStepConfiguration
 
@@ -222,6 +225,7 @@ def __init__(
             retry=retry,
             substitutions=substitutions,
             cache_policy=cache_policy,
+            in_process=in_process,
         )
 
         notebook_utils.try_to_save_notebook_cell_code(self.source_object)
@@ -661,6 +665,7 @@ def configure(
         retry: Optional[StepRetryConfig] = None,
         substitutions: Optional[Dict[str, str]] = None,
         cache_policy: Optional[CachePolicyOrString] = None,
+        in_process: Optional[bool] = None,
         merge: bool = True,
     ) -> T:
         """Configures the step.
@@ -704,6 +709,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.
             merge: If `True`, will merge the given dictionary configurations
                 like `parameters` and `settings` with existing
                 configurations. If `False` the given configurations will
@@ -782,6 +789,7 @@ def _convert_to_tuple(value: Any) -> Tuple[Source, ...]:
                 "retry": retry,
                 "substitutions": substitutions,
                 "cache_policy": cache_policy,
+                "in_process": in_process,
             }
         )
         config = StepConfigurationUpdate(**values)
@@ -810,6 +818,7 @@ def with_options(
         retry: Optional[StepRetryConfig] = None,
         substitutions: Optional[Dict[str, str]] = None,
         cache_policy: Optional[CachePolicyOrString] = None,
+        in_process: Optional[bool] = None,
         merge: bool = True,
     ) -> "BaseStep":
         """Copies the step and applies the given configurations.
@@ -843,6 +852,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.
             merge: If `True`, will merge the given dictionary configurations
                 like `parameters` and `settings` with existing
                 configurations. If `False` the given configurations will
@@ -872,6 +883,7 @@ def with_options(
             retry=retry,
             substitutions=substitutions,
             cache_policy=cache_policy,
+            in_process=in_process,
             merge=merge,
         )
         return step_copy

src/zenml/steps/step_decorator.py

@@ -80,6 +80,7 @@ def step(
     retry: Optional["StepRetryConfig"] = None,
     substitutions: Optional[Dict[str, str]] = None,
     cache_policy: Optional["CachePolicyOrString"] = None,
+    in_process: Optional[bool] = None,
 ) -> Callable[["F"], "BaseStep"]: ...
 
 
@@ -104,6 +105,7 @@ def step(
     retry: Optional["StepRetryConfig"] = None,
     substitutions: Optional[Dict[str, str]] = None,
     cache_policy: Optional["CachePolicyOrString"] = None,
+    in_process: Optional[bool] = None,
 ) -> Union["BaseStep", Callable[["F"], "BaseStep"]]:
     """Decorator to create a ZenML step.
 
@@ -139,7 +141,8 @@ 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.
     Returns:
         The step instance.
     """
@@ -176,6 +179,7 @@ def inner_decorator(func: "F") -> "BaseStep":
             retry=retry,
             substitutions=substitutions,
             cache_policy=cache_policy,
+            in_process=in_process,
         )
 
         return step_instance