Repository: sql-machine-learning/elasticdl

Issue:
[PS-1] RPC service `push_gradient` async-SGD implementation
Async apply gradients to non-embedding parameters. Updating embedding parameters will be supported in issue #1306 .


Discussion:

Resolving patch:
diff --git a/elasticdl/python/ps/parameters.py b/elasticdl/python/ps/parameters.py
--- a/elasticdl/python/ps/parameters.py
+++ b/elasticdl/python/ps/parameters.py
@@ -23,6 +23,9 @@ def __init__(self):
         self.non_embedding_params = {}
         self.embedding_params = {}
 
+    def get_non_embedding_param(self, name, default_value=None):
+        return self.non_embedding_params.get(name, default_value)
+
     def get_embedding_param(self, name, indices):
         if name not in self.embedding_params:
             raise ValueError(
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.log_utils import default_logger as logger
 from elasticdl.python.common.tensor import Tensor, serialize_tensor
 
 
@@ -23,7 +24,7 @@ def __init__(
         self._optimizer = optimizer
         self._lr_staleness_modulation = lr_staleness_modulation
         self._use_async = use_async
-        self._version = 0
+        self._version_lock = threading.Lock()
         self._lock = threading.Lock()
 
     def pull_variable(self, request, _):
@@ -69,5 +70,29 @@ def push_model(self, request, _):
         return empty_pb2.Empty()
 
     def push_gradient(self, request, _):
-        # TODO: implement this RPC service
+        if self._use_async:
+            grad_vars = []
+            for pb in request.gradients:
+                tensor = Tensor.from_tensor_pb(pb)
+                var = self._parameters.get_non_embedding_param(tensor.name)
+                if var is None:
+                    logger.warning(
+                        "Gradients with invalid name %s" % tensor.name
+                    )
+                    continue
+                grad = tensor.to_tf_tensor()
+                grad_vars.append((grad, var))
+
+            self._optimizer.apply_gradients(grad_vars)
+            with self._version_lock:
+                self._parameters.version += 1
+
+            res = elasticdl_pb2.PushGradientResponse()
+            res.accepted = True
+            res.model_version = self._parameters.version
+            return res
+
+        raise NotImplementedError(
+            "Updating parameters synchronously is not implemented."
+        )
         return elasticdl_pb2.PushGradientResponse()


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
@@ -3,6 +3,7 @@
 
 import grpc
 import numpy as np
+import tensorflow as tf
 from google.protobuf import empty_pb2
 
 from elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc
@@ -78,6 +79,8 @@ def create_server_and_stub(
         self._server = pserver.server
         self._stub = elasticdl_pb2_grpc.PserverStub(self._channel)
 
+        self._lr = 0.1
+
     def create_default_server_and_stub(self):
         grads_to_wait = 8
         lr_staleness_modulation = False
@@ -97,15 +100,6 @@ def get_embedding_vectors(self, name, ids):
         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.PushGradientRequest()
-        res = self._stub.push_gradient(req)
-        self.assertEqual(res, elasticdl_pb2.PushGradientResponse())
-
     def testPushModel(self):
         self.create_default_server_and_stub()
         param0 = {
@@ -232,3 +226,66 @@ def testPullEmbeddingVector(self):
 
         vectors = self.get_embedding_vectors("layer_a", [])
         self.assertEqual(vectors, None)
+
+    def testAsyncPushGradient(self):
+        self.create_default_server_and_stub()
+
+        var_name = ["test_1", "test_2"]
+        var_value = [
+            np.array([10.0, 20.0, 30.0], np.float32),
+            np.array([20.0, 40.0, 60.0], np.float32),
+        ]
+        grad_value = [
+            np.array([1.0, 2.0, 3.0], np.float32),
+            np.array([2.0, 4.0, 6.0], np.float32),
+        ]
+        non_embedding_params = self._parameters.non_embedding_params
+        for i in range(2):
+            non_embedding_params[var_name[i]] = tf.Variable(var_value[i])
+
+        # Ignore grad with wrong name, still return accepted and add version
+        req = elasticdl_pb2.PushGradientRequest()
+        emplace_tensor_pb_from_ndarray(
+            req.gradients, grad_value[0], name="wrong_name"
+        )
+        res = self._stub.push_gradient(req)
+        self.assertEqual(res.accepted, True)
+        self.assertEqual(res.model_version, 1)
+        for idx, name in enumerate(var_name):
+            expected_value = var_value[idx]
+            self.assertTrue(
+                np.allclose(expected_value, non_embedding_params[name].numpy())
+            )
+
+        # grads with right name should be applied to variables
+        req = elasticdl_pb2.PushGradientRequest()
+        for g, name in zip(grad_value, var_name):
+            emplace_tensor_pb_from_ndarray(req.gradients, g, name=name)
+        res = self._stub.push_gradient(req)
+        self.assertEqual(res.accepted, True)
+        self.assertEqual(res.model_version, 2)
+        for idx, name in enumerate(var_name):
+            expected_value = var_value[idx] - self._lr * grad_value[idx]
+            self.assertTrue(
+                np.allclose(expected_value, non_embedding_params[name].numpy())
+            )
+
+        # Test applying gradients with same name
+        req = elasticdl_pb2.PushGradientRequest()
+        for g in grad_value:
+            emplace_tensor_pb_from_ndarray(req.gradients, g, name=var_name[0])
+        res = self._stub.push_gradient(req)
+        self.assertEqual(res.accepted, True)
+        self.assertEqual(res.model_version, 3)
+        expected_values = [
+            var_value[0]
+            - self._lr * grad_value[0] * 2
+            - self._lr * grad_value[1],
+            var_value[1] - self._lr * grad_value[1],
+        ]
+        for expected_value, name in zip(expected_values, var_name):
+            self.assertTrue(
+                np.allclose(expected_value, non_embedding_params[name].numpy())
+            )
+
+        # TODO: test applying gradients with embedding params