Merge pull request #36 from codeflare/pickle

Pickle

Author: yuanchi2807

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,64 @@ def get_terminal_nodes(self):
                 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
+
+    def get_pre_nodes(self, node):
+        return self.__pre_graph__[node]
+
+    def get_post_nodes(self, node):
+        return self.__post_graph__[node]
+
+    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

@@ -58,11 +58,11 @@ def execute_or_node_remote(node: dm.EstimatorNode, mode: ExecutionType, xy_ref:
     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 = estimator.predict(X)
+            res_Xref = ray.put(estimator.predict(X))
             result = dm.XYRef(res_Xref, xy_ref.get_yref())
             return result
         else:
-            res_Xref = estimator.transform(X)
+            res_Xref = ray.put(estimator.transform(X))
             result = dm.XYRef(res_Xref, xy_ref.get_yref())
             return result
 
@@ -265,3 +265,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,136 @@
+import codeflare.pipelines.Datamodel as dm
+import codeflare.pipelines.Runtime as rt
+
+import numpy as np
+from sklearn.preprocessing import MinMaxScaler
+import os
+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
+
+import ray
+
+
+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:
+    """
+    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)
+    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']).drop(['Loan_Status'], axis=1).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)
+
+    import ray
+    ray.shutdown()
+    ray.init()
+    pipeline_input = dm.PipelineInput()
+    xy = dm.Xy(X_train, y_train)
+    pipeline_input.add_xy_arg(node_pre, xy)
+
+    pipeline_output = rt.execute_pipeline(pipeline, rt.ExecutionType.FIT, pipeline_input)
+    node_rf_xyrefs = pipeline_output.get_xyrefs(node_rf)
+
+    # save this pipeline for random forest and load and then predict on test data
+    fname = 'random_forest.cfp'
+    w_fh = open(fname, 'wb')
+    rt.save(pipeline_output, node_rf_xyrefs[0], w_fh)
+    w_fh.close()
+
+    # load it
+    r_fh = open(fname, 'rb')
+    saved_pipeline = dm.Pipeline.load(r_fh)
+    nodes = saved_pipeline.get_nodes()
+    # this should not exist in the saved pipeline
+    assert(node_gb.get_node_name() not in nodes.keys())
+
+    # should be preditable as well
+    predict_pipeline_input = dm.PipelineInput()
+    predict_pipeline_input.add_xy_arg(node_pre, dm.Xy(X_test, y_test))
+    try:
+        predict_pipeline_output = rt.execute_pipeline(saved_pipeline, rt.ExecutionType.PREDICT, predict_pipeline_input)
+        predict_pipeline_output.get_xyrefs(node_rf)
+    except Exception:
+        assert False
+
+    os.remove(fname)

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