Refactoring

Author: schustmi

src/zenml/config/compiler.py

@@ -489,15 +489,15 @@ def _compile_step_invocation(
         invocation.step = copy.deepcopy(invocation.step)
 
         step = invocation.step
-        with step._skip_dynamic_configuration():
+        with step._suspend_dynamic_configuration():
             if step_config:
                 step._apply_configuration(
                     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()
+            step._merge_dynamic_configuration()
 
             convert_component_shortcut_settings_keys(
                 step.configuration.settings, stack=stack

src/zenml/execution/step/utils.py

@@ -0,0 +1,99 @@
+#  Copyright (c) ZenML GmbH 2025. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Step execution utilities."""
+
+import time
+from typing import (
+    TYPE_CHECKING,
+)
+
+from zenml.config.step_configurations import Step
+from zenml.exceptions import RunStoppedException
+from zenml.logger import get_logger
+from zenml.models import (
+    PipelineSnapshotResponse,
+)
+from zenml.models.v2.core.step_run import StepRunResponse
+from zenml.orchestrators.step_launcher import StepLauncher
+
+if TYPE_CHECKING:
+    from zenml.config.step_configurations import Step
+
+
+logger = get_logger(__name__)
+
+
+def launch_step(
+    snapshot: "PipelineSnapshotResponse",
+    step: "Step",
+    orchestrator_run_id: str,
+    retry: bool = False,
+) -> StepRunResponse:
+    """Launch a step.
+
+    Args:
+        snapshot: The snapshot.
+        step: The step to run.
+        orchestrator_run_id: The orchestrator run ID.
+        retry: Whether to retry the step if it fails.
+
+    Returns:
+        The step run response.
+    """
+
+    def _launch_without_retry() -> StepRunResponse:
+        launcher = StepLauncher(
+            snapshot=snapshot,
+            step=step,
+            orchestrator_run_id=orchestrator_run_id,
+        )
+        return launcher.launch()
+
+    if not retry:
+        step_run = _launch_without_retry()
+    else:
+        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:
+                step_run = _launch_without_retry()
+            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
+
+    return step_run

src/zenml/orchestrators/base_orchestrator.py

@@ -213,7 +213,18 @@ def submit_dynamic_pipeline(
         environment: Dict[str, str],
         placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Optional[SubmissionResult]:
-        """Submits a dynamic pipeline to the orchestrator."""
+        """Submits a dynamic pipeline to the orchestrator.
+
+        Args:
+            snapshot: The pipeline snapshot to submit.
+            stack: The stack the pipeline will run on.
+            environment: Environment variables to set in the orchestration
+                environment.
+            placeholder_run: An optional placeholder run.
+
+        Returns:
+            Optional submission result.
+        """
         return None
 
     def prepare_or_run_pipeline(
@@ -427,9 +438,11 @@ 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
+        from zenml.execution.step.utils import launch_step
+
+        assert self._active_snapshot
 
-        _run_step_sync(
+        launch_step(
             snapshot=self._active_snapshot,
             step=step,
             orchestrator_run_id=self.get_orchestrator_run_id(),
@@ -438,23 +451,41 @@ def run_step(
 
     @property
     def supports_dynamic_pipelines(self) -> bool:
+        """Whether the orchestrator supports dynamic pipelines.
+
+        Returns:
+            Whether the orchestrator supports dynamic pipelines.
+        """
         return (
             getattr(self.submit_dynamic_pipeline, "__func__", None)
             is not BaseOrchestrator.submit_dynamic_pipeline
         )
 
     @property
-    def supports_dynamic_out_of_process_steps(self) -> bool:
+    def can_launch_dynamic_steps(self) -> bool:
+        """Whether the orchestrator can launch dynamic steps.
+
+        Returns:
+            Whether the orchestrator can launch dynamic steps.
+        """
         return (
-            getattr(self.run_dynamic_out_of_process_step, "__func__", None)
-            is not BaseOrchestrator.run_dynamic_out_of_process_step
+            getattr(self.launch_dynamic_step, "__func__", None)
+            is not BaseOrchestrator.launch_dynamic_step
         )
 
-    def run_dynamic_out_of_process_step(
+    def launch_dynamic_step(
         self, step_run_info: "StepRunInfo", environment: Dict[str, str]
     ) -> None:
+        """Launch a dynamic step.
+
+        Args:
+            step_run_info: The step run information.
+            environment: The environment variables to set in the execution
+                environment.
+        """
         raise NotImplementedError(
-            "Running dynamic out of process steps is not implemented for the orchestrator."
+            "Launching dynamic steps is not implemented for "
+            f"the {self.__class__.__name__} orchestrator."
         )
 
     @staticmethod

src/zenml/orchestrators/step_launcher.py

@@ -481,8 +481,7 @@ def _run_step(
                         # the orchestrator doesn't support it.
                         logger.warning(
                             "The %s does not support running dynamic out of "
-                            "process steps. Running step `%s` in current "
-                            "thread instead.",
+                            "process steps. Running step `%s` locally instead.",
                             self._stack.orchestrator.__class__.__name__,
                             self._invocation_id,
                         )
@@ -574,7 +573,7 @@ def _run_step_with_dynamic_orchestrator(
                 stack=self._stack,
             )
         )
-        self._stack.orchestrator.run_dynamic_out_of_process_step(
+        self._stack.orchestrator.launch_dynamic_step(
             step_run_info=step_run_info,
             environment=environment,
         )