We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
support blocking pull in PS so client don't need to retry in a loop


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `tensorflow/ps/ps.py`
Content:
```
1 import tensorflow as tf
2 tf.enable_eager_execution()
3 import tensorflow.contrib.eager as tfe
4 import numpy as np
5 import queue
6 import threading
7 
8 
9 class ParameterServer(object):
10     def __init__(self, optimizer, vars):
11         self._opt = optimizer
12         self._vars = {}
13         for k, v in vars.items():
14             if (not isinstance(v, np.ndarray)
15                     or v.dtype not in (np.float32, np.float64)):
16                 raise ValueError(
17                     'Initial value for variable %s is not of float type ndarray' %
18                     k)
19             self._vars[k] = tfe.Variable(v, name=k)
20         self._step = 0
21         self._grad_q = queue.Queue()
22         self._lock = threading.Lock()
23         self._runner = threading.Thread(target=self._run, name='ps-runner')
24         self._exiting = False
25 
26     def pull(self, min_step=0, names=None):
27         with self._lock:
28             if min_step > self._step:
29                 raise LookupError('Required step is not ready yet: %s' % min_step)
30             if names:
31                 res = {k: self._vars[k].numpy() for k in names}
32             else:
33                 res = {k: v.numpy() for k, v in self._vars.items()}
34             return self._step, res
35 
36     def push(self, base_step, sub_step, grads):
37         with self._lock:
38             if base_step > self._step:
39                 raise ValueError(
40                     'Illegal base step %s, parameter server step is %s' %
41                     (base_step, self._step))
42 
43         if sub_step < 0:
44             raise ValueError('Illegal sub step %s' % sub_step)
45 
46         for k, g in grads.items():
47             v = self._vars[k]
48             if g.dtype != v.dtype.as_numpy_dtype or g.shape != v.shape:
49                 raise ValueError('Incompatible gradient for variable %s' % k)
50         # TODO(l.zou): use @dataclass when python 3.7 is available.
51         self._grad_q.put((base_step, sub_step, grads))
52 
53     def _compute(self, grads):
54         grads_vars = [(g, self._vars[k]) for k, g in grads.items()]
55         with self._lock:
56             self._opt.apply_gradients(grads_vars)
57             self._step += 1
58 
59     def _run(self):
60         while not self._exiting:
61             # TODO(l.zou): How to properly accumulate and decay grads?
62             try:
63                 base_step, sub_step, grads = self._grad_q.get(timeout=1.0)
64                 self._compute(grads)
65             except queue.Empty:
66                 pass
67 
68     def start(self):
69         self._runner.start()
70 
71     def join(self):
72         self._exiting = True
73         self._runner.join()
74
```
--- 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/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 @@
         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 @@
         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:

{"golden_diff": "diff --git a/tensorflow/ps/ps.py b/tensorflow/ps/ps.py\n--- a/tensorflow/ps/ps.py\n+++ b/tensorflow/ps/ps.py\n@@ -1,9 +1,9 @@\n+import threading\n+import queue\n+import numpy as np\n+import tensorflow.contrib.eager as tfe\n import tensorflow as tf\n tf.enable_eager_execution()\n-import tensorflow.contrib.eager as tfe\n-import numpy as np\n-import queue\n-import threading\n \n \n class ParameterServer(object):\n@@ -22,11 +22,18 @@\n         self._lock = threading.Lock()\n         self._runner = threading.Thread(target=self._run, name='ps-runner')\n         self._exiting = False\n+        self._min_step_cv = threading.Condition()\n \n-    def pull(self, min_step=0, names=None):\n+    def pull(self, names=None, min_step=0, blocking=True, timeout=None):\n+        with self._min_step_cv:\n+            self._min_step_cv.wait_for(\n+                lambda: not blocking or min_step <= self._step,\n+                timeout=timeout)\n         with self._lock:\n             if min_step > self._step:\n-                raise LookupError('Required step is not ready yet: %s' % min_step)\n+                raise LookupError(\n+                    'Required step is not ready yet: %s' %\n+                    min_step)\n             if names:\n                 res = {k: self._vars[k].numpy() for k in names}\n             else:\n@@ -54,7 +61,9 @@\n         grads_vars = [(g, self._vars[k]) for k, g in grads.items()]\n         with self._lock:\n             self._opt.apply_gradients(grads_vars)\n+        with self._min_step_cv:\n             self._step += 1\n+            self._min_step_cv.notify_all()\n \n     def _run(self):\n         while not self._exiting:\n", "issue": "support blocking pull in PS so client don't need to retry in a loop\n\n", "before_files": [{"content": "import tensorflow as tf\ntf.enable_eager_execution()\nimport tensorflow.contrib.eager as tfe\nimport numpy as np\nimport queue\nimport threading\n\n\nclass ParameterServer(object):\n    def __init__(self, optimizer, vars):\n        self._opt = optimizer\n        self._vars = {}\n        for k, v in vars.items():\n            if (not isinstance(v, np.ndarray)\n                    or v.dtype not in (np.float32, np.float64)):\n                raise ValueError(\n                    'Initial value for variable %s is not of float type ndarray' %\n                    k)\n            self._vars[k] = tfe.Variable(v, name=k)\n        self._step = 0\n        self._grad_q = queue.Queue()\n        self._lock = threading.Lock()\n        self._runner = threading.Thread(target=self._run, name='ps-runner')\n        self._exiting = False\n\n    def pull(self, min_step=0, names=None):\n        with self._lock:\n            if min_step > self._step:\n                raise LookupError('Required step is not ready yet: %s' % min_step)\n            if names:\n                res = {k: self._vars[k].numpy() for k in names}\n            else:\n                res = {k: v.numpy() for k, v in self._vars.items()}\n            return self._step, res\n\n    def push(self, base_step, sub_step, grads):\n        with self._lock:\n            if base_step > self._step:\n                raise ValueError(\n                    'Illegal base step %s, parameter server step is %s' %\n                    (base_step, self._step))\n\n        if sub_step < 0:\n            raise ValueError('Illegal sub step %s' % sub_step)\n\n        for k, g in grads.items():\n            v = self._vars[k]\n            if g.dtype != v.dtype.as_numpy_dtype or g.shape != v.shape:\n                raise ValueError('Incompatible gradient for variable %s' % k)\n        # TODO(l.zou): use @dataclass when python 3.7 is available.\n        self._grad_q.put((base_step, sub_step, grads))\n\n    def _compute(self, grads):\n        grads_vars = [(g, self._vars[k]) for k, g in grads.items()]\n        with self._lock:\n            self._opt.apply_gradients(grads_vars)\n            self._step += 1\n\n    def _run(self):\n        while not self._exiting:\n            # TODO(l.zou): How to properly accumulate and decay grads?\n            try:\n                base_step, sub_step, grads = self._grad_q.get(timeout=1.0)\n                self._compute(grads)\n            except queue.Empty:\n                pass\n\n    def start(self):\n        self._runner.start()\n\n    def join(self):\n        self._exiting = True\n        self._runner.join()\n", "path": "tensorflow/ps/ps.py"}], "after_files": [{"content": "import threading\nimport queue\nimport numpy as np\nimport tensorflow.contrib.eager as tfe\nimport tensorflow as tf\ntf.enable_eager_execution()\n\n\nclass ParameterServer(object):\n    def __init__(self, optimizer, vars):\n        self._opt = optimizer\n        self._vars = {}\n        for k, v in vars.items():\n            if (not isinstance(v, np.ndarray)\n                    or v.dtype not in (np.float32, np.float64)):\n                raise ValueError(\n                    'Initial value for variable %s is not of float type ndarray' %\n                    k)\n            self._vars[k] = tfe.Variable(v, name=k)\n        self._step = 0\n        self._grad_q = queue.Queue()\n        self._lock = threading.Lock()\n        self._runner = threading.Thread(target=self._run, name='ps-runner')\n        self._exiting = False\n        self._min_step_cv = threading.Condition()\n\n    def pull(self, names=None, min_step=0, blocking=True, timeout=None):\n        with self._min_step_cv:\n            self._min_step_cv.wait_for(\n                lambda: not blocking or min_step <= self._step,\n                timeout=timeout)\n        with self._lock:\n            if min_step > self._step:\n                raise LookupError(\n                    'Required step is not ready yet: %s' %\n                    min_step)\n            if names:\n                res = {k: self._vars[k].numpy() for k in names}\n            else:\n                res = {k: v.numpy() for k, v in self._vars.items()}\n            return self._step, res\n\n    def push(self, base_step, sub_step, grads):\n        with self._lock:\n            if base_step > self._step:\n                raise ValueError(\n                    'Illegal base step %s, parameter server step is %s' %\n                    (base_step, self._step))\n\n        if sub_step < 0:\n            raise ValueError('Illegal sub step %s' % sub_step)\n\n        for k, g in grads.items():\n            v = self._vars[k]\n            if g.dtype != v.dtype.as_numpy_dtype or g.shape != v.shape:\n                raise ValueError('Incompatible gradient for variable %s' % k)\n        # TODO(l.zou): use @dataclass when python 3.7 is available.\n        self._grad_q.put((base_step, sub_step, grads))\n\n    def _compute(self, grads):\n        grads_vars = [(g, self._vars[k]) for k, g in grads.items()]\n        with self._lock:\n            self._opt.apply_gradients(grads_vars)\n        with self._min_step_cv:\n            self._step += 1\n            self._min_step_cv.notify_all()\n\n    def _run(self):\n        while not self._exiting:\n            # TODO(l.zou): How to properly accumulate and decay grads?\n            try:\n                base_step, sub_step, grads = self._grad_q.get(timeout=1.0)\n                self._compute(grads)\n            except queue.Empty:\n                pass\n\n    def start(self):\n        self._runner.start()\n\n    def join(self):\n        self._exiting = True\n        self._runner.join()\n", "path": "tensorflow/ps/ps.py"}]}