Repository: sql-machine-learning/elasticdl

Issue:
Optimizer for Embedding Layer Step 3: Apply Gradients and Report to Kv Store
According to [design doc](https://github.com/wangkuiyi/elasticdl/blob/develop/elasticdl/doc/distributed_embedding_layer_design.md#optimizer), call `apply_gradients` function of TensorFlow optimizers to update model parameters and slots, and report updated embedding vectors and slots to kv store.



Discussion:

Resolving patch:
diff --git a/elasticdl/python/master/optimizer_wrapper.py b/elasticdl/python/master/optimizer_wrapper.py
--- a/elasticdl/python/master/optimizer_wrapper.py
+++ b/elasticdl/python/master/optimizer_wrapper.py
@@ -168,9 +168,11 @@ def apply_gradients(self, grads_and_vars):
         )
         self._set_slot_values_to_variables(slot_values)
 
-        # TODO: implement the following logic:
-        # * call self._opt.apply_gradients
-        # * report updated values to Redis
+        self._opt.apply_gradients(
+            grads_and_vars_local + grads_and_vars_kv_store
+        )
+
+        self._report_to_kv_store()
 
     def _lookup_embeddings_and_slots(self, grads_and_vars):
         """Look up embedding vectors and slot values form kv store.
@@ -253,6 +255,7 @@ def _generate_lookup_keys(self, grads_and_vars):
         embed_key_index = {}
         slot_keys = []
         slot_key_index = {}
+        self._unique_ids_all_layers = {}
 
         # generate keys
         for it, (grad, layer_name) in enumerate(grads_and_vars):
@@ -325,6 +328,8 @@ def _get_embedding_variable(self, layer_name):
         """Get the variable for the specified ElasticDL embedding layer."""
         return self._embed_variables.get(layer_name, None)
 
+    # TODO: refactor _create_slot_variable and _create_embedding_variable
+    # into one function
     def _create_embedding_variable(self, layer_name, initial_value):
         """Create a variable for an ElasticDL embedding layer."""
         dim = self._embed_dims[layer_name]
@@ -407,6 +412,26 @@ def _create_slot_variable_in_optimizer(
                 "be in self._opt." % (var_key, slot_name)
             )
 
+    def _report_to_kv_store(self):
+        """Report updated embedding vectors and slots to kv store."""
+        keys = []
+        values = []
+        for layer, ids in self._unique_ids_all_layers.items():
+            value = self._get_embedding_variable(layer).numpy()
+            for id, v in zip(ids, value):
+                keys.append(Embedding.get_key([layer, id]))
+                values.append(v)
+
+            for slot in self._allowed_slot_names:
+                value = self._get_slot_variable(layer, slot).numpy()
+                for id, v in zip(ids, value):
+                    keys.append(Embedding.get_key([layer, slot, id]))
+                    values.append(v)
+
+        EmbeddingService.update_embedding(
+            keys, values, self._kv_store_endpoint
+        )
+
     @property
     def allowed_slot_names(self):
         return self._allowed_slot_names


Test patch:
diff --git a/elasticdl/python/tests/optimizer_wrapper_test.py b/elasticdl/python/tests/optimizer_wrapper_test.py
--- a/elasticdl/python/tests/optimizer_wrapper_test.py
+++ b/elasticdl/python/tests/optimizer_wrapper_test.py
@@ -283,6 +283,51 @@ def test_set_embedding_values_to_variables(self):
             (opt_wrapper._embed_variables["test-1"].numpy() < 0.0001).all()
         )
 
+    def test_report_to_kv_store(self):
+        opt = SGD(momentum=0.1)
+        opt_wrapper = OptimizerWrapper(opt, None, {})
+
+        ids_list = [[1, 5], [10]]
+        opt_wrapper._unique_ids_all_layers = {
+            "test_1": np.array(ids_list[0]),
+            "test_2": np.array(ids_list[1]),
+        }
+        t = np.array([1.0, 1.0, 1.0])
+        opt_wrapper._embed_variables = {
+            "test_1": tf.Variable([t, t * 5]),
+            "test_2": tf.Variable([t * 10]),
+        }
+        opt_wrapper._slot_variables = {
+            "test_1": {"momentum": tf.Variable([t / 10.0, t / 2.0])},
+            "test_2": {"momentum": tf.Variable([t])},
+        }
+
+        mock_kv_store = MockKvStore({})
+        with mock.patch.object(
+            EmbeddingService, "update_embedding", mock_kv_store.update
+        ):
+            opt_wrapper._report_to_kv_store()
+
+        expected_mock_kv_store = MockKvStore({})
+        expected_mock_kv_store.update(
+            keys=["test_1-1", "test_1-5", "test_2-10"],
+            values=[t, t * 5.0, t * 10.0],
+        )
+        expected_mock_kv_store.update(
+            keys=[
+                "test_1-momentum-1",
+                "test_1-momentum-5",
+                "test_2-momentum-10",
+            ],
+            values=[t / 10.0, t / 2.0, t],
+        )
+        for k, ids in zip(["test_1", "test_2"], ids_list):
+            for id in ids:
+                key = Embedding.get_key([k, id])
+                v, _ = mock_kv_store.lookup([key])
+                expected_v, _ = expected_mock_kv_store.lookup([key])
+                self.assertTrue((v[0] == expected_v[0]).all())
+
 
 if __name__ == "__main__":
     unittest.main()