Refactored capsule server and fixed bug

build_and_run.sh

@@ -1,3 +1,3 @@
 pip install -r requirements.txt
 python setup.py install
-python capsule_zmq/local_server.py
+python capsule_server/local_server.py --type=zmq --port=5002

build_and_run_mpc.sh

@@ -1,3 +1,3 @@
 pip install -r requirements.txt
 python setup.py install
-python capsule_mpc_server/local_server.py
+python capsule_server/local_server.py --type=mpc --port=5003

capsule/zmq_client.py

@@ -5,13 +5,33 @@
 import pickle
 
 
+class TaskSocket(object):
+    def __init__(self, context):
+        self.context = context
+        self.socket = context.socket(zmq.REQ)
+
+    def connect(self, server):
+        return self.socket.connect(server)
+
+    def send_string(self, string):
+        return self.socket.send_string(string)
+
+    def recv(self):
+        response = self.socket.recv()
+        output = pickle.loads(response)
+        if output.get('success'):
+            return output.get('payload')
+
+        raise Exception(output.get('payload'))
+
+
 class LocalCapsuleClient():
 
     def __init__(self, host='127.0.0.1', port='5000'):
         self.host = host
         self.port = port
         ctx = zmq.Context()
-        self.task_socket = ctx.socket(zmq.REQ)
+        self.task_socket = TaskSocket(ctx)
         self.task_socket.connect('tcp://127.0.0.1:5002')
 
     def keygen(self, scheme='paillier'):
@@ -58,10 +78,8 @@ def decrypt(self, x, id=None):
         }))
         r = self.task_socket.recv()
         try:
-            # print("Hello.........................")
             out = syft.tensor.TensorBase.deserialize(r)
         except Exception as e:
-            # print("Hello.........................")
             print(e)
             out = float(r)
         return out
@@ -72,7 +90,7 @@ def __init__(self, repo):
         self.id = str(random.randint(0, 2**32))
         self.repo = repo
         ctx = zmq.Context()
-        self.task_socket = ctx.socket(zmq.REQ)
+        self.task_socket = TaskSocket(ctx)
         self.task_socket.connect('tcp://127.0.0.1:5003')
 
     def create_siblings(self):

capsule_server/local_server.py

@@ -0,0 +1,56 @@
+import zmq
+from ast import literal_eval
+import pickle
+import argparse
+import importlib
+
+
+def create_payload(data, success=True):
+    output = {
+        'success': success,
+        'payload': data
+    }
+    return pickle.dumps(output)
+
+
+def main(args):
+    tasks = importlib.import_module(args.type + '_' + 'tasks')
+
+    context = zmq.Context()
+    socket = context.socket(zmq.REP)
+    socket.bind('tcp://127.0.0.1:' + args.port)
+    print(args.type.upper() + " Server is listening on 127.0.0.1:" + args.port)
+
+    while True:
+        try:
+            task_data = socket.recv()
+            task_data = literal_eval(task_data.decode('utf-8'))
+            task = task_data.pop('task')
+            task_kwargs = task_data.pop('task_kwargs')
+            server_data = getattr(tasks, task)(**task_kwargs)
+
+            socket.send(create_payload(server_data))
+
+        except Exception as e:
+            print("Encountered exception: ", e)
+            socket.send(create_payload(e, False))
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description="Local Capsule Server")
+    parser.add_argument(
+        '--type',
+        required=True,
+        type=str,
+        help="Type of tasks to be processed.",
+        choices=['mpc', 'zmq']
+    )
+    parser.add_argument(
+        '--port',
+        help="Port for the listening socket.",
+        default='5002'
+    )
+
+    args = parser.parse_args()
+
+    main(args)

capsule_mpc_server/tasks.py → capsule_server/mpc_tasks.py

@@ -1,4 +1,3 @@
-# from flask import Flask, request, Response
 import os
 import pickle
 import redis

capsule_zmq/tasks.py → capsule_server/zmq_tasks.py

@@ -1,14 +1,12 @@
 from syft.he.paillier.keys import KeyPair
 from syft.he.keys import Paillier
 import syft as sy
-# from flask import Flask, request, Response
 import os
 import redis
 
 
 redis_url = os.getenv('REDISTOGO_URL', 'redis://localhost:6379')
 conn = redis.from_url(redis_url)
-# app = Flask(__name__)
 
 
 def create_keys(id, scheme):