fix(dsl): mismatched output names when using subkeys of a nested dag (#…

…48)
larribas · Sep 30, 2021 · c9ae0ef · c9ae0ef
1 parent 581eddd
commit c9ae0ef
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 14 deletions.
diff --git a/dagger/dsl/build.py b/dagger/dsl/build.py
@@ -154,7 +154,7 @@ def _build_dag_outputs(
     if isinstance(dag_output, NodeOutputReference):
         outputs_by_name = {"return_value": dag_output}
     elif isinstance(dag_output, Mapping):
-        outputs_by_name = dag_output
+        outputs_by_name = {f"key_{key}": value for key, value in dag_output.items()}
     elif dag_output is not None:
         raise TypeError(
             f"This DAG returned a value of type {type(dag_output).__name__}. Functions decorated with `dsl.DAG` may only return two types of values: The output of another node or a mapping of [str, the output of another node]"

diff --git a/tests/dsl/test_black_box.py b/tests/dsl/test_black_box.py
@@ -376,28 +376,39 @@ def generate_number() -> int:
         return 1
 
     @dsl.DAG()
-    def dag():
+    def inner_dag():
         return {
             "a": generate_number(),
             "b": generate_number(),
         }
 
+    @dsl.DAG()
+    def outer_dag():
+        return inner_dag()["a"]
+
     verify_dags_are_equivalent(
-        dsl.build(dag),
+        dsl.build(outer_dag),
         DAG(
             nodes={
-                "generate-number-1": Task(
-                    generate_number.func,
-                    outputs={"return_value": FromReturnValue()},
-                ),
-                "generate-number-2": Task(
-                    generate_number.func,
-                    outputs={"return_value": FromReturnValue()},
+                "inner-dag": DAG(
+                    nodes={
+                        "generate-number-1": Task(
+                            generate_number.func,
+                            outputs={"return_value": FromReturnValue()},
+                        ),
+                        "generate-number-2": Task(
+                            generate_number.func,
+                            outputs={"return_value": FromReturnValue()},
+                        ),
+                    },
+                    outputs={
+                        "key_a": FromNodeOutput("generate-number-1", "return_value"),
+                        "key_b": FromNodeOutput("generate-number-2", "return_value"),
+                    },
                 ),
             },
             outputs={
-                "a": FromNodeOutput("generate-number-1", "return_value"),
-                "b": FromNodeOutput("generate-number-2", "return_value"),
+                "return_value": FromNodeOutput("inner-dag", "key_a"),
             },
         ),
     )

diff --git a/tests/dsl/test_build.py b/tests/dsl/test_build.py
@@ -257,12 +257,12 @@ def test__build_dag_outputs__when_multiple_outputs_are_returned():
             "id-2": "name-2",
         },
     ) == {
-        "x": FromNodeOutput(
+        "key_x": FromNodeOutput(
             "name-1",
             "output",
             serializer=AsPickle(),
         ),
-        "y": FromNodeOutput(
+        "key_y": FromNodeOutput(
             "name-2",
             "return_value",
             serializer=AsJSON(indent=1),