Repository: sql-machine-learning/elasticdl

Issue:
support blocking pull in PS so client don't need to retry in a loop



Discussion:

Resolving patch:
diff --git a/tensorflow/ps/ps.py b/tensorflow/ps/ps.py
--- a/tensorflow/ps/ps.py
+++ b/tensorflow/ps/ps.py
@@ -1,9 +1,9 @@
+import threading
+import queue
+import numpy as np
+import tensorflow.contrib.eager as tfe
 import tensorflow as tf
 tf.enable_eager_execution()
-import tensorflow.contrib.eager as tfe
-import numpy as np
-import queue
-import threading
 
 
 class ParameterServer(object):
@@ -22,11 +22,18 @@ def __init__(self, optimizer, vars):
         self._lock = threading.Lock()
         self._runner = threading.Thread(target=self._run, name='ps-runner')
         self._exiting = False
+        self._min_step_cv = threading.Condition()
 
-    def pull(self, min_step=0, names=None):
+    def pull(self, names=None, min_step=0, blocking=True, timeout=None):
+        with self._min_step_cv:
+            self._min_step_cv.wait_for(
+                lambda: not blocking or min_step <= self._step,
+                timeout=timeout)
         with self._lock:
             if min_step > self._step:
-                raise LookupError('Required step is not ready yet: %s' % min_step)
+                raise LookupError(
+                    'Required step is not ready yet: %s' %
+                    min_step)
             if names:
                 res = {k: self._vars[k].numpy() for k in names}
             else:
@@ -54,7 +61,9 @@ 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)
+        with self._min_step_cv:
             self._step += 1
+            self._min_step_cv.notify_all()
 
     def _run(self):
         while not self._exiting:


Test patch:
diff --git a/tensorflow/ps/ps_test.py b/tensorflow/ps/ps_test.py
--- a/tensorflow/ps/ps_test.py
+++ b/tensorflow/ps/ps_test.py
@@ -1,17 +1,20 @@
 import unittest
 import tensorflow as tf
 import numpy as np
+import threading
+import time
 from ps import ParameterServer
+from timeit import default_timer as timer
 
 
 class ParameterServerTestCase(unittest.TestCase):
     def setUp(self):
         self.ps1 = ParameterServer(tf.train.GradientDescentOptimizer(0.1),
-            {'x1': np.array([0.0, 0.0], dtype='float32'),
-            'y1': np.array([1.0, 1.0], dtype='float32')})
+                                   {'x1': np.array([0.0, 0.0], dtype='float32'),
+                                    'y1': np.array([1.0, 1.0], dtype='float32')})
         self.ps2 = ParameterServer(tf.train.GradientDescentOptimizer(0.1),
-            {'x2': np.array([0.0, 0.0], dtype='float32'),
-            'y2': np.array([1.0, 1.0], dtype='float32')})
+                                   {'x2': np.array([0.0, 0.0], dtype='float32'),
+                                    'y2': np.array([1.0, 1.0], dtype='float32')})
         self.ps1.start()
         self.ps2.start()
 
@@ -20,8 +23,8 @@ def tearDown(self):
         self.ps2.join()
 
     def test_all(self):
-        # First pull before any update.
-        step, params = self.ps1.pull()
+        # First non-blocking pull before any update.
+        step, params = self.ps1.pull(blocking=False)
         self.assertEqual(step, 0)
         np.testing.assert_array_almost_equal(
             params['x1'], np.array([0.0, 0.0], dtype='float32'))
@@ -29,8 +32,8 @@ def test_all(self):
             params['y1'], np.array([1.0, 1.0], dtype='float32'))
 
         self.ps1.push(0, 0,
-            {'x1': np.array([0.1, 0.1], dtype='float32'),
-            'y1': np.array([0.1, 0.1], dtype='float32')})
+                      {'x1': np.array([0.1, 0.1], dtype='float32'),
+                       'y1': np.array([0.1, 0.1], dtype='float32')})
 
         # Pull the new version
         while step < 1:
@@ -41,20 +44,48 @@ def test_all(self):
         np.testing.assert_array_almost_equal(
             params['y1'], np.array([0.99, 0.99], dtype='float32'))
 
-        # Pull a version in future should raise exception
-        with self.assertRaises(LookupError):
-            self.ps1.pull(min_step=2)
-        
         # Pull by names
-        step, params = self.ps1.pull(names=['x1',])
+        step, params = self.ps1.pull(names=['x1', ])
         self.assertEqual(step, 1)
         np.testing.assert_array_almost_equal(
             params['x1'], np.array([-0.01, -0.01], dtype='float32'))
-        
+
         # Pull by unknown name
         with self.assertRaises(LookupError):
-            self.ps1.pull(names=['z1',])
+            self.ps1.pull(names=['z1', ])
+        # Pull a version non-blocking in future should raise exception
+        with self.assertRaises(LookupError):
+            self.ps1.pull(min_step=2, blocking=False)
+
+        # Pull a version blocking in future should raise exception and blocks
+        # for ~1s
+        start = timer()
+        with self.assertRaises(LookupError):
+            self.ps1.pull(min_step=2, timeout=1.0)
+        end = timer()
+        self.assertGreater(end - start, 0.5)
 
+        # Block indefinitely and push a new step from another thread with some
+        # delay
+        def waited_push():
+            time.sleep(1)
+            self.ps1.push(0, 0,
+                          {'x1': np.array([0.1, 0.1], dtype='float32'),
+                           'y1': np.array([0.1, 0.1], dtype='float32')})
+        t = threading.Thread(target=waited_push)
+        t.start()
+
+        start = timer()
+        step, params = self.ps1.pull(min_step=2)
+        end = timer()
+        self.assertGreater(end - start, 0.5)
+
+        self.assertEqual(step, 2)
+        np.testing.assert_array_almost_equal(
+            params['x1'], np.array([-0.02, -0.02], dtype='float32'))
+        np.testing.assert_array_almost_equal(
+            params['y1'], np.array([0.98, 0.98], dtype='float32'))
+        t.join()
 
 
 if __name__ == '__main__':