Grid search implementation start, adding:

1. Support for getting pipeline input on a given pipeline and the chosen xyref
2. Adding some new methods for easier access to pipeline internals
3. Refactoring pre_image/post_image -- older names
4. Added a test for pipeline_input check

Author: raghukiran1224

codeflare/pipelines/Datamodel.py

@@ -301,13 +301,15 @@ def __init__(self):
         self.__post_graph__ = {}
         self.__node_levels__ = None
         self.__level_nodes__ = None
+        self.__node_name_map__ = {}
 
     def add_node(self, node: Node):
         self.__node_levels__ = None
         self.__level_nodes__ = None
         if node not in self.__pre_graph__.keys():
             self.__pre_graph__[node] = []
             self.__post_graph__[node] = []
+            self.__node_name_map__[node.get_node_name()] = node
 
     def __str__(self):
         res = ''
@@ -333,23 +335,17 @@ def add_edge(self, from_node: Node, to_node: Node):
         self.__pre_graph__[to_node].append(from_node)
         self.__post_graph__[from_node].append(to_node)
 
-    def get_preimage(self, node: Node):
-        return self.__pre_graph__[node]
-
-    def get_postimage(self, node: Node):
-        return self.__post_graph__[node]
-
     def compute_node_level(self, node: Node, result: dict):
         if node in result:
             return result[node]
 
-        node_preimage = self.get_preimage(node)
-        if not node_preimage:
+        pre_nodes = self.get_pre_nodes(node)
+        if not pre_nodes:
             result[node] = 0
             return 0
 
         max_level = 0
-        for p_node in node_preimage:
+        for p_node in pre_nodes:
             level = self.compute_node_level(p_node, result)
             max_level = max(level, max_level)
 
@@ -369,6 +365,10 @@ def compute_node_levels(self):
 
         return self.__node_levels__
 
+    def get_node_level(self, node: Node):
+        self.compute_node_levels()
+        return self.__node_levels__[node]
+
     def compute_max_level(self):
         levels = self.compute_node_levels()
         max_level = 0
@@ -423,30 +423,42 @@ def get_post_edges(self, node: Node):
             post_edges.append(Edge(node, post_node))
         return post_edges
 
-    def is_terminal(self, node: Node):
-        post_nodes = self.__post_graph__[node]
+    def is_output(self, node: Node):
+        post_nodes = self.get_post_nodes(node)
         return not post_nodes
 
-    def get_terminal_nodes(self):
+    def get_output_nodes(self):
         # dict from level to nodes
         terminal_nodes = []
         for node in self.__pre_graph__.keys():
-            if self.is_terminal(node):
+            if self.is_output(node):
                 terminal_nodes.append(node)
         return terminal_nodes
 
     def get_nodes(self):
-        nodes = {}
-        for node in self.__pre_graph__.keys():
-            nodes[node.get_node_name()] = node
-        return nodes
+        return self.__node_name_map__
 
     def get_pre_nodes(self, node):
         return self.__pre_graph__[node]
 
     def get_post_nodes(self, node):
         return self.__post_graph__[node]
 
+    def is_input(self, node: Node):
+        pre_nodes = self.get_pre_nodes(node)
+        return not pre_nodes
+
+    def get_input_nodes(self):
+        input_nodes = []
+        for node in self.__node_name_map__.values():
+            if self.get_node_level() == 0:
+                input_nodes.append(node)
+
+        return input_nodes
+
+    def get_node(self, node_name: str) -> Node:
+        return self.__node_name_map__[node_name]
+
     def save(self, filehandle):
         nodes = {}
         edges = []

codeflare/pipelines/Runtime.py

@@ -153,7 +153,7 @@ def execute_pipeline(pipeline: dm.Pipeline, mode: ExecutionType, pipeline_input:
                 execute_and_node(node, pre_edges, edge_args, post_edges)
 
     out_args = {}
-    terminal_nodes = pipeline.get_terminal_nodes()
+    terminal_nodes = pipeline.get_output_nodes()
     for terminal_node in terminal_nodes:
         edge = dm.Edge(terminal_node, None)
         out_args[terminal_node] = edge_args[edge]
@@ -184,6 +184,32 @@ def select_pipeline(pipeline_output: dm.PipelineOutput, chosen_xyref: dm.XYRef):
     return pipeline
 
 
+def get_pipeline_input(pipeline: dm.Pipeline, pipeline_output: dm.PipelineOutput, chosen_xyref: dm.XYRef):
+    pipeline_input = dm.PipelineInput()
+
+    xyref_queue = SimpleQueue()
+    xyref_queue.put(chosen_xyref)
+    while not xyref_queue.empty():
+        curr_xyref = xyref_queue.get()
+        curr_node_state_ptr = curr_xyref.get_curr_node_state_ref()
+        curr_node = ray.get(curr_node_state_ptr)
+        curr_node_level = pipeline.get_node_level(curr_node)
+        prev_xyrefs = curr_xyref.get_prev_xyrefs()
+
+        if curr_node_level == 0:
+            # This is an input node
+            for prev_xyref in prev_xyrefs:
+                pipeline_input.add_xyref_arg(curr_node, prev_xyref)
+
+        for prev_xyref in prev_xyrefs:
+            prev_node_state_ptr = prev_xyref.get_curr_node_state_ref()
+            if prev_node_state_ptr is None:
+                continue
+            xyref_queue.put(prev_xyref)
+
+    return pipeline_input
+
+
 @ray.remote(num_returns=2)
 def split(cross_validator: BaseCrossValidator, xy_ref):
     x = ray.get(xy_ref.get_Xref())
@@ -220,7 +246,7 @@ def cross_validate(cross_validator: BaseCrossValidator, pipeline: dm.Pipeline, p
 
     in_args = pipeline_input.get_in_args()
     for node, xyref_ptrs in in_args.items():
-        # NOTE: The assumption is that this node has only one input, the check earlier will ensure this!
+        # NOTE: The assumption is that this node has only one input!
         xyref_ptr = xyref_ptrs[0]
         xy_train_refs_ptr, xy_test_refs_ptr = split.remote(cross_validator, xyref_ptr)
         xy_train_refs = ray.get(xy_train_refs_ptr)
@@ -238,7 +264,7 @@ def cross_validate(cross_validator: BaseCrossValidator, pipeline: dm.Pipeline, p
     pipeline_output_train = execute_pipeline(pipeline, ExecutionType.FIT, pipeline_input_train)
 
     # Now we can choose the pipeline and then score for each of the chosen pipelines
-    out_nodes = pipeline.get_terminal_nodes()
+    out_nodes = pipeline.get_output_nodes()
     if len(out_nodes) > 1:
         raise pe.PipelineException("Cannot cross validate as output is not a single node")
 
@@ -267,6 +293,42 @@ def cross_validate(cross_validator: BaseCrossValidator, pipeline: dm.Pipeline, p
     return result_scores
 
 
+def grid_search(cross_validator: BaseCrossValidator, pipeline: dm.Pipeline, pipeline_input: dm.PipelineInput):
+    pipeline_input_train = dm.PipelineInput()
+
+    pipeline_input_test = []
+    k = cross_validator.get_n_splits()
+    # add k pipeline inputs for testing
+    for i in range(k):
+        pipeline_input_test.append(dm.PipelineInput())
+
+    in_args = pipeline_input.get_in_args()
+    for node, xyref_ptrs in in_args.items():
+        # NOTE: The assumption is that this node has only one input!
+        xyref_ptr = xyref_ptrs[0]
+        if len(xyref_ptrs) > 1:
+            raise pe.PipelineException("Input to grid search is multiple objects, re-run with only single object")
+
+        xy_train_refs_ptr, xy_test_refs_ptr = split.remote(cross_validator, xyref_ptr)
+        xy_train_refs = ray.get(xy_train_refs_ptr)
+        xy_test_refs = ray.get(xy_test_refs_ptr)
+
+        for xy_train_ref in xy_train_refs:
+            pipeline_input_train.add_xyref_arg(node, xy_train_ref)
+
+        # for testing, add only to the specific input
+        for i in range(k):
+            pipeline_input_test[i].add_xyref_arg(node, xy_test_refs[i])
+
+    # Ready for execution now that data has been prepared! This execution happens in parallel
+    # because of the underlying pipeline graph and multiple input objects
+    pipeline_output_train = execute_pipeline(pipeline, ExecutionType.FIT, pipeline_input_train)
+
+    # For grid search, we will have multiple output nodes that need to be iterated on and select the pipeline
+    # that is "best"
+    out_nodes = pipeline.get_output_nodes()
+
+
 def save(pipeline_output: dm.PipelineOutput, xy_ref: dm.XYRef, filehandle):
     pipeline = select_pipeline(pipeline_output, xy_ref)
     pipeline.save(filehandle)

codeflare/pipelines/test_Datamodel.py

@@ -76,7 +76,7 @@ def test_multibranch(self):
         pipeline_input = dm.PipelineInput()
         pipeline_input.add_xy_arg(node_a, dm.Xy(X_train, y_train))
 
-        terminal_nodes = pipeline.get_terminal_nodes()
+        terminal_nodes = pipeline.get_output_nodes()
         assert len(terminal_nodes) == 4
 
         ## execute the codeflare pipeline

codeflare/pipelines/tests/test_helper.py

@@ -0,0 +1,56 @@
+import codeflare.pipelines.Datamodel as dm
+
+import pandas as pd
+from sklearn.pipeline import Pipeline
+from sklearn.impute import SimpleImputer
+from sklearn.preprocessing import StandardScaler, OneHotEncoder
+from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
+
+
+def get_pipeline(train) -> dm.Pipeline:
+    imputer = SimpleImputer(strategy='median')
+    scaler = StandardScaler()
+
+    numeric_transformer = Pipeline(steps=[
+        ('imputer', imputer),
+        ('scaler', scaler)])
+
+    cat_imputer = SimpleImputer(strategy='constant', fill_value='missing')
+    cat_onehot = OneHotEncoder(handle_unknown='ignore')
+
+    categorical_transformer = Pipeline(steps=[
+        ('imputer', cat_imputer),
+        ('onehot', cat_onehot)])
+    numeric_features = train.select_dtypes(include=['int64', 'float64']).columns
+    categorical_features = train.select_dtypes(include=['object']).columns
+    from sklearn.compose import ColumnTransformer
+    preprocessor = ColumnTransformer(
+        transformers=[
+            ('num', numeric_transformer, numeric_features),
+            ('cat', categorical_transformer, categorical_features)])
+
+    classifiers = [
+        RandomForestClassifier(),
+        GradientBoostingClassifier()
+    ]
+    pipeline = dm.Pipeline()
+    node_pre = dm.EstimatorNode('preprocess', preprocessor)
+    node_rf = dm.EstimatorNode('random_forest', classifiers[0])
+    node_gb = dm.EstimatorNode('gradient_boost', classifiers[1])
+
+    pipeline.add_edge(node_pre, node_rf)
+    pipeline.add_edge(node_pre, node_gb)
+
+    return pipeline
+
+
+def get_data():
+    train = pd.read_csv('../../../resources/data/train_ctrUa4K.csv')
+    train = train.drop('Loan_ID', axis=1)
+
+    X = train.drop('Loan_Status', axis=1)
+    y = train['Loan_Status']
+    from sklearn.model_selection import train_test_split
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
+
+    return X_train, X_test, y_train, y_test

codeflare/pipelines/tests/test_runtime.py

@@ -0,0 +1,43 @@
+from codeflare.pipelines.tests import test_helper
+
+import codeflare.pipelines.Datamodel as dm
+import codeflare.pipelines.Runtime as rt
+
+
+def test_runtime_pipeline_input_getter():
+    """
+    A test to get the pipeline inputs after a selection is done
+    :return:
+    """
+
+    import ray
+    ray.shutdown()
+    ray.init()
+    X_train, X_test, y_train, y_test = test_helper.get_data()
+    pipeline = test_helper.get_pipeline(X_train)
+
+    node_rf = pipeline.get_node('random_forest')
+    node_gb = pipeline.get_node('gradient_boost')
+    input_node = pipeline.get_node('preprocess')
+
+    pipeline_input = dm.PipelineInput()
+    xy = dm.Xy(X_train, y_train)
+    pipeline_input.add_xy_arg(input_node, xy)
+
+    pipeline_output = rt.execute_pipeline(pipeline, rt.ExecutionType.FIT, pipeline_input)
+    node_rf_xyrefs = pipeline_output.get_xyrefs(node_rf)
+
+    selected_pipeline_input = rt.get_pipeline_input(pipeline, pipeline_output, node_rf_xyrefs[0])
+    in_args = selected_pipeline_input.get_in_args()
+    is_input_node_present = (input_node in in_args.keys())
+    assert is_input_node_present
+
+    # check if the XYref is the same
+    xyref_ptrs = in_args[input_node]
+    xyref_ptr = xyref_ptrs[0]
+    xyref = ray.get(xyref_ptr)
+
+    input_xyref = ray.get(pipeline_input.get_in_args()[input_node][0])
+    assert xyref.get_Xref() == input_xyref.get_Xref()
+    assert xyref.get_yref() == input_xyref.get_yref()
+