Merge pull request #15 from codeflare/develop

Author: chcost

com/__init__.py renamed to codeflare/__init__.py

@@ -3,40 +3,61 @@
 
 
 class Xy:
-    __X__ = None
-    __y__ = None
+    """
+    Holder class for Xy, where X is array-like and y is array-like. This is the base
+    data structure for fully materialized X and y.
+    """
 
     def __init__(self, X, y):
         self.__X__ = X
         self.__y__ = y
 
+    """
+    Returns the holder value of X
+    """
+
     def get_x(self):
         return self.__X__
 
+    """
+    Returns the holder value of y
+    """
+
     def get_y(self):
         return self.__y__
 
 
 class XYRef:
+    """
+    Holder class that maintains a pointer/reference to X and y. The goal of this is to provide
+    a holder to the object references of Ray. This is used for passing outputs from a transform/fit
+    to the next stage of the pipeline. Since the references can be potentially in flight (or being
+    computed), these holders are essential to the pipeline constructs.
+    """
+
     def __init__(self, Xref, yref):
-        self.Xref = Xref
-        self.yref = yref
+        self.__Xref__ = Xref
+        self.__yref__ = yref
 
     def get_Xref(self):
-        return self.Xref
+        """
+            Returns the object reference to X
+        """
+        return self.__Xref__
 
     def get_yref(self):
-        return self.yref
-
-
-class AndFunc(ABC):
-    @abstractmethod
-    def eval(self, xy_list: list) -> Xy:
-        raise NotImplementedError("Please implement this method")
+        """
+            Returns the object reference to y
+        """
+        return self.__yref__
 
 
 class Node(ABC):
-    __node_name__ = None
+    """
+    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 __str__(self):
         return self.__node_name__
@@ -46,29 +67,71 @@ def get_and_flag(self):
         raise NotImplementedError("Please implement this method")
 
     def __hash__(self):
+        """
+        Hash code, defined as the hash code of the node name
+
+        :return: Hash code
+        """
         return self.__node_name__.__hash__()
 
     def __eq__(self, other):
+        """
+        Equality with another node, defined as the class names match and the
+        node names match
+
+        :param other: Node to compare with
+        :return: True if nodes are equal, else False
+        """
         return (
                 self.__class__ == other.__class__ and
                 self.__node_name__ == other.__node_name__
         )
 
 
 class OrNode(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.
+    """
     __estimator__ = None
 
     def __init__(self, node_name: str, estimator: BaseEstimator):
+        """
+        Init the OrNode with the name of the node and the etimator.
+
+        :param node_name: Name of the node
+        :param estimator: The base estimator
+        """
         self.__node_name__ = node_name
         self.__estimator__ = estimator
 
     def get_estimator(self) -> BaseEstimator:
+        """
+        Return the estimator that this was initialize with
+
+        :return: Estimator
+        """
         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
 
 
+class AndFunc(ABC):
+    """
+    Or nodes are init-ed from the
+    """
+
+    @abstractmethod
+    def eval(self, xy_list: list) -> Xy:
+        raise NotImplementedError("Please implement this method")
+
+
 class AndNode(Node):
     __andfunc__ = None
 
@@ -127,10 +190,9 @@ def get_object_ref(self):
 
 
 class Pipeline:
-    __pre_graph__ = {}
-    __post_graph__ = {}
-    __node_levels__ = None
-    __level_nodes__ = None
+    """
+    The pipeline class that defines the DAG structure composed of Node(s). The
+    """
 
     def __init__(self):
         self.__pre_graph__ = {}

com/ibm/research/ray/graph/Datamodel.py renamed to codeflare/pipelines/Datamodel.py

@@ -1,18 +1,20 @@
 import ray
 
-from com.ibm.research.ray.graph.Datamodel import OrNode
-from com.ibm.research.ray.graph.Datamodel import AndNode
-from com.ibm.research.ray.graph.Datamodel import Edge
-from com.ibm.research.ray.graph.Datamodel import Pipeline
-from com.ibm.research.ray.graph.Datamodel import XYRef
+from codeflare.pipelines.Datamodel import OrNode
+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
 
 import sklearn.base as base
 from enum import Enum
 
 
 class ExecutionType(Enum):
-    TRAIN = 0,
-    TEST = 1
+    FIT = 0,
+    PREDICT = 1,
+    SCORE = 2
 
 
 @ray.remote
@@ -22,87 +24,84 @@ def execute_or_node_inner(node: OrNode, train_mode: ExecutionType, Xy: XYRef):
     X = ray.get(Xy.get_Xref())
     y = ray.get(Xy.get_yref())
 
-    if train_mode == ExecutionType.TRAIN:
+    if train_mode == ExecutionType.FIT:
         if base.is_classifier(estimator) or base.is_regressor(estimator):
             # Always clone before fit, else fit is invalid
             cloned_estimator = base.clone(estimator)
             cloned_estimator.fit(X, y)
             # 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.get_yref())]
+            result = XYRef(res_Xref, Xy.get_yref())
             return result
         else:
             # No need to clone as it is a transform pass through on the fitted estimator
-            res_Xref = ray.put(estimator.fit_transform(X))
-            result = [XYRef(res_Xref, Xy.get_yref())]
+            res_Xref = ray.put(estimator.fit_transform(X, y))
+            result = XYRef(res_Xref, Xy.get_yref())
             return result
-    elif train_mode == ExecutionType.TEST:
+    elif train_mode == ExecutionType.SCORE:
+        if base.is_classifier(estimator) or base.is_regressor(estimator):
+            cloned_estimator = base.clone(estimator)
+            cloned_estimator.fit(X, y)
+            res_Xref = ray.put(cloned_estimator.score(X, y))
+            result = XYRef(res_Xref, Xy.get_yref())
+            return result
+        else:
+            # No need to clone as it is a transform pass through on the fitted estimator
+            res_Xref = ray.put(estimator.fit_transform(X, y))
+            result = XYRef(res_Xref, Xy.get_yref())
+            return result
+    elif train_mode == ExecutionType.PREDICT:
         # Test mode does not clone as it is a simple predict or transform
         if base.is_classifier(estimator) or base.is_regressor(estimator):
             res_Xref = estimator.predict(X)
-            result = [XYRef(res_Xref, Xy.get_yref())]
+            result = XYRef(res_Xref, Xy.get_yref())
             return result
         else:
             res_Xref = estimator.transform(X)
-            result = [XYRef(res_Xref, Xy.get_yref())]
+            result = XYRef(res_Xref, Xy.get_yref())
             return result
 
 
-###
-# in_args is a dict from Node to list of XYRefs
-###
-def execute_pipeline(pipeline: Pipeline, mode: ExecutionType, in_args: dict):
-    nodes_by_level = pipeline.get_nodes_by_level()
-
-    # track args per edge
-    edge_args = {}
-    for node, node_in_args in in_args.items():
-        pre_edges = pipeline.get_pre_edges(node)
-        for pre_edge in pre_edges:
-            edge_args[pre_edge] = node_in_args
-
-    for nodes in nodes_by_level:
-        for node in nodes:
-            pre_edges = pipeline.get_pre_edges(node)
-            post_edges = pipeline.get_post_edges(node)
-            if not node.get_and_flag():
-                execute_or_node(node, pre_edges, edge_args, post_edges, mode)
-            else:
-                cross_product = execute_and_node(node, pre_edges, edge_args, post_edges)
-                for element in cross_product:
-                    print(element)
-
-    out_args = {}
-    last_level_nodes = nodes_by_level[pipeline.compute_max_level()]
-    for last_level_node in last_level_nodes:
-        edge = Edge(last_level_node, None)
-        out_args[last_level_node] = edge_args[edge]
+def execute_or_node(node, pre_edges, edge_args, post_edges, mode: ExecutionType):
+    for pre_edge in pre_edges:
+        Xyref_ptrs = edge_args[pre_edge]
+        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)
+            exec_xyrefs.append(inner_result)
 
-    return out_args
+        for post_edge in post_edges:
+            if post_edge not in edge_args.keys():
+                edge_args[post_edge] = []
+            edge_args[post_edge].extend(exec_xyrefs)
 
 
 @ray.remote
-def and_node_eval(and_func, xy_list):
-    Xy = and_func.eval(xy_list)
-    res_Xref = ray.put(Xy.get_x())
-    res_yref = ray.put(Xy.get_y())
+def and_node_eval(and_func, Xyref_list):
+    xy_list = []
+    for Xyref in Xyref_list:
+        X = ray.get(Xyref.get_Xref())
+        y = ray.get(Xyref.get_yref())
+        xy_list.append(Xy(X, y))
+
+    res_Xy = and_func.eval(xy_list)
+    res_Xref = ray.put(res_Xy.get_x())
+    res_yref = ray.put(res_Xy.get_y())
     return XYRef(res_Xref, res_yref)
 
 
-def execute_and_node_inner(node: AndNode, elements):
+def execute_and_node_inner(node: AndNode, Xyref_ptrs):
     and_func = node.get_and_func()
     result = []
 
-    for element in elements:
-        xy_list = []
-        for Xy in element:
-            X = ray.get(Xy.get_Xref())
-            y = ray.get(Xy.get_yref())
+    Xyref_list = []
+    for Xyref_ptr in Xyref_ptrs:
+        Xyref = ray.get(Xyref_ptr)
+        Xyref_list.append(Xyref)
 
-            Xy = Xy(X, y)
-            xy_list.append(Xy)
-        Xyref = and_node_eval(and_func, xy_list)
-        result.append(Xyref)
+    Xyref_ptr = and_node_eval.remote(and_func, Xyref_list)
+    result.append(Xyref_ptr)
     return result
 
 
@@ -116,24 +115,36 @@ def execute_and_node(node, pre_edges, edge_args, post_edges):
     cross_product = itertools.product(*edge_args_lists)
 
     for element in cross_product:
-        exec_xyrefs = execute_and_node_inner(node, element)
+        exec_xyref_ptrs = execute_and_node_inner(node, element)
         for post_edge in post_edges:
             if post_edge not in edge_args.keys():
                 edge_args[post_edge] = []
-            edge_args[post_edge].extend(exec_xyrefs)
+            edge_args[post_edge].extend(exec_xyref_ptrs)
 
 
-def execute_or_node(node, pre_edges, edge_args, post_edges, mode: ExecutionType):
-    for pre_edge in pre_edges:
-        Xyrefs = edge_args[pre_edge]
-        exec_xyrefs = []
-        for xy_ref in Xyrefs:
-            xy_ref_list = ray.get(xy_ref)
-            for xy_ref in xy_ref_list:
-                inner_result = execute_or_node_inner.remote(node, mode, xy_ref)
-                exec_xyrefs.append(inner_result)
+def execute_pipeline(pipeline: Pipeline, mode: ExecutionType, in_args: dict):
+    nodes_by_level = pipeline.get_nodes_by_level()
 
-        for post_edge in post_edges:
-            if post_edge not in edge_args.keys():
-                edge_args[post_edge] = []
-            edge_args[post_edge].extend(exec_xyrefs)
+    # track args per edge
+    edge_args = {}
+    for node, node_in_args in in_args.items():
+        pre_edges = pipeline.get_pre_edges(node)
+        for pre_edge in pre_edges:
+            edge_args[pre_edge] = node_in_args
+
+    for nodes in nodes_by_level:
+        for node in nodes:
+            pre_edges = pipeline.get_pre_edges(node)
+            post_edges = pipeline.get_post_edges(node)
+            if not node.get_and_flag():
+                execute_or_node(node, pre_edges, edge_args, post_edges, mode)
+            elif node.get_and_flag():
+                execute_and_node(node, pre_edges, edge_args, post_edges)
+
+    out_args = {}
+    last_level_nodes = nodes_by_level[pipeline.compute_max_level()]
+    for last_level_node in last_level_nodes:
+        edge = Edge(last_level_node, None)
+        out_args[last_level_node] = edge_args[edge]
+
+    return out_args

com/ibm/research/ray/graph/Runtime.py renamed to codeflare/pipelines/Runtime.py

@@ -1,10 +1,10 @@
 Metadata-Version: 1.0
-Name: ray-graphs
+Name: codeflare-pipelines
 Version: 1.0.0
-Summary: Ray
+Summary: Codeflare pipelines
 Home-page: UNKNOWN
-Author: rganti
+Author: Raghu Ganti, Mudhakar Srivatsa
 Author-email: rganti@us.ibm.com
-License: UNKNOWN
+License: Apache v2.0
 Description: UNKNOWN
 Platform: UNKNOWN