make dag info available through argo workflow template annotation

Netflix · Dec 11, 2024 · 3904fe8 · 3904fe8
1 parent b3b109a
commit 3904fe8
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/metaflow/plugins/argo/argo_workflows.py b/metaflow/plugins/argo/argo_workflows.py
@@ -744,6 +744,18 @@ def _compile_workflow_template(self):
                     )
                 }
             )
+        try:
+            # Build the DAG based on the DAGNodes given by the FlowGraph for the found FlowSpec class.
+            steps_info, graph_structure = self.graph.output_steps()
+            graph_info = {
+                "steps": steps_info,
+                "graph_structure": graph_structure,
+                "doc": self.graph.doc,
+            }
+        except Exception:
+            graph_info = None
+
+        dag_annotation = {"metaflow/dag": json.dumps(graph_info)}
 
         return (
             WorkflowTemplate()
@@ -758,6 +770,7 @@ def _compile_workflow_template(self):
                 .label("app.kubernetes.io/name", "metaflow-flow")
                 .label("app.kubernetes.io/part-of", "metaflow")
                 .annotations(annotations)
+                .annotations(dag_annotation)
             )
             .spec(
                 WorkflowSpec()