Adding grid search in runtime that provides a similar API as GridSearchCV, but a richer one that can explore non-gridded spaces as well.

Tests have been added to cover this

Modifications to datamodel have been done to accommodate for grid search

Author: raghukiran1224

.gitignore

@@ -5,3 +5,4 @@ dist/
 ray-graphs.egginfo/
 .idea/
 .ipynb_checkpoints/
+*__pycache*

codeflare/pipelines/Datamodel.py

@@ -2,8 +2,6 @@
 from enum import Enum
 
 import sklearn.base as base
-from codeflare.pipelines.Datamodel import PipelineParam
-from sklearn.base import TransformerMixin
 from sklearn.base import BaseEstimator
 from sklearn.model_selection import ParameterGrid
 
@@ -159,16 +157,17 @@ class Node(ABC):
     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):
-        if '__' in node_name:
-            raise pe.PipelineException("Node name cannot have __, please rename")
+    def __init__(self, node_name, estimator: BaseEstimator, node_input_type: NodeInputType, node_firing_type: NodeFiringType, node_state_type: NodeStateType):
         self.__node_name__ = node_name
+        self.__estimator__ = estimator
         self.__node_input_type__ = node_input_type
         self.__node_firing_type__ = node_firing_type
         self.__node_state_type__ = node_state_type
 
     def __str__(self):
-        return self.__node_name__
+        estimator_params_str = str(self.get_estimator().get_params())
+        retval = self.__node_name__ + estimator_params_str
+        return retval
 
     def get_node_name(self):
         return self.__node_name__
@@ -182,6 +181,16 @@ def get_node_firing_type(self):
     def get_node_state_type(self):
         return self.__node_state_type__
 
+    def get_estimator(self):
+        return self.__estimator__
+
+    def get_parameterized_node(self, node_name, **params):
+        cloned_node = self.clone()
+        cloned_node.__node_name__ = node_name
+        estimator = cloned_node.get_estimator()
+        estimator.set_params(**params)
+        return cloned_node
+
     @abstractmethod
     def clone(self):
         raise NotImplementedError("Please implement the clone method")
@@ -222,44 +231,50 @@ def __init__(self, node_name: str, estimator: BaseEstimator):
         :param estimator: The base estimator
         """
 
-        super().__init__(node_name, NodeInputType.OR, NodeFiringType.ANY, NodeStateType.IMMUTABLE)
-        self.__estimator__ = estimator
-
-    def get_estimator(self) -> BaseEstimator:
-        """
-        Return the estimator that this was initialize with
-
-        :return: Estimator
-        """
-        return self.__estimator__
+        super().__init__(node_name, estimator, NodeInputType.OR, NodeFiringType.ANY, NodeStateType.IMMUTABLE)
 
     def clone(self):
         cloned_estimator = base.clone(self.__estimator__)
         return EstimatorNode(self.__node_name__, cloned_estimator)
 
 
-class AndTransform(TransformerMixin, BaseEstimator):
+class AndEstimator(BaseEstimator):
     @abstractmethod
     def transform(self, xy_list: list) -> Xy:
-        raise NotImplementedError("Please implement this method")
+        raise NotImplementedError("And estimator needs to implement a transform method")
+
+    @abstractmethod
+    def fit(self, xy_list: list):
+        raise NotImplementedError("And estimator needs to implement a fit method")
 
+    @abstractmethod
+    def fit_transform(self, xy_list: list):
+        raise NotImplementedError("And estimator needs to implement a fit method")
 
-class GeneralTransform(TransformerMixin, BaseEstimator):
     @abstractmethod
-    def transform(self, xy: Xy) -> Xy:
-        raise NotImplementedError("Please implement this method")
+    def predict(self, xy_list: list) -> Xy:
+        raise NotImplementedError("And classifier needs to implement the predict method")
 
+    @abstractmethod
+    def score(self, xy_list: list) -> Xy:
+        raise NotImplementedError("And classifier needs to implement the score method")
 
-class AndNode(Node):
-    def __init__(self, node_name: str, and_func: AndTransform):
-        super().__init__(node_name, NodeInputType.AND, NodeFiringType.ANY, NodeStateType.STATELESS)
-        self.__andfunc__ = and_func
+    @abstractmethod
+    def get_estimator_type(self):
+        raise NotImplementedError("And classifier needs to implement the get_estimator_type method")
+
+    @abstractmethod
+    def clone(self):
+        raise NotImplementedError("And estimator needs to implement a clone method")
 
-    def get_and_func(self) -> AndTransform:
-        return self.__andfunc__
+
+class AndNode(Node):
+    def __init__(self, node_name: str, and_estimator: AndEstimator):
+        super().__init__(node_name, and_estimator, NodeInputType.AND, NodeFiringType.ANY, NodeStateType.STATELESS)
 
     def clone(self):
-        return AndNode(self.__node_name__, self.__andfunc__)
+        cloned_estimator = self.__estimator__.clone()
+        return AndNode(self.__node_name__, cloned_estimator)
 
 
 class Edge:
@@ -477,7 +492,7 @@ def is_input(self, node: Node):
     def get_input_nodes(self):
         input_nodes = []
         for node in self.__node_name_map__.values():
-            if self.get_node_level() == 0:
+            if self.get_node_level(node) == 0:
                 input_nodes.append(node)
 
         return input_nodes
@@ -514,7 +529,7 @@ def save(self, filehandle):
         saved_pipeline = _SavedPipeline(nodes, edges)
         pickle.dump(saved_pipeline, filehandle)
 
-    def set_param_grid(self, pipeline_param: PipelineParam):
+    def get_parameterized_pipeline(self, pipeline_param):
         result = Pipeline()
         pipeline_params = pipeline_param.get_all_params()
         parameterized_nodes = {}
@@ -523,14 +538,27 @@ def set_param_grid(self, pipeline_param: PipelineParam):
             if node_name_part not in parameterized_nodes.keys():
                 parameterized_nodes[node_name_part] = []
             node = self.__node_name_map__[node_name_part]
-            estimator = node.get_estimator()
-            cloned_estimator = estimator.clone()
-            cloned_estimator.set_params(**params)
+            parameterized_node = node.get_parameterized_node(node_name, **params)
+            parameterized_nodes[node_name_part].append(parameterized_node)
 
-            parameterized_nodes[node_name_part].append()
-            result.add_node()
+        # update parameterized nodes with missing non-parameterized nodes for completeness
+        for node in self.__pre_graph__.keys():
+            node_name = node.get_node_name()
+            if node_name not in parameterized_nodes.keys():
+                parameterized_nodes[node_name] = [node]
+
+        # loop through the graph and add edges
+        for node, pre_nodes in self.__pre_graph__.items():
+            node_name = node.get_node_name()
+            expanded_nodes = parameterized_nodes[node_name]
+            for pre_node in pre_nodes:
+                pre_node_name = pre_node.get_node_name()
+                expanded_pre_nodes = parameterized_nodes[pre_node_name]
+                for expanded_pre_node in expanded_pre_nodes:
+                    for expanded_node in expanded_nodes:
+                        result.add_edge(expanded_pre_node, expanded_node)
 
-        # construct nodes
+        return result
 
     @staticmethod
     def load(filehandle):
@@ -617,9 +645,27 @@ def add_xy_arg(self, node: Node, xy: Xy):
         xyref = XYRef(x_ref, y_ref)
         self.add_xyref_arg(node, xyref)
 
+    def add_all(self, node, node_inargs):
+        self.__in_args__[node] = node_inargs
+
     def get_in_args(self):
         return self.__in_args__
 
+    def get_parameterized_input(self, pipeline: Pipeline, parameterized_pipeline: Pipeline):
+        input_nodes = parameterized_pipeline.get_input_nodes()
+        parameterized_pipeline_input = PipelineInput()
+        for input_node in input_nodes:
+            input_node_name = input_node.get_node_name()
+            if '__' not in input_node_name:
+                node_name = input_node_name
+            else:
+                node_name, param = input_node.get_node_name().split('__', 1)
+
+            pipeline_node = pipeline.get_node(node_name)
+            if pipeline_node in self.__in_args__:
+                parameterized_pipeline_input.add_all(input_node, self.__in_args__[pipeline_node])
+        return parameterized_pipeline_input
+
 
 class PipelineParam:
     def __init__(self):

codeflare/pipelines/Runtime.py

@@ -23,11 +23,11 @@ def execute_or_node_remote(node: dm.EstimatorNode, mode: ExecutionType, xy_ref:
     # Blocking operation -- not avoidable
     X = ray.get(xy_ref.get_Xref())
     y = ray.get(xy_ref.get_yref())
+    prev_node_ptr = ray.put(node)
 
     # TODO: Can optimize the node pointers without replicating them
     if mode == ExecutionType.FIT:
         cloned_node = node.clone()
-        prev_node_ptr = ray.put(node)
 
         if base.is_classifier(estimator) or base.is_regressor(estimator):
             # Always clone before fit, else fit is invalid
@@ -49,22 +49,22 @@ def execute_or_node_remote(node: dm.EstimatorNode, mode: ExecutionType, xy_ref:
         if base.is_classifier(estimator) or base.is_regressor(estimator):
             estimator = node.get_estimator()
             res_Xref = ray.put(estimator.score(X, y))
-            result = dm.XYRef(res_Xref, xy_ref.get_yref())
+            result = dm.XYRef(res_Xref, xy_ref.get_yref(), prev_node_ptr, prev_node_ptr, [xy_ref])
             return result
         else:
             res_Xref = ray.put(estimator.transform(X))
-            result = dm.XYRef(res_Xref, xy_ref.get_yref())
+            result = dm.XYRef(res_Xref, xy_ref.get_yref(), prev_node_ptr, prev_node_ptr, [xy_ref])
 
             return result
     elif 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 = ray.put(estimator.predict(X))
-            result = dm.XYRef(res_Xref, xy_ref.get_yref())
+            result = dm.XYRef(res_Xref, xy_ref.get_yref(), prev_node_ptr, prev_node_ptr, [xy_ref])
             return result
         else:
             res_Xref = ray.put(estimator.transform(X))
-            result = dm.XYRef(res_Xref, xy_ref.get_yref())
+            result = dm.XYRef(res_Xref, xy_ref.get_yref(), prev_node_ptr, prev_node_ptr, [xy_ref])
             return result
 
 
@@ -84,38 +84,88 @@ def execute_or_node(node, pre_edges, edge_args, post_edges, mode: ExecutionType)
 
 
 @ray.remote
-def execute_and_node_remote(node: dm.AndNode, Xyref_list):
+def execute_and_node_remote(node: dm.AndNode, mode: ExecutionType, 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(dm.Xy(X, y))
 
-    cloned_node = node.clone()
-    curr_node_ptr = ray.put(cloned_node)
+    estimator = node.get_estimator()
+
+    # TODO: Can optimize the node pointers without replicating them
+    if mode == ExecutionType.FIT:
+        cloned_node = node.clone()
+
+        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(xy_list)
 
-    cloned_and_func = cloned_node.get_and_func()
-    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 dm.XYRef(res_Xref, res_yref, prev_node_ptr, curr_node_ptr, Xyref_list)
+            curr_node_ptr = ray.put(cloned_node)
+            res_xy = cloned_estimator.predict(xy_list)
+            res_xref = ray.put(res_xy.get_x())
+            res_yref = ray.put(res_xy.get_y())
 
+            result = dm.XYRef(res_xref, res_yref, prev_node_ptr, curr_node_ptr, Xyref_list)
+            return result
+        else:
+            cloned_estimator = cloned_node.get_estimator()
+            res_xy = cloned_estimator.fit_transform(xy_list)
+            res_xref = ray.put(res_xy.get_x())
+            res_yref = ray.put(res_xy.get_y())
 
-def execute_and_node_inner(node: dm.AndNode, Xyref_ptrs):
+            curr_node_ptr = ray.put(cloned_node)
+            result = dm.XYRef(res_xref, res_yref, prev_node_ptr, curr_node_ptr, Xyref_list)
+            return result
+    elif mode == ExecutionType.SCORE:
+        if base.is_classifier(estimator) or base.is_regressor(estimator):
+            estimator = node.get_estimator()
+            res_xy = estimator.score(xy_list)
+            res_xref = ray.put(res_xy.get_x())
+            res_yref = ray.put(res_xy.get_y())
+
+            result = dm.XYRef(res_xref, res_yref, prev_node_ptr, prev_node_ptr, Xyref_list)
+            return result
+        else:
+            res_xy = estimator.transform(xy_list)
+            res_xref = ray.put(res_xy.get_x())
+            res_yref = ray.put(res_xy.get_y())
+            result = dm.XYRef(res_xref, res_yref, prev_node_ptr, prev_node_ptr, Xyref_list)
+
+            return result
+    elif 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_xy = estimator.predict(xy_list)
+            res_xref = ray.put(res_xy.get_x())
+            res_yref = ray.put(res_xy.get_y())
+
+            result = dm.XYRef(res_xref, res_yref, prev_node_ptr, prev_node_ptr, Xyref_list)
+            return result
+        else:
+            res_xy = estimator.transform(xy_list)
+            res_xref = ray.put(res_xy.get_x())
+            res_yref = ray.put(res_xy.get_y())
+            result = dm.XYRef(res_xref, res_yref, prev_node_ptr, prev_node_ptr, Xyref_list)
+            return result
+
+
+def execute_and_node_inner(node: dm.AndNode, mode: ExecutionType, Xyref_ptrs):
     result = []
 
     Xyref_list = []
     for Xyref_ptr in Xyref_ptrs:
         Xyref = ray.get(Xyref_ptr)
         Xyref_list.append(Xyref)
 
-    Xyref_ptr = execute_and_node_remote.remote(node, Xyref_list)
+    Xyref_ptr = execute_and_node_remote.remote(node, mode, Xyref_list)
     result.append(Xyref_ptr)
     return result
 
 
-def execute_and_node(node, pre_edges, edge_args, post_edges):
+def execute_and_node(node, pre_edges, edge_args, post_edges, mode: ExecutionType):
     edge_args_lists = list()
     for pre_edge in pre_edges:
         edge_args_lists.append(edge_args[pre_edge])
@@ -125,7 +175,7 @@ def execute_and_node(node, pre_edges, edge_args, post_edges):
     cross_product = itertools.product(*edge_args_lists)
 
     for element in cross_product:
-        exec_xyref_ptrs = execute_and_node_inner(node, element)
+        exec_xyref_ptrs = execute_and_node_inner(node, mode, element)
         for post_edge in post_edges:
             if post_edge not in edge_args.keys():
                 edge_args[post_edge] = []
@@ -151,7 +201,7 @@ def execute_pipeline(pipeline: dm.Pipeline, mode: ExecutionType, pipeline_input:
             if node.get_node_input_type() == dm.NodeInputType.OR:
                 execute_or_node(node, pre_edges, edge_args, post_edges, mode)
             elif node.get_node_input_type() == dm.NodeInputType.AND:
-                execute_and_node(node, pre_edges, edge_args, post_edges)
+                execute_and_node(node, pre_edges, edge_args, post_edges, mode)
 
     out_args = {}
     terminal_nodes = pipeline.get_output_nodes()
@@ -249,7 +299,7 @@ def cross_validate(cross_validator: BaseCrossValidator, pipeline: dm.Pipeline, p
         raise pe.PipelineException("Cross validation can only be done on pipelines with single estimator, "
                                    "use grid_search_cv instead")
 
-    result_grid_search_cv = grid_search_cv(cross_validator, pipeline, pipeline_input)
+    result_grid_search_cv = _grid_search_cv(cross_validator, pipeline, pipeline_input)
     # only one output here
     result_scores = None
     for scores in result_grid_search_cv.values():
@@ -259,7 +309,13 @@ def cross_validate(cross_validator: BaseCrossValidator, pipeline: dm.Pipeline, p
     return result_scores
 
 
-def grid_search_cv(cross_validator: BaseCrossValidator, pipeline: dm.Pipeline, pipeline_input: dm.PipelineInput):
+def grid_search_cv(cross_validator: BaseCrossValidator, pipeline: dm.Pipeline, pipeline_input: dm.PipelineInput, pipeline_params: dm.PipelineParam):
+    parameterized_pipeline = pipeline.get_parameterized_pipeline(pipeline_params)
+    parameterized_pipeline_input = pipeline_input.get_parameterized_input(pipeline, parameterized_pipeline)
+    return _grid_search_cv(cross_validator, parameterized_pipeline, parameterized_pipeline_input)
+
+
+def _grid_search_cv(cross_validator: BaseCrossValidator, pipeline: dm.Pipeline, pipeline_input: dm.PipelineInput):
     pipeline_input_train = dm.PipelineInput()
 
     pipeline_input_test = []