Repository: sql-machine-learning/elasticdl

Issue:
Master aggregates gradients and updates embedding table
This issue is for [Master update embedding table with embedding vectors' gradients](https://github.com/wangkuiyi/elasticdl/blob/develop/elasticdl/doc/distributed_embedding_layer_design.md#master-update-embedding-table-with-embedding-vectors-gradients)  in design doc. 
Master aggregates gradients from all workers, updates embedding table through `optimizer.apply_gradient`, and reports embedding table to distributed storage.


Discussion:

Resolving patch:
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
@@ -5,6 +5,7 @@
 from google.protobuf import empty_pb2
 
 from elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc
+from elasticdl.python.common.embedding_service import EmbeddingService
 from elasticdl.python.common.file_helper import copy_if_not_exists
 from elasticdl.python.common.log_util import default_logger as logger
 from elasticdl.python.common.model_helper import load_from_checkpoint_file
@@ -13,6 +14,7 @@
     tensor_to_ndarray,
 )
 from elasticdl.python.common.tensor_helper import merge_indexed_slices
+from elasticdl.python.elasticdl.layers.embedding import Embedding
 from elasticdl.python.master.checkpoint_service import CheckpointService
 
 
@@ -30,6 +32,7 @@ def __init__(
         checkpoint_filename_for_init,
         checkpoint_service,
         evaluation_service,
+        embedding_service_address=None
     ):
         # TODO: group params together into a single object.
         self._opt = optimizer
@@ -47,6 +50,7 @@ def __init__(
         # optimizer's apply_gradients() function.
         self._model = {}
         self._version = 0
+        self.embedding_service_address = embedding_service_address
         self._init_model(checkpoint_filename_for_init, init_var)
 
         self._checkpoint_service = checkpoint_service
@@ -138,17 +142,84 @@ def _update_model_version(self):
         assert self._lock.locked()
         self._version += 1
 
+    def _update_edl_embedding_table(self, name_var_list):
+        """
+            Put updated embedding vectors' ids and values together
+            and use EmbeddingService.update_embedding() to update
+            embedding table in the distributed storage
+        """
+        keys = []
+        embeddings = None
+        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 = [
+                    Embedding.get_key([layer_name, i])
+                    for i in unique_ids.numpy()
+                ]
+                embeddings = embedding_var.numpy()
+
+        if embeddings is not None:
+            EmbeddingService.update_embedding(
+                keys=keys,
+                embeddings=embeddings,
+                embedding_service_address=self.embedding_service_address,
+            )
+
     def _update_model(self):
         assert self._lock.locked()
         grad_var = []
+
+        # (grad, var) pairs excluding keras Embedding layer and
+        # ElasticDL Embedding layer
         for k in self._gradient_sum:
             self._gradient_sum[k] = self._gradient_sum[k] / self._grad_to_wait
             grad_var.append((self._gradient_sum[k], self._model[k]))
 
+        # (grad, var) pair of Keras Embedding layer
         for k in self._gradient_sum_indexed:
             grad_var.append((self._gradient_sum_indexed[k], self._model[k]))
 
+        # (grad, var) pair of ElasticDL Embedding layer
+        edl_embedding_offset = len(grad_var)
+        unique_ids_list = []
+        if self._edl_embedding_gradients:
+            for layer_name, grads in self._edl_embedding_gradients.items():
+                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(
+                    embedding_service_address=self.embedding_service_address,
+                    keys=[
+                        Embedding.get_key([layer_name, i])
+                        for i in unique_ids.numpy()
+                    ],
+                )
+                if embeddings is None:
+                    continue
+                embedding_var = tf.Variable(embeddings)
+                grad_var.append((grads_idx_transformed, embedding_var))
+
+        # TODO: support optimizer with slots such as Adam, FTRL
         self._opt.apply_gradients(grad_var)
+
+        # report updated embedding table to EmbeddingService
+        self._update_edl_embedding_table(
+            zip(
+                self._edl_embedding_gradients.keys(),
+                unique_ids_list,
+                [v for g, v in grad_var[edl_embedding_offset:]],
+            )
+        )
         self._update_model_version()
         self._gradient_sum.clear()
         self._gradient_sum_indexed.clear()
@@ -248,31 +319,22 @@ def ReportGradient(self, request, _):
         with self._lock:
             tmp = {}
             indexed_grads = {}
+            edl_embedding_gradients = {}
             # Do sanity check before accumulating gradients.
             for k, v in request.gradient.items():
                 if k not in self._model:
                     if v.indices:
-                        # IndexedSlices Gradients of elasticdl.layers.embedding
-                        # TODO: Use the gradients to update embedding
-                        # table of elasticdl.layers.embedding
+                        # grads of ElasticDL Embedding layer
+                        # TODO: check arr.shape[1] = embedding_dim of this
+                        # EdlEmbedding layer
                         arr = tensor_to_ndarray(v)
-                        if k in self._edl_embedding_gradients:
-                            self._edl_embedding_gradients[
-                                k
-                            ] = merge_indexed_slices(
-                                self._edl_embedding_gradients[k], arr
-                            )
-                        else:
-                            self._edl_embedding_gradients[k] = arr
-                        logger.warning(
-                            "Update embedding table of "
-                            "elasticdl.layers.embedding is not implemented."
-                        )
+                        edl_embedding_gradients[k] = arr
                         continue
                     else:
                         raise ValueError(
                             "Gradient key: %s is not part of model", k
                         )
+
                 arr = tensor_to_ndarray(v)
                 if isinstance(arr, tf.IndexedSlices):
                     if arr.values.shape[1] != self._model[k].numpy().shape[1]:
@@ -305,16 +367,26 @@ def ReportGradient(self, request, _):
                         )
                     tmp[k] = arr
 
+            # grads of ElasticDL Embedding layer
+            for k, v in edl_embedding_gradients.items():
+                if k in self._edl_embedding_gradients:
+                    self._edl_embedding_gradients[k] = merge_indexed_slices(
+                        self._edl_embedding_gradients[k], v
+                    )
+                else:
+                    self._edl_embedding_gradients[k] = v
+
+            # grads of Keras Embedding layer
             for k, v in indexed_grads.items():
                 if k not in self._gradient_sum_indexed:
                     self._gradient_sum_indexed[k] = v
                 else:
                     grads_s = self._gradient_sum_indexed[k]
-                    self._gradient_sum_indexed[k] = tf.IndexedSlices(
-                        tf.concat([grads_s.values, v.values], axis=0),
-                        tf.concat([grads_s.indices, v.indices], axis=0),
+                    self._gradient_sum_indexed[k] = merge_indexed_slices(
+                        grads_s, v
                     )
 
+            # other grads
             for k, v in tmp.items():
                 if k in self._gradient_sum:
                     self._gradient_sum[k] = self._gradient_sum[k] + v


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
@@ -1,15 +1,54 @@
 import os
 import unittest
 
+import mock
 import numpy as np
 import tensorflow as tf
 
 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.servicer import MasterServicer
 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 = None
+        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
+
+    def mock_update_embedding(self, **kwargs):
+        keys, embeddings = kwargs["keys"], kwargs["embeddings"]
+        if embeddings is None:
+            return
+        for k, emb in zip(keys, embeddings):
+            layer_name, idx = k.split("-")
+            idx = int(idx)
+            self.mock_embedding_table[layer_name][idx] = emb
+
+
 _model_file = get_model_file(
     os.path.dirname(os.path.realpath(__file__)), "test_module.custom_model"
 )
@@ -70,9 +109,15 @@ def _create_master_and_worker(self):
 
         return master, worker
 
-    def test_report_bet_gradients(self):
+    def test_report_bet_gradients_worker_to_master(self):
         master, worker = self._create_master_and_worker()
 
+        mock_embedding_service = MockEmbeddingService()
+        mock_embedding_service.mock_embedding_table = {
+            "test_edlembedding_1": np.zeros((5, 3), dtype=np.float32),
+            "test_edlembedding_2": np.zeros((5, 3), dtype=np.float32),
+        }
+
         layer1 = MockEdlEmbedding("test_edlembedding_1")
         layer1.bet_ids_pair = [
             (None, tf.constant([1, 2])),
@@ -175,6 +220,69 @@ def test_report_bet_gradients(self):
         for i, j in zip(master._model.values(), expected_weights):
             self.assertTrue(np.all(i.numpy() - j < 0.0001))
 
+    def test_report_bet_gradients_master_to_service(self):
+        master, _ = self._create_master_and_worker()
+
+        layer_names = ["test_layer_1", "test_layer_2"]
+
+        mock_embedding_service = MockEmbeddingService()
+        mock_embedding_service.mock_embedding_table = {
+            layer_names[0]: np.zeros((2, 4), dtype=np.float32),
+            layer_names[1]: np.zeros((4, 4), dtype=np.float32),
+        }
+        for i in range(2):
+            mock_embedding_service.mock_embedding_table[layer_names[0]][
+                i
+            ].fill(i)
+        for i in range(4):
+            mock_embedding_service.mock_embedding_table[layer_names[1]][
+                i
+            ].fill(i)
+
+        grads = tf.reshape(tf.range(28, dtype=tf.float32), (7, 4))
+        indices = tf.convert_to_tensor([0, 1, 0, 2, 0, 2, 3])
+
+        master._edl_embedding_gradients = {
+            layer_names[0]: tf.IndexedSlices(grads[:3], indices[:3]),
+            layer_names[1]: tf.IndexedSlices(grads[3:], indices[3:]),
+        }
+
+        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,
+        ):
+            with master._lock:
+                assert master._lock.locked()
+                master._update_model()
+
+        expected_embedding_table = {
+            layer_names[0]: np.array(
+                [[-0.8, -1, -1.2, -1.4], [0.6, 0.5, 0.4, 0.3]]
+            ),
+            layer_names[1]: np.array(
+                [
+                    [-1.6, -1.7, -1.8, -1.9],
+                    [1, 1, 1, 1],
+                    [-1.2, -1.4, -1.6, -1.8],
+                    [0.6, 0.5, 0.4, 0.3],
+                ]
+            ),
+        }
+
+        for layer in layer_names:
+            self.assertTrue(
+                (
+                    expected_embedding_table[layer]
+                    - mock_embedding_service.mock_embedding_table[layer]
+                    < 0.0001
+                ).all()
+            )
+
     def test_get_trainable_variable(self):
         master, worker = self._create_master_and_worker()
         layer = MockEdlEmbedding("test")
@@ -183,7 +291,6 @@ def test_get_trainable_variable(self):
         ]
         worker._embedding_layers = [layer]
         train_vars = worker.get_trainable_items()
-        print(train_vars)
         self.assertTrue("embedding" in train_vars[0].name)
         self.assertTrue("dense" in train_vars[1].name)
         self.assertTrue("test_bet" in train_vars[2].name)