misc

Author: schustmi

src/zenml/config/compiler.py

@@ -35,6 +35,7 @@
 from zenml.config.step_configurations import (
     InputSpec,
     Step,
+    StepConfiguration,
     StepConfigurationUpdate,
     StepSpec,
 )
@@ -47,6 +48,7 @@
 if TYPE_CHECKING:
     from zenml.pipelines.pipeline_definition import Pipeline
     from zenml.stack import Stack, StackComponent
+    from zenml.steps.base_step import BaseStep
     from zenml.steps.step_invocation import StepInvocation
 
 from zenml.logger import get_logger
@@ -126,17 +128,29 @@ def compile(
                 merge=False,
             )
 
-        steps = {
-            invocation_id: self._compile_step_invocation(
-                invocation=invocation,
-                stack=stack,
-                step_config=(run_configuration.steps or {}).get(invocation_id),
-                pipeline_configuration=pipeline.configuration,
-            )
-            for invocation_id, invocation in self._get_sorted_invocations(
-                pipeline=pipeline
-            )
-        }
+        if pipeline.is_dynamic:
+            step_templates = {
+                step.name: self._compile_config_template(
+                    step=step, stack=stack
+                )
+                for step in pipeline.depends_on
+            }
+            steps = {}
+        else:
+            step_templates = None
+            steps = {
+                invocation_id: self._compile_step_invocation(
+                    invocation=invocation,
+                    stack=stack,
+                    step_config=(run_configuration.steps or {}).get(
+                        invocation_id
+                    ),
+                    pipeline_configuration=pipeline.configuration,
+                )
+                for invocation_id, invocation in self._get_sorted_invocations(
+                    pipeline=pipeline
+                )
+            }
 
         self._ensure_required_stack_components_exist(stack=stack, steps=steps)
 
@@ -156,6 +170,7 @@ def compile(
             is_dynamic=pipeline.is_dynamic,
             pipeline_configuration=pipeline.configuration,
             step_configurations=steps,
+            step_configuration_templates=step_templates,
             client_environment=get_run_environment_dict(),
             client_version=client_version,
             server_version=server_version,
@@ -521,6 +536,48 @@ def _compile_step_invocation(
             step_config_overrides=step_configuration_overrides,
         )
 
+    def _compile_config_template(
+        self,
+        step: "BaseStep",
+        stack: "Stack",
+        step_config: Optional["StepConfigurationUpdate"],
+    ) -> StepConfiguration:
+        """Compiles a ZenML step.
+
+        Args:
+            invocation: The step invocation to compile.
+            stack: The stack on which the pipeline will be run.
+            step_config: Run configuration for the step.
+            pipeline_configuration: Configuration for the pipeline.
+
+        Returns:
+            The compiled step.
+        """
+        if step_config:
+            step._apply_configuration(step_config)
+
+        convert_component_shortcut_settings_keys(
+            step.configuration.settings, stack=stack
+        )
+        step_secrets = secret_utils.resolve_and_verify_secrets(
+            step.configuration.secrets
+        )
+        step_settings = self._filter_and_validate_settings(
+            settings=step.configuration.settings,
+            configuration_level=ConfigurationLevel.STEP,
+            stack=stack,
+        )
+        step.configure(
+            secrets=step_secrets,
+            settings=step_settings,
+            merge=False,
+        )
+
+        # TODO: apply pipeline config
+        return StepConfiguration.model_validate(
+            step.configuration.model_dump()
+        )
+
     def _get_sorted_invocations(
         self,
         pipeline: "Pipeline",

src/zenml/constants.py

@@ -246,11 +246,6 @@ def handle_int_env_var(var: str, default: int = 0) -> int:
 # ZenML Analytics Server - URL
 ANALYTICS_SERVER_URL = "https://analytics.zenml.io/"
 
-# Container utils
-SHOULD_PREVENT_PIPELINE_EXECUTION = handle_bool_env_var(
-    ENV_ZENML_PREVENT_PIPELINE_EXECUTION
-)
-
 # Repository and local store directory paths:
 REPOSITORY_DIRECTORY_NAME = ".zen"
 LOCAL_STORES_DIRECTORY_NAME = "local_stores"

src/zenml/entrypoints/entrypoint.py

@@ -17,11 +17,11 @@
 import logging
 import sys
 
-from zenml import constants
 from zenml.entrypoints.base_entrypoint_configuration import (
     ENTRYPOINT_CONFIG_SOURCE_OPTION,
     BaseEntrypointConfiguration,
 )
+from zenml.pipelines.run_utils import prevent_pipeline_execution
 from zenml.utils import source_utils
 
 
@@ -35,23 +35,22 @@ def main() -> None:
     _setup_logging()
 
     # Make sure this entrypoint does not run an entire pipeline when
-    # importing user modules. This could happen if the `pipeline.run()` call
-    # is not wrapped in a function or an `if __name__== "__main__":` check)
-    constants.SHOULD_PREVENT_PIPELINE_EXECUTION = True
-
-    # Read the source for the entrypoint configuration class from the command
-    # line arguments
-    parser = argparse.ArgumentParser()
-    parser.add_argument(f"--{ENTRYPOINT_CONFIG_SOURCE_OPTION}", required=True)
-    args, remaining_args = parser.parse_known_args()
-
-    entrypoint_config_class = source_utils.load_and_validate_class(
-        args.entrypoint_config_source,
-        expected_class=BaseEntrypointConfiguration,
-    )
-    entrypoint_config = entrypoint_config_class(arguments=remaining_args)
-
-    entrypoint_config.run()
+    # importing user modules. This could happen if a pipeline is called in a
+    # module without an `if __name__== "__main__":` check)
+    with prevent_pipeline_execution():
+        parser = argparse.ArgumentParser()
+        parser.add_argument(
+            f"--{ENTRYPOINT_CONFIG_SOURCE_OPTION}", required=True
+        )
+        args, remaining_args = parser.parse_known_args()
+
+        entrypoint_config_class = source_utils.load_and_validate_class(
+            args.entrypoint_config_source,
+            expected_class=BaseEntrypointConfiguration,
+        )
+        entrypoint_config = entrypoint_config_class(arguments=remaining_args)
+
+        entrypoint_config.run()
 
 
 if __name__ == "__main__":

src/zenml/models/v2/core/pipeline_snapshot.py

@@ -31,7 +31,7 @@
 from zenml.config.pipeline_configurations import PipelineConfiguration
 from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
 from zenml.config.pipeline_spec import PipelineSpec
-from zenml.config.step_configurations import Step
+from zenml.config.step_configurations import Step, StepConfiguration
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.enums import ExecutionStatus, StackComponentType
 from zenml.models.v2.base.base import BaseUpdate, BaseZenModel
@@ -82,6 +82,12 @@ class PipelineSnapshotBase(BaseZenModel):
     step_configurations: Dict[str, Step] = Field(
         default={}, title="The step configurations for this snapshot."
     )
+    step_configuration_templates: Optional[Dict[str, StepConfiguration]] = (
+        Field(
+            default=None,
+            title="The step configuration templates of the snapshot.",
+        )
+    )
     client_environment: Dict[str, Any] = Field(
         default={}, title="The client environment for this snapshot."
     )

src/zenml/orchestrators/containerized_orchestrator.py

@@ -120,6 +120,24 @@ def get_docker_builds(
                 builds.append(pipeline_build)
                 included_pipeline_build = True
 
+        for name, step_config in snapshot.step_configuration_templates.items():
+            step_settings = step_config.docker_settings
+
+            if step_settings != pipeline_settings:
+                build = BuildConfiguration(
+                    key=ORCHESTRATOR_DOCKER_IMAGE_KEY,
+                    settings=step_settings,
+                    step_name=name,
+                )
+                builds.append(build)
+            elif not included_pipeline_build:
+                pipeline_build = BuildConfiguration(
+                    key=ORCHESTRATOR_DOCKER_IMAGE_KEY,
+                    settings=pipeline_settings,
+                )
+                builds.append(pipeline_build)
+                included_pipeline_build = True
+
         if not included_pipeline_build and self.should_build_pipeline_image(
             snapshot
         ):

src/zenml/orchestrators/local/local_orchestrator.py

@@ -96,13 +96,6 @@ def submit_pipeline(
                 step.
             RuntimeError: If the pipeline run fails.
         """
-        if snapshot.schedule:
-            logger.warning(
-                "Local Orchestrator currently does not support the "
-                "use of schedules. The `schedule` will be ignored "
-                "and the pipeline will be run immediately."
-            )
-
         self._orchestrator_run_id = str(uuid4())
         start_time = time.time()
 
@@ -203,10 +196,19 @@ def submit_dynamic_pipeline(
         """Submits a dynamic pipeline to the orchestrator."""
         from zenml.pipelines.dynamic.runner import DynamicPipelineRunner
 
+        self._orchestrator_run_id = str(uuid4())
+        start_time = time.time()
+
         runner = DynamicPipelineRunner(snapshot=snapshot, run=placeholder_run)
         with temporary_environment(environment):
             runner.run_pipeline()
 
+        run_duration = time.time() - start_time
+        logger.info(
+            "Pipeline run has finished in `%s`.",
+            string_utils.get_human_readable_time(run_duration),
+        )
+        self._orchestrator_run_id = None
         return None
 
     def get_orchestrator_run_id(self) -> str:

src/zenml/pipelines/dynamic/pipeline_definition.py

@@ -18,22 +18,25 @@
     Any,
     Callable,
     Dict,
+    List,
     Optional,
     Set,
     TypeVar,
     Union,
 )
-from uuid import uuid4
 
 from pydantic import ConfigDict, ValidationError
 
-from zenml import ExternalArtifact, constants
+from zenml import ExternalArtifact
 from zenml.client import Client
 from zenml.logger import get_logger
 from zenml.logging.step_logging import setup_pipeline_logging
 from zenml.models import ArtifactVersionResponse, PipelineRunResponse
 from zenml.pipelines.pipeline_definition import Pipeline
-from zenml.pipelines.run_utils import create_placeholder_run
+from zenml.pipelines.run_utils import (
+    create_placeholder_run,
+    should_prevent_pipeline_execution,
+)
 from zenml.steps.step_invocation import StepInvocation
 from zenml.utils import dict_utils, pydantic_utils
 
@@ -49,6 +52,10 @@
 class DynamicPipeline(Pipeline):
     """ZenML pipeline class."""
 
+    @property
+    def depends_on(self) -> List["BaseStep"]:
+        return []
+
     @property
     def is_dynamic(self) -> bool:
         """If the pipeline is dynamic.
@@ -167,15 +174,15 @@ def add_dynamic_invocation(
     def __call__(
         self, *args: Any, **kwargs: Any
     ) -> Optional[PipelineRunResponse]:
-        should_prevent_execution = constants.SHOULD_PREVENT_PIPELINE_EXECUTION
-
-        if should_prevent_execution:
-            logger.warning("Preventing execution of pipeline '%s'.", self.name)
+        if should_prevent_pipeline_execution():
+            logger.info("Preventing execution of pipeline '%s'.", self.name)
             return
 
-        if not Client().active_stack.orchestrator.supports_dynamic_pipelines:
+        stack = Client().active_stack
+
+        if not stack.orchestrator.supports_dynamic_pipelines:
             raise RuntimeError(
-                f"The {Client().active_stack.orchestrator.__class__.__name__} does not support dynamic pipelines. "
+                f"The {stack.orchestrator.__class__.__name__} does not support dynamic pipelines. "
             )
 
         self.prepare(*args, **kwargs)
@@ -186,12 +193,11 @@ def __call__(
         ) as logs_request:
             run = create_placeholder_run(
                 snapshot=snapshot,
-                orchestrator_run_id=str(uuid4()),
                 logs=logs_request,
             )
-            Client().active_stack.orchestrator.run(
+            stack.orchestrator.run(
                 snapshot=snapshot,
-                stack=Client().active_stack,
+                stack=stack,
                 placeholder_run=run,
             )
         return Client().get_pipeline_run(run.id)