Cleanup

Author: schustmi

src/zenml/config/compiler.py

@@ -42,7 +42,6 @@
 from zenml.exceptions import StackValidationError
 from zenml.models import PipelineSnapshotBase
 from zenml.pipelines.run_utils import get_default_run_name
-from zenml.steps.step_invocation import StepInvocation
 from zenml.utils import pydantic_utils, secret_utils, settings_utils
 
 if TYPE_CHECKING:
@@ -127,12 +126,18 @@ def compile(
                 merge=False,
             )
 
+        # If we're compiling a dynamic pipeline, the steps are only templates
+        # and might not have all inputs defined, so we skip the input
+        # validation.
+        skip_input_validation = pipeline.is_dynamic
+
         steps = {
             invocation_id: self._compile_step_invocation(
                 invocation=invocation,
                 stack=stack,
                 step_config=(run_configuration.steps or {}).get(invocation_id),
                 pipeline_configuration=pipeline.configuration,
+                skip_input_validation=skip_input_validation,
             )
             for invocation_id, invocation in self._get_sorted_invocations(
                 pipeline=pipeline
@@ -465,6 +470,7 @@ def _compile_step_invocation(
         stack: "Stack",
         step_config: Optional["StepConfigurationUpdate"],
         pipeline_configuration: "PipelineConfiguration",
+        skip_input_validation: bool = False,
     ) -> Step:
         """Compiles a ZenML step.
 
@@ -473,6 +479,7 @@ def _compile_step_invocation(
             stack: The stack on which the pipeline will be run.
             step_config: Run configuration for the step.
             pipeline_configuration: Configuration for the pipeline.
+            skip_input_validation: If True, will skip the input validation.
 
         Returns:
             The compiled step.
@@ -487,6 +494,10 @@ def _compile_step_invocation(
                 step_config, runtime_parameters=invocation.parameters
             )
 
+        # Apply the dynamic configuration (which happened while executing the
+        # pipeline function) after all other step-specific configurations.
+        step._apply_dynamic_configuration()
+
         convert_component_shortcut_settings_keys(
             step.configuration.settings, stack=stack
         )
@@ -509,7 +520,8 @@ def _compile_step_invocation(
             set(step_config.parameters or {}) if step_config else set()
         )
         step_configuration_overrides = invocation.finalize(
-            parameters_to_ignore=parameters_to_ignore
+            parameters_to_ignore=parameters_to_ignore,
+            skip_input_validation=skip_input_validation,
         )
         full_step_config = (
             step_configuration_overrides.apply_pipeline_configuration(
@@ -535,9 +547,13 @@ def _get_sorted_invocations(
             pipeline: The pipeline of which to sort the invocations
 
         Returns:
-            The sorted steps.
+            The sorted step invocations.
         """
         if pipeline.is_dynamic:
+            # In dynamic pipelines, we require the static invocations to be
+            # sorted the same way they were passed in `pipeline.depends_on`, as
+            # we index this list later to figure out the correct template for
+            # each step invocation.
             return list(pipeline.invocations.items())
 
         from zenml.orchestrators.dag_runner import reverse_dag

src/zenml/config/step_configurations.py

@@ -261,6 +261,7 @@ class PartialStepConfiguration(StepConfigurationUpdate):
     """Class representing a partial step configuration."""
 
     name: str
+    # TODO: maybe move to spec?
     template: Optional[str] = None
     parameters: Dict[str, Any] = {}
     settings: Dict[str, SerializeAsAny[BaseSettings]] = {}

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

@@ -178,9 +178,9 @@ class PipelineRunUpdate(BaseUpdate):
         max_length=STR_FIELD_MAX_LENGTH,
     )
     end_time: Optional[datetime] = None
-    completed: Optional[bool] = Field(
+    finished: Optional[bool] = Field(
         default=None,
-        title="Whether the pipeline run is completed.",
+        title="Whether the pipeline run is finished.",
     )
     orchestrator_run_id: Optional[str] = None
     # TODO: we should maybe have a different update model here, the upper

src/zenml/orchestrators/base_orchestrator.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Base orchestrator class."""
 
-import time
 from abc import ABC, abstractmethod
 from typing import (
     TYPE_CHECKING,
@@ -40,7 +39,6 @@
     HookExecutionException,
     IllegalOperationError,
     RunMonitoringError,
-    RunStoppedException,
 )
 from zenml.hooks.hook_validators import load_and_run_hook
 from zenml.logger import get_logger
@@ -50,7 +48,6 @@
     publish_pipeline_run_status_update,
     publish_schedule_metadata,
 )
-from zenml.orchestrators.step_launcher import StepLauncher
 from zenml.orchestrators.utils import get_config_environment_vars
 from zenml.stack import Flavor, Stack, StackComponent, StackComponentConfig
 from zenml.steps.step_context import RunContext, get_or_create_run_context
@@ -274,16 +271,6 @@ def run(
         # in the orchestrator environment
         base_environment.update(secrets)
 
-        is_dynamic = True
-        if is_dynamic:
-            submission_result = self.submit_dynamic_pipeline(
-                snapshot=snapshot,
-                stack=stack,
-                environment=base_environment,
-                placeholder_run=placeholder_run,
-            )
-            return
-
         prevent_client_side_caching = handle_bool_env_var(
             ENV_ZENML_PREVENT_CLIENT_SIDE_CACHING, default=False
         )
@@ -292,6 +279,7 @@ def run(
             placeholder_run
             and self.config.supports_client_side_caching
             and not snapshot.schedule
+            and not snapshot.is_dynamic
             and not prevent_client_side_caching
         ):
             from zenml.orchestrators import cache_utils
@@ -310,22 +298,10 @@ def run(
         else:
             logger.debug("Skipping client-side caching.")
 
-        step_environments = {}
-        for invocation_id, step in snapshot.step_configurations.items():
-            from zenml.utils.env_utils import get_step_environment
-
-            step_environment = get_step_environment(
-                step_config=step.config,
-                stack=stack,
-            )
-
-            combined_environment = base_environment.copy()
-            combined_environment.update(step_environment)
-            step_environments[invocation_id] = combined_environment
-
         try:
             if (
-                getattr(self.submit_pipeline, "__func__", None)
+                not snapshot.is_dynamic
+                and getattr(self.submit_pipeline, "__func__", None)
                 is BaseOrchestrator.submit_pipeline
             ):
                 logger.warning(
@@ -357,13 +333,37 @@ def run(
                                 f"run metadata: {e}"
                             )
             else:
-                submission_result = self.submit_pipeline(
-                    snapshot=snapshot,
-                    stack=stack,
-                    base_environment=base_environment,
-                    step_environments=step_environments,
-                    placeholder_run=placeholder_run,
-                )
+                if snapshot.is_dynamic:
+                    submission_result = self.submit_dynamic_pipeline(
+                        snapshot=snapshot,
+                        stack=stack,
+                        environment=base_environment,
+                        placeholder_run=placeholder_run,
+                    )
+                else:
+                    step_environments = {}
+                    for (
+                        invocation_id,
+                        step,
+                    ) in snapshot.step_configurations.items():
+                        from zenml.utils.env_utils import get_step_environment
+
+                        step_environment = get_step_environment(
+                            step_config=step.config,
+                            stack=stack,
+                        )
+
+                        combined_environment = base_environment.copy()
+                        combined_environment.update(step_environment)
+                        step_environments[invocation_id] = combined_environment
+
+                    submission_result = self.submit_pipeline(
+                        snapshot=snapshot,
+                        stack=stack,
+                        base_environment=base_environment,
+                        step_environments=step_environments,
+                        placeholder_run=placeholder_run,
+                    )
                 if placeholder_run:
                     publish_pipeline_run_status_update(
                         pipeline_run_id=placeholder_run.id,
@@ -427,54 +427,14 @@ def run_step(
             RunStoppedException: If the run was stopped.
             BaseException: If the step failed all retries.
         """
+        from zenml.pipelines.dynamic.runner import _run_step_sync
 
-        def _launch_step() -> None:
-            assert self._active_snapshot
-
-            launcher = StepLauncher(
-                snapshot=self._active_snapshot,
-                step=step,
-                orchestrator_run_id=self.get_orchestrator_run_id(),
-            )
-            launcher.launch()
-
-        if self.config.handles_step_retries:
-            _launch_step()
-        else:
-            # The orchestrator subclass doesn't handle step retries, so we
-            # handle it in-process instead
-            retries = 0
-            retry_config = step.config.retry
-            max_retries = retry_config.max_retries if retry_config else 0
-            delay = retry_config.delay if retry_config else 0
-            backoff = retry_config.backoff if retry_config else 1
-
-            while retries <= max_retries:
-                try:
-                    _launch_step()
-                except RunStoppedException:
-                    # Don't retry if the run was stopped
-                    raise
-                except BaseException:
-                    retries += 1
-                    if retries <= max_retries:
-                        logger.info(
-                            "Sleeping for %d seconds before retrying step `%s`.",
-                            delay,
-                            step.config.name,
-                        )
-                        time.sleep(delay)
-                        delay *= backoff
-                    else:
-                        if max_retries > 0:
-                            logger.error(
-                                "Failed to run step `%s` after %d retries.",
-                                step.config.name,
-                                max_retries,
-                            )
-                        raise
-                else:
-                    break
+        _run_step_sync(
+            snapshot=self._active_snapshot,
+            step=step,
+            orchestrator_run_id=self.get_orchestrator_run_id(),
+            retry=not self.config.handles_step_retries,
+        )
 
     @property
     def supports_dynamic_pipelines(self) -> bool:

src/zenml/orchestrators/publish_utils.py

@@ -131,7 +131,7 @@ def publish_successful_pipeline_run(
         run_update=PipelineRunUpdate(
             status=ExecutionStatus.COMPLETED,
             end_time=utc_now(),
-            completed=True,
+            finished=True,
         ),
     )
 
@@ -152,7 +152,7 @@ def publish_failed_pipeline_run(
         run_update=PipelineRunUpdate(
             status=ExecutionStatus.FAILED,
             end_time=utc_now(),
-            completed=True,
+            finished=True,
         ),
     )
 

src/zenml/orchestrators/step_launcher.py

@@ -107,7 +107,6 @@ def __init__(
         snapshot: PipelineSnapshotResponse,
         step: Step,
         orchestrator_run_id: str,
-        dynamic: bool = False,
     ):
         """Initializes the launcher.
 
@@ -122,7 +121,6 @@ def __init__(
         self._snapshot = snapshot
         self._step = step
         self._orchestrator_run_id = orchestrator_run_id
-        self._dynamic = dynamic
 
         if not snapshot.stack:
             raise RuntimeError(
@@ -306,9 +304,10 @@ def launch(self) -> StepRunResponse:
                 pipeline_run=pipeline_run,
                 stack=self._stack,
             )
+            dynamic_config = self._step if self._snapshot.is_dynamic else None
             step_run_request = request_factory.create_request(
                 invocation_id=self._invocation_id,
-                dynamic_config=self._step if self._dynamic else None,
+                dynamic_config=dynamic_config,
             )
             step_run_request.logs = logs_model
 
@@ -461,19 +460,19 @@ def _run_step(
                     step_run_info=step_run_info,
                 )
             else:
+                from zenml.pipelines.dynamic.runner import (
+                    should_run_in_process,
+                )
+
                 should_run_out_of_process = (
                     self._snapshot.is_dynamic
                     and self._step.config.in_process is False
                 )
 
-                if (
-                    should_run_out_of_process
-                    and self._stack.orchestrator.supports_dynamic_out_of_process_steps
+                if should_run_in_process(
+                    self._step,
+                    self._snapshot.pipeline_configuration.docker_settings,
                 ):
-                    self._run_step_with_dynamic_orchestrator(
-                        step_run_info=step_run_info
-                    )
-                else:
                     if should_run_out_of_process:
                         logger.warning(
                             "The %s does not support running dynamic out of "
@@ -490,6 +489,10 @@ def _run_step(
                         input_artifacts=step_run.regular_inputs,
                         output_artifact_uris=output_artifact_uris,
                     )
+                else:
+                    self._run_step_with_dynamic_orchestrator(
+                        step_run_info=step_run_info
+                    )
         except:  # noqa: E722
             output_utils.remove_artifact_dirs(
                 artifact_uris=list(output_artifact_uris.values())