Some cleanup

Author: schustmi

src/zenml/config/compiler.py

@@ -27,6 +27,7 @@
 
 from zenml import __version__
 from zenml.config.base_settings import BaseSettings, ConfigurationLevel
+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.settings_resolver import SettingsResolver
@@ -43,7 +44,6 @@
 from zenml.utils import pydantic_utils, settings_utils
 
 if TYPE_CHECKING:
-    from zenml.config.source import Source
     from zenml.pipelines.pipeline_definition import Pipeline
     from zenml.stack import Stack, StackComponent
     from zenml.steps.step_invocation import StepInvocation
@@ -121,12 +121,9 @@ def compile(
         steps = {
             invocation_id: self._compile_step_invocation(
                 invocation=invocation,
-                pipeline_settings=settings_to_passdown,
-                pipeline_extra=pipeline.configuration.extra,
                 stack=stack,
                 step_config=run_configuration.steps.get(invocation_id),
-                pipeline_failure_hook_source=pipeline.configuration.failure_hook_source,
-                pipeline_success_hook_source=pipeline.configuration.success_hook_source,
+                pipeline_configuration=pipeline.configuration,
             )
             for invocation_id, invocation in self._get_sorted_invocations(
                 pipeline=pipeline
@@ -438,24 +435,17 @@ def _get_step_spec(
     def _compile_step_invocation(
         self,
         invocation: "StepInvocation",
-        pipeline_settings: Dict[str, "BaseSettings"],
-        pipeline_extra: Dict[str, Any],
         stack: "Stack",
         step_config: Optional["StepConfigurationUpdate"],
-        pipeline_failure_hook_source: Optional["Source"] = None,
-        pipeline_success_hook_source: Optional["Source"] = None,
+        pipeline_configuration: "PipelineConfiguration",
     ) -> Step:
         """Compiles a ZenML step.
 
         Args:
             invocation: The step invocation to compile.
-            pipeline_settings: settings configured on the
-                pipeline of the step.
-            pipeline_extra: Extra values configured on the pipeline of the step.
             stack: The stack on which the pipeline will be run.
             step_config: Run configuration for the step.
-            pipeline_failure_hook_source: Source for the failure hook.
-            pipeline_success_hook_source: Source for the success hook.
+            pipeline_configuration: Configuration for the pipeline.
 
         Returns:
             The compiled step.
@@ -479,32 +469,27 @@ 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.configure(
             settings=step_settings,
-            # extra=step_extra,
-            # on_failure=step_on_failure_hook_source,
-            # on_success=step_on_success_hook_source,
             merge=False,
         )
 
         parameters_to_ignore = (
             set(step_config.parameters) if step_config else set()
         )
-        complete_step_configuration = invocation.finalize(
+        step_configuration_overrides = invocation.finalize(
             parameters_to_ignore=parameters_to_ignore
         )
-        return Step(spec=step_spec, config=complete_step_configuration)
+        full_step_config = (
+            step_configuration_overrides.apply_pipeline_configuration(
+                pipeline_configuration=pipeline_configuration
+            )
+        )
+        return Step(
+            spec=step_spec,
+            config=full_step_config,
+            step_config_overrides=step_configuration_overrides,
+        )
 
     def _get_sorted_invocations(
         self,

src/zenml/config/step_configurations.py

@@ -344,3 +344,21 @@ class Step(StrictBaseModel):
 
     spec: StepSpec
     config: StepConfiguration
+    step_config_overrides: StepConfiguration
+
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _autocomplete_step_config_overrides(cls, data: Any) -> Any:
+        """Autocompletes the step config overrides.
+
+        Args:
+            data: The values dict used to instantiate the model.
+
+        Returns:
+            The values dict with the step config overrides autocompleted.
+        """
+        if "step_config_overrides" not in data:
+            data["step_config_overrides"] = data["config"]
+
+        return data

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

@@ -60,9 +60,6 @@ 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."
     )
@@ -146,9 +143,6 @@ 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."
     )
@@ -247,15 +241,6 @@ 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/zen_server/template_execution/utils.py

@@ -371,7 +371,7 @@ def deployment_request_from_template(
     )
 
     steps = {}
-    for invocation_id, step in deployment.raw_step_configurations.items():
+    for invocation_id, step in deployment.step_configurations.items():
         step_update = config.steps.get(
             invocation_id, StepConfigurationUpdate()
         ).model_dump(
@@ -380,10 +380,12 @@ def deployment_request_from_template(
             exclude={"name"},
             exclude_unset=True,
         )
-        step_config = pydantic_utils.update_model(step.config, step_update)
-        # step_config = step_config.apply_pipeline_configuration(
-        #     pipeline_configuration
-        # )
+        step_config = pydantic_utils.update_model(
+            step.step_config_overrides, step_update
+        )
+        merged_step_config = step_config.apply_pipeline_configuration(
+            pipeline_configuration
+        )
 
         required_parameters = set(step.config.parameters)
         configured_parameters = set(step_config.parameters)
@@ -402,7 +404,11 @@ def deployment_request_from_template(
                 f"parameters for step {invocation_id}: {missing_parameters}."
             )
 
-        steps[invocation_id] = Step(spec=step.spec, config=step_config)
+        steps[invocation_id] = Step(
+            spec=step.spec,
+            config=merged_step_config,
+            step_config_overrides=step_config,
+        )
 
     code_reference_request = None
     if deployment.code_reference:
@@ -420,8 +426,7 @@ 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,
-        raw_step_configurations=steps,
+        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.raw_step_configurations,
+                request.step_configurations,
                 sort_keys=False,
                 default=pydantic_encoder,
             ),
@@ -249,16 +249,7 @@ def to_model(
                 run_name_template=self.run_name_template,
                 pipeline_configuration=pipeline_configuration,
                 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()
-                },
+                step_configurations=step_configurations,
                 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,9 +254,6 @@ 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
             )
@@ -272,14 +269,9 @@ 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": merged_config.model_copy(
+                        "config": full_step_config.config.model_copy(
                             update={"substitutions": new_substitutions}
                         )
                     }

src/zenml/zen_stores/template_utils.py

@@ -96,12 +96,12 @@ def generate_config_template(
     deployment_model = deployment.to_model(include_metadata=True)
 
     steps_configs = {
-        name: step.config.model_dump(
+        name: step.step_config_overrides.model_dump(
             include=set(StepConfigurationUpdate.model_fields),
             exclude={"name", "outputs"},
             exclude_none=True,
         )
-        for name, step in deployment_model.raw_step_configurations.items()
+        for name, step in deployment_model.step_configurations.items()
     }
 
     for config in steps_configs.values():