Repository: sql-machine-learning/elasticdl

Issue:
[PS-1] Initialize PS RPC servicer
Create a Servicer class, and `ps.main` start the RPC service.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/ps/main.py b/elasticdl/python/ps/main.py
--- a/elasticdl/python/ps/main.py
+++ b/elasticdl/python/ps/main.py
@@ -1,19 +1,13 @@
 from elasticdl.python.common.args import parse_ps_args
-from elasticdl.python.common.log_utils import get_logger
-from elasticdl.python.common.model_utils import (
-    get_module_file_path,
-    load_module,
-)
+from elasticdl.python.ps.parameter_server import ParameterServer
 
 
 def main():
     args = parse_ps_args()
-    logger = get_logger("PS", level=args.log_level.upper())
+    pserver = ParameterServer(args)
+    pserver.prepare()
+    pserver.run()
 
-    model_module = load_module(
-        get_module_file_path(args.model_zoo, args.model_def)
-    ).__dict__
-    optimizer = model_module[args.optimizer]()
 
-    logger.info("Starting PS pod with optimizer as %s", optimizer._name)
-    # TODO: create `Parameters` class instance, and start PS RPC servicer.
+if __name__ == "__main__":
+    main()
diff --git a/elasticdl/python/ps/parameter_server.py b/elasticdl/python/ps/parameter_server.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/ps/parameter_server.py
@@ -0,0 +1,67 @@
+import time
+from concurrent import futures
+
+import grpc
+
+from elasticdl.proto import elasticdl_pb2_grpc
+from elasticdl.python.common.constants import GRPC
+from elasticdl.python.common.log_utils import get_logger
+from elasticdl.python.common.model_utils import (
+    get_module_file_path,
+    load_module,
+)
+from elasticdl.python.ps.parameters import Parameters
+from elasticdl.python.ps.servicer import PserverServicer
+
+
+class ParameterServer(object):
+    def __init__(self, args):
+        self.logger = get_logger("PS", level=args.log_level.upper())
+
+        self.grads_to_wait = args.grads_to_wait
+        self.lr_staleness_modulation = args.lr_staleness_modulation
+        self.use_async = args.use_async
+        self.port = args.port
+        model_module = load_module(
+            get_module_file_path(args.model_zoo, args.model_def)
+        ).__dict__
+        self.optimizer = model_module[args.optimizer]()
+        # Create Parameters instance
+        self.parameters = Parameters()
+
+    def prepare(self):
+        server = grpc.server(
+            futures.ThreadPoolExecutor(max_workers=64),
+            options=[
+                ("grpc.max_send_message_length", GRPC.MAX_SEND_MESSAGE_LENGTH),
+                (
+                    "grpc.max_receive_message_length",
+                    GRPC.MAX_RECEIVE_MESSAGE_LENGTH,
+                ),
+            ],
+        )
+        pserver_servicer = PserverServicer(
+            self.parameters,
+            self.grads_to_wait,
+            self.optimizer,
+            lr_staleness_modulation=self.lr_staleness_modulation,
+            use_async=self.use_async,
+        )
+        elasticdl_pb2_grpc.add_PserverServicer_to_server(
+            pserver_servicer, server
+        )
+        server.add_insecure_port("[::]:{}".format(self.port))
+        server.start()
+        self.server = server
+        self.logger.info("RPC Server started at port: %d", self.port)
+
+    def run(self):
+        try:
+            while True:
+                # TODO: add loop break condition
+                time.sleep(30)
+        except KeyboardInterrupt:
+            self.logger.warning("Server stopping")
+
+        self.server.stop(0)
+        self.logger.info("RPC server stopped")
diff --git a/elasticdl/python/ps/servicer.py b/elasticdl/python/ps/servicer.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/ps/servicer.py
@@ -0,0 +1,38 @@
+from google.protobuf import empty_pb2
+
+from elasticdl.proto import elasticdl_pb2_grpc
+
+
+class PserverServicer(elasticdl_pb2_grpc.PserverServicer):
+    """PS service implementation"""
+
+    def __init__(
+        self,
+        parameters,
+        grads_to_wait,
+        optimizer,
+        lr_staleness_modulation=False,
+        use_async=False,
+    ):
+        self._parameters = parameters
+        self._grads_to_wait = grads_to_wait
+        self._optimizer = optimizer
+        self._lr_staleness_modulation = lr_staleness_modulation
+        self._use_async = use_async
+        self._version = 0
+
+    def pull_variable(self, request, _):
+        # TODO: implement this RPC service
+        return empty_pb2.Empty()
+
+    def pull_embedding_vector(self, request, _):
+        # TODO: implement this RPC service
+        return empty_pb2.Empty()
+
+    def push_model(self, request, _):
+        # TODO: implement this RPC service
+        return empty_pb2.Empty()
+
+    def push_gradient(self, request, _):
+        # TODO: implement this RPC service
+        return empty_pb2.Empty()


Test patch:
diff --git a/elasticdl/python/tests/pserver_servicer_test.py b/elasticdl/python/tests/pserver_servicer_test.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/tests/pserver_servicer_test.py
@@ -0,0 +1,94 @@
+import os
+import unittest
+
+import grpc
+from google.protobuf import empty_pb2
+
+from elasticdl.proto import elasticdl_pb2_grpc
+from elasticdl.python.common.constants import GRPC
+from elasticdl.python.ps.parameter_server import ParameterServer
+from elasticdl.python.ps.parameters import Parameters
+
+_test_model_zoo_path = os.path.dirname(os.path.realpath(__file__))
+
+
+class PserverArgs(object):
+    def __init__(
+        self,
+        grads_to_wait=8,
+        lr_staleness_modulation=0,
+        use_async=False,
+        model_zoo=_test_model_zoo_path,
+        model_def="test_module.custom_model",
+        optimizer="optimizer",
+        port=9999,
+        log_level="INFO",
+    ):
+        self.grads_to_wait = grads_to_wait
+        self.lr_staleness_modulation = lr_staleness_modulation
+        self.use_async = use_async
+        self.model_zoo = model_zoo
+        self.model_def = model_def
+        self.optimizer = optimizer
+        self.port = port
+        self.log_level = log_level
+
+
+class PserverServicerTest(unittest.TestCase):
+    def setUp(self):
+        self._port = 9999
+        addr = "localhost:%d" % self._port
+        self._channel = grpc.insecure_channel(
+            addr,
+            options=[
+                ("grpc.max_send_message_length", GRPC.MAX_SEND_MESSAGE_LENGTH),
+                (
+                    "grpc.max_receive_message_length",
+                    GRPC.MAX_RECEIVE_MESSAGE_LENGTH,
+                ),
+            ],
+        )
+        self._server = None
+
+    def tearDown(self):
+        if self._server:
+            self._server.stop(0)
+
+    def create_server_and_stub(
+        self, parameters, grads_to_wait, lr_staleness_modulation, use_async
+    ):
+        args = PserverArgs(
+            grads_to_wait=grads_to_wait,
+            lr_staleness_modulation=lr_staleness_modulation,
+            use_async=use_async,
+            port=self._port,
+        )
+        pserver = ParameterServer(args)
+        pserver.prepare()
+        self._server = pserver.server
+        self._stub = elasticdl_pb2_grpc.PserverStub(self._channel)
+
+    def testServicer(self):
+        parameters = Parameters()
+        grads_to_wait = 8
+        lr_staleness_modulation = False
+        use_async = True
+
+        self.create_server_and_stub(
+            parameters, grads_to_wait, lr_staleness_modulation, use_async
+        )
+
+        # TODO: replace the section below with real RPC service tests
+        # after service implementation
+        try:
+            req = empty_pb2.Empty()
+            res = self._stub.pull_variable(req)
+            self.assertEqual(res, req)
+            res = self._stub.pull_embedding_vector(req)
+            self.assertEqual(res, req)
+            res = self._stub.push_model(req)
+            self.assertEqual(res, req)
+            res = self._stub.push_gradient(req)
+            self.assertEqual(res, req)
+        except Exception:
+            self.assertTrue(False)