Repository: sql-machine-learning/elasticdl

Issue:
[PS-1] RPC service `pull_embedding_vector` 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
@@ -4,6 +4,7 @@
 
 from elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc
 from elasticdl.python.common.dtypes import dtype_numpy_to_tensor
+from elasticdl.python.common.tensor import Tensor, serialize_tensor
 
 
 class PserverServicer(elasticdl_pb2_grpc.PserverServicer):
@@ -52,8 +53,15 @@ def pull_variable(self, request, _):
         return res
 
     def pull_embedding_vector(self, request, _):
-        # TODO: implement this RPC service
-        return elasticdl_pb2.Tensor()
+        ret = elasticdl_pb2.Tensor()
+        if not request.ids:
+            return ret
+        embedding_vectors = self._parameters.get_embedding_param(
+            request.name, request.ids
+        )
+        tensor = Tensor(values=embedding_vectors)
+        serialize_tensor(tensor, ret)
+        return ret
 
     def push_model(self, request, _):
         with self._lock:


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
@@ -52,6 +52,11 @@ def setUp(self):
                 ),
             ],
         )
+        embedding_info = elasticdl_pb2.EmbeddingTableInfo()
+        embedding_info.name = "layer_a"
+        embedding_info.dim = 32
+        embedding_info.initializer = "normal"
+        self._embedding_info = embedding_info
         self._server = None
 
     def tearDown(self):
@@ -82,15 +87,21 @@ def create_default_server_and_stub(self):
             grads_to_wait, lr_staleness_modulation, use_async
         )
 
+    def get_embedding_vectors(self, name, ids):
+        pull_req = elasticdl_pb2.PullEmbeddingVectorRequest()
+        pull_req.name = name
+        pull_req.ids.extend(ids)
+        res = self._stub.pull_embedding_vector(pull_req)
+        if res.content:
+            return tensor_pb_to_ndarray(res)
+        else:
+            return None
+
     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.PullEmbeddingVectorRequest()
-        res = self._stub.pull_embedding_vector(req)
-        self.assertEqual(res, elasticdl_pb2.Tensor())
-
         req = elasticdl_pb2.PushGradientRequest()
         res = self._stub.push_gradient(req)
         self.assertEqual(res, elasticdl_pb2.PushGradientResponse())
@@ -105,10 +116,6 @@ def testPushModel(self):
             "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]
 
@@ -119,7 +126,7 @@ def testPushModel(self):
                 emplace_tensor_pb_from_ndarray(
                     req.param, model[name], name=name
                 )
-            req.embedding_table_info.append(embedding_info)
+            req.embedding_table_info.append(self._embedding_info)
             res = self._stub.push_model(req)
             self.assertEqual(res, empty_pb2.Empty())
             # self._parameters is initialized with the first push_model call
@@ -133,17 +140,21 @@ def testPushModel(self):
                     )
                 )
             self.assertEqual(
-                embedding_info.name,
-                self._parameters.embedding_params[embedding_info.name].name,
+                self._embedding_info.name,
+                self._parameters.embedding_params[
+                    self._embedding_info.name
+                ].name,
             )
             self.assertEqual(
-                embedding_info.dim,
-                self._parameters.embedding_params[embedding_info.name].dim,
+                self._embedding_info.dim,
+                self._parameters.embedding_params[
+                    self._embedding_info.name
+                ].dim,
             )
             self.assertEqual(
-                embedding_info.initializer,
+                self._embedding_info.initializer,
                 self._parameters.embedding_params[
-                    embedding_info.name
+                    self._embedding_info.name
                 ].initializer,
             )
 
@@ -175,3 +186,49 @@ def testPullVariable(self):
             name = param.name
             tensor = tensor_pb_to_ndarray(param)
             self.assertTrue(np.allclose(param0[name], tensor))
+
+    def testPullEmbeddingVector(self):
+        self.create_default_server_and_stub()
+
+        id_list_0 = [1, 3, 9, 6]
+        id_list_1 = [8, 9, 1, 0, 6]
+
+        req = elasticdl_pb2.Model()
+        req.version = 1
+        req.embedding_table_info.append(self._embedding_info)
+        another_embedding_info = elasticdl_pb2.EmbeddingTableInfo()
+        another_embedding_info.name = "layer_b"
+        another_embedding_info.dim = 16
+        another_embedding_info.initializer = "normal"
+        req.embedding_table_info.append(another_embedding_info)
+        res = self._stub.push_model(req)
+        self.assertEqual(res, empty_pb2.Empty())
+
+        vectors_a_0 = self.get_embedding_vectors("layer_a", id_list_0)
+        self.assertEqual(vectors_a_0.shape[0], len(id_list_0))
+        self.assertEqual(vectors_a_0.shape[1], 32)
+
+        vectors_a_1 = self.get_embedding_vectors("layer_a", id_list_1)
+        self.assertEqual(vectors_a_1.shape[0], len(id_list_1))
+        self.assertEqual(vectors_a_1.shape[1], 32)
+
+        vectors_b_1 = self.get_embedding_vectors("layer_b", id_list_1)
+        self.assertEqual(vectors_b_1.shape[0], len(id_list_1))
+        self.assertEqual(vectors_b_1.shape[1], 16)
+
+        vectors_b_0 = self.get_embedding_vectors("layer_b", id_list_0)
+        self.assertEqual(vectors_b_0.shape[0], len(id_list_0))
+        self.assertEqual(vectors_b_0.shape[1], 16)
+
+        for idx0, id0 in enumerate(id_list_0):
+            for idx1, id1 in enumerate(id_list_1):
+                if id0 == id1:
+                    self.assertTrue(
+                        np.array_equal(vectors_a_0[idx0], vectors_a_1[idx1])
+                    )
+                    self.assertTrue(
+                        np.array_equal(vectors_b_0[idx0], vectors_b_1[idx1])
+                    )
+
+        vectors = self.get_embedding_vectors("layer_a", [])
+        self.assertEqual(vectors, None)