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/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/step_operators/step_operator_entrypoint_configuration.py

@@ -13,10 +13,11 @@
 #  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
+from zenml.config.step_configurations import Step
 from zenml.config.step_run_info import StepRunInfo
 from zenml.entrypoints.step_entrypoint_configuration import (
     STEP_NAME_OPTION,
@@ -26,15 +27,18 @@
 from zenml.orchestrators.step_runner import StepRunner
 
 if TYPE_CHECKING:
-    from zenml.config.step_configurations import Step
-    from zenml.models import PipelineSnapshotResponse
+    from zenml.models import PipelineSnapshotResponse, StepRunResponse
 
 STEP_RUN_ID_OPTION = "step_run_id"
 
 
 class StepOperatorEntrypointConfiguration(StepEntrypointConfiguration):
     """Base class for step operator entrypoint configurations."""
 
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._step_run: Optional["StepRunResponse"] = None
+
     @classmethod
     def get_entrypoint_options(cls) -> Dict[str, bool]:
         """Gets all options required for running with this configuration.
@@ -64,6 +68,27 @@ def get_entrypoint_arguments(
             kwargs[STEP_RUN_ID_OPTION],
         ]
 
+    @property
+    def step_run(self) -> "StepRunResponse":
+        """The step run configured for this entrypoint configuration.
+
+        Returns:
+            The step run.
+        """
+        if self._step_run is None:
+            step_run_id = UUID(self.entrypoint_args[STEP_RUN_ID_OPTION])
+            self._step_run = Client().zen_store.get_run_step(step_run_id)
+        return self._step_run
+
+    @property
+    def step(self) -> "Step":
+        """The step configured for this entrypoint configuration.
+
+        Returns:
+            The step.
+        """
+        return Step(spec=self.step_run.spec, config=self.step_run.config)
+
     def _run_step(
         self,
         step: "Step",
@@ -75,8 +100,7 @@ def _run_step(
             step: The step to run.
             snapshot: The snapshot configuration.
         """
-        step_run_id = UUID(self.entrypoint_args[STEP_RUN_ID_OPTION])
-        step_run = Client().zen_store.get_run_step(step_run_id)
+        step_run = self.step_run
         pipeline_run = Client().get_pipeline_run(step_run.pipeline_run_id)
 
         step_run_info = StepRunInfo(
@@ -87,7 +111,7 @@ def _run_step(
             run_name=pipeline_run.name,
             pipeline_step_name=self.entrypoint_args[STEP_NAME_OPTION],
             run_id=pipeline_run.id,
-            step_run_id=step_run_id,
+            step_run_id=step_run.id,
             force_write_logs=lambda: None,
         )