WIP integrate obliv-c into tests

benchmarks/aspirin/workload.py

@@ -129,22 +129,7 @@ def run_mpc(pid: str, data_root: str, mpc_backend: str):
     conclave_config = CodeGenConfig(workflow_name, int(pid))
     conclave_config.use_leaky_ops = False
 
-    if mpc_backend == "sharemind":
-        sharemind_conf = SharemindCodeGenConfig("/mnt/shared", use_docker=True, use_hdfs=False)
-        conclave_config.with_sharemind_config(sharemind_conf)
-    elif mpc_backend == "obliv-c":
-        conclave_config.all_pids = [1, 2]
-        net_conf = [
-            {"host": "ca-spark-node-0", "port": 8001},
-            {"host": "cb-spark-node-0", "port": 8002}
-        ]
-        net = NetworkConfig(net_conf, int(pid))
-        conclave_config.with_network_config(net)
-
-        oc_conf = OblivcConfig("/obliv-c/bin/oblivcc", "ca-spark-node-0:9000")
-        conclave_config.with_oc_config(oc_conf)
-    else:
-        raise Exception("Unknown MPC backend {}".format(mpc_backend))
+
 
     conclave_config.code_path = os.path.join("/mnt/shared", workflow_name)
     conclave_config.input_path = os.path.join("/mnt/shared", data_root)

conclave/config/__init__.py

@@ -106,6 +106,29 @@ def __init__(self, job_name: [str, None] = None, pid: int = 1):
             }
         }
 
+    def with_default_mpc_config(self, mpc_backend: str):
+        """
+        Boiler plate code for configuring MPC backend.
+        """
+        if mpc_backend == "sharemind":
+            sharemind_conf = SharemindCodeGenConfig("/mnt/shared", use_docker=True, use_hdfs=False)
+            self.with_sharemind_config(sharemind_conf)
+            return self
+        elif mpc_backend == "obliv-c":
+            self.all_pids = [1, 2]
+            net_conf = [
+                {"host": "ca-spark-node-0", "port": 8001},
+                {"host": "cb-spark-node-0", "port": 8002}
+            ]
+            net = NetworkConfig(net_conf, self.pid)
+            self.with_network_config(net)
+
+            oc_conf = OblivcConfig("/obliv-c/bin/oblivcc", "ca-spark-node-0:9000")
+            self.with_oc_config(oc_conf)
+            return self
+        else:
+            raise Exception("Unknown MPC backend {}".format(mpc_backend))
+
     def with_pid(self, pid: int):
         """ Change pid of this party (default is 1). """
 

tests/integration/distinct_count/real.py

@@ -3,7 +3,7 @@
 
 import conclave.lang as cc
 from conclave import generate_code, dispatch_jobs
-from conclave.config import CodeGenConfig, SharemindCodeGenConfig
+from conclave.config import CodeGenConfig
 from conclave.utils import defCol
 
 
@@ -28,15 +28,12 @@ def protocol():
 
 if __name__ == "__main__":
     pid = sys.argv[1]
+    mpc_backend = sys.argv[2]
     # define name for the workflow
     workflow_name = "real-distinct-count-test-" + pid
     # configure conclave
-    conclave_config = CodeGenConfig(workflow_name, int(pid))
-    conclave_config.all_pids = [1, 2, 3]
-    sharemind_conf = SharemindCodeGenConfig("/mnt/shared",
-                                            use_docker=True,
-                                            use_hdfs=False)
-    conclave_config.with_sharemind_config(sharemind_conf)
+    conclave_config = CodeGenConfig(workflow_name, int(pid)).with_default_mpc_config(mpc_backend)
+
     current_dir = os.path.dirname(os.path.realpath(__file__))
     # point conclave to the directory where the generated code should be stored/ read from
     conclave_config.code_path = os.path.join("/mnt/shared", workflow_name)
@@ -45,5 +42,5 @@ def protocol():
     # and written to
     conclave_config.output_path = os.path.join(current_dir, "data")
     # define this party's unique ID (in this demo there is only one party)
-    job_queue = generate_code(protocol, conclave_config, ["sharemind"], ["python"], apply_optimizations=True)
+    job_queue = generate_code(protocol, conclave_config, [mpc_backend], ["python"], apply_optimizations=True)
     dispatch_jobs(job_queue, conclave_config)

tests/integration/single_test.sh

@@ -12,6 +12,12 @@ trap kill_c_procs INT
 
 TEST_SUB_DIR=$1
 DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"/${TEST_SUB_DIR}
+BACKEND=$2
+
+if [ -z ${BACKEND} ]
+then
+    BACKEND=sharemind
+fi
 
 # set up data
 mkdir -p ${DIR}/data
@@ -35,7 +41,7 @@ python ${DIR}/simple.py 1
 # run real workflow
 for i in 1 2 3;
 do
-    python ${DIR}/real.py ${i} &
+    python ${DIR}/real.py ${i} ${BACKEND} &
 done
 wait