Adding node semantics: input type, firing type, state type; Adding lineage: previous and current state as well as objects

Author: raghukiran1224

codeflare/pipelines/Datamodel.py

@@ -1,8 +1,11 @@
-from sklearn.base import BaseEstimator
 from abc import ABC, abstractmethod
+import uuid
+from enum import Enum
+
 
 import sklearn.base as base
-import uuid
+from sklearn.base import TransformerMixin
+from sklearn.base import BaseEstimator
 
 
 class Xy:
@@ -38,10 +41,11 @@ class XYRef:
     computed), these holders are essential to the pipeline constructs.
     """
 
-    def __init__(self, Xref, yref, node=None, prev_Xyrefs = None):
+    def __init__(self, Xref, yref, prev_noderef=None, curr_noderef=None, prev_Xyrefs = None):
         self.__Xref__ = Xref
         self.__yref__ = yref
-        self.__noderef__ = node
+        self.__prevnoderef__ = prev_noderef
+        self.__currnoderef__ = curr_noderef
         self.__prev_Xyrefs__ = prev_Xyrefs
 
     def get_Xref(self):
@@ -56,29 +60,60 @@ def get_yref(self):
         """
         return self.__yref__
 
-    def get_noderef(self):
-        return self.__noderef__
+    def get_prevnoderef(self):
+        return self.__prevnoderef__
+
+    def get_currnoderef(self):
+        return self.__currnoderef__
 
     def get_prev_xyrefs(self):
         return self.__prev_Xyrefs__
 
 
+class NodeInputType(Enum):
+    OR = 0,
+    AND = 1
+
+
+class NodeFiringType(Enum):
+    ANY = 0,
+    ALL = 1
+
+
+class NodeStateType(Enum):
+    STATELESS = 0,
+    IMMUTABLE = 1,
+    MUTABLE_SEQUENTIAL = 2,
+    MUTABLE_AGGREGATE = 3
+
+
 class Node(ABC):
     """
     A node class that is an abstract one, this is capturing basic info re the Node.
     The hash code of this node is the name of the node and equality is defined if the
     node name and the type of the node match.
     """
+    def __init__(self, node_name, node_input_type: NodeInputType, node_firing_type: NodeFiringType, node_state_type: NodeStateType):
+        self.__node_name__ = node_name
+        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__
 
-    @abstractmethod
-    def get_and_flag(self):
-        raise NotImplementedError("Please implement this method")
+    def get_node_input_type(self):
+        return self.__node_input_type__
+
+    def get_node_firing_type(self):
+        return self.__node_firing_type__
+
+    def get_node_state_type(self):
+        return self.__node_state_type__
 
     @abstractmethod
     def clone(self):
@@ -107,7 +142,7 @@ def __eq__(self, other):
         )
 
 
-class OrNode(Node):
+class EstimatorNode(Node):
     """
     Or node, which is the basic node that would be the equivalent of any SKlearn pipeline
     stage. This node is initialized with an estimator that needs to extend sklearn.BaseEstimator.
@@ -120,9 +155,8 @@ def __init__(self, node_name: str, estimator: BaseEstimator):
         :param node_name: Name of the node
         :param estimator: The base estimator
         """
-        self.__node_name__ = node_name
+        super().__init__(node_name, NodeInputType.OR, NodeFiringType.ANY, NodeStateType.IMMUTABLE)
         self.__estimator__ = estimator
-        self.__id__ = uuid.uuid4()
 
     def get_estimator(self) -> BaseEstimator:
         """
@@ -132,41 +166,31 @@ def get_estimator(self) -> BaseEstimator:
         """
         return self.__estimator__
 
-    def get_and_flag(self):
-        """
-        A flag to check if node is AND or not. By definition, this is NOT
-        an AND node.
-        :return: False, always
-        """
-        return False
-
     def clone(self):
         cloned_estimator = base.clone(self.__estimator__)
-        return OrNode(self.__node_name__, cloned_estimator)
+        return EstimatorNode(self.__node_name__, cloned_estimator)
 
 
-class AndFunc(ABC):
-    """
-    Or nodes are init-ed from the
-    """
+class AndTransform(TransformerMixin, BaseEstimator):
+    @abstractmethod
+    def transform(self, xy_list: list) -> Xy:
+        raise NotImplementedError("Please implement this method")
+
 
+class GeneralTransform(TransformerMixin, BaseEstimator):
     @abstractmethod
-    def eval(self, xy_list: list) -> Xy:
+    def transform(self, xy: Xy) -> Xy:
         raise NotImplementedError("Please implement this method")
 
 
 class AndNode(Node):
-    def __init__(self, node_name: str, and_func: AndFunc):
-        self.__node_name__ = node_name
+    def __init__(self, node_name: str, and_func: AndTransform):
+        super().__init__(node_name, NodeInputType.AND, NodeFiringType.ANY, NodeStateType.STATELESS)
         self.__andfunc__ = and_func
-        self.__id__ = uuid.uuid4()
 
-    def get_and_func(self) -> AndFunc:
+    def get_and_func(self) -> AndTransform:
         return self.__andfunc__
 
-    def get_and_flag(self):
-        return True
-
     def clone(self):
         return AndNode(self.__node_name__, self.__andfunc__)
 

codeflare/pipelines/Runtime.py

@@ -1,11 +1,14 @@
 import ray
 
-from codeflare.pipelines.Datamodel import OrNode
+from codeflare.pipelines.Datamodel import EstimatorNode
 from codeflare.pipelines.Datamodel import AndNode
 from codeflare.pipelines.Datamodel import Edge
 from codeflare.pipelines.Datamodel import Pipeline
 from codeflare.pipelines.Datamodel import XYRef
 from codeflare.pipelines.Datamodel import Xy
+from codeflare.pipelines.Datamodel import NodeInputType
+from codeflare.pipelines.Datamodel import NodeStateType
+from codeflare.pipelines.Datamodel import NodeFiringType
 
 import sklearn.base as base
 from enum import Enum
@@ -18,43 +21,49 @@ class ExecutionType(Enum):
 
 
 @ray.remote
-def execute_or_node_inner(node: OrNode, train_mode: ExecutionType, xy_ref: XYRef):
+def execute_or_node_remote(node: EstimatorNode, train_mode: ExecutionType, xy_ref: XYRef):
     estimator = node.get_estimator()
     # Blocking operation -- not avoidable
     X = ray.get(xy_ref.get_Xref())
     y = ray.get(xy_ref.get_yref())
 
+    # TODO: Can optimize the node pointers without replicating them
     if train_mode == ExecutionType.FIT:
         cloned_node = node.clone()
-        node_ptr = ray.put(cloned_node)
+        prev_node_ptr = ray.put(node)
 
         if base.is_classifier(estimator) or base.is_regressor(estimator):
             # Always clone before fit, else fit is invalid
             cloned_estimator = cloned_node.get_estimator()
             cloned_estimator.fit(X, y)
+
+            curr_node_ptr = ray.put(cloned_node)
             # TODO: For now, make yref passthrough - this has to be fixed more comprehensively
             res_Xref = ray.put(cloned_estimator.predict(X))
-            result = XYRef(res_Xref, xy_ref.get_yref(), node_ptr, [xy_ref])
+            result = XYRef(res_Xref, xy_ref.get_yref(), prev_node_ptr, curr_node_ptr, [xy_ref])
             return result
         else:
             cloned_estimator = cloned_node.get_estimator()
             res_Xref = ray.put(cloned_estimator.fit_transform(X, y))
-            result = XYRef(res_Xref, xy_ref.get_yref(), node_ptr, [xy_ref])
+            curr_node_ptr = ray.put(cloned_node)
+            result = XYRef(res_Xref, xy_ref.get_yref(), prev_node_ptr, curr_node_ptr, [xy_ref])
             return result
     elif train_mode == ExecutionType.SCORE:
         cloned_node = node.clone()
-        node_ptr = ray.put(cloned_node)
+        prev_node_ptr = ray.put(node)
 
         if base.is_classifier(estimator) or base.is_regressor(estimator):
             cloned_estimator = cloned_node.get_estimator()
             cloned_estimator.fit(X, y)
+            curr_node_ptr = ray.put(cloned_node)
             res_Xref = ray.put(cloned_estimator.score(X, y))
-            result = XYRef(res_Xref, xy_ref.get_yref(), node_ptr, [xy_ref])
+            result = XYRef(res_Xref, xy_ref.get_yref(), prev_node_ptr, curr_node_ptr, [xy_ref])
             return result
         else:
             cloned_estimator = cloned_node.get_estimator()
             res_Xref = ray.put(cloned_estimator.fit_transform(X, y))
-            result = XYRef(res_Xref, xy_ref.get_yref(), node_ptr, [xy_ref])
+            curr_node_ptr = ray.put(cloned_node)
+            result = XYRef(res_Xref, xy_ref.get_yref(), prev_node_ptr, curr_node_ptr, [xy_ref])
             return result
     elif train_mode == ExecutionType.PREDICT:
         # Test mode does not clone as it is a simple predict or transform
@@ -74,7 +83,7 @@ def execute_or_node(node, pre_edges, edge_args, post_edges, mode: ExecutionType)
         exec_xyrefs = []
         for xy_ref_ptr in Xyref_ptrs:
             xy_ref = ray.get(xy_ref_ptr)
-            inner_result = execute_or_node_inner.remote(node, mode, xy_ref)
+            inner_result = execute_or_node_remote.remote(node, mode, xy_ref)
             exec_xyrefs.append(inner_result)
 
         for post_edge in post_edges:
@@ -84,21 +93,22 @@ def execute_or_node(node, pre_edges, edge_args, post_edges, mode: ExecutionType)
 
 
 @ray.remote
-def and_node_eval(node: AndNode, Xyref_list):
+def execute_and_node_remote(node: AndNode, Xyref_list):
     xy_list = []
+    prev_node_ptr = ray.put(node)
     for Xyref in Xyref_list:
         X = ray.get(Xyref.get_Xref())
         y = ray.get(Xyref.get_yref())
         xy_list.append(Xy(X, y))
 
     cloned_node = node.clone()
-    node_ptr = ray.put(cloned_node)
+    curr_node_ptr = ray.put(cloned_node)
 
     cloned_and_func = cloned_node.get_and_func()
-    res_Xy = cloned_and_func.eval(xy_list)
+    res_Xy = cloned_and_func.transform(xy_list)
     res_Xref = ray.put(res_Xy.get_x())
     res_yref = ray.put(res_Xy.get_y())
-    return XYRef(res_Xref, res_yref, node_ptr, Xyref_list)
+    return XYRef(res_Xref, res_yref, prev_node_ptr, curr_node_ptr, Xyref_list)
 
 
 def execute_and_node_inner(node: AndNode, Xyref_ptrs):
@@ -109,7 +119,7 @@ def execute_and_node_inner(node: AndNode, Xyref_ptrs):
         Xyref = ray.get(Xyref_ptr)
         Xyref_list.append(Xyref)
 
-    Xyref_ptr = and_node_eval.remote(node, Xyref_list)
+    Xyref_ptr = execute_and_node_remote.remote(node, Xyref_list)
     result.append(Xyref_ptr)
     return result
 
@@ -145,9 +155,9 @@ def execute_pipeline(pipeline: Pipeline, mode: ExecutionType, in_args: dict):
         for node in nodes:
             pre_edges = pipeline.get_pre_edges(node)
             post_edges = pipeline.get_post_edges(node)
-            if not node.get_and_flag():
+            if node.get_node_input_type() == NodeInputType.OR:
                 execute_or_node(node, pre_edges, edge_args, post_edges, mode)
-            elif node.get_and_flag():
+            elif node.get_node_input_type() == NodeInputType.AND:
                 execute_and_node(node, pre_edges, edge_args, post_edges)
 
     out_args = {}