Repository: sql-machine-learning/elasticdl

Issue:
worker.embedding_lookup(ids)
As [description in design doc](https://github.com/wangkuiyi/elasticdl/blob/develop/elasticdl/doc/distributed_embedding_layer_design.md#initialization-of-embedding-vectors), worker will get embedding vectors from embeddding table by
 `lookup_embedding`. If embedding vectors do not exist, worker will initialize it.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/worker/worker.py b/elasticdl/python/worker/worker.py
--- a/elasticdl/python/worker/worker.py
+++ b/elasticdl/python/worker/worker.py
@@ -1,9 +1,11 @@
 import traceback
 
+import numpy as np
 import tensorflow as tf
 
 from elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc
 from elasticdl.python.common.constants import JobType, Mode
+from elasticdl.python.common.embedding_service import EmbeddingService
 from elasticdl.python.common.log_util import default_logger as logger
 from elasticdl.python.common.model_helper import (
     find_layer,
@@ -39,6 +41,7 @@ def __init__(
         optimizer="optimizer",
         eval_metrics_fn="eval_metrics_fn",
         channel=None,
+        embedding_service_endpoint=None,
         model_def=None,
         model_params="",
         prediction_outputs_processor="PredictionOutputsProcessor",
@@ -69,6 +72,7 @@ def __init__(
             self._stub = None
         else:
             self._stub = elasticdl_pb2_grpc.MasterStub(channel)
+        self._embedding_service_endpoint = embedding_service_endpoint
         self._max_minibatch_retry_num = max_minibatch_retry_num
         self._model_version = -1
         self._task_data_service = TaskDataService(
@@ -132,6 +136,51 @@ def get_model(self, version, method):
             var.assign(tensor_to_ndarray(model.param[var.name]))
         self._model_version = model.version
 
+    def lookup_embedding(
+        self, ids, layer_name, initializer="uniform", embedding_table_dim=128
+    ):
+        keys = [Embedding.get_key([layer_name, id]) for id in ids]
+        ES_lookup_embedding = EmbeddingService.lookup_embedding
+        embedding_vectors, unknown_keys_index = ES_lookup_embedding(
+            keys=keys,
+            embedding_service_endpoint=self._embedding_service_endpoint,
+        )
+        if unknown_keys_index:
+            # Initialize unknown_keys' embedding vectors and write into Redis.
+            unknown_keys = [keys[index] for index in unknown_keys_index]
+            initializer = tf.keras.initializers.get(initializer)
+            embedding_vector_init = [
+                initializer(shape=[1, embedding_table_dim]).numpy()
+                for _ in unknown_keys
+            ]
+            embedding_vector_init = np.concatenate(
+                embedding_vector_init, axis=0
+            )
+            EmbeddingService.update_embedding(
+                keys=unknown_keys,
+                embedding_vectors=embedding_vector_init,
+                embedding_service_endpoint=self._embedding_service_endpoint,
+                set_if_not_exist=True,
+            )
+            # Lookup unknown_keys' embedding vectors
+            embedding_vectors_new, unknown_keys_idx_new = ES_lookup_embedding(
+                keys=unknown_keys,
+                embedding_service_endpoint=self._embedding_service_endpoint,
+            )
+            if unknown_keys_idx_new:
+                raise Exception(
+                    "Update embedding vector: %s failed."
+                    % str(
+                        [unknown_keys[index] for index in unknown_keys_idx_new]
+                    )
+                )
+            for key_index, vector in zip(
+                unknown_keys_index, embedding_vectors_new
+            ):
+                embedding_vectors[key_index] = vector
+        embedding_vectors = np.concatenate(embedding_vectors, axis=0)
+        return embedding_vectors.reshape((len(keys), embedding_table_dim))
+
     def report_task_result(self, task_id, err_msg):
         """
         report task result to master


Test patch:
diff --git a/elasticdl/python/tests/worker_test.py b/elasticdl/python/tests/worker_test.py
--- a/elasticdl/python/tests/worker_test.py
+++ b/elasticdl/python/tests/worker_test.py
@@ -3,12 +3,14 @@
 import unittest
 from contextlib import closing
 
+import mock
 import numpy as np
 import recordio
 import tensorflow as tf
 
 from elasticdl.proto import elasticdl_pb2
 from elasticdl.python.common.constants import JobType
+from elasticdl.python.common.embedding_service import EmbeddingService
 from elasticdl.python.common.model_helper import get_model_file, load_module
 from elasticdl.python.master.evaluation_service import EvaluationService
 from elasticdl.python.master.servicer import MasterServicer
@@ -16,6 +18,33 @@
 from elasticdl.python.tests.in_process_master import InProcessMaster
 from elasticdl.python.worker.worker import Worker
 
+
+class MockEmbeddingService:
+    def __init__(self):
+        self.mock_embedding_table = None
+
+    def mock_lookup_embedding(self, **kwargs):
+        keys = kwargs["keys"]
+        embeddings = []
+        unknown_index = []
+        for index, k in enumerate(keys):
+            if k in self.mock_embedding_table:
+                embeddings.append(
+                    self.mock_embedding_table[k].reshape((1, -1))
+                )
+            else:
+                unknown_index.append(index)
+                embeddings.append(None)
+        return embeddings, unknown_index
+
+    def mock_update_embedding(self, **kwargs):
+        keys, embeddings = kwargs["keys"], kwargs["embedding_vectors"]
+        if embeddings is None:
+            return
+        for k, emb in zip(keys, embeddings):
+            self.mock_embedding_table[k] = emb
+
+
 _model_file = get_model_file(
     os.path.dirname(os.path.realpath(__file__)), "test_module.custom_model"
 )
@@ -139,6 +168,53 @@ def test_embedding_layer(self):
         )
         self.assertTrue(len(worker._embedding_layers) == 2)
 
+    def test_lookup_embedding(self):
+        mock_embedding_service = MockEmbeddingService()
+
+        ids = [1, 2, 3, 4, 5, 6]
+        layer_name = "test_edlembedding"
+        embedding_table_dim = 10
+        mock_embedding_service.mock_embedding_table = {
+            "test_edlembedding-1": np.zeros(
+                (1, embedding_table_dim), dtype=np.float32
+            ),
+            "test_edlembedding-2": np.zeros(
+                (1, embedding_table_dim), dtype=np.float32
+            ),
+            "test_edlembedding-3": np.zeros(
+                (1, embedding_table_dim), dtype=np.float32
+            ),
+        }
+        worker = Worker(
+            1,
+            JobType.TRAINING_ONLY,
+            32,
+            _embedding_model_file,
+            model_def="embedding_test_module.CustomModel",
+            channel=None,
+        )
+        with mock.patch.object(
+            EmbeddingService,
+            "lookup_embedding",
+            mock_embedding_service.mock_lookup_embedding,
+        ), mock.patch.object(
+            EmbeddingService,
+            "update_embedding",
+            mock_embedding_service.mock_update_embedding,
+        ):
+            e_lookup, e_unknown = EmbeddingService.lookup_embedding(
+                keys=["-".join([layer_name, str(id)]) for id in ids]
+            )
+            lookup_result = worker.lookup_embedding(
+                ids=ids,
+                layer_name=layer_name,
+                embedding_table_dim=embedding_table_dim,
+            )
+            self.assertTrue(len(e_lookup) == 6)
+            self.assertTrue(len(e_unknown) == 3)
+            self.assertTrue(len(lookup_result) == 6)
+            self.assertFalse(None in lookup_result)
+
 
 if __name__ == "__main__":
     unittest.main()