Repository: sql-machine-learning/elasticdl

Issue:
[PS-1] RPC service `push_model` implementation



Discussion:

Resolving patch:
diff --git a/elasticdl/python/ps/parameters.py b/elasticdl/python/ps/parameters.py
--- a/elasticdl/python/ps/parameters.py
+++ b/elasticdl/python/ps/parameters.py
@@ -18,6 +18,7 @@ class Parameters(object):
     """
 
     def __init__(self):
+        self.version = 0
         self.init_status = False
         self.non_embedding_params = {}
         self.embedding_params = {}
@@ -42,6 +43,7 @@ def init_from_model_pb(self, model_pb):
             embeddings_pb = model_pb.embedding_table_info
             self._init_non_embedding_params(tensors_pb)
             self._init_embedding_params(embeddings_pb)
+            self.version = model_pb.version
             self.init_status = True
 
     def _init_non_embedding_params(self, tensors_pb):
diff --git a/elasticdl/python/ps/servicer.py b/elasticdl/python/ps/servicer.py
--- a/elasticdl/python/ps/servicer.py
+++ b/elasticdl/python/ps/servicer.py
@@ -1,3 +1,5 @@
+import threading
+
 from google.protobuf import empty_pb2
 
 from elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc
@@ -20,6 +22,7 @@ def __init__(
         self._lr_staleness_modulation = lr_staleness_modulation
         self._use_async = use_async
         self._version = 0
+        self._lock = threading.Lock()
 
     def pull_variable(self, request, _):
         # TODO: implement this RPC service
@@ -30,7 +33,8 @@ def pull_embedding_vector(self, request, _):
         return elasticdl_pb2.Tensor()
 
     def push_model(self, request, _):
-        # TODO: implement this RPC service
+        with self._lock:
+            self._parameters.init_from_model_pb(request)
         return empty_pb2.Empty()
 
     def push_gradient(self, request, _):


Test patch:
diff --git a/elasticdl/python/tests/pserver_servicer_test.py b/elasticdl/python/tests/pserver_servicer_test.py
--- a/elasticdl/python/tests/pserver_servicer_test.py
+++ b/elasticdl/python/tests/pserver_servicer_test.py
@@ -2,12 +2,13 @@
 import unittest
 
 import grpc
+import numpy as np
 from google.protobuf import empty_pb2
 
 from elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc
 from elasticdl.python.common.constants import GRPC
+from elasticdl.python.common.ndarray import ndarray_to_tensor
 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__))
 
@@ -55,7 +56,7 @@ def tearDown(self):
             self._server.stop(0)
 
     def create_server_and_stub(
-        self, parameters, grads_to_wait, lr_staleness_modulation, use_async
+        self, grads_to_wait, lr_staleness_modulation, use_async
     ):
         args = PserverArgs(
             grads_to_wait=grads_to_wait,
@@ -65,19 +66,22 @@ def create_server_and_stub(
         )
         pserver = ParameterServer(args)
         pserver.prepare()
+        self._parameters = pserver.parameters
         self._server = pserver.server
         self._stub = elasticdl_pb2_grpc.PserverStub(self._channel)
 
-    def testServicer(self):
-        parameters = Parameters()
+    def create_default_server_and_stub(self):
         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
+            grads_to_wait, lr_staleness_modulation, use_async
         )
 
+    def testServicer(self):
+        self.create_default_server_and_stub()
+
         # TODO: replace the section below with real RPC service tests
         # after service implementation
         req = elasticdl_pb2.PullVariableRequest()
@@ -88,10 +92,56 @@ def testServicer(self):
         res = self._stub.pull_embedding_vector(req)
         self.assertEqual(res, elasticdl_pb2.Tensor())
 
-        req = elasticdl_pb2.Model()
-        res = self._stub.push_model(req)
-        self.assertEqual(res, empty_pb2.Empty())
-
         req = elasticdl_pb2.PushGradientRequest()
         res = self._stub.push_gradient(req)
         self.assertEqual(res, elasticdl_pb2.PushGradientResponse())
+
+    def testPushModel(self):
+        self.create_default_server_and_stub()
+        param0 = {
+            "v0": np.random.rand(3, 2).astype(np.float32),
+            "v1": np.random.rand(10, 32).astype(np.float32),
+        }
+        param1 = {
+            "v0": np.ones([3, 2], dtype=np.float32),
+            "v1": np.ones([10, 32], dtype=np.float32),
+        }
+        embedding_info = elasticdl_pb2.EmbeddingTableInfo()
+        embedding_info.name = "layer0"
+        embedding_info.dim = 32
+        embedding_info.initializer = "normal"
+
+        models = [param0, param1]
+
+        for idx, model in enumerate(models):
+            req = elasticdl_pb2.Model()
+            req.version = idx + 1
+            for name in model:
+                req.param.append(ndarray_to_tensor(model[name], name))
+            req.embedding_table_info.append(embedding_info)
+            res = self._stub.push_model(req)
+            self.assertEqual(res, empty_pb2.Empty())
+            # self._parameters is initialized with the first push_model call
+            # and the second push_model has no effect
+            self.assertEqual(self._parameters.version, 1)
+            for name in param0:
+                self.assertTrue(
+                    np.allclose(
+                        param0[name],
+                        self._parameters.non_embedding_params[name].numpy(),
+                    )
+                )
+            self.assertEqual(
+                embedding_info.name,
+                self._parameters.embedding_params[embedding_info.name].name,
+            )
+            self.assertEqual(
+                embedding_info.dim,
+                self._parameters.embedding_params[embedding_info.name].dim,
+            )
+            self.assertEqual(
+                embedding_info.initializer,
+                self._parameters.embedding_params[
+                    embedding_info.name
+                ].initializer,
+            )