WIP

Author: schustmi

src/zenml/config/compiler.py

@@ -149,7 +149,17 @@ def compile(
         deployment = PipelineDeploymentBase(
             run_name_template=run_name,
             pipeline_configuration=pipeline.configuration,
-            step_configurations=steps,
+            step_configurations={
+                k: v.model_copy(
+                    update={
+                        "config": v.config.apply_pipeline_configuration(
+                            pipeline.configuration
+                        )
+                    }
+                )
+                for k, v in steps.items()
+            },
+            raw_step_configurations=steps,
             client_environment=get_run_environment_dict(),
             client_version=client_version,
             server_version=server_version,
@@ -469,23 +479,23 @@ def _compile_step_invocation(
             configuration_level=ConfigurationLevel.STEP,
             stack=stack,
         )
-        step_extra = step.configuration.extra
-        step_on_failure_hook_source = step.configuration.failure_hook_source
-        step_on_success_hook_source = step.configuration.success_hook_source
-
-        step.configure(
-            settings=pipeline_settings,
-            extra=pipeline_extra,
-            on_failure=pipeline_failure_hook_source,
-            on_success=pipeline_success_hook_source,
-            merge=False,
-        )
+        # step_extra = step.configuration.extra
+        # step_on_failure_hook_source = step.configuration.failure_hook_source
+        # step_on_success_hook_source = step.configuration.success_hook_source
+
+        # step.configure(
+        #     settings=pipeline_settings,
+        #     extra=pipeline_extra,
+        #     on_failure=pipeline_failure_hook_source,
+        #     on_success=pipeline_success_hook_source,
+        #     merge=False,
+        # )
         step.configure(
             settings=step_settings,
-            extra=step_extra,
-            on_failure=step_on_failure_hook_source,
-            on_success=step_on_success_hook_source,
-            merge=True,
+            # extra=step_extra,
+            # on_failure=step_on_failure_hook_source,
+            # on_success=step_on_success_hook_source,
+            merge=False,
         )
 
         parameters_to_ignore = (

src/zenml/config/step_configurations.py

@@ -46,7 +46,7 @@
 from zenml.model.lazy_load import ModelVersionDataLazyLoader
 from zenml.model.model import Model
 from zenml.utils import deprecation_utils
-from zenml.utils.pydantic_utils import before_validator_handler
+from zenml.utils.pydantic_utils import before_validator_handler, update_model
 
 if TYPE_CHECKING:
     from zenml.config import DockerSettings, ResourceSettings
@@ -259,6 +259,39 @@ def _get_full_substitutions(
         ret.update(self.substitutions)
         return ret
 
+    def apply_pipeline_configuration(
+        self, pipeline_configuration: "PipelineConfiguration"
+    ) -> "StepConfiguration":
+        """Apply the pipeline configuration to this step configuration.
+
+        Args:
+            pipeline_configuration: The pipeline configuration to apply.
+        """
+        pipeline_values = pipeline_configuration.model_dump(
+            include={
+                "settings",
+                "extra",
+                "failure_hook_source",
+                "success_hook_source",
+            },
+            exclude_none=True,
+        )
+        if pipeline_values:
+            original_values = self.model_dump(
+                include={
+                    "settings",
+                    "extra",
+                    "failure_hook_source",
+                    "success_hook_source",
+                },
+                exclude_none=True,
+            )
+
+            updated_config = self.model_copy(update=pipeline_values, deep=True)
+            return update_model(updated_config, original_values)
+        else:
+            return self.model_copy(deep=True)
+
 
 class InputSpec(StrictBaseModel):
     """Step input specification."""

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

@@ -60,6 +60,9 @@ class PipelineDeploymentBase(BaseZenModel):
     step_configurations: Dict[str, Step] = Field(
         default={}, title="The step configurations for this deployment."
     )
+    raw_step_configurations: Dict[str, Step] = Field(
+        default={}, title="The raw step configurations for this deployment."
+    )
     client_environment: Dict[str, str] = Field(
         default={}, title="The client environment for this deployment."
     )
@@ -143,6 +146,9 @@ class PipelineDeploymentResponseMetadata(ProjectScopedResponseMetadata):
     step_configurations: Dict[str, Step] = Field(
         default={}, title="The step configurations for this deployment."
     )
+    raw_step_configurations: Dict[str, Step] = Field(
+        default={}, title="The raw step configurations for this deployment."
+    )
     client_environment: Dict[str, str] = Field(
         default={}, title="The client environment for this deployment."
     )
@@ -241,6 +247,15 @@ def step_configurations(self) -> Dict[str, Step]:
         """
         return self.get_metadata().step_configurations
 
+    @property
+    def raw_step_configurations(self) -> Dict[str, Step]:
+        """The `raw_step_configurations` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().raw_step_configurations
+
     @property
     def client_environment(self) -> Dict[str, str]:
         """The `client_environment` property.

src/zenml/pipelines/run_utils.py

@@ -113,6 +113,7 @@ def get_placeholder_run(
         size=1,
         deployment_id=deployment_id,
         status=ExecutionStatus.INITIALIZING,
+        hydrate=True,
     )
     if len(runs.items) == 0:
         return None

src/zenml/zen_server/template_execution/utils.py

@@ -1,6 +1,5 @@
 """Utility functions to run a pipeline from the server."""
 
-import copy
 import hashlib
 import sys
 from typing import Any, Dict, List, Optional
@@ -16,7 +15,7 @@
 from zenml.config.pipeline_run_configuration import (
     PipelineRunConfiguration,
 )
-from zenml.config.step_configurations import Step, StepConfiguration
+from zenml.config.step_configurations import Step, StepConfigurationUpdate
 from zenml.constants import (
     ENV_ZENML_ACTIVE_PROJECT_ID,
     ENV_ZENML_ACTIVE_STACK_ID,
@@ -43,7 +42,7 @@
     validate_stack_is_runnable_from_server,
 )
 from zenml.stack.flavor import Flavor
-from zenml.utils import dict_utils, requirements_utils, settings_utils
+from zenml.utils import pydantic_utils, requirements_utils, settings_utils
 from zenml.zen_server.auth import AuthContext, generate_access_token
 from zenml.zen_server.template_execution.runner_entrypoint_configuration import (
     RunnerEntrypointConfiguration,
@@ -106,7 +105,6 @@ def run_template(
     deployment_request = deployment_request_from_template(
         template=template,
         config=run_config or PipelineRunConfiguration(),
-        user_id=auth_context.user.id,
     )
 
     ensure_async_orchestrator(deployment=deployment_request, stack=stack)
@@ -345,14 +343,12 @@ def generate_dockerfile(
 def deployment_request_from_template(
     template: RunTemplateResponse,
     config: PipelineRunConfiguration,
-    user_id: UUID,
 ) -> "PipelineDeploymentRequest":
     """Generate a deployment request from a template.
 
     Args:
         template: The template from which to create the deployment request.
         config: The run configuration.
-        user_id: ID of the user that is trying to run the template.
 
     Raises:
         ValueError: If there are missing/extra step parameters in the run
@@ -363,49 +359,34 @@ def deployment_request_from_template(
     """
     deployment = template.source_deployment
     assert deployment
-    pipeline_configuration = PipelineConfiguration(
-        **config.model_dump(
-            include=set(PipelineConfiguration.model_fields),
-            exclude={"name", "parameters"},
-        ),
-        name=deployment.pipeline_configuration.name,
-        parameters=deployment.pipeline_configuration.parameters,
-    )
 
-    step_config_dict_base = pipeline_configuration.model_dump(
-        exclude={"name", "parameters", "tags", "enable_pipeline_logs"}
+    pipeline_update = config.model_dump(
+        include=set(PipelineConfiguration.model_fields),
+        exclude={"name", "parameters"},
+        # TODO: Make sure all unset values are actually passed as unset
+        exclude_unset=True,
+    )
+    pipeline_configuration = pydantic_utils.update_model(
+        deployment.pipeline_configuration, pipeline_update
     )
-    steps = {}
-    for invocation_id, step in deployment.step_configurations.items():
-        step_config_dict = {
-            **copy.deepcopy(step_config_dict_base),
-            **step.config.model_dump(
-                # TODO: Maybe we need to make some of these configurable via
-                # yaml as well, e.g. the lazy loaders?
-                include={
-                    "name",
-                    "caching_parameters",
-                    "external_input_artifacts",
-                    "model_artifacts_or_metadata",
-                    "client_lazy_loaders",
-                    "substitutions",
-                    "outputs",
-                }
-            ),
-        }
-
-        required_parameters = set(step.config.parameters)
-        configured_parameters = set()
 
-        if update := config.steps.get(invocation_id):
-            update_dict = update.model_dump()
+    steps = {}
+    for invocation_id, step in deployment.raw_step_configurations.items():
+        step_update = config.steps.get(
+            invocation_id, StepConfigurationUpdate()
+        ).model_dump(
             # Get rid of deprecated name to prevent overriding the step name
             # with `None`.
-            update_dict.pop("name", None)
-            configured_parameters = set(update.parameters)
-            step_config_dict = dict_utils.recursive_update(
-                step_config_dict, update=update_dict
-            )
+            exclude={"name"},
+            exclude_unset=True,
+        )
+        step_config = pydantic_utils.update_model(step.config, step_update)
+        # step_config = step_config.apply_pipeline_configuration(
+        #     pipeline_configuration
+        # )
+
+        required_parameters = set(step.config.parameters)
+        configured_parameters = set(step_config.parameters)
 
         unknown_parameters = configured_parameters - required_parameters
         if unknown_parameters:
@@ -421,7 +402,6 @@ def deployment_request_from_template(
                 f"parameters for step {invocation_id}: {missing_parameters}."
             )
 
-        step_config = StepConfiguration.model_validate(step_config_dict)
         steps[invocation_id] = Step(spec=step.spec, config=step_config)
 
     code_reference_request = None
@@ -440,7 +420,8 @@ def deployment_request_from_template(
         run_name_template=config.run_name
         or get_default_run_name(pipeline_name=pipeline_configuration.name),
         pipeline_configuration=pipeline_configuration,
-        step_configurations=steps,
+        # step_configurations=steps,
+        raw_step_configurations=steps,
         client_environment={},
         client_version=zenml_version,
         server_version=zenml_version,

src/zenml/zen_stores/schemas/pipeline_deployment_schemas.py

@@ -197,7 +197,7 @@ def from_request(
             run_name_template=request.run_name_template,
             pipeline_configuration=request.pipeline_configuration.model_dump_json(),
             step_configurations=json.dumps(
-                request.step_configurations,
+                request.raw_step_configurations,
                 sort_keys=False,
                 default=pydantic_encoder,
             ),
@@ -248,7 +248,17 @@ def to_model(
                 project=self.project.to_model(),
                 run_name_template=self.run_name_template,
                 pipeline_configuration=pipeline_configuration,
-                step_configurations=step_configurations,
+                raw_step_configurations=step_configurations,
+                step_configurations={
+                    k: v.model_copy(
+                        update={
+                            "config": v.config.apply_pipeline_configuration(
+                                pipeline_configuration
+                            )
+                        }
+                    )
+                    for k, v in step_configurations.items()
+                },
                 client_environment=json.loads(self.client_environment),
                 client_version=self.client_version,
                 server_version=self.server_version,

src/zenml/zen_stores/schemas/step_run_schemas.py

@@ -254,6 +254,9 @@ def to_model(
 
         full_step_config = None
         if self.deployment is not None:
+            pipeline_configuration = PipelineConfiguration.model_validate_json(
+                self.deployment.pipeline_configuration
+            )
             step_configuration = json.loads(
                 self.deployment.step_configurations
             )
@@ -269,9 +272,14 @@ def to_model(
                         self.pipeline_run.start_time,
                     )
                 )
+                merged_config = (
+                    full_step_config.config.apply_pipeline_configuration(
+                        pipeline_configuration
+                    )
+                )
                 full_step_config = full_step_config.model_copy(
                     update={
-                        "config": full_step_config.config.model_copy(
+                        "config": merged_config.model_copy(
                             update={"substitutions": new_substitutions}
                         )
                     }

src/zenml/zen_stores/template_utils.py

@@ -101,7 +101,7 @@ def generate_config_template(
             exclude={"name", "outputs"},
             exclude_none=True,
         )
-        for name, step in deployment_model.step_configurations.items()
+        for name, step in deployment_model.raw_step_configurations.items()
     }
 
     for config in steps_configs.values():

tests/integration/functional/zen_server/template_execution/test_template_execution_utils.py

@@ -69,5 +69,4 @@ def test_creating_deployment_request_from_template(
         deployment_request_from_template(
             template=template_response,
             config=PipelineRunConfiguration(),
-            user_id=deployment.user.id,
         )