Adding save/load feature for pipeline, semantics are as follows:

1. Save load on an untrained pipeline will save the graph
2. On a trained/selected pipeline, will save the node state
3. Tests for simple graph save have been added

Author: raghukiran1224

codeflare/pipelines/Datamodel.py

@@ -6,6 +6,7 @@
 from sklearn.base import BaseEstimator
 
 import ray
+import pickle5 as pickle
 import codeflare.pipelines.Exceptions as pe
 
 class Xy:
@@ -103,6 +104,9 @@ def __init__(self, node_name, node_input_type: NodeInputType, node_firing_type:
     def __str__(self):
         return self.__node_name__
 
+    def get_node_name(self):
+        return self.__node_name__
+
     def get_node_input_type(self):
         return self.__node_input_type__
 
@@ -379,6 +383,52 @@ def get_terminal_nodes(self):
                 terminal_nodes.append(node)
         return terminal_nodes
 
+    def save(self, filehandle):
+        nodes = {}
+        edges = []
+
+        for node in self.__pre_graph__.keys():
+            nodes[node.get_node_name()] = node
+            pre_edges = self.get_pre_edges(node)
+            for edge in pre_edges:
+                # Since we are iterating on pre_edges, to_node cannot be None
+                from_node = edge.get_from_node()
+                if from_node is not None:
+                    to_node = edge.get_to_node()
+                    edge_tuple = (from_node.get_node_name(), to_node.get_node_name())
+                    edges.append(edge_tuple)
+        saved_pipeline = _SavedPipeline(nodes, edges)
+        pickle.dump(saved_pipeline, filehandle)
+
+    @staticmethod
+    def load(filehandle):
+        saved_pipeline = pickle.load(filehandle)
+        if not isinstance(saved_pipeline, _SavedPipeline):
+            raise pe.PipelineException("Filehandle is not a saved pipeline instance")
+
+        nodes = saved_pipeline.get_nodes()
+        edges = saved_pipeline.get_edges()
+
+        pipeline = Pipeline()
+        for edge in edges:
+            (from_node_str, to_node_str) = edge
+            from_node = nodes[from_node_str]
+            to_node = nodes[to_node_str]
+            pipeline.add_edge(from_node, to_node)
+        return pipeline
+
+
+class _SavedPipeline:
+    def __init__(self, nodes, edges):
+        self.__nodes__ = nodes
+        self.__edges__ = edges
+
+    def get_nodes(self):
+        return self.__nodes__
+
+    def get_edges(self):
+        return self.__edges__
+
 
 class PipelineOutput:
     """

codeflare/pipelines/Runtime.py

@@ -8,6 +8,7 @@
 from enum import Enum
 
 from queue import SimpleQueue
+import pickle5 as pickle
 
 
 class ExecutionType(Enum):
@@ -265,3 +266,8 @@ def cross_validate(cross_validator: BaseCrossValidator, pipeline: dm.Pipeline, p
         result_scores.append(out_x)
 
     return result_scores
+
+
+def save(pipeline_output: dm.PipelineOutput, xy_ref: dm.XYRef, filehandle):
+    pipeline = select_pipeline(pipeline_output, xy_ref)
+    pipeline.save(filehandle)

codeflare/pipelines/tests/__init__.py

@@ -0,0 +1,61 @@
+import pytest
+
+import codeflare.pipelines.Datamodel as dm
+import codeflare.pipelines.Runtime as rt
+
+import numpy as np
+from sklearn.preprocessing import FunctionTransformer
+from sklearn.preprocessing import MinMaxScaler
+import os
+
+
+class FeatureUnion(dm.AndTransform):
+    def __init__(self):
+        pass
+
+    def transform(self, xy_list):
+        X_list = []
+        y_list = []
+
+        for xy in xy_list:
+            X_list.append(xy.get_x())
+        X_concat = np.concatenate(X_list, axis=0)
+
+        return dm.Xy(X_concat, None)
+
+
+def test_save_load():
+    """
+    A simple save load test for a pipeline graph
+    :return:
+    """
+    pipeline = dm.Pipeline()
+    minmax_scaler = MinMaxScaler()
+
+    node_a = dm.EstimatorNode('a', minmax_scaler)
+    node_b = dm.EstimatorNode('b', minmax_scaler)
+    node_c = dm.AndNode('c', FeatureUnion())
+
+    pipeline.add_edge(node_a, node_c)
+    pipeline.add_edge(node_b, node_c)
+
+    fname = 'save_pipeline.cfp'
+    fh = open(fname, 'wb')
+    pipeline.save(fh)
+    fh.close()
+
+    r_fh = open(fname, 'rb')
+    saved_pipeline = dm.Pipeline.load(r_fh)
+    pre_edges = saved_pipeline.get_pre_edges(node_c)
+    assert(len(pre_edges) == 2)
+
+    os.remove(fname)
+
+
+def test_runtime_save_load():
+    """
+    Tests for selecting a pipeline and save/load it, we also test the predict to ensure state is
+    captured accurately
+    :return:
+    """
+    

codeflare/pipelines/tests/test_save_load.py

@@ -2,4 +2,7 @@ ray~=1.3.0
 setuptools~=52.0.0
 sklearn~=0.0
 scikit-learn~=0.24.1
-pandas~=1.2.4
+pandas~=1.2.4
+pytest~=6.2.4
+numpy~=1.18.5
+pickle5~=0.0.11