Repository: sql-machine-learning/elasticdl

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



Discussion:

Resolving patch:
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
@@ -3,6 +3,7 @@
 from google.protobuf import empty_pb2
 
 from elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc
+from elasticdl.python.common.dtypes import dtype_numpy_to_tensor
 
 
 class PserverServicer(elasticdl_pb2_grpc.PserverServicer):
@@ -25,8 +26,30 @@ def __init__(
         self._lock = threading.Lock()
 
     def pull_variable(self, request, _):
-        # TODO: implement this RPC service
-        return elasticdl_pb2.PullVariableResponse()
+        """
+        Response with all non-embedding parameters if initialized.
+        """
+        res = elasticdl_pb2.PullVariableResponse()
+        if not self._parameters.init_status:
+            res.model_init_status = False
+            return res
+
+        # Only sync-SGD needs lock
+        # TODO: use a read-write lock to support multiple concurrent reads
+        if not self._use_async:
+            self._lock.acquire()
+        res.model.version = self._parameters.version
+        for name, var in self._parameters.non_embedding_params.items():
+            tensor = res.model.param.add()
+            tensor.name = name
+            tensor.dim.extend(var.shape.as_list())
+            var_values = var.numpy()
+            tensor.content = var_values.tobytes()
+            tensor.dtype = dtype_numpy_to_tensor(var_values.dtype)
+        if not self._use_async:
+            self._lock.release()
+        res.model_init_status = True
+        return res
 
     def pull_embedding_vector(self, request, _):
         # TODO: implement this RPC service


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
@@ -7,7 +7,10 @@
 
 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.common.ndarray import (
+    ndarray_to_tensor,
+    tensor_to_ndarray,
+)
 from elasticdl.python.ps.parameter_server import ParameterServer
 
 _test_model_zoo_path = os.path.dirname(os.path.realpath(__file__))
@@ -84,10 +87,6 @@ def testServicer(self):
 
         # TODO: replace the section below with real RPC service tests
         # after service implementation
-        req = elasticdl_pb2.PullVariableRequest()
-        res = self._stub.pull_variable(req)
-        self.assertEqual(res, elasticdl_pb2.PullVariableResponse())
-
         req = elasticdl_pb2.PullEmbeddingVectorRequest()
         res = self._stub.pull_embedding_vector(req)
         self.assertEqual(res, elasticdl_pb2.Tensor())
@@ -145,3 +144,32 @@ def testPushModel(self):
                     embedding_info.name
                 ].initializer,
             )
+
+    def testPullVariable(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),
+        }
+        pull_req = empty_pb2.Empty()
+        # try to pull variable
+        res = self._stub.pull_variable(pull_req)
+        # not initialized
+        self.assertFalse(res.model_init_status)
+
+        # init variable
+        req = elasticdl_pb2.Model()
+        req.version = 1
+        for name in param0:
+            req.param.append(ndarray_to_tensor(param0[name], name))
+        res = self._stub.push_model(req)
+        self.assertEqual(res, empty_pb2.Empty())
+
+        # pull variable back
+        res = self._stub.pull_variable(pull_req)
+        self.assertTrue(res.model_init_status)
+        self.assertEqual(res.model.version, req.version)
+        for param in res.model.param:
+            name = param.name
+            tensor = tensor_to_ndarray(param)
+            self.assertTrue(np.allclose(param0[name], tensor))