fix: multi-merge case when key is not used

Signed-off-by: zjgemi <[email protected]>
deepmodeling · Mar 28, 2024 · 74e1bd8 · 74e1bd8
1 parent 0f2f823
commit 74e1bd8
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/src/dflow/step.py b/src/dflow/step.py
@@ -590,6 +590,19 @@ def merge_step_output_artifact(art, parent):
                         step.inputs.parameters["dflow_artifact_key"] = \
                             InputParameter(value="{{inputs.parameters."
                                            "dflow_artifact_key}}")
+                        if step.prepare_step is not None and art.name in \
+                                step.prepare_step.outputs.artifacts:
+                            # for multi-merge, save output artifact of inner
+                            # slices all together
+                            step.prepare_step.template.inputs.parameters[
+                                "dflow_artifact_key"] = InputParameter()
+                            step.prepare_step.inputs.parameters[
+                                "dflow_artifact_key"] = InputParameter(
+                                value="{{inputs.parameters."
+                                "dflow_artifact_key}}")
+                            merge_output_artifact(
+                                step.prepare_step.template.outputs.artifacts[
+                                    art.name], None)
                         merge_output_artifact(
                             template.outputs.artifacts[art.name], template)