Misc fixes

Author: schustmi

src/zenml/deployers/server/entrypoint_configuration.py

@@ -113,7 +113,7 @@ def run(self) -> None:
             raise RuntimeError(f"Deployment {deployment.id} has no snapshot")
 
         # Download code if necessary (for remote execution environments)
-        self.download_code_if_necessary(snapshot=deployment.snapshot)
+        self.download_code_if_necessary()
 
         app_runner = BaseDeploymentAppRunner.load_app_runner(deployment)
         app_runner.run()

src/zenml/entrypoints/base_entrypoint_configuration.py

@@ -34,6 +34,7 @@
 
 if TYPE_CHECKING:
     from zenml.artifact_stores import BaseArtifactStore
+    from zenml.config import DockerSettings
     from zenml.models import CodeReferenceResponse, PipelineSnapshotResponse
 
 logger = get_logger(__name__)
@@ -64,6 +65,7 @@ def __init__(self, arguments: List[str]):
             arguments: Command line arguments to configure this object.
         """
         self.entrypoint_args = self._parse_arguments(arguments)
+        self._snapshot: Optional["PipelineSnapshotResponse"] = None
 
     @classmethod
     def get_entrypoint_command(cls) -> List[str]:
@@ -189,7 +191,48 @@ def error(self, message: str) -> NoReturn:
         result, _ = parser.parse_known_args(arguments)
         return vars(result)
 
-    def load_snapshot(self) -> "PipelineSnapshotResponse":
+    @property
+    def snapshot(self) -> "PipelineSnapshotResponse":
+        """The snapshot configured for this entrypoint configuration.
+
+        Returns:
+            The snapshot.
+        """
+        if self._snapshot is None:
+            self._snapshot = self._load_snapshot()
+        return self._snapshot
+
+    @property
+    def docker_settings(self) -> "DockerSettings":
+        """The Docker settings configured for this entrypoint configuration.
+
+        Returns:
+            The Docker settings.
+        """
+        return self.snapshot.pipeline_configuration.docker_settings
+
+    @property
+    def should_download_code(self) -> bool:
+        """Whether code should be downloaded.
+
+        Returns:
+            Whether code should be downloaded.
+        """
+        if (
+            self.snapshot.code_reference
+            and self.docker_settings.allow_download_from_code_repository
+        ):
+            return True
+
+        if (
+            self.snapshot.code_path
+            and self.docker_settings.allow_download_from_artifact_store
+        ):
+            return True
+
+        return False
+
+    def _load_snapshot(self) -> "PipelineSnapshotResponse":
         """Loads the snapshot.
 
         Returns:
@@ -198,34 +241,19 @@ def load_snapshot(self) -> "PipelineSnapshotResponse":
         snapshot_id = UUID(self.entrypoint_args[SNAPSHOT_ID_OPTION])
         return Client().zen_store.get_snapshot(snapshot_id=snapshot_id)
 
-    def download_code_if_necessary(
-        self,
-        snapshot: "PipelineSnapshotResponse",
-        step_name: Optional[str] = None,
-    ) -> None:
+    def download_code_if_necessary(self) -> None:
         """Downloads user code if necessary.
 
-        Args:
-            snapshot: The snapshot for which to download the code.
-            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.
-
         Raises:
             CustomFlavorImportError: If the artifact store flavor can't be
                 imported.
             RuntimeError: If the current environment requires code download
                 but the snapshot does not have a reference to any code.
         """
-        should_download_code = self._should_download_code(
-            snapshot=snapshot, step_name=step_name
-        )
-
-        if not should_download_code:
+        if not self.should_download_code:
             return
 
-        if code_path := snapshot.code_path:
+        if code_path := self.snapshot.code_path:
             # Load the artifact store not from the active stack but separately.
             # This is required in case the stack has custom flavor components
             # (other than the artifact store) for which the flavor
@@ -247,7 +275,7 @@ def download_code_if_necessary(
             code_utils.download_code_from_artifact_store(
                 code_path=code_path, artifact_store=artifact_store
             )
-        elif code_reference := snapshot.code_reference:
+        elif code_reference := self.snapshot.code_reference:
             # TODO: This might fail if the code repository had unpushed changes
             # at the time the pipeline run was started.
             self.download_code_from_code_repository(
@@ -294,43 +322,6 @@ def download_code_from_code_repository(
         sys.path.insert(0, download_dir)
         os.chdir(download_dir)
 
-    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.
-        """
-        docker_settings = (
-            snapshot.step_configurations[step_name].config.docker_settings
-            if step_name
-            else snapshot.pipeline_configuration.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 _load_active_artifact_store(self) -> "BaseArtifactStore":
         """Load the active artifact store.
 

src/zenml/entrypoints/pipeline_entrypoint_configuration.py

@@ -25,13 +25,13 @@ class PipelineEntrypointConfiguration(BaseEntrypointConfiguration):
 
     def run(self) -> None:
         """Prepares the environment and runs the configured pipeline."""
-        snapshot = self.load_snapshot()
+        snapshot = self.snapshot
 
         # Activate all the integrations. This makes sure that all materializers
         # and stack component flavors are registered.
         integration_registry.activate_integrations()
 
-        self.download_code_if_necessary(snapshot=snapshot)
+        self.download_code_if_necessary()
 
         orchestrator = Client().active_stack.orchestrator
         orchestrator._prepare_run(snapshot=snapshot)

src/zenml/entrypoints/step_entrypoint_configuration.py

@@ -27,6 +27,7 @@
 from zenml.logger import get_logger
 
 if TYPE_CHECKING:
+    from zenml.config import DockerSettings
     from zenml.config.step_configurations import Step
     from zenml.models import PipelineSnapshotResponse
 
@@ -149,7 +150,26 @@ def get_entrypoint_arguments(
             kwargs[STEP_NAME_OPTION],
         ]
 
-    def load_snapshot(self) -> "PipelineSnapshotResponse":
+    @property
+    def docker_settings(self) -> "DockerSettings":
+        """The Docker settings configured for this entrypoint configuration.
+
+        Returns:
+            The Docker settings.
+        """
+        return self.step.config.docker_settings
+
+    @property
+    def step(self) -> "Step":
+        """The step configured for this entrypoint configuration.
+
+        Returns:
+            The step.
+        """
+        step_name = self.entrypoint_args[STEP_NAME_OPTION]
+        return self.snapshot.step_configurations[step_name]
+
+    def _load_snapshot(self) -> "PipelineSnapshotResponse":
         """Loads the snapshot.
 
         Returns:
@@ -163,7 +183,7 @@ def load_snapshot(self) -> "PipelineSnapshotResponse":
 
     def run(self) -> None:
         """Prepares the environment and runs the configured step."""
-        snapshot = self.load_snapshot()
+        snapshot = self.snapshot
 
         # Activate all the integrations. This makes sure that all materializers
         # and stack component flavors are registered.
@@ -178,7 +198,7 @@ def run(self) -> None:
         os.makedirs("/app", exist_ok=True)
         os.chdir("/app")
 
-        self.download_code_if_necessary(snapshot=snapshot, step_name=step_name)
+        self.download_code_if_necessary()
 
         # If the working directory is not in the sys.path, we include it to make
         # sure user code gets correctly imported.
@@ -188,8 +208,7 @@ def run(self) -> None:
 
         pipeline_name = snapshot.pipeline_configuration.name
 
-        step = snapshot.step_configurations[step_name]
-        self._run_step(step, snapshot=snapshot)
+        self._run_step(step=self.step, snapshot=snapshot)
 
         self.post_run(
             pipeline_name=pipeline_name,

src/zenml/orchestrators/publish_utils.py

@@ -238,11 +238,10 @@ def get_pipeline_run_status(
         or ExecutionStatus.RETRYING in step_statuses
     ):
         return ExecutionStatus.RUNNING
-
-    # If there are less steps than the total number of steps, it is running
-    elif is_dynamic_pipeline or len(step_statuses) < num_steps:
+    elif is_dynamic_pipeline:
+        return run_status
+    elif len(step_statuses) < num_steps:
         return ExecutionStatus.RUNNING
-
     # Any other state is completed
     else:
         return ExecutionStatus.COMPLETED

src/zenml/pipelines/dynamic/entrypoint_configuration.py

@@ -58,13 +58,13 @@ def get_entrypoint_arguments(
 
     def run(self) -> None:
         """Prepares the environment and runs the configured dynamic pipeline."""
-        snapshot = self.load_snapshot()
+        snapshot = self.snapshot
 
         # Activate all the integrations. This makes sure that all materializers
         # and stack component flavors are registered.
         integration_registry.activate_integrations()
 
-        self.download_code_if_necessary(snapshot=snapshot)
+        self.download_code_if_necessary()
 
         run = None
         if run_id := self.entrypoint_args.get(RUN_ID_OPTION, None):

src/zenml/pipelines/dynamic/runner.py

@@ -7,6 +7,7 @@
     Any,
     Dict,
     Iterator,
+    List,
     Optional,
     Sequence,
     Set,
@@ -65,7 +66,7 @@ def __init__(self, wrapped: Future[StepRunOutputs], invocation_id: str):
         self._invocation_id = invocation_id
 
     def wait(self) -> None:
-        self._wrapped.wait()
+        self._wrapped.result()
 
     def result(self) -> StepRunOutputs:
         return self._wrapped.result()
@@ -103,12 +104,13 @@ def __init__(
     ) -> None:
         self._snapshot = snapshot
         self._run = run
-        self._executor = ThreadPoolExecutor(max_workers=3)
+        self._executor = ThreadPoolExecutor(max_workers=10)
         self._pipeline: Optional["DynamicPipeline"] = None
         self._orchestrator = Stack.from_model(snapshot.stack).orchestrator
         self._orchestrator_run_id = (
             self._orchestrator.get_orchestrator_run_id()
         )
+        self._futures: List[StepRunOutputsFuture] = []
 
     @property
     def pipeline(self) -> "DynamicPipeline":
@@ -159,6 +161,7 @@ def run_pipeline(self) -> None:
                     self._orchestrator.run_cleanup_hook(
                         snapshot=self._snapshot
                     )
+                self.await_all_futures()
                 publish_successful_pipeline_run(run.id)
 
     def run_step_sync(
@@ -217,9 +220,16 @@ def _run() -> StepRunOutputs:
 
         ctx = contextvars.copy_context()
         future = self._executor.submit(ctx.run, _run)
-        return StepRunOutputsFuture(
+        step_run_future = StepRunOutputsFuture(
             wrapped=future, invocation_id=invocation_id
         )
+        self._futures.append(step_run_future)
+        return step_run_future
+
+    def await_all_futures(self) -> None:
+        for future in self._futures:
+            future.wait()
+        self._futures = []
 
 
 def _prepare_step_run(
@@ -419,7 +429,7 @@ def should_run_in_process(
         return False
 
     if not Client().active_stack.orchestrator.supports_dynamic_out_of_process_steps:
-        return False
+        return True
 
     if step.config.in_process is False:
         return False