Repository: sql-machine-learning/elasticdl

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.  


Discussion:
I think this kind of wrapper is a little confusing. Can we just provide our optimizers to users instead of silently using them? For example:

 ```python
from elasticdl.optimizer.adam import AdamOptimizer
```

This way we do not have to detect whether users have used the embedding layer. We just tell users they need to use our optimizers if they use certain functions like the embedding layer.
1. Whether or not we implement ElasticDL's optimizers, we need to discuss whether we need to use locks in optimizers.
2. Some users told us they want to use native TensorFlow api.
3. Through optimizer wrapper, we don't need to implement every optimizer in ElasticDL, which would cause a lot of repeat code.
Just a reference, Horovod uses:
> Wrap optimizer in `hvd.DistributedOptimizer`. The distributed optimizer delegates gradient computation to the original optimizer, averages gradients using allreduce or allgather, and then applies those averaged gradients.

which can be used as the following:
```
opt = tf.train.AdagradOptimizer(0.01 * hvd.size())
# Add Horovod Distributed Optimizer
opt = hvd.DistributedOptimizer(opt)
```
`OptimizerWrapper` in ElasticDL behaves like `hvd.DistributedOptimizer`. It also delegates gradient computation to the original optimizer and does following things in brief: 
* pull embedding vectors and slot variables from distributed kv store.
* call `apply_gradient` function of TensorFlow optimizer.
* push updated embedding vectors and slot variables to distributed kv store.

This issue aims to describe the following problem:
*In async mode, many threads wants to update embedding vectors at the same time. However, these embedding vectors are saved in distributed kv store. If we need to ensure that updates from all threads do not conflict, we need to add a lock here.*

Currently we use Redis as our distributed kv store, which complicates the code of `OptimizerWrapper`. After implementing PS, `OptimizerWrapper` will be refactored as a part of PS and the code can be greatly simplified. So if we decide to add a lock, we can implement this after finishing PS.
"The entire optimizer is currently thread compatible, not thread-safe. The user needs to perform synchronization if necessary." https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Optimizer?version=stable#thread_compatibility_2
I think we need to use a lock for async Optimizer Wrapper to apply gradient.

Resolving patch:
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 @@ def __init__(
         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 @@ def apply_gradients(self, grads_and_vars):
                 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 @@ def _delete_variables(self):
             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]
 


Test patch: