Repository: sql-machine-learning/elasticdl

Issue:
Change PS to run in graph mode.
在multi-thread下面， 如果ps enable eager mode, 在同一个process下面，所有thread都是在eager mode下面了。 eager mode 可以run graph, 反之则不行。


Discussion:

Resolving patch:
diff --git a/python/elasticdl/tflib/ps/ps.py b/python/elasticdl/tflib/ps/ps.py
--- a/python/elasticdl/tflib/ps/ps.py
+++ b/python/elasticdl/tflib/ps/ps.py
@@ -1,22 +1,12 @@
 import threading
 import queue
 import numpy as np
-import tensorflow.contrib.eager as tfe
 import tensorflow as tf
-tf.enable_eager_execution()
+from tensorflow.python.ops import array_ops
 
 
 class ParameterServer(object):
     def __init__(self, optimizer, vars):
-        self._opt = optimizer
-        self._vars = {}
-        for k, v in vars.items():
-            if (not isinstance(v, np.ndarray)
-                    or v.dtype not in (np.float32, np.float64)):
-                raise ValueError(
-                    'Initial value for variable %s is not of float type ndarray' %
-                    k)
-            self._vars[k] = tfe.Variable(v, name=k)
         self._step = 0
         self._grad_q = queue.Queue()
         self._lock = threading.Lock()
@@ -24,6 +14,23 @@ def __init__(self, optimizer, vars):
         self._exiting = False
         self._min_step_cv = threading.Condition()
 
+        self._grads_vars = {}
+        for k, v in vars.items():
+            if (not isinstance(v, np.ndarray)
+                    or v.dtype not in (np.float32, np.float64)):
+                raise ValueError(
+                    'Initial value for variable %s is not of float type ndarray' %
+                    k)
+            # TODO: In graph mode we don't need to keep track of variables by ourselves.
+            self._grads_vars[k] = (array_ops.placeholder(dtype=v.dtype), tf.Variable(v, name=k))
+
+        self._opt = optimizer
+        self._apply_grad_op = self._opt.apply_gradients(self._grads_vars.values())
+
+        self._sess = tf.Session()
+        init_op = tf.global_variables_initializer()
+        self._sess.run(init_op)
+
     def pull(self, names=None, min_step=0, blocking=True, timeout=None):
         with self._min_step_cv:
             self._min_step_cv.wait_for(
@@ -35,9 +42,9 @@ def pull(self, names=None, min_step=0, blocking=True, timeout=None):
                     'Required step is not ready yet: %s' %
                     min_step)
             if names:
-                res = {k: self._vars[k].numpy() for k in names}
+                res = {k: self._grads_vars[k][1].eval(self._sess) for k in names}
             else:
-                res = {k: v.numpy() for k, v in self._vars.items()}
+                res = {k: v[1].eval(self._sess) for k, v in self._grads_vars.items()}
             return self._step, res
 
     def push(self, base_step, sub_step, grads):
@@ -51,16 +58,16 @@ def push(self, base_step, sub_step, grads):
             raise ValueError('Illegal sub step %s' % sub_step)
 
         for k, g in grads.items():
-            v = self._vars[k]
+            v = self._grads_vars[k][1]
             if g.dtype != v.dtype.as_numpy_dtype or g.shape != v.shape:
                 raise ValueError('Incompatible gradient for variable %s' % k)
         # TODO(l.zou): use @dataclass when python 3.7 is available.
         self._grad_q.put((base_step, sub_step, grads))
 
     def _compute(self, grads):
-        grads_vars = [(g, self._vars[k]) for k, g in grads.items()]
         with self._lock:
-            self._opt.apply_gradients(grads_vars)
+            feed_dict = {self._grads_vars[k][0]:v for k, v in grads.items()}
+            self._sess.run(self._apply_grad_op, feed_dict=feed_dict)
         with self._min_step_cv:
             self._step += 1
             self._min_step_cv.notify_all()
@@ -80,3 +87,4 @@ def start(self):
     def join(self):
         self._exiting = True
         self._runner.join()
+        self._sess.close()


Test patch: