We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
May need to use lock in async Optimizer Wrapper
Currently, `OptimizerWrapper` is not thread-safe for embedding table. However, in async mode, many threads call `OptimizerWrapper.apply_gradients` without using lock. 

After implementing PS, we should discuss whether we need to add a lock here.  

--- 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/optimizer_wrapper.py`
Content:
```
1 """Optimizer Wrapper for ElasticDL"""
2 
3 import threading
4 
5 import tensorflow as tf
6 from tensorflow.keras.optimizers import (
7     SGD,
8     Adadelta,
9     Adagrad,
10     Adam,
11     Adamax,
12     Ftrl,
13     Nadam,
14     RMSprop,
15 )
16 
17 from elasticdl.python.common.log_utils import default_logger as logger
18 from elasticdl.python.ps.embedding_table import get_slot_table_name
19 
20 
21 def _get_embedding_layer_name_from_var(var):
22     """Get name for ElasticDL embedding layer from variable."""
23     # Assumes that for ElasticDL embedding layer, variable will be a
24     # string representing its layer name
25     if isinstance(var, str):
26         return var
27     return None
28 
29 
30 # This function is taken from `tensorflow.keras.optimizers.Optimizer._var_key`.
31 # https://github.com/tensorflow/tensorflow/blob/71d73e56a2e66e4a6805d967cfa48ea
32 # 594f8c54e/tensorflow/python/keras/optimizer_v2/optimizer_v2.py#L1033
33 def _var_key(var):
34     """Key for representing a primary variable, for looking up slots.
35 
36     In graph mode the name is derived from the var shared name.
37     In eager mode the name is derived from the var unique id.
38     If distribution strategy exists, get the primary variable first.
39 
40     Arguments:
41         var: the variable.
42 
43     Returns:
44         the unique name of the variable.
45     """
46 
47     # pylint: disable=protected-access
48     # Get the distributed variable if it exists.
49     if getattr(var, "_distributed_container", None) is not None:
50         var = var._distributed_container()
51     if var._in_graph_mode:
52         return var._shared_name
53     return var._unique_id
54 
55 
56 class OptimizerWrapper(object):
57     """ ElasticDL optimizer wrapper.
58 
59     If model does not use ElasticDL embedding layer, `OptimizerWrapper`
60     does nothing but calls `apply_gradients` function of TensorFlow optimizer.
61     Otherwise, `OptimizerWrapper` looks up embedding vectors and slot values
62     from external kv store before updating variables, and updates embedding
63     vectors and slot values in kv store after updating variables.
64     """
65 
66     def __init__(
67         self,
68         opt,
69         use_async=False,
70         lookup_embedding_func=None,
71         update_embedding_func=None,
72     ):
73         """
74         Note:
75             We need to support Redis and ElasticDL parameter server at the
76             same time. If `lookup_embedding_func`/`update_embedding_func`
77             is not None, use parameter server to lookup/update embedding.
78             Otherwise use Redis.
79 
80         Arguments:
81             opt: A TensorFlow optimizer instance.
82             kv_store_endpoint: The endpoint to kv store.
83             embedding_dims: A python dictionary of
84                 {layer name: `embedding_dim`} where layer name is the
85                 name of ElasticDL embedding layer and `embedding_dim`
86                 is the output dimension of corresponding embedding layer.
87             use_async: A python bool. True if using asynchronous updates. When
88                 using asynchronoues updates, `OptimizerWrapper` is thread-safe
89                 for non-embedding variables and is not thread-safe for
90                 embedding table.
91             lookup_embedding_func: The function to lookup embeddings. The
92                 argument of this function is a list of keys.
93             update_embedding_func: The function to update embeddings. The
94                 arguments of this function is a key list and a value list.
95         """
96         self._opt = opt
97         self._use_async = use_async
98         self._lookup_embedding_func = lookup_embedding_func
99         self._update_embedding_func = update_embedding_func
100         self._slot_initial_value = {}
101 
102         self._opt_weights_delete_lock = threading.Lock()
103         self._tls = threading.local()
104         self._init_thread_local()
105 
106         # "-" in slot name is not supported
107         if isinstance(opt, SGD):
108             self._allowed_slot_names = []
109             if opt._momentum:
110                 self._allowed_slot_names.append("momentum")
111 
112         elif isinstance(opt, (Adam, Adamax, Nadam)):
113             self._allowed_slot_names = ["m", "v"]
114             if isinstance(opt, Adam) and self._opt.amsgrad:
115                 self._allowed_slot_names.append("vhat")
116 
117         elif isinstance(opt, Adadelta):
118             self._allowed_slot_names = ["accum_grad", "accum_var"]
119 
120         elif isinstance(opt, (Adagrad, Ftrl)):
121             self._allowed_slot_names = ["accumulator"]
122             if isinstance(opt, Ftrl):
123                 self._allowed_slot_names.append("linear")
124             accumu_init = opt._initial_accumulator_value
125             self._slot_initial_value["accumulator"] = accumu_init
126 
127         elif isinstance(opt, RMSprop):
128             self._allowed_slot_names = ["rms"]
129             if self._opt._momentum:
130                 self._allowed_slot_names.append("momentum")
131             if self._opt.centered:
132                 self._allowed_slot_names.append("mg")
133 
134         else:
135             raise NotImplementedError(
136                 "Optimizer %s is not supported in ElasticDL." % type(opt)
137             )
138 
139         for slot in self._allowed_slot_names:
140             self._slot_initial_value.setdefault(slot, 0.0)
141 
142     def _init_thread_local(self):
143         self._tls._unique_ids_all_layers = {}
144         self._tls._embed_variables = {}
145         self._tls._slot_variables = {}
146 
147     def apply_gradients(self, grads_and_vars):
148         """Update variable values.
149 
150         Args:
151             grads_and_vars: A list of (gradient, variable) pairs. If the
152                 variable is from ElasticDL embedding layer, it should be a
153                 ElasticDL `Tensor` object. Otherwise it is a TensorFlow
154                 variable.
155         """
156         # TODO (#1255): Discuss whether `OptimizerWrapper` needs a lock after
157         # implementing PS.
158         self._init_thread_local()
159 
160         grads_and_vars_new = []
161         for grad, var in grads_and_vars:
162             # If var is a string, this grad var pair is for ElasticDL embedding
163             if isinstance(var, str):
164                 grads_and_vars_new.append(
165                     self._get_embedding_var_and_grad(grad, var)
166                 )
167             else:
168                 grads_and_vars_new.append((grad, var))
169 
170         self._opt.apply_gradients(grads_and_vars_new)
171         self._update_embedding_param()
172         self._delete_variables()
173 
174     def _get_embedding_var_and_grad(self, grad, layer_name):
175         unique_ids, indices = tf.unique(grad.indices)
176         unique_ids = unique_ids.numpy()
177         if layer_name in self._tls._unique_ids_all_layers:
178             # TODO: support grads_and_vars with duplicated layer name
179             logger.warning(
180                 "grads_and_vars has duplicated layer name %s." % layer_name
181             )
182         self._tls._unique_ids_all_layers[layer_name] = unique_ids
183         new_grad = tf.IndexedSlices(values=grad.values, indices=indices)
184 
185         embed_value = self._lookup_embedding_func(layer_name, unique_ids)
186         embed_var = self._create_embedding_variable(layer_name, embed_value)
187         self._get_slot_and_set_to_optimizer(layer_name)
188         return new_grad, embed_var
189 
190     def _create_embedding_variable(self, name, initial_value):
191         """Creates a TensorFlow variable using given initial value.
192 
193         Note that this function saves the created variable to
194         `self._tls._embed_variables`.
195         """
196         embed_var = tf.Variable(
197             initial_value,
198             name=name + str(threading.get_ident()),
199             shape=initial_value.shape,
200             dtype=tf.float32,
201             trainable=False,
202         )
203         self._tls._embed_variables[name] = embed_var
204         return embed_var
205 
206     def _get_slot_and_set_to_optimizer(self, layer_name):
207         """Looks up slot value and set it to TensorFlow optimizer."""
208         for slot_name in self._allowed_slot_names:
209             param_name = get_slot_table_name(layer_name, slot_name)
210             indices = self._tls._unique_ids_all_layers[layer_name]
211             slot_value = self._lookup_embedding_func(param_name, indices)
212             # self._create_slot_variable creates a slot variable in tf
213             # optimizer and set slot_value to it.
214             self._create_slot_variable(layer_name, slot_name, slot_value)
215 
216     def _get_slot_variable(self, layer_name, slot_name):
217         """Get the variable for specified slot."""
218         return self._tls._slot_variables.get(layer_name, {}).get(
219             slot_name, None
220         )
221 
222     def _get_embedding_variable(self, layer_name):
223         """Get the variable for the specified ElasticDL embedding layer."""
224         return self._tls._embed_variables.get(layer_name, None)
225 
226     def _create_slot_variable(self, layer_name, slot_name, initial_value):
227         """Creates a slot variable in TensorFlow optimizer using given
228         value.
229         """
230         embed_var = self._get_embedding_variable(layer_name)
231         if embed_var is None:
232             raise RuntimeError(
233                 "Embedding variable for layer %s should be already created."
234                 % (layer_name)
235             )
236         slot_var = self._create_slot_variable_in_optimizer(
237             embed_var, slot_name, initial_value.shape, initial_value
238         )
239         slot_variables_dict = self._tls._slot_variables.setdefault(
240             layer_name, {}
241         )
242         slot_variables_dict[slot_name] = slot_var
243         return slot_var
244 
245     # This is a function modified from TensorFlow optimizers.
246     # https://github.com/tensorflow/tensorflow/blob/
247     # 69b1feac62276edcc509ac88af229c6236e645fe/tensorflow/python
248     # /keras/optimizer_v2/optimizer_v2.py#L567
249     def _create_slot_variable_in_optimizer(
250         self, embed_var, slot_name, shape, initial_value
251     ):
252         """Create variable for a slot and save it in TensorFlow optimizer."""
253         if slot_name not in self._opt._slot_names:
254             self._opt._slot_names.append(slot_name)
255         var_key = _var_key(embed_var)
256         slot_dict = self._opt._slots.setdefault(var_key, {})
257         slot_var = slot_dict.get(slot_name, None)
258         if slot_var is None:
259             slot_var_name = "%s/%s" % (embed_var._shared_name, slot_name)
260             slot_var = self._opt.add_weight(
261                 name=slot_var_name,
262                 shape=shape,
263                 dtype=embed_var.dtype,
264                 initializer=initial_value,
265                 trainable=False,
266             )
267             slot_dict[slot_name] = slot_var
268             self._opt._weights.append(slot_var)
269             return slot_var
270         else:
271             raise RuntimeError(
272                 "Variable with var_key %s and slot_name %s is not expected to "
273                 "be in self._opt." % (var_key, slot_name)
274             )
275 
276     def _update_embedding_param(self):
277         """Report updated embedding vectors and slots to kv store."""
278         for layer, ids in self._tls._unique_ids_all_layers.items():
279             value = self._get_embedding_variable(layer).numpy()
280             self._update_embedding_func(layer, ids, value)
281 
282             for slot in self._allowed_slot_names:
283                 value = self._get_slot_variable(layer, slot).numpy()
284                 slot_table_name = get_slot_table_name(layer, slot)
285                 self._update_embedding_func(slot_table_name, ids, value)
286 
287     def _delete_variables(self):
288         # Slot variable access in optimizer requires corresponding embedding
289         # variable information. Delete slot variables first.
290         for layer_name, slots in self._tls._slot_variables.items():
291             embed_var = self._get_embedding_variable(layer_name)
292             embed_var_key = _var_key(embed_var)
293             del self._opt._slots[embed_var_key]
294             for _, var in slots.items():
295                 opt_weight_iter = 0
296                 with self._opt_weights_delete_lock:
297                     while opt_weight_iter < len(self._opt._weights):
298                         if var is self._opt._weights[opt_weight_iter]:
299                             self._opt._weights.pop(opt_weight_iter)
300                             break
301                         else:
302                             opt_weight_iter += 1
303         for key in list(self._tls._slot_variables.keys()):
304             del self._tls._slot_variables[key]
305 
306         # Delete variables in embed_variables.
307         for key in list(self._tls._embed_variables.keys()):
308             del self._tls._embed_variables[key]
309 
310         # Delete variables in unique_ids_all_layers.
311         for key in list(self._tls._unique_ids_all_layers.keys()):
312             del self._tls._unique_ids_all_layers[key]
313 
314     @property
315     def allowed_slot_names(self):
316         return self._allowed_slot_names
317 
318     # TODO(yunjian.lmh): Do not need to save slot_initial_value in
319     #     optimizer wrapper after we do not need to support Redis.
320     @property
321     def slot_initial_value(self):
322         return self._slot_initial_value
323
```
--- 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/optimizer_wrapper.py b/elasticdl/python/ps/optimizer_wrapper.py
--- a/elasticdl/python/ps/optimizer_wrapper.py
+++ b/elasticdl/python/ps/optimizer_wrapper.py
@@ -99,7 +99,7 @@
         self._update_embedding_func = update_embedding_func
         self._slot_initial_value = {}
 
-        self._opt_weights_delete_lock = threading.Lock()
+        self._update_gradient_lock = threading.Lock()
         self._tls = threading.local()
         self._init_thread_local()
 
@@ -153,23 +153,23 @@
                 ElasticDL `Tensor` object. Otherwise it is a TensorFlow
                 variable.
         """
-        # TODO (#1255): Discuss whether `OptimizerWrapper` needs a lock after
-        # implementing PS.
         self._init_thread_local()
 
-        grads_and_vars_new = []
-        for grad, var in grads_and_vars:
-            # If var is a string, this grad var pair is for ElasticDL embedding
-            if isinstance(var, str):
-                grads_and_vars_new.append(
-                    self._get_embedding_var_and_grad(grad, var)
-                )
-            else:
-                grads_and_vars_new.append((grad, var))
-
-        self._opt.apply_gradients(grads_and_vars_new)
-        self._update_embedding_param()
-        self._delete_variables()
+        with self._update_gradient_lock:
+            grads_and_vars_new = []
+            for grad, var in grads_and_vars:
+                # If var is a string, create the grad var pair for
+                # ElasticDL embedding
+                if isinstance(var, str):
+                    grads_and_vars_new.append(
+                        self._get_embedding_var_and_grad(grad, var)
+                    )
+                else:
+                    grads_and_vars_new.append((grad, var))
+
+            self._opt.apply_gradients(grads_and_vars_new)
+            self._update_embedding_param()
+            self._delete_variables()
 
     def _get_embedding_var_and_grad(self, grad, layer_name):
         unique_ids, indices = tf.unique(grad.indices)
@@ -293,13 +293,12 @@
             del self._opt._slots[embed_var_key]
             for _, var in slots.items():
                 opt_weight_iter = 0
-                with self._opt_weights_delete_lock:
-                    while opt_weight_iter < len(self._opt._weights):
-                        if var is self._opt._weights[opt_weight_iter]:
-                            self._opt._weights.pop(opt_weight_iter)
-                            break
-                        else:
-                            opt_weight_iter += 1
+                while opt_weight_iter < len(self._opt._weights):
+                    if var is self._opt._weights[opt_weight_iter]:
+                        self._opt._weights.pop(opt_weight_iter)
+                        break
+                    else:
+                        opt_weight_iter += 1
         for key in list(self._tls._slot_variables.keys()):
             del self._tls._slot_variables[key]

{"golden_diff": "diff --git a/elasticdl/python/ps/optimizer_wrapper.py b/elasticdl/python/ps/optimizer_wrapper.py\n--- a/elasticdl/python/ps/optimizer_wrapper.py\n+++ b/elasticdl/python/ps/optimizer_wrapper.py\n@@ -99,7 +99,7 @@\n         self._update_embedding_func = update_embedding_func\n         self._slot_initial_value = {}\n \n-        self._opt_weights_delete_lock = threading.Lock()\n+        self._update_gradient_lock = threading.Lock()\n         self._tls = threading.local()\n         self._init_thread_local()\n \n@@ -153,23 +153,23 @@\n                 ElasticDL `Tensor` object. Otherwise it is a TensorFlow\n                 variable.\n         \"\"\"\n-        # TODO (#1255): Discuss whether `OptimizerWrapper` needs a lock after\n-        # implementing PS.\n         self._init_thread_local()\n \n-        grads_and_vars_new = []\n-        for grad, var in grads_and_vars:\n-            # If var is a string, this grad var pair is for ElasticDL embedding\n-            if isinstance(var, str):\n-                grads_and_vars_new.append(\n-                    self._get_embedding_var_and_grad(grad, var)\n-                )\n-            else:\n-                grads_and_vars_new.append((grad, var))\n-\n-        self._opt.apply_gradients(grads_and_vars_new)\n-        self._update_embedding_param()\n-        self._delete_variables()\n+        with self._update_gradient_lock:\n+            grads_and_vars_new = []\n+            for grad, var in grads_and_vars:\n+                # If var is a string, create the grad var pair for\n+                # ElasticDL embedding\n+                if isinstance(var, str):\n+                    grads_and_vars_new.append(\n+                        self._get_embedding_var_and_grad(grad, var)\n+                    )\n+                else:\n+                    grads_and_vars_new.append((grad, var))\n+\n+            self._opt.apply_gradients(grads_and_vars_new)\n+            self._update_embedding_param()\n+            self._delete_variables()\n \n     def _get_embedding_var_and_grad(self, grad, layer_name):\n         unique_ids, indices = tf.unique(grad.indices)\n@@ -293,13 +293,12 @@\n             del self._opt._slots[embed_var_key]\n             for _, var in slots.items():\n                 opt_weight_iter = 0\n-                with self._opt_weights_delete_lock:\n-                    while opt_weight_iter < len(self._opt._weights):\n-                        if var is self._opt._weights[opt_weight_iter]:\n-                            self._opt._weights.pop(opt_weight_iter)\n-                            break\n-                        else:\n-                            opt_weight_iter += 1\n+                while opt_weight_iter < len(self._opt._weights):\n+                    if var is self._opt._weights[opt_weight_iter]:\n+                        self._opt._weights.pop(opt_weight_iter)\n+                        break\n+                    else:\n+                        opt_weight_iter += 1\n         for key in list(self._tls._slot_variables.keys()):\n             del self._tls._slot_variables[key]\n", "issue": "May need to use lock in async Optimizer Wrapper\nCurrently, `OptimizerWrapper` is not thread-safe for embedding table. However, in async mode, many threads call `OptimizerWrapper.apply_gradients` without using lock. \r\n\r\nAfter implementing PS, we should discuss whether we need to add a lock here.  \n", "before_files": [{"content": "\"\"\"Optimizer Wrapper for ElasticDL\"\"\"\n\nimport threading\n\nimport tensorflow as tf\nfrom tensorflow.keras.optimizers import (\n    SGD,\n    Adadelta,\n    Adagrad,\n    Adam,\n    Adamax,\n    Ftrl,\n    Nadam,\n    RMSprop,\n)\n\nfrom elasticdl.python.common.log_utils import default_logger as logger\nfrom elasticdl.python.ps.embedding_table import get_slot_table_name\n\n\ndef _get_embedding_layer_name_from_var(var):\n    \"\"\"Get name for ElasticDL embedding layer from variable.\"\"\"\n    # Assumes that for ElasticDL embedding layer, variable will be a\n    # string representing its layer name\n    if isinstance(var, str):\n        return var\n    return None\n\n\n# This function is taken from `tensorflow.keras.optimizers.Optimizer._var_key`.\n# https://github.com/tensorflow/tensorflow/blob/71d73e56a2e66e4a6805d967cfa48ea\n# 594f8c54e/tensorflow/python/keras/optimizer_v2/optimizer_v2.py#L1033\ndef _var_key(var):\n    \"\"\"Key for representing a primary variable, for looking up slots.\n\n    In graph mode the name is derived from the var shared name.\n    In eager mode the name is derived from the var unique id.\n    If distribution strategy exists, get the primary variable first.\n\n    Arguments:\n        var: the variable.\n\n    Returns:\n        the unique name of the variable.\n    \"\"\"\n\n    # pylint: disable=protected-access\n    # Get the distributed variable if it exists.\n    if getattr(var, \"_distributed_container\", None) is not None:\n        var = var._distributed_container()\n    if var._in_graph_mode:\n        return var._shared_name\n    return var._unique_id\n\n\nclass OptimizerWrapper(object):\n    \"\"\" ElasticDL optimizer wrapper.\n\n    If model does not use ElasticDL embedding layer, `OptimizerWrapper`\n    does nothing but calls `apply_gradients` function of TensorFlow optimizer.\n    Otherwise, `OptimizerWrapper` looks up embedding vectors and slot values\n    from external kv store before updating variables, and updates embedding\n    vectors and slot values in kv store after updating variables.\n    \"\"\"\n\n    def __init__(\n        self,\n        opt,\n        use_async=False,\n        lookup_embedding_func=None,\n        update_embedding_func=None,\n    ):\n        \"\"\"\n        Note:\n            We need to support Redis and ElasticDL parameter server at the\n            same time. If `lookup_embedding_func`/`update_embedding_func`\n            is not None, use parameter server to lookup/update embedding.\n            Otherwise use Redis.\n\n        Arguments:\n            opt: A TensorFlow optimizer instance.\n            kv_store_endpoint: The endpoint to kv store.\n            embedding_dims: A python dictionary of\n                {layer name: `embedding_dim`} where layer name is the\n                name of ElasticDL embedding layer and `embedding_dim`\n                is the output dimension of corresponding embedding layer.\n            use_async: A python bool. True if using asynchronous updates. When\n                using asynchronoues updates, `OptimizerWrapper` is thread-safe\n                for non-embedding variables and is not thread-safe for\n                embedding table.\n            lookup_embedding_func: The function to lookup embeddings. The\n                argument of this function is a list of keys.\n            update_embedding_func: The function to update embeddings. The\n                arguments of this function is a key list and a value list.\n        \"\"\"\n        self._opt = opt\n        self._use_async = use_async\n        self._lookup_embedding_func = lookup_embedding_func\n        self._update_embedding_func = update_embedding_func\n        self._slot_initial_value = {}\n\n        self._opt_weights_delete_lock = threading.Lock()\n        self._tls = threading.local()\n        self._init_thread_local()\n\n        # \"-\" in slot name is not supported\n        if isinstance(opt, SGD):\n            self._allowed_slot_names = []\n            if opt._momentum:\n                self._allowed_slot_names.append(\"momentum\")\n\n        elif isinstance(opt, (Adam, Adamax, Nadam)):\n            self._allowed_slot_names = [\"m\", \"v\"]\n            if isinstance(opt, Adam) and self._opt.amsgrad:\n                self._allowed_slot_names.append(\"vhat\")\n\n        elif isinstance(opt, Adadelta):\n            self._allowed_slot_names = [\"accum_grad\", \"accum_var\"]\n\n        elif isinstance(opt, (Adagrad, Ftrl)):\n            self._allowed_slot_names = [\"accumulator\"]\n            if isinstance(opt, Ftrl):\n                self._allowed_slot_names.append(\"linear\")\n            accumu_init = opt._initial_accumulator_value\n            self._slot_initial_value[\"accumulator\"] = accumu_init\n\n        elif isinstance(opt, RMSprop):\n            self._allowed_slot_names = [\"rms\"]\n            if self._opt._momentum:\n                self._allowed_slot_names.append(\"momentum\")\n            if self._opt.centered:\n                self._allowed_slot_names.append(\"mg\")\n\n        else:\n            raise NotImplementedError(\n                \"Optimizer %s is not supported in ElasticDL.\" % type(opt)\n            )\n\n        for slot in self._allowed_slot_names:\n            self._slot_initial_value.setdefault(slot, 0.0)\n\n    def _init_thread_local(self):\n        self._tls._unique_ids_all_layers = {}\n        self._tls._embed_variables = {}\n        self._tls._slot_variables = {}\n\n    def apply_gradients(self, grads_and_vars):\n        \"\"\"Update variable values.\n\n        Args:\n            grads_and_vars: A list of (gradient, variable) pairs. If the\n                variable is from ElasticDL embedding layer, it should be a\n                ElasticDL `Tensor` object. Otherwise it is a TensorFlow\n                variable.\n        \"\"\"\n        # TODO (#1255): Discuss whether `OptimizerWrapper` needs a lock after\n        # implementing PS.\n        self._init_thread_local()\n\n        grads_and_vars_new = []\n        for grad, var in grads_and_vars:\n            # If var is a string, this grad var pair is for ElasticDL embedding\n            if isinstance(var, str):\n                grads_and_vars_new.append(\n                    self._get_embedding_var_and_grad(grad, var)\n                )\n            else:\n                grads_and_vars_new.append((grad, var))\n\n        self._opt.apply_gradients(grads_and_vars_new)\n        self._update_embedding_param()\n        self._delete_variables()\n\n    def _get_embedding_var_and_grad(self, grad, layer_name):\n        unique_ids, indices = tf.unique(grad.indices)\n        unique_ids = unique_ids.numpy()\n        if layer_name in self._tls._unique_ids_all_layers:\n            # TODO: support grads_and_vars with duplicated layer name\n            logger.warning(\n                \"grads_and_vars has duplicated layer name %s.\" % layer_name\n            )\n        self._tls._unique_ids_all_layers[layer_name] = unique_ids\n        new_grad = tf.IndexedSlices(values=grad.values, indices=indices)\n\n        embed_value = self._lookup_embedding_func(layer_name, unique_ids)\n        embed_var = self._create_embedding_variable(layer_name, embed_value)\n        self._get_slot_and_set_to_optimizer(layer_name)\n        return new_grad, embed_var\n\n    def _create_embedding_variable(self, name, initial_value):\n        \"\"\"Creates a TensorFlow variable using given initial value.\n\n        Note that this function saves the created variable to\n        `self._tls._embed_variables`.\n        \"\"\"\n        embed_var = tf.Variable(\n            initial_value,\n            name=name + str(threading.get_ident()),\n            shape=initial_value.shape,\n            dtype=tf.float32,\n            trainable=False,\n        )\n        self._tls._embed_variables[name] = embed_var\n        return embed_var\n\n    def _get_slot_and_set_to_optimizer(self, layer_name):\n        \"\"\"Looks up slot value and set it to TensorFlow optimizer.\"\"\"\n        for slot_name in self._allowed_slot_names:\n            param_name = get_slot_table_name(layer_name, slot_name)\n            indices = self._tls._unique_ids_all_layers[layer_name]\n            slot_value = self._lookup_embedding_func(param_name, indices)\n            # self._create_slot_variable creates a slot variable in tf\n            # optimizer and set slot_value to it.\n            self._create_slot_variable(layer_name, slot_name, slot_value)\n\n    def _get_slot_variable(self, layer_name, slot_name):\n        \"\"\"Get the variable for specified slot.\"\"\"\n        return self._tls._slot_variables.get(layer_name, {}).get(\n            slot_name, None\n        )\n\n    def _get_embedding_variable(self, layer_name):\n        \"\"\"Get the variable for the specified ElasticDL embedding layer.\"\"\"\n        return self._tls._embed_variables.get(layer_name, None)\n\n    def _create_slot_variable(self, layer_name, slot_name, initial_value):\n        \"\"\"Creates a slot variable in TensorFlow optimizer using given\n        value.\n        \"\"\"\n        embed_var = self._get_embedding_variable(layer_name)\n        if embed_var is None:\n            raise RuntimeError(\n                \"Embedding variable for layer %s should be already created.\"\n                % (layer_name)\n            )\n        slot_var = self._create_slot_variable_in_optimizer(\n            embed_var, slot_name, initial_value.shape, initial_value\n        )\n        slot_variables_dict = self._tls._slot_variables.setdefault(\n            layer_name, {}\n        )\n        slot_variables_dict[slot_name] = slot_var\n        return slot_var\n\n    # This is a function modified from TensorFlow optimizers.\n    # https://github.com/tensorflow/tensorflow/blob/\n    # 69b1feac62276edcc509ac88af229c6236e645fe/tensorflow/python\n    # /keras/optimizer_v2/optimizer_v2.py#L567\n    def _create_slot_variable_in_optimizer(\n        self, embed_var, slot_name, shape, initial_value\n    ):\n        \"\"\"Create variable for a slot and save it in TensorFlow optimizer.\"\"\"\n        if slot_name not in self._opt._slot_names:\n            self._opt._slot_names.append(slot_name)\n        var_key = _var_key(embed_var)\n        slot_dict = self._opt._slots.setdefault(var_key, {})\n        slot_var = slot_dict.get(slot_name, None)\n        if slot_var is None:\n            slot_var_name = \"%s/%s\" % (embed_var._shared_name, slot_name)\n            slot_var = self._opt.add_weight(\n                name=slot_var_name,\n                shape=shape,\n                dtype=embed_var.dtype,\n                initializer=initial_value,\n                trainable=False,\n            )\n            slot_dict[slot_name] = slot_var\n            self._opt._weights.append(slot_var)\n            return slot_var\n        else:\n            raise RuntimeError(\n                \"Variable with var_key %s and slot_name %s is not expected to \"\n                \"be in self._opt.\" % (var_key, slot_name)\n            )\n\n    def _update_embedding_param(self):\n        \"\"\"Report updated embedding vectors and slots to kv store.\"\"\"\n        for layer, ids in self._tls._unique_ids_all_layers.items():\n            value = self._get_embedding_variable(layer).numpy()\n            self._update_embedding_func(layer, ids, value)\n\n            for slot in self._allowed_slot_names:\n                value = self._get_slot_variable(layer, slot).numpy()\n                slot_table_name = get_slot_table_name(layer, slot)\n                self._update_embedding_func(slot_table_name, ids, value)\n\n    def _delete_variables(self):\n        # Slot variable access in optimizer requires corresponding embedding\n        # variable information. Delete slot variables first.\n        for layer_name, slots in self._tls._slot_variables.items():\n            embed_var = self._get_embedding_variable(layer_name)\n            embed_var_key = _var_key(embed_var)\n            del self._opt._slots[embed_var_key]\n            for _, var in slots.items():\n                opt_weight_iter = 0\n                with self._opt_weights_delete_lock:\n                    while opt_weight_iter < len(self._opt._weights):\n                        if var is self._opt._weights[opt_weight_iter]:\n                            self._opt._weights.pop(opt_weight_iter)\n                            break\n                        else:\n                            opt_weight_iter += 1\n        for key in list(self._tls._slot_variables.keys()):\n            del self._tls._slot_variables[key]\n\n        # Delete variables in embed_variables.\n        for key in list(self._tls._embed_variables.keys()):\n            del self._tls._embed_variables[key]\n\n        # Delete variables in unique_ids_all_layers.\n        for key in list(self._tls._unique_ids_all_layers.keys()):\n            del self._tls._unique_ids_all_layers[key]\n\n    @property\n    def allowed_slot_names(self):\n        return self._allowed_slot_names\n\n    # TODO(yunjian.lmh): Do not need to save slot_initial_value in\n    #     optimizer wrapper after we do not need to support Redis.\n    @property\n    def slot_initial_value(self):\n        return self._slot_initial_value\n", "path": "elasticdl/python/ps/optimizer_wrapper.py"}], "after_files": [{"content": "\"\"\"Optimizer Wrapper for ElasticDL\"\"\"\n\nimport threading\n\nimport tensorflow as tf\nfrom tensorflow.keras.optimizers import (\n    SGD,\n    Adadelta,\n    Adagrad,\n    Adam,\n    Adamax,\n    Ftrl,\n    Nadam,\n    RMSprop,\n)\n\nfrom elasticdl.python.common.log_utils import default_logger as logger\nfrom elasticdl.python.ps.embedding_table import get_slot_table_name\n\n\ndef _get_embedding_layer_name_from_var(var):\n    \"\"\"Get name for ElasticDL embedding layer from variable.\"\"\"\n    # Assumes that for ElasticDL embedding layer, variable will be a\n    # string representing its layer name\n    if isinstance(var, str):\n        return var\n    return None\n\n\n# This function is taken from `tensorflow.keras.optimizers.Optimizer._var_key`.\n# https://github.com/tensorflow/tensorflow/blob/71d73e56a2e66e4a6805d967cfa48ea\n# 594f8c54e/tensorflow/python/keras/optimizer_v2/optimizer_v2.py#L1033\ndef _var_key(var):\n    \"\"\"Key for representing a primary variable, for looking up slots.\n\n    In graph mode the name is derived from the var shared name.\n    In eager mode the name is derived from the var unique id.\n    If distribution strategy exists, get the primary variable first.\n\n    Arguments:\n        var: the variable.\n\n    Returns:\n        the unique name of the variable.\n    \"\"\"\n\n    # pylint: disable=protected-access\n    # Get the distributed variable if it exists.\n    if getattr(var, \"_distributed_container\", None) is not None:\n        var = var._distributed_container()\n    if var._in_graph_mode:\n        return var._shared_name\n    return var._unique_id\n\n\nclass OptimizerWrapper(object):\n    \"\"\" ElasticDL optimizer wrapper.\n\n    If model does not use ElasticDL embedding layer, `OptimizerWrapper`\n    does nothing but calls `apply_gradients` function of TensorFlow optimizer.\n    Otherwise, `OptimizerWrapper` looks up embedding vectors and slot values\n    from external kv store before updating variables, and updates embedding\n    vectors and slot values in kv store after updating variables.\n    \"\"\"\n\n    def __init__(\n        self,\n        opt,\n        use_async=False,\n        lookup_embedding_func=None,\n        update_embedding_func=None,\n    ):\n        \"\"\"\n        Note:\n            We need to support Redis and ElasticDL parameter server at the\n            same time. If `lookup_embedding_func`/`update_embedding_func`\n            is not None, use parameter server to lookup/update embedding.\n            Otherwise use Redis.\n\n        Arguments:\n            opt: A TensorFlow optimizer instance.\n            kv_store_endpoint: The endpoint to kv store.\n            embedding_dims: A python dictionary of\n                {layer name: `embedding_dim`} where layer name is the\n                name of ElasticDL embedding layer and `embedding_dim`\n                is the output dimension of corresponding embedding layer.\n            use_async: A python bool. True if using asynchronous updates. When\n                using asynchronoues updates, `OptimizerWrapper` is thread-safe\n                for non-embedding variables and is not thread-safe for\n                embedding table.\n            lookup_embedding_func: The function to lookup embeddings. The\n                argument of this function is a list of keys.\n            update_embedding_func: The function to update embeddings. The\n                arguments of this function is a key list and a value list.\n        \"\"\"\n        self._opt = opt\n        self._use_async = use_async\n        self._lookup_embedding_func = lookup_embedding_func\n        self._update_embedding_func = update_embedding_func\n        self._slot_initial_value = {}\n\n        self._update_gradient_lock = threading.Lock()\n        self._tls = threading.local()\n        self._init_thread_local()\n\n        # \"-\" in slot name is not supported\n        if isinstance(opt, SGD):\n            self._allowed_slot_names = []\n            if opt._momentum:\n                self._allowed_slot_names.append(\"momentum\")\n\n        elif isinstance(opt, (Adam, Adamax, Nadam)):\n            self._allowed_slot_names = [\"m\", \"v\"]\n            if isinstance(opt, Adam) and self._opt.amsgrad:\n                self._allowed_slot_names.append(\"vhat\")\n\n        elif isinstance(opt, Adadelta):\n            self._allowed_slot_names = [\"accum_grad\", \"accum_var\"]\n\n        elif isinstance(opt, (Adagrad, Ftrl)):\n            self._allowed_slot_names = [\"accumulator\"]\n            if isinstance(opt, Ftrl):\n                self._allowed_slot_names.append(\"linear\")\n            accumu_init = opt._initial_accumulator_value\n            self._slot_initial_value[\"accumulator\"] = accumu_init\n\n        elif isinstance(opt, RMSprop):\n            self._allowed_slot_names = [\"rms\"]\n            if self._opt._momentum:\n                self._allowed_slot_names.append(\"momentum\")\n            if self._opt.centered:\n                self._allowed_slot_names.append(\"mg\")\n\n        else:\n            raise NotImplementedError(\n                \"Optimizer %s is not supported in ElasticDL.\" % type(opt)\n            )\n\n        for slot in self._allowed_slot_names:\n            self._slot_initial_value.setdefault(slot, 0.0)\n\n    def _init_thread_local(self):\n        self._tls._unique_ids_all_layers = {}\n        self._tls._embed_variables = {}\n        self._tls._slot_variables = {}\n\n    def apply_gradients(self, grads_and_vars):\n        \"\"\"Update variable values.\n\n        Args:\n            grads_and_vars: A list of (gradient, variable) pairs. If the\n                variable is from ElasticDL embedding layer, it should be a\n                ElasticDL `Tensor` object. Otherwise it is a TensorFlow\n                variable.\n        \"\"\"\n        self._init_thread_local()\n\n        with self._update_gradient_lock:\n            grads_and_vars_new = []\n            for grad, var in grads_and_vars:\n                # If var is a string, create the grad var pair for\n                # ElasticDL embedding\n                if isinstance(var, str):\n                    grads_and_vars_new.append(\n                        self._get_embedding_var_and_grad(grad, var)\n                    )\n                else:\n                    grads_and_vars_new.append((grad, var))\n\n            self._opt.apply_gradients(grads_and_vars_new)\n            self._update_embedding_param()\n            self._delete_variables()\n\n    def _get_embedding_var_and_grad(self, grad, layer_name):\n        unique_ids, indices = tf.unique(grad.indices)\n        unique_ids = unique_ids.numpy()\n        if layer_name in self._tls._unique_ids_all_layers:\n            # TODO: support grads_and_vars with duplicated layer name\n            logger.warning(\n                \"grads_and_vars has duplicated layer name %s.\" % layer_name\n            )\n        self._tls._unique_ids_all_layers[layer_name] = unique_ids\n        new_grad = tf.IndexedSlices(values=grad.values, indices=indices)\n\n        embed_value = self._lookup_embedding_func(layer_name, unique_ids)\n        embed_var = self._create_embedding_variable(layer_name, embed_value)\n        self._get_slot_and_set_to_optimizer(layer_name)\n        return new_grad, embed_var\n\n    def _create_embedding_variable(self, name, initial_value):\n        \"\"\"Creates a TensorFlow variable using given initial value.\n\n        Note that this function saves the created variable to\n        `self._tls._embed_variables`.\n        \"\"\"\n        embed_var = tf.Variable(\n            initial_value,\n            name=name + str(threading.get_ident()),\n            shape=initial_value.shape,\n            dtype=tf.float32,\n            trainable=False,\n        )\n        self._tls._embed_variables[name] = embed_var\n        return embed_var\n\n    def _get_slot_and_set_to_optimizer(self, layer_name):\n        \"\"\"Looks up slot value and set it to TensorFlow optimizer.\"\"\"\n        for slot_name in self._allowed_slot_names:\n            param_name = get_slot_table_name(layer_name, slot_name)\n            indices = self._tls._unique_ids_all_layers[layer_name]\n            slot_value = self._lookup_embedding_func(param_name, indices)\n            # self._create_slot_variable creates a slot variable in tf\n            # optimizer and set slot_value to it.\n            self._create_slot_variable(layer_name, slot_name, slot_value)\n\n    def _get_slot_variable(self, layer_name, slot_name):\n        \"\"\"Get the variable for specified slot.\"\"\"\n        return self._tls._slot_variables.get(layer_name, {}).get(\n            slot_name, None\n        )\n\n    def _get_embedding_variable(self, layer_name):\n        \"\"\"Get the variable for the specified ElasticDL embedding layer.\"\"\"\n        return self._tls._embed_variables.get(layer_name, None)\n\n    def _create_slot_variable(self, layer_name, slot_name, initial_value):\n        \"\"\"Creates a slot variable in TensorFlow optimizer using given\n        value.\n        \"\"\"\n        embed_var = self._get_embedding_variable(layer_name)\n        if embed_var is None:\n            raise RuntimeError(\n                \"Embedding variable for layer %s should be already created.\"\n                % (layer_name)\n            )\n        slot_var = self._create_slot_variable_in_optimizer(\n            embed_var, slot_name, initial_value.shape, initial_value\n        )\n        slot_variables_dict = self._tls._slot_variables.setdefault(\n            layer_name, {}\n        )\n        slot_variables_dict[slot_name] = slot_var\n        return slot_var\n\n    # This is a function modified from TensorFlow optimizers.\n    # https://github.com/tensorflow/tensorflow/blob/\n    # 69b1feac62276edcc509ac88af229c6236e645fe/tensorflow/python\n    # /keras/optimizer_v2/optimizer_v2.py#L567\n    def _create_slot_variable_in_optimizer(\n        self, embed_var, slot_name, shape, initial_value\n    ):\n        \"\"\"Create variable for a slot and save it in TensorFlow optimizer.\"\"\"\n        if slot_name not in self._opt._slot_names:\n            self._opt._slot_names.append(slot_name)\n        var_key = _var_key(embed_var)\n        slot_dict = self._opt._slots.setdefault(var_key, {})\n        slot_var = slot_dict.get(slot_name, None)\n        if slot_var is None:\n            slot_var_name = \"%s/%s\" % (embed_var._shared_name, slot_name)\n            slot_var = self._opt.add_weight(\n                name=slot_var_name,\n                shape=shape,\n                dtype=embed_var.dtype,\n                initializer=initial_value,\n                trainable=False,\n            )\n            slot_dict[slot_name] = slot_var\n            self._opt._weights.append(slot_var)\n            return slot_var\n        else:\n            raise RuntimeError(\n                \"Variable with var_key %s and slot_name %s is not expected to \"\n                \"be in self._opt.\" % (var_key, slot_name)\n            )\n\n    def _update_embedding_param(self):\n        \"\"\"Report updated embedding vectors and slots to kv store.\"\"\"\n        for layer, ids in self._tls._unique_ids_all_layers.items():\n            value = self._get_embedding_variable(layer).numpy()\n            self._update_embedding_func(layer, ids, value)\n\n            for slot in self._allowed_slot_names:\n                value = self._get_slot_variable(layer, slot).numpy()\n                slot_table_name = get_slot_table_name(layer, slot)\n                self._update_embedding_func(slot_table_name, ids, value)\n\n    def _delete_variables(self):\n        # Slot variable access in optimizer requires corresponding embedding\n        # variable information. Delete slot variables first.\n        for layer_name, slots in self._tls._slot_variables.items():\n            embed_var = self._get_embedding_variable(layer_name)\n            embed_var_key = _var_key(embed_var)\n            del self._opt._slots[embed_var_key]\n            for _, var in slots.items():\n                opt_weight_iter = 0\n                while opt_weight_iter < len(self._opt._weights):\n                    if var is self._opt._weights[opt_weight_iter]:\n                        self._opt._weights.pop(opt_weight_iter)\n                        break\n                    else:\n                        opt_weight_iter += 1\n        for key in list(self._tls._slot_variables.keys()):\n            del self._tls._slot_variables[key]\n\n        # Delete variables in embed_variables.\n        for key in list(self._tls._embed_variables.keys()):\n            del self._tls._embed_variables[key]\n\n        # Delete variables in unique_ids_all_layers.\n        for key in list(self._tls._unique_ids_all_layers.keys()):\n            del self._tls._unique_ids_all_layers[key]\n\n    @property\n    def allowed_slot_names(self):\n        return self._allowed_slot_names\n\n    # TODO(yunjian.lmh): Do not need to save slot_initial_value in\n    #     optimizer wrapper after we do not need to support Redis.\n    @property\n    def slot_initial_value(self):\n        return self._slot_initial_value\n", "path": "elasticdl/python/ps/optimizer_wrapper.py"}]}