We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Handle gradients of embedding parameters in PS servicer's `push_gradient`


--- 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 (
5     EmbeddingTable,
6     create_embedding_table,
7     get_slot_table_name,
8 )
9 
10 
11 class Parameters(object):
12     """
13     There are two kinds of parameters:
14 
15     1. non-embedding parameters, or dense tensor parameters. We save it
16        in a hashmap `non-embedding_params`, the key is the parameter name,
17        the value is a tf.Variable` object.
18     2. embedding parameters, or row-sparse parameters. We save it in a
19        hashmap `embedding_params`, the key is the embedding layer name,
20        the value is an `EmbeddingTable` object.
21 
22     """
23 
24     def __init__(self):
25         self.version = 0
26         self.init_status = False
27         self.non_embedding_params = {}
28         self.embedding_params = {}
29 
30     def get_non_embedding_param(self, name, default_value=None):
31         return self.non_embedding_params.get(name, default_value)
32 
33     def get_embedding_param(self, name, indices):
34         if name not in self.embedding_params:
35             raise ValueError(
36                 "Please initialize embedding param %s first!", name
37             )
38         return self.embedding_params[name].get(indices)
39 
40     def set_embedding_param(self, name, indices, values):
41         if name not in self.embedding_params:
42             raise ValueError(
43                 "Please initialize embedding param %s first!", name
44             )
45         self.embedding_params[name].set(indices, values)
46 
47     def check_grad(self, grad):
48         name = grad.name
49         if name in self.non_embedding_params:
50             param_shape = tuple(
51                 self.non_embedding_params[name].get_shape().as_list()
52             )
53             if grad.is_indexed_slices():
54                 dim0 = tf.math.reduce_max(grad.indices).numpy()
55                 dim1 = grad.values.shape[1]
56                 if dim0 > param_shape[0] or dim1 != param_shape[1]:
57                     raise ValueError(
58                         "Keras embedding param error: \
59                         the shape of gradient %s is (%d, %d), \
60                         the shape of parameter %s is (%d, %d), \
61                         which is incompatible",
62                         name,
63                         dim0,
64                         dim1,
65                         name,
66                         param_shape[0],
67                         param_shape[1],
68                     )
69             else:
70                 if grad.values.shape != param_shape:
71                     raise ValueError(
72                         "Non embedding param error: \
73                         the shape of gradient %s is %s, \
74                         the shape of parameter %s is %s, \
75                         which is incompatible",
76                         name,
77                         str(grad.values.shape),
78                         name,
79                         str(param_shape),
80                     )
81         elif name in self.embedding_params:
82             if grad.values.shape[1] != self.embedding_params[name].dim:
83                 raise ValueError(
84                     "ElasticDL embedding param error: \
85                     the shape of gradient %s is (None, %d), \
86                     the shape of parameter %s is (None, %d), \
87                     which is incompatible",
88                     name,
89                     grad.values.shape[1],
90                     name,
91                     self.embedding_params[name].dim,
92                 )
93         else:
94             raise ValueError(
95                 "Name error: Gradient %s is not in Parameters", name
96             )
97 
98     def init_from_model_pb(self, model_pb):
99         """Initializes `Parameters` with model protocol buffer.
100 
101         The `Parameters` accepts model pb and initialize only when it is
102         not initialized. Otherwise, it ignores the model pb.
103 
104         Args:
105             model_pb: The model protocol buffer used for initialization.
106 
107         Returns:
108             A bool indicates whether `Parameters` accepts this model pb or not.
109         """
110         if not self.init_status:
111             tensors_pb = model_pb.param
112             embeddings_pb = model_pb.embedding_table_info
113             self._init_non_embedding_params(tensors_pb)
114             self._init_embedding_params(embeddings_pb)
115             self.version = model_pb.version
116             self.init_status = True
117             return True
118         return False
119 
120     def _init_non_embedding_params(self, tensors_pb):
121         for pb in tensors_pb:
122             name = pb.name
123             arr = tensor_pb_to_ndarray(pb)
124             # Please note that `tf.Variable` will do something with magic.
125             # If you pass a name "somename" to a `tf.Variable`, the final
126             # variable name will be "somename:0". So the `tf.Variable.name`
127             # is meaningless, we must avoid use it in PS side.
128             var = tf.Variable(initial_value=arr, trainable=True)
129             self.non_embedding_params[name] = var
130 
131     def _init_embedding_params(self, embeddings_pb):
132         for pb in embeddings_pb:
133             self.embedding_params[pb.name] = create_embedding_table(pb)
134 
135     def has_embedding_params(self):
136         return len(self.embedding_params) > 0
137 
138     def create_slot_params(self, slot_names, init_values):
139         embed_layer_names = list(self.embedding_params.keys())
140         for layer_name in embed_layer_names:
141             for slot_name in slot_names:
142                 key = get_slot_table_name(layer_name, slot_name)
143                 if key in self.embedding_params:
144                     raise ValueError(
145                         "An embedding layer has unexpected name %s" % key
146                     )
147                 self.embedding_params[key] = EmbeddingTable(
148                     key,
149                     self.embedding_params[layer_name].dim,
150                     init_values[slot_name],
151                     True,
152                 )
153
```

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.tensor import (
7     Tensor,
8     emplace_tensor_pb_from_ndarray,
9     serialize_tensor,
10 )
11 from elasticdl.python.master.optimizer_wrapper import OptimizerWrapper
12 
13 
14 class PserverServicer(elasticdl_pb2_grpc.PserverServicer):
15     """PS service implementation"""
16 
17     def __init__(
18         self,
19         parameters,
20         grads_to_wait,
21         optimizer,
22         lr_staleness_modulation=False,
23         use_async=False,
24     ):
25         self._parameters = parameters
26         self._grads_to_wait = grads_to_wait
27         self._optimizer = optimizer
28         self._lr_staleness_modulation = lr_staleness_modulation
29         self._use_async = use_async
30         self._version_lock = threading.Lock()
31         self._lock = threading.Lock()
32 
33         self._grads_n = 0
34         self._grads_buffer = {}
35 
36     def pull_variable(self, request, _):
37         """
38         Response with all non-embedding parameters if initialized.
39         """
40         res = elasticdl_pb2.PullVariableResponse()
41         if not self._parameters.init_status:
42             res.model_init_status = False
43             return res
44 
45         # Only sync-SGD needs lock
46         # TODO: use a read-write lock to support multiple concurrent reads
47         if not self._use_async:
48             self._lock.acquire()
49         res.model.version = self._parameters.version
50         for name, var in self._parameters.non_embedding_params.items():
51             emplace_tensor_pb_from_ndarray(
52                 res.model.param, var.numpy(), name=name
53             )
54         if not self._use_async:
55             self._lock.release()
56         res.model_init_status = True
57         return res
58 
59     def pull_embedding_vector(self, request, _):
60         ret = elasticdl_pb2.Tensor()
61         if not request.ids:
62             return ret
63         embedding_vectors = self._parameters.get_embedding_param(
64             request.name, request.ids
65         )
66         tensor = Tensor(values=embedding_vectors)
67         serialize_tensor(tensor, ret)
68         return ret
69 
70     def push_model(self, request, _):
71         with self._lock:
72             accepted = self._parameters.init_from_model_pb(request)
73         if accepted and self._parameters.has_embedding_params():
74             self.wrap_optimizer()
75             self._parameters.create_slot_params(
76                 self._optimizer.allowed_slot_names,
77                 self._optimizer.slot_initial_value,
78             )
79         return empty_pb2.Empty()
80 
81     def push_gradient(self, request, _):
82         res = elasticdl_pb2.PushGradientResponse()
83         if self._use_async:
84             grad_vars = []
85             for pb in request.gradients:
86                 grad = Tensor.from_tensor_pb(pb)
87                 self._parameters.check_grad(grad)
88                 var = self._parameters.get_non_embedding_param(grad.name)
89                 grad = grad.to_tf_tensor()
90                 grad_vars.append((grad, var))
91 
92             self._optimizer.apply_gradients(grad_vars)
93             with self._version_lock:
94                 self._parameters.version += 1
95 
96             res.accepted = True
97             res.model_version = self._parameters.version
98             return res
99         else:
100             if request.model_version != self._parameters.version:
101                 res.accepted = False
102                 res.model_version = self._parameters.version
103                 return res
104 
105             with self._lock:
106                 for pb in request.gradients:
107                     grad = Tensor.from_tensor_pb(pb)
108                     self._parameters.check_grad(grad)
109                     if grad.name in self._grads_buffer:
110                         self._grads_buffer[grad.name] = (
111                             self._grads_buffer[grad.name] + grad
112                         )
113                     else:
114                         self._grads_buffer[grad.name] = grad
115 
116                 self._grads_n += 1
117                 res.accepted = True
118 
119                 if self._grads_n == self._grads_to_wait:
120                     grad_vars = []
121                     for name, grad in self._grads_buffer.items():
122                         # Dense gradients are averaged,
123                         # while sparse gradients are summed
124                         if not grad.is_indexed_slices():
125                             grad.values = grad.values / self._grads_to_wait
126                         var = self._parameters.get_non_embedding_param(name)
127                         grad_vars.append((grad.to_tf_tensor(), var))
128 
129                     self._optimizer.apply_gradients(grad_vars)
130                     self._grads_n = 0
131                     self._grads_buffer.clear()
132                     self._parameters.version += 1
133 
134                 res.model_version = self._parameters.version
135                 return res
136 
137     def wrap_optimizer(self):
138         # TODO(yunjian.lmh): refine these arguments when we don't need
139         # to support using Redis as distributed KV storage.
140         embedding_dims = None
141         embedding_service_endpoint = None
142 
143         def lookup_embedding_func(keys):
144             embeddings = []
145             for key in keys:
146                 arrs = key.split("-")
147                 layer_name = "-".join(arrs[:-1])
148                 id = int(arrs[-1])
149                 embedding = self._parameters.get_embedding_param(
150                     layer_name, [id]
151                 )
152                 embeddings.append(embedding.flatten())
153             return embeddings, []
154 
155         def update_embedding_func(keys, values):
156             for key, value in zip(keys, values):
157                 arrs = key.split("-")
158                 layer_name = "-".join(arrs[:-1])
159                 id = int(arrs[-1])
160                 self._params.set_embedding_param(layer_name, [id], [value])
161 
162         self._optimizer = OptimizerWrapper(
163             self._optimizer,
164             embedding_service_endpoint,
165             embedding_dims,
166             self._use_async,
167             lookup_embedding_func,
168             update_embedding_func,
169         )
170
```
--- 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
@@ -55,10 +55,10 @@
                 dim1 = grad.values.shape[1]
                 if dim0 > param_shape[0] or dim1 != param_shape[1]:
                     raise ValueError(
-                        "Keras embedding param error: \
-                        the shape of gradient %s is (%d, %d), \
-                        the shape of parameter %s is (%d, %d), \
-                        which is incompatible",
+                        "Keras embedding param error: "
+                        "the shape of gradient %s is (%d, %d), "
+                        "the shape of parameter %s is (%d, %d), "
+                        "which is incompatible",
                         name,
                         dim0,
                         dim1,
@@ -69,10 +69,10 @@
             else:
                 if grad.values.shape != param_shape:
                     raise ValueError(
-                        "Non embedding param error: \
-                        the shape of gradient %s is %s, \
-                        the shape of parameter %s is %s, \
-                        which is incompatible",
+                        "Non embedding param error: "
+                        "the shape of gradient %s is %s, "
+                        "the shape of parameter %s is %s, "
+                        "which is incompatible",
                         name,
                         str(grad.values.shape),
                         name,
@@ -81,10 +81,10 @@
         elif name in self.embedding_params:
             if grad.values.shape[1] != self.embedding_params[name].dim:
                 raise ValueError(
-                    "ElasticDL embedding param error: \
-                    the shape of gradient %s is (None, %d), \
-                    the shape of parameter %s is (None, %d), \
-                    which is incompatible",
+                    "ElasticDL embedding param error: "
+                    "the shape of gradient %s is (None, %d), "
+                    "the shape of parameter %s is (None, %d), "
+                    "which is incompatible",
                     name,
                     grad.values.shape[1],
                     name,
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
@@ -85,9 +85,13 @@
             for pb in request.gradients:
                 grad = Tensor.from_tensor_pb(pb)
                 self._parameters.check_grad(grad)
-                var = self._parameters.get_non_embedding_param(grad.name)
+                name = grad.name
+                var = self._parameters.get_non_embedding_param(name)
                 grad = grad.to_tf_tensor()
-                grad_vars.append((grad, var))
+                if var is None:
+                    grad_vars.append((grad, name))
+                else:
+                    grad_vars.append((grad, var))
 
             self._optimizer.apply_gradients(grad_vars)
             with self._version_lock:
@@ -124,7 +128,11 @@
                         if not grad.is_indexed_slices():
                             grad.values = grad.values / self._grads_to_wait
                         var = self._parameters.get_non_embedding_param(name)
-                        grad_vars.append((grad.to_tf_tensor(), var))
+                        grad = grad.to_tf_tensor()
+                        if var is None:
+                            grad_vars.append((grad, name))
+                        else:
+                            grad_vars.append((grad, var))
 
                     self._optimizer.apply_gradients(grad_vars)
                     self._grads_n = 0
@@ -137,7 +145,9 @@
     def wrap_optimizer(self):
         # TODO(yunjian.lmh): refine these arguments when we don't need
         # to support using Redis as distributed KV storage.
-        embedding_dims = None
+        embedding_dims = {}
+        for table in self._parameters.embedding_params.values():
+            embedding_dims[table.name] = table.dim
         embedding_service_endpoint = None
 
         def lookup_embedding_func(keys):
@@ -157,7 +167,7 @@
                 arrs = key.split("-")
                 layer_name = "-".join(arrs[:-1])
                 id = int(arrs[-1])
-                self._params.set_embedding_param(layer_name, [id], [value])
+                self._parameters.set_embedding_param(layer_name, [id], [value])
 
         self._optimizer = OptimizerWrapper(
             self._optimizer,

{"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@@ -55,10 +55,10 @@\n                 dim1 = grad.values.shape[1]\n                 if dim0 > param_shape[0] or dim1 != param_shape[1]:\n                     raise ValueError(\n-                        \"Keras embedding param error: \\\n-                        the shape of gradient %s is (%d, %d), \\\n-                        the shape of parameter %s is (%d, %d), \\\n-                        which is incompatible\",\n+                        \"Keras embedding param error: \"\n+                        \"the shape of gradient %s is (%d, %d), \"\n+                        \"the shape of parameter %s is (%d, %d), \"\n+                        \"which is incompatible\",\n                         name,\n                         dim0,\n                         dim1,\n@@ -69,10 +69,10 @@\n             else:\n                 if grad.values.shape != param_shape:\n                     raise ValueError(\n-                        \"Non embedding param error: \\\n-                        the shape of gradient %s is %s, \\\n-                        the shape of parameter %s is %s, \\\n-                        which is incompatible\",\n+                        \"Non embedding param error: \"\n+                        \"the shape of gradient %s is %s, \"\n+                        \"the shape of parameter %s is %s, \"\n+                        \"which is incompatible\",\n                         name,\n                         str(grad.values.shape),\n                         name,\n@@ -81,10 +81,10 @@\n         elif name in self.embedding_params:\n             if grad.values.shape[1] != self.embedding_params[name].dim:\n                 raise ValueError(\n-                    \"ElasticDL embedding param error: \\\n-                    the shape of gradient %s is (None, %d), \\\n-                    the shape of parameter %s is (None, %d), \\\n-                    which is incompatible\",\n+                    \"ElasticDL embedding param error: \"\n+                    \"the shape of gradient %s is (None, %d), \"\n+                    \"the shape of parameter %s is (None, %d), \"\n+                    \"which is incompatible\",\n                     name,\n                     grad.values.shape[1],\n                     name,\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@@ -85,9 +85,13 @@\n             for pb in request.gradients:\n                 grad = Tensor.from_tensor_pb(pb)\n                 self._parameters.check_grad(grad)\n-                var = self._parameters.get_non_embedding_param(grad.name)\n+                name = grad.name\n+                var = self._parameters.get_non_embedding_param(name)\n                 grad = grad.to_tf_tensor()\n-                grad_vars.append((grad, var))\n+                if var is None:\n+                    grad_vars.append((grad, name))\n+                else:\n+                    grad_vars.append((grad, var))\n \n             self._optimizer.apply_gradients(grad_vars)\n             with self._version_lock:\n@@ -124,7 +128,11 @@\n                         if not grad.is_indexed_slices():\n                             grad.values = grad.values / self._grads_to_wait\n                         var = self._parameters.get_non_embedding_param(name)\n-                        grad_vars.append((grad.to_tf_tensor(), var))\n+                        grad = grad.to_tf_tensor()\n+                        if var is None:\n+                            grad_vars.append((grad, name))\n+                        else:\n+                            grad_vars.append((grad, var))\n \n                     self._optimizer.apply_gradients(grad_vars)\n                     self._grads_n = 0\n@@ -137,7 +145,9 @@\n     def wrap_optimizer(self):\n         # TODO(yunjian.lmh): refine these arguments when we don't need\n         # to support using Redis as distributed KV storage.\n-        embedding_dims = None\n+        embedding_dims = {}\n+        for table in self._parameters.embedding_params.values():\n+            embedding_dims[table.name] = table.dim\n         embedding_service_endpoint = None\n \n         def lookup_embedding_func(keys):\n@@ -157,7 +167,7 @@\n                 arrs = key.split(\"-\")\n                 layer_name = \"-\".join(arrs[:-1])\n                 id = int(arrs[-1])\n-                self._params.set_embedding_param(layer_name, [id], [value])\n+                self._parameters.set_embedding_param(layer_name, [id], [value])\n \n         self._optimizer = OptimizerWrapper(\n             self._optimizer,\n", "issue": "Handle gradients of embedding parameters in PS servicer's `push_gradient`\n\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 (\n    EmbeddingTable,\n    create_embedding_table,\n    get_slot_table_name,\n)\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 check_grad(self, grad):\n        name = grad.name\n        if name in self.non_embedding_params:\n            param_shape = tuple(\n                self.non_embedding_params[name].get_shape().as_list()\n            )\n            if grad.is_indexed_slices():\n                dim0 = tf.math.reduce_max(grad.indices).numpy()\n                dim1 = grad.values.shape[1]\n                if dim0 > param_shape[0] or dim1 != param_shape[1]:\n                    raise ValueError(\n                        \"Keras embedding param error: \\\n                        the shape of gradient %s is (%d, %d), \\\n                        the shape of parameter %s is (%d, %d), \\\n                        which is incompatible\",\n                        name,\n                        dim0,\n                        dim1,\n                        name,\n                        param_shape[0],\n                        param_shape[1],\n                    )\n            else:\n                if grad.values.shape != param_shape:\n                    raise ValueError(\n                        \"Non embedding param error: \\\n                        the shape of gradient %s is %s, \\\n                        the shape of parameter %s is %s, \\\n                        which is incompatible\",\n                        name,\n                        str(grad.values.shape),\n                        name,\n                        str(param_shape),\n                    )\n        elif name in self.embedding_params:\n            if grad.values.shape[1] != self.embedding_params[name].dim:\n                raise ValueError(\n                    \"ElasticDL embedding param error: \\\n                    the shape of gradient %s is (None, %d), \\\n                    the shape of parameter %s is (None, %d), \\\n                    which is incompatible\",\n                    name,\n                    grad.values.shape[1],\n                    name,\n                    self.embedding_params[name].dim,\n                )\n        else:\n            raise ValueError(\n                \"Name error: Gradient %s is not in Parameters\", name\n            )\n\n    def init_from_model_pb(self, model_pb):\n        \"\"\"Initializes `Parameters` with model protocol buffer.\n\n        The `Parameters` accepts model pb and initialize only when it is\n        not initialized. Otherwise, it ignores the model pb.\n\n        Args:\n            model_pb: The model protocol buffer used for initialization.\n\n        Returns:\n            A bool indicates whether `Parameters` accepts this model pb or not.\n        \"\"\"\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            return True\n        return False\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            # Please note that `tf.Variable` will do something with magic.\n            # If you pass a name \"somename\" to a `tf.Variable`, the final\n            # variable name will be \"somename:0\". So the `tf.Variable.name`\n            # is meaningless, we must avoid use it in PS side.\n            var = tf.Variable(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\n    def has_embedding_params(self):\n        return len(self.embedding_params) > 0\n\n    def create_slot_params(self, slot_names, init_values):\n        embed_layer_names = list(self.embedding_params.keys())\n        for layer_name in embed_layer_names:\n            for slot_name in slot_names:\n                key = get_slot_table_name(layer_name, slot_name)\n                if key in self.embedding_params:\n                    raise ValueError(\n                        \"An embedding layer has unexpected name %s\" % key\n                    )\n                self.embedding_params[key] = EmbeddingTable(\n                    key,\n                    self.embedding_params[layer_name].dim,\n                    init_values[slot_name],\n                    True,\n                )\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.tensor import (\n    Tensor,\n    emplace_tensor_pb_from_ndarray,\n    serialize_tensor,\n)\nfrom elasticdl.python.master.optimizer_wrapper import OptimizerWrapper\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        self._grads_n = 0\n        self._grads_buffer = {}\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            emplace_tensor_pb_from_ndarray(\n                res.model.param, var.numpy(), name=name\n            )\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            accepted = self._parameters.init_from_model_pb(request)\n        if accepted and self._parameters.has_embedding_params():\n            self.wrap_optimizer()\n            self._parameters.create_slot_params(\n                self._optimizer.allowed_slot_names,\n                self._optimizer.slot_initial_value,\n            )\n        return empty_pb2.Empty()\n\n    def push_gradient(self, request, _):\n        res = elasticdl_pb2.PushGradientResponse()\n        if self._use_async:\n            grad_vars = []\n            for pb in request.gradients:\n                grad = Tensor.from_tensor_pb(pb)\n                self._parameters.check_grad(grad)\n                var = self._parameters.get_non_embedding_param(grad.name)\n                grad = grad.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.accepted = True\n            res.model_version = self._parameters.version\n            return res\n        else:\n            if request.model_version != self._parameters.version:\n                res.accepted = False\n                res.model_version = self._parameters.version\n                return res\n\n            with self._lock:\n                for pb in request.gradients:\n                    grad = Tensor.from_tensor_pb(pb)\n                    self._parameters.check_grad(grad)\n                    if grad.name in self._grads_buffer:\n                        self._grads_buffer[grad.name] = (\n                            self._grads_buffer[grad.name] + grad\n                        )\n                    else:\n                        self._grads_buffer[grad.name] = grad\n\n                self._grads_n += 1\n                res.accepted = True\n\n                if self._grads_n == self._grads_to_wait:\n                    grad_vars = []\n                    for name, grad in self._grads_buffer.items():\n                        # Dense gradients are averaged,\n                        # while sparse gradients are summed\n                        if not grad.is_indexed_slices():\n                            grad.values = grad.values / self._grads_to_wait\n                        var = self._parameters.get_non_embedding_param(name)\n                        grad_vars.append((grad.to_tf_tensor(), var))\n\n                    self._optimizer.apply_gradients(grad_vars)\n                    self._grads_n = 0\n                    self._grads_buffer.clear()\n                    self._parameters.version += 1\n\n                res.model_version = self._parameters.version\n                return res\n\n    def wrap_optimizer(self):\n        # TODO(yunjian.lmh): refine these arguments when we don't need\n        # to support using Redis as distributed KV storage.\n        embedding_dims = None\n        embedding_service_endpoint = None\n\n        def lookup_embedding_func(keys):\n            embeddings = []\n            for key in keys:\n                arrs = key.split(\"-\")\n                layer_name = \"-\".join(arrs[:-1])\n                id = int(arrs[-1])\n                embedding = self._parameters.get_embedding_param(\n                    layer_name, [id]\n                )\n                embeddings.append(embedding.flatten())\n            return embeddings, []\n\n        def update_embedding_func(keys, values):\n            for key, value in zip(keys, values):\n                arrs = key.split(\"-\")\n                layer_name = \"-\".join(arrs[:-1])\n                id = int(arrs[-1])\n                self._params.set_embedding_param(layer_name, [id], [value])\n\n        self._optimizer = OptimizerWrapper(\n            self._optimizer,\n            embedding_service_endpoint,\n            embedding_dims,\n            self._use_async,\n            lookup_embedding_func,\n            update_embedding_func,\n        )\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 (\n    EmbeddingTable,\n    create_embedding_table,\n    get_slot_table_name,\n)\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 check_grad(self, grad):\n        name = grad.name\n        if name in self.non_embedding_params:\n            param_shape = tuple(\n                self.non_embedding_params[name].get_shape().as_list()\n            )\n            if grad.is_indexed_slices():\n                dim0 = tf.math.reduce_max(grad.indices).numpy()\n                dim1 = grad.values.shape[1]\n                if dim0 > param_shape[0] or dim1 != param_shape[1]:\n                    raise ValueError(\n                        \"Keras embedding param error: \"\n                        \"the shape of gradient %s is (%d, %d), \"\n                        \"the shape of parameter %s is (%d, %d), \"\n                        \"which is incompatible\",\n                        name,\n                        dim0,\n                        dim1,\n                        name,\n                        param_shape[0],\n                        param_shape[1],\n                    )\n            else:\n                if grad.values.shape != param_shape:\n                    raise ValueError(\n                        \"Non embedding param error: \"\n                        \"the shape of gradient %s is %s, \"\n                        \"the shape of parameter %s is %s, \"\n                        \"which is incompatible\",\n                        name,\n                        str(grad.values.shape),\n                        name,\n                        str(param_shape),\n                    )\n        elif name in self.embedding_params:\n            if grad.values.shape[1] != self.embedding_params[name].dim:\n                raise ValueError(\n                    \"ElasticDL embedding param error: \"\n                    \"the shape of gradient %s is (None, %d), \"\n                    \"the shape of parameter %s is (None, %d), \"\n                    \"which is incompatible\",\n                    name,\n                    grad.values.shape[1],\n                    name,\n                    self.embedding_params[name].dim,\n                )\n        else:\n            raise ValueError(\n                \"Name error: Gradient %s is not in Parameters\", name\n            )\n\n    def init_from_model_pb(self, model_pb):\n        \"\"\"Initializes `Parameters` with model protocol buffer.\n\n        The `Parameters` accepts model pb and initialize only when it is\n        not initialized. Otherwise, it ignores the model pb.\n\n        Args:\n            model_pb: The model protocol buffer used for initialization.\n\n        Returns:\n            A bool indicates whether `Parameters` accepts this model pb or not.\n        \"\"\"\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            return True\n        return False\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            # Please note that `tf.Variable` will do something with magic.\n            # If you pass a name \"somename\" to a `tf.Variable`, the final\n            # variable name will be \"somename:0\". So the `tf.Variable.name`\n            # is meaningless, we must avoid use it in PS side.\n            var = tf.Variable(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\n    def has_embedding_params(self):\n        return len(self.embedding_params) > 0\n\n    def create_slot_params(self, slot_names, init_values):\n        embed_layer_names = list(self.embedding_params.keys())\n        for layer_name in embed_layer_names:\n            for slot_name in slot_names:\n                key = get_slot_table_name(layer_name, slot_name)\n                if key in self.embedding_params:\n                    raise ValueError(\n                        \"An embedding layer has unexpected name %s\" % key\n                    )\n                self.embedding_params[key] = EmbeddingTable(\n                    key,\n                    self.embedding_params[layer_name].dim,\n                    init_values[slot_name],\n                    True,\n                )\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.tensor import (\n    Tensor,\n    emplace_tensor_pb_from_ndarray,\n    serialize_tensor,\n)\nfrom elasticdl.python.master.optimizer_wrapper import OptimizerWrapper\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        self._grads_n = 0\n        self._grads_buffer = {}\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            emplace_tensor_pb_from_ndarray(\n                res.model.param, var.numpy(), name=name\n            )\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            accepted = self._parameters.init_from_model_pb(request)\n        if accepted and self._parameters.has_embedding_params():\n            self.wrap_optimizer()\n            self._parameters.create_slot_params(\n                self._optimizer.allowed_slot_names,\n                self._optimizer.slot_initial_value,\n            )\n        return empty_pb2.Empty()\n\n    def push_gradient(self, request, _):\n        res = elasticdl_pb2.PushGradientResponse()\n        if self._use_async:\n            grad_vars = []\n            for pb in request.gradients:\n                grad = Tensor.from_tensor_pb(pb)\n                self._parameters.check_grad(grad)\n                name = grad.name\n                var = self._parameters.get_non_embedding_param(name)\n                grad = grad.to_tf_tensor()\n                if var is None:\n                    grad_vars.append((grad, name))\n                else:\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.accepted = True\n            res.model_version = self._parameters.version\n            return res\n        else:\n            if request.model_version != self._parameters.version:\n                res.accepted = False\n                res.model_version = self._parameters.version\n                return res\n\n            with self._lock:\n                for pb in request.gradients:\n                    grad = Tensor.from_tensor_pb(pb)\n                    self._parameters.check_grad(grad)\n                    if grad.name in self._grads_buffer:\n                        self._grads_buffer[grad.name] = (\n                            self._grads_buffer[grad.name] + grad\n                        )\n                    else:\n                        self._grads_buffer[grad.name] = grad\n\n                self._grads_n += 1\n                res.accepted = True\n\n                if self._grads_n == self._grads_to_wait:\n                    grad_vars = []\n                    for name, grad in self._grads_buffer.items():\n                        # Dense gradients are averaged,\n                        # while sparse gradients are summed\n                        if not grad.is_indexed_slices():\n                            grad.values = grad.values / self._grads_to_wait\n                        var = self._parameters.get_non_embedding_param(name)\n                        grad = grad.to_tf_tensor()\n                        if var is None:\n                            grad_vars.append((grad, name))\n                        else:\n                            grad_vars.append((grad, var))\n\n                    self._optimizer.apply_gradients(grad_vars)\n                    self._grads_n = 0\n                    self._grads_buffer.clear()\n                    self._parameters.version += 1\n\n                res.model_version = self._parameters.version\n                return res\n\n    def wrap_optimizer(self):\n        # TODO(yunjian.lmh): refine these arguments when we don't need\n        # to support using Redis as distributed KV storage.\n        embedding_dims = {}\n        for table in self._parameters.embedding_params.values():\n            embedding_dims[table.name] = table.dim\n        embedding_service_endpoint = None\n\n        def lookup_embedding_func(keys):\n            embeddings = []\n            for key in keys:\n                arrs = key.split(\"-\")\n                layer_name = \"-\".join(arrs[:-1])\n                id = int(arrs[-1])\n                embedding = self._parameters.get_embedding_param(\n                    layer_name, [id]\n                )\n                embeddings.append(embedding.flatten())\n            return embeddings, []\n\n        def update_embedding_func(keys, values):\n            for key, value in zip(keys, values):\n                arrs = key.split(\"-\")\n                layer_name = \"-\".join(arrs[:-1])\n                id = int(arrs[-1])\n                self._parameters.set_embedding_param(layer_name, [id], [value])\n\n        self._optimizer = OptimizerWrapper(\n            self._optimizer,\n            embedding_service_endpoint,\n            embedding_dims,\n            self._use_async,\n            lookup_embedding_func,\n            update_embedding_func,\n        )\n", "path": "elasticdl/python/ps/servicer.py"}]}