We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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 .

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `elasticdl/python/ps/parameters.py`
Content:
```
1 import tensorflow as tf
2 
3 from elasticdl.python.common.tensor import tensor_pb_to_ndarray
4 from elasticdl.python.ps.embedding_table import create_embedding_table
5 
6 
7 class Parameters(object):
8     """
9     There are two kinds of parameters:
10 
11     1. non-embedding parameters, or dense tensor parameters. We save it
12        in a hashmap `non-embedding_params`, the key is the parameter name,
13        the value is a tf.Variable` object.
14     2. embedding parameters, or row-sparse parameters. We save it in a
15        hashmap `embedding_params`, the key is the embedding layer name,
16        the value is an `EmbeddingTable` object.
17 
18     """
19 
20     def __init__(self):
21         self.version = 0
22         self.init_status = False
23         self.non_embedding_params = {}
24         self.embedding_params = {}
25 
26     def get_embedding_param(self, name, indices):
27         if name not in self.embedding_params:
28             raise ValueError(
29                 "Please initialize embedding param %s first!", name
30             )
31         return self.embedding_params[name].get(indices)
32 
33     def set_embedding_param(self, name, indices, values):
34         if name not in self.embedding_params:
35             raise ValueError(
36                 "Please initialize embedding param %s first!", name
37             )
38         self.embedding_params[name].set(indices, values)
39 
40     def init_from_model_pb(self, model_pb):
41         if not self.init_status:
42             tensors_pb = model_pb.param
43             embeddings_pb = model_pb.embedding_table_info
44             self._init_non_embedding_params(tensors_pb)
45             self._init_embedding_params(embeddings_pb)
46             self.version = model_pb.version
47             self.init_status = True
48 
49     def _init_non_embedding_params(self, tensors_pb):
50         for pb in tensors_pb:
51             name = pb.name
52             arr = tensor_pb_to_ndarray(pb)
53             var = tf.Variable(name=name, initial_value=arr, trainable=True)
54             self.non_embedding_params[name] = var
55 
56     def _init_embedding_params(self, embeddings_pb):
57         for pb in embeddings_pb:
58             self.embedding_params[pb.name] = create_embedding_table(pb)
59
```

Path: `elasticdl/python/ps/servicer.py`
Content:
```
1 import threading
2 
3 from google.protobuf import empty_pb2
4 
5 from elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc
6 from elasticdl.python.common.dtypes import dtype_numpy_to_tensor
7 from elasticdl.python.common.tensor import Tensor, serialize_tensor
8 
9 
10 class PserverServicer(elasticdl_pb2_grpc.PserverServicer):
11     """PS service implementation"""
12 
13     def __init__(
14         self,
15         parameters,
16         grads_to_wait,
17         optimizer,
18         lr_staleness_modulation=False,
19         use_async=False,
20     ):
21         self._parameters = parameters
22         self._grads_to_wait = grads_to_wait
23         self._optimizer = optimizer
24         self._lr_staleness_modulation = lr_staleness_modulation
25         self._use_async = use_async
26         self._version = 0
27         self._lock = threading.Lock()
28 
29     def pull_variable(self, request, _):
30         """
31         Response with all non-embedding parameters if initialized.
32         """
33         res = elasticdl_pb2.PullVariableResponse()
34         if not self._parameters.init_status:
35             res.model_init_status = False
36             return res
37 
38         # Only sync-SGD needs lock
39         # TODO: use a read-write lock to support multiple concurrent reads
40         if not self._use_async:
41             self._lock.acquire()
42         res.model.version = self._parameters.version
43         for name, var in self._parameters.non_embedding_params.items():
44             tensor = res.model.param.add()
45             tensor.name = name
46             tensor.dim.extend(var.shape.as_list())
47             var_values = var.numpy()
48             tensor.content = var_values.tobytes()
49             tensor.dtype = dtype_numpy_to_tensor(var_values.dtype)
50         if not self._use_async:
51             self._lock.release()
52         res.model_init_status = True
53         return res
54 
55     def pull_embedding_vector(self, request, _):
56         ret = elasticdl_pb2.Tensor()
57         if not request.ids:
58             return ret
59         embedding_vectors = self._parameters.get_embedding_param(
60             request.name, request.ids
61         )
62         tensor = Tensor(values=embedding_vectors)
63         serialize_tensor(tensor, ret)
64         return ret
65 
66     def push_model(self, request, _):
67         with self._lock:
68             self._parameters.init_from_model_pb(request)
69         return empty_pb2.Empty()
70 
71     def push_gradient(self, request, _):
72         # TODO: implement this RPC service
73         return elasticdl_pb2.PushGradientResponse()
74
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
         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 @@
         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 @@
         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()

{"golden_diff": "diff --git a/elasticdl/python/ps/parameters.py b/elasticdl/python/ps/parameters.py\n--- a/elasticdl/python/ps/parameters.py\n+++ b/elasticdl/python/ps/parameters.py\n@@ -23,6 +23,9 @@\n         self.non_embedding_params = {}\n         self.embedding_params = {}\n \n+    def get_non_embedding_param(self, name, default_value=None):\n+        return self.non_embedding_params.get(name, default_value)\n+\n     def get_embedding_param(self, name, indices):\n         if name not in self.embedding_params:\n             raise ValueError(\ndiff --git a/elasticdl/python/ps/servicer.py b/elasticdl/python/ps/servicer.py\n--- a/elasticdl/python/ps/servicer.py\n+++ b/elasticdl/python/ps/servicer.py\n@@ -4,6 +4,7 @@\n \n from elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc\n from elasticdl.python.common.dtypes import dtype_numpy_to_tensor\n+from elasticdl.python.common.log_utils import default_logger as logger\n from elasticdl.python.common.tensor import Tensor, serialize_tensor\n \n \n@@ -23,7 +24,7 @@\n         self._optimizer = optimizer\n         self._lr_staleness_modulation = lr_staleness_modulation\n         self._use_async = use_async\n-        self._version = 0\n+        self._version_lock = threading.Lock()\n         self._lock = threading.Lock()\n \n     def pull_variable(self, request, _):\n@@ -69,5 +70,29 @@\n         return empty_pb2.Empty()\n \n     def push_gradient(self, request, _):\n-        # TODO: implement this RPC service\n+        if self._use_async:\n+            grad_vars = []\n+            for pb in request.gradients:\n+                tensor = Tensor.from_tensor_pb(pb)\n+                var = self._parameters.get_non_embedding_param(tensor.name)\n+                if var is None:\n+                    logger.warning(\n+                        \"Gradients with invalid name %s\" % tensor.name\n+                    )\n+                    continue\n+                grad = tensor.to_tf_tensor()\n+                grad_vars.append((grad, var))\n+\n+            self._optimizer.apply_gradients(grad_vars)\n+            with self._version_lock:\n+                self._parameters.version += 1\n+\n+            res = elasticdl_pb2.PushGradientResponse()\n+            res.accepted = True\n+            res.model_version = self._parameters.version\n+            return res\n+\n+        raise NotImplementedError(\n+            \"Updating parameters synchronously is not implemented.\"\n+        )\n         return elasticdl_pb2.PushGradientResponse()\n", "issue": "[PS-1] RPC service `push_gradient` async-SGD implementation\nAsync apply gradients to non-embedding parameters. Updating embedding parameters will be supported in issue #1306 .\n", "before_files": [{"content": "import tensorflow as tf\n\nfrom elasticdl.python.common.tensor import tensor_pb_to_ndarray\nfrom elasticdl.python.ps.embedding_table import create_embedding_table\n\n\nclass Parameters(object):\n    \"\"\"\n    There are two kinds of parameters:\n\n    1. non-embedding parameters, or dense tensor parameters. We save it\n       in a hashmap `non-embedding_params`, the key is the parameter name,\n       the value is a tf.Variable` object.\n    2. embedding parameters, or row-sparse parameters. We save it in a\n       hashmap `embedding_params`, the key is the embedding layer name,\n       the value is an `EmbeddingTable` object.\n\n    \"\"\"\n\n    def __init__(self):\n        self.version = 0\n        self.init_status = False\n        self.non_embedding_params = {}\n        self.embedding_params = {}\n\n    def get_embedding_param(self, name, indices):\n        if name not in self.embedding_params:\n            raise ValueError(\n                \"Please initialize embedding param %s first!\", name\n            )\n        return self.embedding_params[name].get(indices)\n\n    def set_embedding_param(self, name, indices, values):\n        if name not in self.embedding_params:\n            raise ValueError(\n                \"Please initialize embedding param %s first!\", name\n            )\n        self.embedding_params[name].set(indices, values)\n\n    def init_from_model_pb(self, model_pb):\n        if not self.init_status:\n            tensors_pb = model_pb.param\n            embeddings_pb = model_pb.embedding_table_info\n            self._init_non_embedding_params(tensors_pb)\n            self._init_embedding_params(embeddings_pb)\n            self.version = model_pb.version\n            self.init_status = True\n\n    def _init_non_embedding_params(self, tensors_pb):\n        for pb in tensors_pb:\n            name = pb.name\n            arr = tensor_pb_to_ndarray(pb)\n            var = tf.Variable(name=name, initial_value=arr, trainable=True)\n            self.non_embedding_params[name] = var\n\n    def _init_embedding_params(self, embeddings_pb):\n        for pb in embeddings_pb:\n            self.embedding_params[pb.name] = create_embedding_table(pb)\n", "path": "elasticdl/python/ps/parameters.py"}, {"content": "import threading\n\nfrom google.protobuf import empty_pb2\n\nfrom elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc\nfrom elasticdl.python.common.dtypes import dtype_numpy_to_tensor\nfrom elasticdl.python.common.tensor import Tensor, serialize_tensor\n\n\nclass PserverServicer(elasticdl_pb2_grpc.PserverServicer):\n    \"\"\"PS service implementation\"\"\"\n\n    def __init__(\n        self,\n        parameters,\n        grads_to_wait,\n        optimizer,\n        lr_staleness_modulation=False,\n        use_async=False,\n    ):\n        self._parameters = parameters\n        self._grads_to_wait = grads_to_wait\n        self._optimizer = optimizer\n        self._lr_staleness_modulation = lr_staleness_modulation\n        self._use_async = use_async\n        self._version = 0\n        self._lock = threading.Lock()\n\n    def pull_variable(self, request, _):\n        \"\"\"\n        Response with all non-embedding parameters if initialized.\n        \"\"\"\n        res = elasticdl_pb2.PullVariableResponse()\n        if not self._parameters.init_status:\n            res.model_init_status = False\n            return res\n\n        # Only sync-SGD needs lock\n        # TODO: use a read-write lock to support multiple concurrent reads\n        if not self._use_async:\n            self._lock.acquire()\n        res.model.version = self._parameters.version\n        for name, var in self._parameters.non_embedding_params.items():\n            tensor = res.model.param.add()\n            tensor.name = name\n            tensor.dim.extend(var.shape.as_list())\n            var_values = var.numpy()\n            tensor.content = var_values.tobytes()\n            tensor.dtype = dtype_numpy_to_tensor(var_values.dtype)\n        if not self._use_async:\n            self._lock.release()\n        res.model_init_status = True\n        return res\n\n    def pull_embedding_vector(self, request, _):\n        ret = elasticdl_pb2.Tensor()\n        if not request.ids:\n            return ret\n        embedding_vectors = self._parameters.get_embedding_param(\n            request.name, request.ids\n        )\n        tensor = Tensor(values=embedding_vectors)\n        serialize_tensor(tensor, ret)\n        return ret\n\n    def push_model(self, request, _):\n        with self._lock:\n            self._parameters.init_from_model_pb(request)\n        return empty_pb2.Empty()\n\n    def push_gradient(self, request, _):\n        # TODO: implement this RPC service\n        return elasticdl_pb2.PushGradientResponse()\n", "path": "elasticdl/python/ps/servicer.py"}], "after_files": [{"content": "import tensorflow as tf\n\nfrom elasticdl.python.common.tensor import tensor_pb_to_ndarray\nfrom elasticdl.python.ps.embedding_table import create_embedding_table\n\n\nclass Parameters(object):\n    \"\"\"\n    There are two kinds of parameters:\n\n    1. non-embedding parameters, or dense tensor parameters. We save it\n       in a hashmap `non-embedding_params`, the key is the parameter name,\n       the value is a tf.Variable` object.\n    2. embedding parameters, or row-sparse parameters. We save it in a\n       hashmap `embedding_params`, the key is the embedding layer name,\n       the value is an `EmbeddingTable` object.\n\n    \"\"\"\n\n    def __init__(self):\n        self.version = 0\n        self.init_status = False\n        self.non_embedding_params = {}\n        self.embedding_params = {}\n\n    def get_non_embedding_param(self, name, default_value=None):\n        return self.non_embedding_params.get(name, default_value)\n\n    def get_embedding_param(self, name, indices):\n        if name not in self.embedding_params:\n            raise ValueError(\n                \"Please initialize embedding param %s first!\", name\n            )\n        return self.embedding_params[name].get(indices)\n\n    def set_embedding_param(self, name, indices, values):\n        if name not in self.embedding_params:\n            raise ValueError(\n                \"Please initialize embedding param %s first!\", name\n            )\n        self.embedding_params[name].set(indices, values)\n\n    def init_from_model_pb(self, model_pb):\n        if not self.init_status:\n            tensors_pb = model_pb.param\n            embeddings_pb = model_pb.embedding_table_info\n            self._init_non_embedding_params(tensors_pb)\n            self._init_embedding_params(embeddings_pb)\n            self.version = model_pb.version\n            self.init_status = True\n\n    def _init_non_embedding_params(self, tensors_pb):\n        for pb in tensors_pb:\n            name = pb.name\n            arr = tensor_pb_to_ndarray(pb)\n            var = tf.Variable(name=name, initial_value=arr, trainable=True)\n            self.non_embedding_params[name] = var\n\n    def _init_embedding_params(self, embeddings_pb):\n        for pb in embeddings_pb:\n            self.embedding_params[pb.name] = create_embedding_table(pb)\n", "path": "elasticdl/python/ps/parameters.py"}, {"content": "import threading\n\nfrom google.protobuf import empty_pb2\n\nfrom elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc\nfrom elasticdl.python.common.dtypes import dtype_numpy_to_tensor\nfrom elasticdl.python.common.log_utils import default_logger as logger\nfrom elasticdl.python.common.tensor import Tensor, serialize_tensor\n\n\nclass PserverServicer(elasticdl_pb2_grpc.PserverServicer):\n    \"\"\"PS service implementation\"\"\"\n\n    def __init__(\n        self,\n        parameters,\n        grads_to_wait,\n        optimizer,\n        lr_staleness_modulation=False,\n        use_async=False,\n    ):\n        self._parameters = parameters\n        self._grads_to_wait = grads_to_wait\n        self._optimizer = optimizer\n        self._lr_staleness_modulation = lr_staleness_modulation\n        self._use_async = use_async\n        self._version_lock = threading.Lock()\n        self._lock = threading.Lock()\n\n    def pull_variable(self, request, _):\n        \"\"\"\n        Response with all non-embedding parameters if initialized.\n        \"\"\"\n        res = elasticdl_pb2.PullVariableResponse()\n        if not self._parameters.init_status:\n            res.model_init_status = False\n            return res\n\n        # Only sync-SGD needs lock\n        # TODO: use a read-write lock to support multiple concurrent reads\n        if not self._use_async:\n            self._lock.acquire()\n        res.model.version = self._parameters.version\n        for name, var in self._parameters.non_embedding_params.items():\n            tensor = res.model.param.add()\n            tensor.name = name\n            tensor.dim.extend(var.shape.as_list())\n            var_values = var.numpy()\n            tensor.content = var_values.tobytes()\n            tensor.dtype = dtype_numpy_to_tensor(var_values.dtype)\n        if not self._use_async:\n            self._lock.release()\n        res.model_init_status = True\n        return res\n\n    def pull_embedding_vector(self, request, _):\n        ret = elasticdl_pb2.Tensor()\n        if not request.ids:\n            return ret\n        embedding_vectors = self._parameters.get_embedding_param(\n            request.name, request.ids\n        )\n        tensor = Tensor(values=embedding_vectors)\n        serialize_tensor(tensor, ret)\n        return ret\n\n    def push_model(self, request, _):\n        with self._lock:\n            self._parameters.init_from_model_pb(request)\n        return empty_pb2.Empty()\n\n    def push_gradient(self, request, _):\n        if self._use_async:\n            grad_vars = []\n            for pb in request.gradients:\n                tensor = Tensor.from_tensor_pb(pb)\n                var = self._parameters.get_non_embedding_param(tensor.name)\n                if var is None:\n                    logger.warning(\n                        \"Gradients with invalid name %s\" % tensor.name\n                    )\n                    continue\n                grad = tensor.to_tf_tensor()\n                grad_vars.append((grad, var))\n\n            self._optimizer.apply_gradients(grad_vars)\n            with self._version_lock:\n                self._parameters.version += 1\n\n            res = elasticdl_pb2.PushGradientResponse()\n            res.accepted = True\n            res.model_version = self._parameters.version\n            return res\n\n        raise NotImplementedError(\n            \"Updating parameters synchronously is not implemented.\"\n        )\n        return elasticdl_pb2.PushGradientResponse()\n", "path": "elasticdl/python/ps/servicer.py"}]}