Merge pull request #27 from codeflare/pipeline_terminal_bug

Pipeline terminal bug

Author: yuanchi2807

codeflare/pipelines/Datamodel.py

@@ -1,12 +1,13 @@
 from abc import ABC, abstractmethod
-import uuid
 from enum import Enum
 
-
 import sklearn.base as base
 from sklearn.base import TransformerMixin
 from sklearn.base import BaseEstimator
 
+import ray
+import codeflare.pipelines.Exceptions as pe
+
 class Xy:
     """
     Holder class for Xy, where X is array-like and y is array-like. This is the base
@@ -40,11 +41,11 @@ class XYRef:
     computed), these holders are essential to the pipeline constructs.
     """
 
-    def __init__(self, Xref, yref, prev_noderef=None, curr_noderef=None, prev_Xyrefs = None):
+    def __init__(self, Xref, yref, prev_node_state_ref=None, curr_node_state_ref=None, prev_Xyrefs = None):
         self.__Xref__ = Xref
         self.__yref__ = yref
-        self.__prevnoderef__ = prev_noderef
-        self.__currnoderef__ = curr_noderef
+        self.__prev_node_state_ref__ = prev_node_state_ref
+        self.__curr_node_state_ref__ = curr_node_state_ref
         self.__prev_Xyrefs__ = prev_Xyrefs
 
     def get_Xref(self):
@@ -59,11 +60,11 @@ def get_yref(self):
         """
         return self.__yref__
 
-    def get_prevnoderef(self):
-        return self.__prevnoderef__
+    def get_prev_node_state_ref(self):
+        return self.__prev_node_state_ref__
 
-    def get_currnoderef(self):
-        return self.__currnoderef__
+    def get_curr_node_state_ref(self):
+        return self.__curr_node_state_ref__
 
     def get_prev_xyrefs(self):
         return self.__prev_Xyrefs__
@@ -98,14 +99,10 @@ def __init__(self, node_name, node_input_type: NodeInputType, node_firing_type:
         self.__node_input_type__ = node_input_type
         self.__node_firing_type__ = node_firing_type
         self.__node_state_type__ = node_state_type
-        self.__id__ = uuid.uuid4()
 
     def __str__(self):
         return self.__node_name__
 
-    def get_id(self):
-        return self.__id__
-
     def get_node_input_type(self):
         return self.__node_input_type__
 
@@ -125,8 +122,7 @@ def __hash__(self):
 
         :return: Hash code
         """
-
-        return self.__id__.__hash__()
+        return self.__node_name__.__hash__()
 
     def __eq__(self, other):
         """
@@ -138,7 +134,6 @@ def __eq__(self, other):
         """
         return (
                 self.__class__ == other.__class__ and
-                self.__id__ == other.__id__ and
                 self.__node_name__ == other.__node_name__
         )
 
@@ -373,5 +368,69 @@ def get_post_edges(self, node: Node):
         return post_edges
 
     def is_terminal(self, node: Node):
-        node_post_edges = self.get_post_edges(node)
-        return len(node_post_edges) == 0
+        post_nodes = self.__post_graph__[node]
+        return not post_nodes
+
+    def get_terminal_nodes(self):
+        # dict from level to nodes
+        terminal_nodes = []
+        for node in self.__pre_graph__.keys():
+            if self.is_terminal(node):
+                terminal_nodes.append(node)
+        return terminal_nodes
+
+
+class PipelineOutput:
+    """
+    Pipeline output to keep reference counters so that pipelines can be materialized
+    """
+    def __init__(self, out_args, edge_args):
+        self.__out_args__ = out_args
+        self.__edge_args__ = edge_args
+
+    def get_xyrefs(self, node: Node):
+        if node in self.__out_args__:
+            xyrefs_ptr = self.__out_args__[node]
+        elif node in self.__edge_args__:
+            xyrefs_ptr = self.__edge_args__[node]
+        else:
+            raise pe.PipelineNodeNotFoundException("Node " + str(node) + " not found")
+
+        xyrefs = ray.get(xyrefs_ptr)
+        return xyrefs
+
+    def get_edge_args(self):
+        return self.__edge_args__
+
+
+class PipelineInput:
+    """
+    in_args is a dict from a node -> [Xy]
+    """
+    def __init__(self):
+        self.__in_args__ = {}
+
+    def add_xyref_ptr_arg(self, node: Node, xyref_ptr):
+        if node not in self.__in_args__:
+            self.__in_args__[node] = []
+
+        self.__in_args__[node].append(xyref_ptr)
+
+    def add_xyref_arg(self, node: Node, xyref: XYRef):
+        if node not in self.__in_args__:
+            self.__in_args__[node] = []
+
+        xyref_ptr = ray.put(xyref)
+        self.__in_args__[node].append(xyref_ptr)
+
+    def add_xy_arg(self, node: Node, xy: Xy):
+        if node not in self.__in_args__:
+            self.__in_args__[node] = []
+
+        x_ref = ray.put(xy.get_x())
+        y_ref = ray.put(xy.get_y())
+        xyref = XYRef(x_ref, y_ref)
+        self.add_xyref_arg(node, xyref)
+
+    def get_in_args(self):
+        return self.__in_args__

codeflare/pipelines/Exceptions.py

@@ -0,0 +1,17 @@
+class BasePipelineException(Exception):
+    pass
+
+
+class PipelineSaveException(BasePipelineException):
+    def __init__(self, message):
+        self.message = message
+
+
+class PipelineNodeNotFoundException(BasePipelineException):
+    def __init__(self, message):
+        self.message = message
+
+
+class PipelineException(BasePipelineException):
+    def __init__(self, message):
+        self.message = message