Add a function to get the list of flow graph nodes in strict left-to-right order

Author: cstenac

dataikuapi/dss/flow.py

@@ -2,7 +2,7 @@
 from .dataset import DSSDataset
 from .recipe import DSSRecipe, DSSRecipeDefinitionAndPayload
 from .future import DSSFuture
-import logging
+import logging, json
 
 class DSSProjectFlow(object):
     def __init__(self, client, project):
@@ -90,6 +90,7 @@ class DSSProjectFlowGraph(object):
     def __init__(self, flow, data):
         self.flow = flow
         self.data = data
+        self.nodes = data["nodes"]
 
     def get_source_computables(self, as_type="dict"):
         """
@@ -102,17 +103,31 @@ def get_source_computables(self, as_type="dict"):
                     :class:`dataikuapi.dss.savedmodel.DSSSavedModel`, or streaming endpoint
         """
         ret = []
-        for computable in self.data["computables"].values():
-            if len(computable["predecessors"]) == 0:
-                ret.append(computable)
-        return self._convert_computables_list(ret, as_type)
+        for node in self.nodes.values():
+            if len(node["predecessors"]) == 0 and node["type"].startswith("COMPUTABLE"):
+                ret.append(node)
+        return self._convert_nodes_list(ret, as_type)
+
+    def get_source_recipes(self, as_type="dict"):
+        """
+        Returns the list of source recipes.
+        :param as_type: How to return the source recipes. Possible values are "dict" and "object"
+        
+        :return: if as_type=dict, each recipes is returned as a dict containing at least "ref" and "type". 
+                 if as_type=object, each computable is returned as a  :class:`dataikuapi.dss.recipe.DSSRecipe`,
+        """
+        ret = []
+        for node in self.nodes.values():
+            if len(node["predecessors"]) == 0 and node["type"] == "RUNNABLE_RECIPE":
+                ret.append(node)
+        return self._convert_nodes_list(ret, as_type)
 
     def get_source_datasets(self):
         """
         Returns the list of source datasets for this project.
         :rtype list of :class:`dataikuapi.dss.dataset.DSSDataset`
         """
-        return [self._get_object_from_computable(x) for x in self.get_source_computables() if x["type"] == "COMPUTABLE_DATASET"]
+        return [self._get_object_from_graph_node(x) for x in self.get_source_computables() if x["type"] == "COMPUTABLE_DATASET"]
 
     def get_successor_recipes(self, node, as_type="name"):
         """
@@ -125,43 +140,82 @@ def get_successor_recipes(self, node, as_type="name"):
         if isinstance(node, DSSDataset):
             node = node.dataset_name
 
-        computable = self.data["computables"].get(node, None)
+        computable = self.nodes.get(node, None)
         if computable is None:
             raise ValueError("Computable %s not found in Flow graph" % node)
 
-        names= computable["successors"]
-
-        if as_type == "object":
-            return [DSSRecipe(self.flow.client, self.flow.project.project_key, x) for x in names]
-        else:
-            return names
-
+        runnables = [self.nodes[x] for x in computable["successors"]]
+        return self._convert_nodes_list(runnables, as_type)
+        
     def get_successor_computables(self, node, as_type="dict"):
         """
         Returns a list of computables that are a successor of a given graph node
         Each computable is returned as a dict containing at least "ref" and "type"
         """
         if isinstance(node, DSSRecipe):
             node = node.recipe_name
-        runnable = self.data["runnables"].get(node, None)
+        runnable = self.nodes.get(node, None)
         if runnable is None:
             raise ValueError("Runnable %s not found in Flow graph" % node)
 
-        computables = [self.data["computables"][x] for x in runnable["successors"]]
-        return self._convert_computables_list(computables, as_type)
+        computables = [self.nodes[x] for x in runnable["successors"]]
+        return self._convert_nodes_list(computables, as_type)
 
-    def _convert_computables_list(self, computables, as_type):
+    def _convert_nodes_list(self, nodes, as_type):
         if as_type == "object":
-            return [self._get_object_from_computable(computable) for computable in computables]
+            return [self._get_object_from_graph_node(node) for node in nodes]
         else:
-            return computables
+            return nodes
 
-    def _get_object_from_computable(self, computable):
-        if computable["type"] == "COMPUTABLE_DATASET":
-            return DSSDataset(self.flow.client, self.flow.project.project_key, computable["ref"])
+    def _get_object_from_graph_node(self, node):
+        if node["type"] == "COMPUTABLE_DATASET":
+            return DSSDataset(self.flow.client, self.flow.project.project_key, node["ref"])
+        elif node["type"] == "RUNNABLE_RECIPE":
+            return DSSRecipe(self.flow.client, self.flow.project.project_key, node["ref"])
         else:
             # TODO
-            raise Exception("unsupported %s" % computable["type"])
+            raise Exception("unsupported node type  %s" % node["type"])
+
+    def get_items_in_traversal_order(self, as_type="dict"):
+        ret = []
+        def add_to_set(node):
+            print("*** Adding: %s" % node["ref"])
+            ret.append(node)
+        def in_set(obj):
+            for candidate in ret:
+                if candidate["type"] == obj["type"] and candidate["ref"] == obj["ref"]:
+                    return True
+            return False
+
+        def add_from(graph_node):
+            #print("Add from %s" % graph_node["ref"])
+            # To keep traversal order, we add predecessors before successors. And we recurse before adding
+            # our own predecessors
+            for predecessor_ref in graph_node["predecessors"]:
+                #print("  Pred = %s " % predecessor_ref)
+                predecessor_node = self.nodes[predecessor_ref]
+                if not in_set(predecessor_node):
+                    add_from(predecessor_node)
+
+            # Then ourselves
+            #print("  Itself %s" % graph_node["ref"])
+            if not in_set(graph_node):
+                add_to_set(graph_node)
+
+            # For successors, we first add ourselves before successors
+            for successor_ref in graph_node["successors"]:
+                #print("  Succ = %s " % successor_ref)
+                successor_node = self.nodes[successor_ref]
+                if not in_set(successor_node):
+                    add_from(successor_node)
+
+        for source_computable in self.get_source_computables():
+            add_from(source_computable)
+
+        for source_recipe in self.get_source_recipes():
+            add_from(source_recipe)
+
+        return self._convert_nodes_list(ret, as_type)
 
 class DSSFlowTool(object):
     """