Adding a test for save/load of a selected pipeline; this fixed a bug in the PREDICT pipeline, which was incorrect earlier. The issue was that a X object was

being treated as XRef, leading to pointer messup... that is fixed now!

Author: raghukiran1224

codeflare/pipelines/Datamodel.py

@@ -383,6 +383,18 @@ 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 = []

codeflare/pipelines/Runtime.py

@@ -8,7 +8,6 @@
 from enum import Enum
 
 from queue import SimpleQueue
-import pickle5 as pickle
 
 
 class ExecutionType(Enum):
@@ -59,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
 

codeflare/pipelines/tests/test_save_load.py

@@ -1,12 +1,16 @@
-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
+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):
@@ -47,8 +51,7 @@ def test_save_load():
     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)
-
+    assert (len(pre_edges) == 2)
     os.remove(fname)
 
 
@@ -58,4 +61,76 @@ def test_runtime_save_load():
     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)