Maybe solution for step configs

Author: schustmi

src/zenml/config/compiler.py

@@ -42,6 +42,7 @@
 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:
@@ -536,6 +537,9 @@ def _get_sorted_invocations(
         Returns:
             The sorted steps.
         """
+        if pipeline.is_dynamic:
+            return list(pipeline.invocations.items())
+
         from zenml.orchestrators.dag_runner import reverse_dag
         from zenml.orchestrators.topsort import topsorted_layers
 

src/zenml/config/step_configurations.py

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

src/zenml/config/step_run_info.py

@@ -59,18 +59,8 @@ def get_image(self, key: str) -> str:
             )
 
         if self.snapshot.is_dynamic:
-            # TODO: better way for this
-            for (
-                invocation_id,
-                compiled_step,
-            ) in self.snapshot.step_configurations.values():
-                if (
-                    compiled_step.spec.source.import_path
-                    == self.spec.source.import_path
-                ):
-                    step_key = invocation_id
-                    break
-            else:
+            step_key = self.config.template
+            if not step_key:
                 logger.warning(
                     "Unable to find config template for step %s. Falling "
                     "back to the pipeline image.",

src/zenml/pipelines/dynamic/runner.py

@@ -286,8 +286,9 @@ def _compile_step(
         else:
             external_artifacts[name] = ExternalArtifact(value=value)
 
-    if template := get_static_step_template(snapshot, step):
+    if template := get_static_step_template(snapshot, step, pipeline):
         step._configuration = template.config
+        step._configuration.template = template.spec.invocation_id
 
     step._apply_dynamic_configuration()
     invocation_id = pipeline.add_dynamic_invocation(
@@ -417,11 +418,12 @@ def _runs_in_process(step: "Step") -> bool:
 def get_static_step_template(
     snapshot: "PipelineSnapshotResponse",
     step: "BaseStep",
+    pipeline: "DynamicPipeline",
 ) -> Optional["Step"]:
-    step_source = step.resolve().import_path
-
-    for compiled_step in snapshot.step_configurations.values():
-        if compiled_step.spec.source.import_path == step_source:
-            return compiled_step
+    for index, step_ in enumerate(pipeline.depends_on):
+        if step_._static_id == step._static_id:
+            break
+    else:
+        return None
 
-    return None
+    return list(snapshot.step_configurations.values())[index]

src/zenml/steps/base_step.py

@@ -166,6 +166,7 @@ def __init__(
             reserved_arguments=["after", "id"],
         )
 
+        self._static_id = id(self)
         name = name or self.__class__.__name__
 
         logger.debug(
@@ -881,7 +882,14 @@ def copy(self) -> "BaseStep":
         Returns:
             The step copy.
         """
-        return copy.deepcopy(self)
+        copy_ = copy.deepcopy(self)
+
+        from zenml.pipelines.dynamic.context import DynamicPipelineRunContext
+
+        if not DynamicPipelineRunContext.get():
+            copy_._static_id = id(copy_)
+
+        return copy_
 
     def _apply_configuration(
         self,