Repository: sql-machine-learning/elasticdl

Issue:
The interface between master and embedding service is wrong
Currently, 
* `EmbeddingService.lookup_embedding` returns a list of embedding vectors while master assumes that `lookup_embedding` returns a 2D numpy array. 
* `EmbeddingService.update_embedding` should take argument `embedding_vectors` as a list of embedding vectors while master passes in a 2D numpy array.



Discussion:

Resolving patch:
diff --git a/elasticdl/python/common/embedding_service.py b/elasticdl/python/common/embedding_service.py
--- a/elasticdl/python/common/embedding_service.py
+++ b/elasticdl/python/common/embedding_service.py
@@ -296,7 +296,7 @@ def update_embedding(
         if (
             keys is None
             or embedding_vectors is None
-            or len(keys) != embedding_vectors.shape[0]
+            or len(keys) != len(embedding_vectors)
         ):
             raise Exception(
                 "keys and embedding_vectors can not be 'None'. "
diff --git a/elasticdl/python/master/servicer.py b/elasticdl/python/master/servicer.py
--- a/elasticdl/python/master/servicer.py
+++ b/elasticdl/python/master/servicer.py
@@ -149,29 +149,20 @@ def _update_edl_embedding_table(self, name_var_list):
             embedding table in the distributed storage
         """
         keys = []
-        embeddings = None
+        embeddings = []
         for layer_name, unique_ids, embedding_var in name_var_list:
-            if keys:
-                keys.extend(
-                    [
-                        Embedding.get_key([layer_name, i])
-                        for i in unique_ids.numpy()
-                    ]
-                )
-                embeddings = np.concatenate(
-                    (embeddings, embedding_var.numpy()), axis=0
-                )
-            else:
-                keys = [
+            keys.extend(
+                [
                     Embedding.get_key([layer_name, i])
                     for i in unique_ids.numpy()
                 ]
-                embeddings = embedding_var.numpy()
+            )
+            embeddings.extend([i for i in embedding_var.numpy()])
 
-        if embeddings is not None:
+        if embeddings:
             EmbeddingService.update_embedding(
                 keys=keys,
-                embeddings=embeddings,
+                embedding_vectors=embeddings,
                 embedding_service_endpoint=self._embedding_service_endpoint,
             )
 
@@ -197,17 +188,26 @@ def _update_model(self):
                 unique_ids, idx = tf.unique(grads.indices)
                 unique_ids_list.append(unique_ids)
                 grads_idx_transformed = tf.IndexedSlices(grads.values, idx)
-                embeddings = EmbeddingService.lookup_embedding(
+                keys = [
+                    Embedding.get_key([layer_name, i])
+                    for i in unique_ids.numpy()
+                ]
+                embeddings, unknown_keys = EmbeddingService.lookup_embedding(
                     embedding_service_endpoint=(
                         self._embedding_service_endpoint
                     ),
-                    keys=[
-                        Embedding.get_key([layer_name, i])
-                        for i in unique_ids.numpy()
-                    ],
+                    keys=keys,
                 )
-                if embeddings is None:
+                if unknown_keys:
+                    raise RuntimeError(
+                        "Master reviced %d unknown embedding keys: %s ..."
+                        % (len(unknown_keys), str(unknown_keys[0]))
+                    )
+                if not embeddings:
                     continue
+                embeddings = np.concatenate(embeddings, axis=0).reshape(
+                    len(keys), -1
+                )
                 embedding_var = tf.Variable(embeddings)
                 grad_var.append((grads_idx_transformed, embedding_var))
 


Test patch:
diff --git a/elasticdl/python/tests/report_gradients_of_bet_test.py b/elasticdl/python/tests/report_gradients_of_bet_test.py
--- a/elasticdl/python/tests/report_gradients_of_bet_test.py
+++ b/elasticdl/python/tests/report_gradients_of_bet_test.py
@@ -19,28 +19,15 @@ def __init__(self):
 
     def mock_lookup_embedding(self, **kwargs):
         keys = kwargs["keys"]
-        embeddings = None
+        embeddings = []
         for k in keys:
             layer_name, idx = k.split("-")
             idx = int(idx)
-            if embeddings is None:
-                embeddings = self.mock_embedding_table[layer_name][
-                    idx
-                ].reshape((1, -1))
-            else:
-                embeddings = np.concatenate(
-                    [
-                        embeddings,
-                        self.mock_embedding_table[layer_name][idx].reshape(
-                            (1, -1)
-                        ),
-                    ],
-                    axis=0,
-                )
-        return embeddings
+            embeddings.append(self.mock_embedding_table[layer_name][idx])
+        return embeddings, []
 
     def mock_update_embedding(self, **kwargs):
-        keys, embeddings = kwargs["keys"], kwargs["embeddings"]
+        keys, embeddings = kwargs["keys"], kwargs["embedding_vectors"]
         if embeddings is None:
             return
         for k, emb in zip(keys, embeddings):