updated api call interfaces

Author: yuanchi2807

codeflare/tests/test_and.py

@@ -30,13 +30,7 @@ def test_and():
 
     ## prepare the data
     X = np.random.randint(0,100,size=(10000, 4))
-
-    X_ref = ray.put(X)
-    y_ref = ray.put(None)
-
-    Xy_ref = XYRef(X_ref, y_ref)
-    Xy_ref_ptr = ray.put(Xy_ref)
-    Xy_ref_ptrs = [Xy_ref_ptr]
+    y = np.random.randint(0,2,size=(10000, 1))
 
     ## initialize codeflare pipeline by first creating the nodes
     pipeline = dm.Pipeline()
@@ -48,19 +42,17 @@ def test_and():
     pipeline.add_edge(node_a, node_c)
     pipeline.add_edge(node_b, node_c)
 
-    in_args={node_a: Xy_ref_ptrs, node_b: Xy_ref_ptrs}
+    pipeline_input = dm.PipelineInput()
+    xy = dm.Xy(X,y)
+    pipeline_input.add_xy_arg(node_a, xy)
+    pipeline_input.add_xy_arg(node_b, xy)
+
     ## execute the codeflare pipeline
-    out_args = rt.execute_pipeline(pipeline, ExecutionType.FIT, in_args)
+    pipeline_output = rt.execute_pipeline(pipeline, ExecutionType.FIT, pipeline_input)
 
     ## retrieve node c
-    out_Xyrefs = ray.get(out_args[node_c])
-    assert out_Xyrefs
-
-    for out_xyref in out_Xyrefs:
-        x = ray.get(out_xyref.get_Xref())
-        and_func = ray.get(out_xyref.get_currnoderef()).get_and_func()
-        assert x.any()
-        print(x)
+    node_c_output = pipeline_output.get_xyrefs(node_c)
+    assert node_c_output
 
     ray.shutdown()
 

codeflare/tests/test_multibranch.py

@@ -48,13 +48,6 @@ def test_multibranch():
 
     X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
 
-    X_ref = ray.put(X_train)
-    y_ref = ray.put(y_train)
-
-    Xy_ref = XYRef(X_ref, y_ref)
-    Xy_ref_ptr = ray.put(Xy_ref)
-    Xy_ref_ptrs = [Xy_ref_ptr]
-
     ## create two decision tree classifiers with different depth limit
     c_a = DecisionTreeClassifier(max_depth=3)
     c_b = DecisionTreeClassifier(max_depth=5)
@@ -78,23 +71,17 @@ def test_multibranch():
     pipeline.add_edge(node_d, node_f)
     pipeline.add_edge(node_e, node_f)
 
-    in_args={node_a: Xy_ref_ptrs}
+    pipeline_input = dm.PipelineInput()
+    xy = dm.Xy(X_train, y_train)
+    pipeline_input.add_xy_arg(node_a, xy)
+
     ## execute the codeflare pipeline
-    out_args = rt.execute_pipeline(pipeline, ExecutionType.FIT, in_args)
-    assert out_args
+    pipeline_output = rt.execute_pipeline(pipeline, ExecutionType.FIT, pipeline_input)
+    assert pipeline_output
 
     ## retrieve node b
-    node_b_out_args = ray.get(out_args[node_b])
-    b_out_xyref = node_b_out_args[0]
-    ray.get(b_out_xyref.get_Xref())
-    b_out_node = ray.get(b_out_xyref.get_currnoderef())
-    sct_b = b_out_node.get_estimator()
-    assert sct_b
-    print(sct_b.feature_importances_)
-
-	## retrieve node f
-    out_Xyrefs_f = ray.get(out_args[node_f])
-    assert out_Xyrefs_f
+    node_b_output = pipeline_output.get_xyrefs(node_b)
+    assert node_b_output
 
     ray.shutdown()
 

codeflare/tests/test_or.py

@@ -33,14 +33,7 @@ def test_or():
             ])
 
     X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
-
-    X_ref = ray.put(X_train)
-    y_ref = ray.put(y_train)
-
-    Xy_ref = XYRef(X_ref, y_ref)
-    Xy_ref_ptr = ray.put(Xy_ref)
-    Xy_ref_ptrs = [Xy_ref_ptr]
-
+    
     ## create two decision tree classifiers with different depth limit
     c_a = DecisionTreeClassifier(max_depth=3)
     c_b = DecisionTreeClassifier(max_depth=5)
@@ -55,19 +48,16 @@ def test_or():
     pipeline.add_edge(node_a, node_b)
     pipeline.add_edge(node_a, node_c)
 
-    in_args={node_a: Xy_ref_ptrs}
-    ## execute the codeflare pipeline
-    out_args = rt.execute_pipeline(pipeline, ExecutionType.FIT, in_args)
-    assert out_args
+    pipeline_input = dm.PipelineInput()
+    xy = dm.Xy(X_train, y_train)
+    pipeline_input.add_xy_arg(node_a, xy)
+
+    pipeline_output = rt.execute_pipeline(pipeline, ExecutionType.FIT, pipeline_input)
+
+    node_b_output = pipeline_output.get_xyrefs(node_b)
+    node_c_output = pipeline_output.get_xyrefs(node_c)
 
-    ## retrieve node b
-    node_b_out_args = ray.get(out_args[node_b])
-    b_out_xyref = node_b_out_args[0]
-    ray.get(b_out_xyref.get_Xref())
-    b_out_node = ray.get(b_out_xyref.get_currnoderef())
-    sct_b = b_out_node.get_estimator()
-    assert sct_b
-    print(sct_b.feature_importances_)
+    assert node_b_output
 
     ray.shutdown()