We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Remove duplicate feed dict constructing
in `python/ray/experimental/sgd/sgd_worker.py` 

```python
    def compute_gradients(self):
        start = time.time()
        feed_dict = self._grad_feed_dict()
        # Aggregate feed dicts for each model on this worker.
        for model in self.models:
            feed_dict.update(model.get_feed_dict())
```

with `_grad_feed_dict` definitions:

```python
    def _grad_feed_dict(self):
        # Aggregate feed dicts for each model on this worker.
        feed_dict = {}
        for model in self.models:
            feed_dict.update(model.get_feed_dict())
        return feed_dict
```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `python/ray/experimental/sgd/sgd_worker.py`
Content:
```
1 from __future__ import absolute_import
2 from __future__ import division
3 from __future__ import print_function
4 
5 import logging
6 import time
7 
8 import pyarrow.plasma as plasma
9 import tensorflow as tf
10 
11 import ray
12 from ray.experimental.sgd.util import fetch, run_timeline, warmup
13 from ray.experimental.sgd.modified_allreduce import sum_gradients_all_reduce, \
14     unpack_small_tensors
15 
16 logger = logging.getLogger(__name__)
17 
18 
19 class SGDWorker(object):
20     """Helper class for ray.experimental.sgd.DistributedSGD."""
21 
22     def __init__(self,
23                  worker_index,
24                  model_creator,
25                  all_reduce_alg="simple",
26                  num_devices=1,
27                  gpu=False,
28                  max_bytes=10000000,
29                  plasma_op=False):
30         self.worker_index = worker_index
31         assert num_devices > 0
32 
33         # TODO(ekl) support custom session
34         tf_session_args = {
35             "device_count": {
36                 "CPU": num_devices
37             },
38             "log_device_placement": False,
39             "gpu_options": tf.GPUOptions(force_gpu_compatible=True),
40             "inter_op_parallelism_threads": 128,
41         }
42         config_proto = tf.ConfigProto(**tf_session_args)
43         self.sess = tf.Session(config=config_proto)
44         self.models = []
45         grad_ops = []
46 
47         if gpu:
48             device_tmpl = "/gpu:%d"
49         else:
50             device_tmpl = "/cpu:%d"
51         with self.sess.as_default():
52             for device_idx in range(num_devices):
53                 device = device_tmpl % device_idx
54                 with tf.device(device):
55                     with tf.variable_scope("device_%d" % device_idx):
56                         model = model_creator(worker_index, device_idx)
57                         self.models.append(model)
58                         grads = [
59                             t for t in model.optimizer.compute_gradients(
60                                 model.loss) if t[0] is not None
61                         ]
62                         grad_ops.append(grads)
63 
64         if num_devices == 1:
65             if max_bytes:
66                 raise ValueError(
67                     "Implementation limitation: grad_shard_bytes > 0 "
68                     "({}) currently requires > 1 device".format(max_bytes))
69             self.packed_grads_and_vars = grad_ops
70         else:
71             if max_bytes:
72                 self.packed_grads_and_vars, packing_vals = (
73                     sum_gradients_all_reduce(
74                         "",
75                         grad_ops,
76                         1,
77                         all_reduce_alg,
78                         1,
79                         list(range(num_devices)),
80                         agg_small_grads_max_bytes=max_bytes))
81             else:
82                 self.packed_grads_and_vars, _ = (sum_gradients_all_reduce(
83                     "",
84                     grad_ops,
85                     1,
86                     all_reduce_alg,
87                     1,
88                     list(range(num_devices)),
89                     agg_small_grads_max_bytes=0))
90         self.per_device_grads = [
91             list(zip(*dev_gv))[0] for dev_gv in self.packed_grads_and_vars
92         ]
93         assert (len(self.per_device_grads) == num_devices)
94         self.num_grads = num_grads = len(self.packed_grads_and_vars[0])
95         if max_bytes:
96             logger.info("Packed grads => {} tensors".format(num_grads))
97 
98         # Ops for reading grads with the right control deps
99         nccl_noops = []
100         for j in range(num_grads)[::-1]:
101             deps = nccl_noops + [
102                 dev_grad[j] for dev_grad in self.per_device_grads
103             ]
104             with tf.control_dependencies(deps):
105                 nccl_noops = [tf.no_op()]
106 
107         # You must fetch this otherwise the NCCL allreduce will hang
108         self.nccl_control_out = tf.group(*nccl_noops)
109 
110         if plasma_op:
111             store_socket = (
112                 ray.worker.global_worker.plasma_client.store_socket_name)
113             manager_socket = (
114                 ray.worker.global_worker.plasma_client.manager_socket_name)
115             if not plasma.tf_plasma_op:
116                 plasma.build_plasma_tensorflow_op()
117 
118             # For fetching grads -> plasma
119             self.plasma_in_grads = []
120             self.plasma_in_grads_oids = [
121                 tf.placeholder(shape=[], dtype=tf.string, name="in_grad_oids")
122                 for _ in range(num_grads)
123             ]
124             for j in range(num_grads):
125                 grad = self.per_device_grads[0][j]
126                 with tf.device(self.models[0].loss.device):
127                     plasma_grad = plasma.tf_plasma_op.tensor_to_plasma(
128                         [grad],
129                         self.plasma_in_grads_oids[j],
130                         plasma_store_socket_name=store_socket,
131                         plasma_manager_socket_name=manager_socket)
132                 self.plasma_in_grads.append(plasma_grad)
133 
134             # For applying grads <- plasma
135             unpacked_gv = []
136             self.plasma_out_grads_oids = [
137                 tf.placeholder(
138                     shape=[], dtype=tf.string, name="grad_out_oids")
139                 for _ in range(num_grads)
140             ]
141             packed_plasma_grads = []
142             for j in range(num_grads):
143                 with tf.device(self.plasma_in_grads[j].device):
144                     with tf.control_dependencies([self.plasma_in_grads[j]]):
145                         grad_ph = plasma.tf_plasma_op.plasma_to_tensor(
146                             self.plasma_out_grads_oids[j],
147                             dtype=tf.float32,
148                             plasma_store_socket_name=store_socket,
149                             plasma_manager_socket_name=manager_socket)
150                 grad_ph = tf.reshape(grad_ph,
151                                      self.packed_grads_and_vars[0][j][0].shape)
152                 logger.debug("Packed tensor {}".format(grad_ph))
153                 packed_plasma_grads.append(grad_ph)
154             for i in range(num_devices):
155                 per_device = []
156                 for j, (g, v) in enumerate(self.packed_grads_and_vars[i]):
157                     grad_ph = packed_plasma_grads[j]
158                     per_device.append((grad_ph, v))
159                 unpacked_gv.append(per_device)
160 
161             if max_bytes:
162                 unpacked_gv = unpack_small_tensors(unpacked_gv, packing_vals)
163 
164         elif max_bytes:
165             unpacked_gv = unpack_small_tensors(self.packed_grads_and_vars,
166                                                packing_vals)
167         else:
168             unpacked_gv = self.packed_grads_and_vars
169 
170         # Same shape as packed_grads_and_vars
171         assert len(unpacked_gv) == num_devices
172         assert len(unpacked_gv[0][0]) == 2
173 
174         apply_ops = []
175         to_apply = unpacked_gv[0]
176         for ix, m in enumerate(self.models):
177             apply_ops.append(
178                 m.optimizer.apply_gradients(
179                     [(g, v)
180                      for ((g, _), (_, v)) in zip(to_apply, unpacked_gv[ix])]))
181         self.apply_op = tf.group(*apply_ops)
182         init_op = tf.group(tf.global_variables_initializer(),
183                            tf.local_variables_initializer())
184         self.sess.run(init_op)
185 
186     def _grad_feed_dict(self):
187         # Aggregate feed dicts for each model on this worker.
188         feed_dict = {}
189         for model in self.models:
190             feed_dict.update(model.get_feed_dict())
191         return feed_dict
192 
193     def foreach_model(self, fn):
194         with self.sess.as_default():
195             return [fn(m) for m in self.models]
196 
197     def foreach_worker(self, fn):
198         with self.sess.as_default():
199             return fn(self)
200 
201     def for_model(self, fn):
202         with self.sess.as_default():
203             return fn(self.models[0])
204 
205     def compute_gradients(self):
206         start = time.time()
207         feed_dict = self._grad_feed_dict()
208         # Aggregate feed dicts for each model on this worker.
209         for model in self.models:
210             feed_dict.update(model.get_feed_dict())
211         # We only need to fetch the first per_device_grad, since they are
212         # averaged across all devices by allreduce.
213         fetches = self.sess.run(
214             [
215                 self.models[0].loss, self.per_device_grads[0],
216                 self.nccl_control_out
217             ],
218             feed_dict=feed_dict)
219         logger.debug(
220             "Compute grad interior time {}".format(time.time() - start))
221         return fetches
222 
223     def apply_gradients(self, avg_grads):
224         start = time.time()
225         result = {
226             g: avg_grads[i]
227             for (i, g) in enumerate(self.per_device_grads[0])
228         }
229         self.sess.run(self.apply_op, feed_dict=result)
230         logger.debug("Apply grad interior time {}".format(time.time() - start))
231 
232     def compute_apply(self):
233         fetches = run_timeline(
234             self.sess,
235             [self.models[0].loss, self.apply_op, self.nccl_control_out],
236             feed_dict=self._grad_feed_dict(),
237             name="compute_apply")
238         return fetches[0]
239 
240     def ps_compute_apply(self,
241                          out_grad_shard_oids,
242                          agg_grad_shard_oids,
243                          tl_name="ps_compute_apply",
244                          write_timeline=False):
245         feed_dict = self._grad_feed_dict()
246         feed_dict.update(
247             dict(zip(self.plasma_in_grads_oids, out_grad_shard_oids)))
248         feed_dict.update(
249             dict(zip(self.plasma_out_grads_oids, agg_grad_shard_oids)))
250         fetch(agg_grad_shard_oids)
251         fetches = run_timeline(
252             self.sess, [
253                 self.models[0].loss, self.plasma_in_grads, self.apply_op,
254                 self.nccl_control_out
255             ],
256             feed_dict=feed_dict,
257             write_timeline=write_timeline)
258         return fetches[0]
259 
260     def num_grad_shards(self):
261         return self.num_grads
262 
263     def shard_shapes(self):
264         main_gv = self.packed_grads_and_vars[0]
265         return [g.shape for g, _ in main_gv]
266 
267     def ip(self):
268         return ray.services.get_node_ip_address()
269 
270     def warmup(self):
271         warmup()
272
```
--- 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/python/ray/experimental/sgd/sgd_worker.py b/python/ray/experimental/sgd/sgd_worker.py
--- a/python/ray/experimental/sgd/sgd_worker.py
+++ b/python/ray/experimental/sgd/sgd_worker.py
@@ -205,9 +205,6 @@
     def compute_gradients(self):
         start = time.time()
         feed_dict = self._grad_feed_dict()
-        # Aggregate feed dicts for each model on this worker.
-        for model in self.models:
-            feed_dict.update(model.get_feed_dict())
         # We only need to fetch the first per_device_grad, since they are
         # averaged across all devices by allreduce.
         fetches = self.sess.run(

{"golden_diff": "diff --git a/python/ray/experimental/sgd/sgd_worker.py b/python/ray/experimental/sgd/sgd_worker.py\n--- a/python/ray/experimental/sgd/sgd_worker.py\n+++ b/python/ray/experimental/sgd/sgd_worker.py\n@@ -205,9 +205,6 @@\n     def compute_gradients(self):\n         start = time.time()\n         feed_dict = self._grad_feed_dict()\n-        # Aggregate feed dicts for each model on this worker.\n-        for model in self.models:\n-            feed_dict.update(model.get_feed_dict())\n         # We only need to fetch the first per_device_grad, since they are\n         # averaged across all devices by allreduce.\n         fetches = self.sess.run(\n", "issue": "Remove duplicate feed dict constructing\nin `python/ray/experimental/sgd/sgd_worker.py` \r\n\r\n```python\r\n    def compute_gradients(self):\r\n        start = time.time()\r\n        feed_dict = self._grad_feed_dict()\r\n        # Aggregate feed dicts for each model on this worker.\r\n        for model in self.models:\r\n            feed_dict.update(model.get_feed_dict())\r\n```\r\n\r\nwith `_grad_feed_dict` definitions:\r\n\r\n```python\r\n    def _grad_feed_dict(self):\r\n        # Aggregate feed dicts for each model on this worker.\r\n        feed_dict = {}\r\n        for model in self.models:\r\n            feed_dict.update(model.get_feed_dict())\r\n        return feed_dict\r\n```\n", "before_files": [{"content": "from __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport logging\nimport time\n\nimport pyarrow.plasma as plasma\nimport tensorflow as tf\n\nimport ray\nfrom ray.experimental.sgd.util import fetch, run_timeline, warmup\nfrom ray.experimental.sgd.modified_allreduce import sum_gradients_all_reduce, \\\n    unpack_small_tensors\n\nlogger = logging.getLogger(__name__)\n\n\nclass SGDWorker(object):\n    \"\"\"Helper class for ray.experimental.sgd.DistributedSGD.\"\"\"\n\n    def __init__(self,\n                 worker_index,\n                 model_creator,\n                 all_reduce_alg=\"simple\",\n                 num_devices=1,\n                 gpu=False,\n                 max_bytes=10000000,\n                 plasma_op=False):\n        self.worker_index = worker_index\n        assert num_devices > 0\n\n        # TODO(ekl) support custom session\n        tf_session_args = {\n            \"device_count\": {\n                \"CPU\": num_devices\n            },\n            \"log_device_placement\": False,\n            \"gpu_options\": tf.GPUOptions(force_gpu_compatible=True),\n            \"inter_op_parallelism_threads\": 128,\n        }\n        config_proto = tf.ConfigProto(**tf_session_args)\n        self.sess = tf.Session(config=config_proto)\n        self.models = []\n        grad_ops = []\n\n        if gpu:\n            device_tmpl = \"/gpu:%d\"\n        else:\n            device_tmpl = \"/cpu:%d\"\n        with self.sess.as_default():\n            for device_idx in range(num_devices):\n                device = device_tmpl % device_idx\n                with tf.device(device):\n                    with tf.variable_scope(\"device_%d\" % device_idx):\n                        model = model_creator(worker_index, device_idx)\n                        self.models.append(model)\n                        grads = [\n                            t for t in model.optimizer.compute_gradients(\n                                model.loss) if t[0] is not None\n                        ]\n                        grad_ops.append(grads)\n\n        if num_devices == 1:\n            if max_bytes:\n                raise ValueError(\n                    \"Implementation limitation: grad_shard_bytes > 0 \"\n                    \"({}) currently requires > 1 device\".format(max_bytes))\n            self.packed_grads_and_vars = grad_ops\n        else:\n            if max_bytes:\n                self.packed_grads_and_vars, packing_vals = (\n                    sum_gradients_all_reduce(\n                        \"\",\n                        grad_ops,\n                        1,\n                        all_reduce_alg,\n                        1,\n                        list(range(num_devices)),\n                        agg_small_grads_max_bytes=max_bytes))\n            else:\n                self.packed_grads_and_vars, _ = (sum_gradients_all_reduce(\n                    \"\",\n                    grad_ops,\n                    1,\n                    all_reduce_alg,\n                    1,\n                    list(range(num_devices)),\n                    agg_small_grads_max_bytes=0))\n        self.per_device_grads = [\n            list(zip(*dev_gv))[0] for dev_gv in self.packed_grads_and_vars\n        ]\n        assert (len(self.per_device_grads) == num_devices)\n        self.num_grads = num_grads = len(self.packed_grads_and_vars[0])\n        if max_bytes:\n            logger.info(\"Packed grads => {} tensors\".format(num_grads))\n\n        # Ops for reading grads with the right control deps\n        nccl_noops = []\n        for j in range(num_grads)[::-1]:\n            deps = nccl_noops + [\n                dev_grad[j] for dev_grad in self.per_device_grads\n            ]\n            with tf.control_dependencies(deps):\n                nccl_noops = [tf.no_op()]\n\n        # You must fetch this otherwise the NCCL allreduce will hang\n        self.nccl_control_out = tf.group(*nccl_noops)\n\n        if plasma_op:\n            store_socket = (\n                ray.worker.global_worker.plasma_client.store_socket_name)\n            manager_socket = (\n                ray.worker.global_worker.plasma_client.manager_socket_name)\n            if not plasma.tf_plasma_op:\n                plasma.build_plasma_tensorflow_op()\n\n            # For fetching grads -> plasma\n            self.plasma_in_grads = []\n            self.plasma_in_grads_oids = [\n                tf.placeholder(shape=[], dtype=tf.string, name=\"in_grad_oids\")\n                for _ in range(num_grads)\n            ]\n            for j in range(num_grads):\n                grad = self.per_device_grads[0][j]\n                with tf.device(self.models[0].loss.device):\n                    plasma_grad = plasma.tf_plasma_op.tensor_to_plasma(\n                        [grad],\n                        self.plasma_in_grads_oids[j],\n                        plasma_store_socket_name=store_socket,\n                        plasma_manager_socket_name=manager_socket)\n                self.plasma_in_grads.append(plasma_grad)\n\n            # For applying grads <- plasma\n            unpacked_gv = []\n            self.plasma_out_grads_oids = [\n                tf.placeholder(\n                    shape=[], dtype=tf.string, name=\"grad_out_oids\")\n                for _ in range(num_grads)\n            ]\n            packed_plasma_grads = []\n            for j in range(num_grads):\n                with tf.device(self.plasma_in_grads[j].device):\n                    with tf.control_dependencies([self.plasma_in_grads[j]]):\n                        grad_ph = plasma.tf_plasma_op.plasma_to_tensor(\n                            self.plasma_out_grads_oids[j],\n                            dtype=tf.float32,\n                            plasma_store_socket_name=store_socket,\n                            plasma_manager_socket_name=manager_socket)\n                grad_ph = tf.reshape(grad_ph,\n                                     self.packed_grads_and_vars[0][j][0].shape)\n                logger.debug(\"Packed tensor {}\".format(grad_ph))\n                packed_plasma_grads.append(grad_ph)\n            for i in range(num_devices):\n                per_device = []\n                for j, (g, v) in enumerate(self.packed_grads_and_vars[i]):\n                    grad_ph = packed_plasma_grads[j]\n                    per_device.append((grad_ph, v))\n                unpacked_gv.append(per_device)\n\n            if max_bytes:\n                unpacked_gv = unpack_small_tensors(unpacked_gv, packing_vals)\n\n        elif max_bytes:\n            unpacked_gv = unpack_small_tensors(self.packed_grads_and_vars,\n                                               packing_vals)\n        else:\n            unpacked_gv = self.packed_grads_and_vars\n\n        # Same shape as packed_grads_and_vars\n        assert len(unpacked_gv) == num_devices\n        assert len(unpacked_gv[0][0]) == 2\n\n        apply_ops = []\n        to_apply = unpacked_gv[0]\n        for ix, m in enumerate(self.models):\n            apply_ops.append(\n                m.optimizer.apply_gradients(\n                    [(g, v)\n                     for ((g, _), (_, v)) in zip(to_apply, unpacked_gv[ix])]))\n        self.apply_op = tf.group(*apply_ops)\n        init_op = tf.group(tf.global_variables_initializer(),\n                           tf.local_variables_initializer())\n        self.sess.run(init_op)\n\n    def _grad_feed_dict(self):\n        # Aggregate feed dicts for each model on this worker.\n        feed_dict = {}\n        for model in self.models:\n            feed_dict.update(model.get_feed_dict())\n        return feed_dict\n\n    def foreach_model(self, fn):\n        with self.sess.as_default():\n            return [fn(m) for m in self.models]\n\n    def foreach_worker(self, fn):\n        with self.sess.as_default():\n            return fn(self)\n\n    def for_model(self, fn):\n        with self.sess.as_default():\n            return fn(self.models[0])\n\n    def compute_gradients(self):\n        start = time.time()\n        feed_dict = self._grad_feed_dict()\n        # Aggregate feed dicts for each model on this worker.\n        for model in self.models:\n            feed_dict.update(model.get_feed_dict())\n        # We only need to fetch the first per_device_grad, since they are\n        # averaged across all devices by allreduce.\n        fetches = self.sess.run(\n            [\n                self.models[0].loss, self.per_device_grads[0],\n                self.nccl_control_out\n            ],\n            feed_dict=feed_dict)\n        logger.debug(\n            \"Compute grad interior time {}\".format(time.time() - start))\n        return fetches\n\n    def apply_gradients(self, avg_grads):\n        start = time.time()\n        result = {\n            g: avg_grads[i]\n            for (i, g) in enumerate(self.per_device_grads[0])\n        }\n        self.sess.run(self.apply_op, feed_dict=result)\n        logger.debug(\"Apply grad interior time {}\".format(time.time() - start))\n\n    def compute_apply(self):\n        fetches = run_timeline(\n            self.sess,\n            [self.models[0].loss, self.apply_op, self.nccl_control_out],\n            feed_dict=self._grad_feed_dict(),\n            name=\"compute_apply\")\n        return fetches[0]\n\n    def ps_compute_apply(self,\n                         out_grad_shard_oids,\n                         agg_grad_shard_oids,\n                         tl_name=\"ps_compute_apply\",\n                         write_timeline=False):\n        feed_dict = self._grad_feed_dict()\n        feed_dict.update(\n            dict(zip(self.plasma_in_grads_oids, out_grad_shard_oids)))\n        feed_dict.update(\n            dict(zip(self.plasma_out_grads_oids, agg_grad_shard_oids)))\n        fetch(agg_grad_shard_oids)\n        fetches = run_timeline(\n            self.sess, [\n                self.models[0].loss, self.plasma_in_grads, self.apply_op,\n                self.nccl_control_out\n            ],\n            feed_dict=feed_dict,\n            write_timeline=write_timeline)\n        return fetches[0]\n\n    def num_grad_shards(self):\n        return self.num_grads\n\n    def shard_shapes(self):\n        main_gv = self.packed_grads_and_vars[0]\n        return [g.shape for g, _ in main_gv]\n\n    def ip(self):\n        return ray.services.get_node_ip_address()\n\n    def warmup(self):\n        warmup()\n", "path": "python/ray/experimental/sgd/sgd_worker.py"}], "after_files": [{"content": "from __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport logging\nimport time\n\nimport pyarrow.plasma as plasma\nimport tensorflow as tf\n\nimport ray\nfrom ray.experimental.sgd.util import fetch, run_timeline, warmup\nfrom ray.experimental.sgd.modified_allreduce import sum_gradients_all_reduce, \\\n    unpack_small_tensors\n\nlogger = logging.getLogger(__name__)\n\n\nclass SGDWorker(object):\n    \"\"\"Helper class for ray.experimental.sgd.DistributedSGD.\"\"\"\n\n    def __init__(self,\n                 worker_index,\n                 model_creator,\n                 all_reduce_alg=\"simple\",\n                 num_devices=1,\n                 gpu=False,\n                 max_bytes=10000000,\n                 plasma_op=False):\n        self.worker_index = worker_index\n        assert num_devices > 0\n\n        # TODO(ekl) support custom session\n        tf_session_args = {\n            \"device_count\": {\n                \"CPU\": num_devices\n            },\n            \"log_device_placement\": False,\n            \"gpu_options\": tf.GPUOptions(force_gpu_compatible=True),\n            \"inter_op_parallelism_threads\": 128,\n        }\n        config_proto = tf.ConfigProto(**tf_session_args)\n        self.sess = tf.Session(config=config_proto)\n        self.models = []\n        grad_ops = []\n\n        if gpu:\n            device_tmpl = \"/gpu:%d\"\n        else:\n            device_tmpl = \"/cpu:%d\"\n        with self.sess.as_default():\n            for device_idx in range(num_devices):\n                device = device_tmpl % device_idx\n                with tf.device(device):\n                    with tf.variable_scope(\"device_%d\" % device_idx):\n                        model = model_creator(worker_index, device_idx)\n                        self.models.append(model)\n                        grads = [\n                            t for t in model.optimizer.compute_gradients(\n                                model.loss) if t[0] is not None\n                        ]\n                        grad_ops.append(grads)\n\n        if num_devices == 1:\n            if max_bytes:\n                raise ValueError(\n                    \"Implementation limitation: grad_shard_bytes > 0 \"\n                    \"({}) currently requires > 1 device\".format(max_bytes))\n            self.packed_grads_and_vars = grad_ops\n        else:\n            if max_bytes:\n                self.packed_grads_and_vars, packing_vals = (\n                    sum_gradients_all_reduce(\n                        \"\",\n                        grad_ops,\n                        1,\n                        all_reduce_alg,\n                        1,\n                        list(range(num_devices)),\n                        agg_small_grads_max_bytes=max_bytes))\n            else:\n                self.packed_grads_and_vars, _ = (sum_gradients_all_reduce(\n                    \"\",\n                    grad_ops,\n                    1,\n                    all_reduce_alg,\n                    1,\n                    list(range(num_devices)),\n                    agg_small_grads_max_bytes=0))\n        self.per_device_grads = [\n            list(zip(*dev_gv))[0] for dev_gv in self.packed_grads_and_vars\n        ]\n        assert (len(self.per_device_grads) == num_devices)\n        self.num_grads = num_grads = len(self.packed_grads_and_vars[0])\n        if max_bytes:\n            logger.info(\"Packed grads => {} tensors\".format(num_grads))\n\n        # Ops for reading grads with the right control deps\n        nccl_noops = []\n        for j in range(num_grads)[::-1]:\n            deps = nccl_noops + [\n                dev_grad[j] for dev_grad in self.per_device_grads\n            ]\n            with tf.control_dependencies(deps):\n                nccl_noops = [tf.no_op()]\n\n        # You must fetch this otherwise the NCCL allreduce will hang\n        self.nccl_control_out = tf.group(*nccl_noops)\n\n        if plasma_op:\n            store_socket = (\n                ray.worker.global_worker.plasma_client.store_socket_name)\n            manager_socket = (\n                ray.worker.global_worker.plasma_client.manager_socket_name)\n            if not plasma.tf_plasma_op:\n                plasma.build_plasma_tensorflow_op()\n\n            # For fetching grads -> plasma\n            self.plasma_in_grads = []\n            self.plasma_in_grads_oids = [\n                tf.placeholder(shape=[], dtype=tf.string, name=\"in_grad_oids\")\n                for _ in range(num_grads)\n            ]\n            for j in range(num_grads):\n                grad = self.per_device_grads[0][j]\n                with tf.device(self.models[0].loss.device):\n                    plasma_grad = plasma.tf_plasma_op.tensor_to_plasma(\n                        [grad],\n                        self.plasma_in_grads_oids[j],\n                        plasma_store_socket_name=store_socket,\n                        plasma_manager_socket_name=manager_socket)\n                self.plasma_in_grads.append(plasma_grad)\n\n            # For applying grads <- plasma\n            unpacked_gv = []\n            self.plasma_out_grads_oids = [\n                tf.placeholder(\n                    shape=[], dtype=tf.string, name=\"grad_out_oids\")\n                for _ in range(num_grads)\n            ]\n            packed_plasma_grads = []\n            for j in range(num_grads):\n                with tf.device(self.plasma_in_grads[j].device):\n                    with tf.control_dependencies([self.plasma_in_grads[j]]):\n                        grad_ph = plasma.tf_plasma_op.plasma_to_tensor(\n                            self.plasma_out_grads_oids[j],\n                            dtype=tf.float32,\n                            plasma_store_socket_name=store_socket,\n                            plasma_manager_socket_name=manager_socket)\n                grad_ph = tf.reshape(grad_ph,\n                                     self.packed_grads_and_vars[0][j][0].shape)\n                logger.debug(\"Packed tensor {}\".format(grad_ph))\n                packed_plasma_grads.append(grad_ph)\n            for i in range(num_devices):\n                per_device = []\n                for j, (g, v) in enumerate(self.packed_grads_and_vars[i]):\n                    grad_ph = packed_plasma_grads[j]\n                    per_device.append((grad_ph, v))\n                unpacked_gv.append(per_device)\n\n            if max_bytes:\n                unpacked_gv = unpack_small_tensors(unpacked_gv, packing_vals)\n\n        elif max_bytes:\n            unpacked_gv = unpack_small_tensors(self.packed_grads_and_vars,\n                                               packing_vals)\n        else:\n            unpacked_gv = self.packed_grads_and_vars\n\n        # Same shape as packed_grads_and_vars\n        assert len(unpacked_gv) == num_devices\n        assert len(unpacked_gv[0][0]) == 2\n\n        apply_ops = []\n        to_apply = unpacked_gv[0]\n        for ix, m in enumerate(self.models):\n            apply_ops.append(\n                m.optimizer.apply_gradients(\n                    [(g, v)\n                     for ((g, _), (_, v)) in zip(to_apply, unpacked_gv[ix])]))\n        self.apply_op = tf.group(*apply_ops)\n        init_op = tf.group(tf.global_variables_initializer(),\n                           tf.local_variables_initializer())\n        self.sess.run(init_op)\n\n    def _grad_feed_dict(self):\n        # Aggregate feed dicts for each model on this worker.\n        feed_dict = {}\n        for model in self.models:\n            feed_dict.update(model.get_feed_dict())\n        return feed_dict\n\n    def foreach_model(self, fn):\n        with self.sess.as_default():\n            return [fn(m) for m in self.models]\n\n    def foreach_worker(self, fn):\n        with self.sess.as_default():\n            return fn(self)\n\n    def for_model(self, fn):\n        with self.sess.as_default():\n            return fn(self.models[0])\n\n    def compute_gradients(self):\n        start = time.time()\n        feed_dict = self._grad_feed_dict()\n        # We only need to fetch the first per_device_grad, since they are\n        # averaged across all devices by allreduce.\n        fetches = self.sess.run(\n            [\n                self.models[0].loss, self.per_device_grads[0],\n                self.nccl_control_out\n            ],\n            feed_dict=feed_dict)\n        logger.debug(\n            \"Compute grad interior time {}\".format(time.time() - start))\n        return fetches\n\n    def apply_gradients(self, avg_grads):\n        start = time.time()\n        result = {\n            g: avg_grads[i]\n            for (i, g) in enumerate(self.per_device_grads[0])\n        }\n        self.sess.run(self.apply_op, feed_dict=result)\n        logger.debug(\"Apply grad interior time {}\".format(time.time() - start))\n\n    def compute_apply(self):\n        fetches = run_timeline(\n            self.sess,\n            [self.models[0].loss, self.apply_op, self.nccl_control_out],\n            feed_dict=self._grad_feed_dict(),\n            name=\"compute_apply\")\n        return fetches[0]\n\n    def ps_compute_apply(self,\n                         out_grad_shard_oids,\n                         agg_grad_shard_oids,\n                         tl_name=\"ps_compute_apply\",\n                         write_timeline=False):\n        feed_dict = self._grad_feed_dict()\n        feed_dict.update(\n            dict(zip(self.plasma_in_grads_oids, out_grad_shard_oids)))\n        feed_dict.update(\n            dict(zip(self.plasma_out_grads_oids, agg_grad_shard_oids)))\n        fetch(agg_grad_shard_oids)\n        fetches = run_timeline(\n            self.sess, [\n                self.models[0].loss, self.plasma_in_grads, self.apply_op,\n                self.nccl_control_out\n            ],\n            feed_dict=feed_dict,\n            write_timeline=write_timeline)\n        return fetches[0]\n\n    def num_grad_shards(self):\n        return self.num_grads\n\n    def shard_shapes(self):\n        main_gv = self.packed_grads_and_vars[0]\n        return [g.shape for g, _ in main_gv]\n\n    def ip(self):\n        return ray.services.get_node_ip_address()\n\n    def warmup(self):\n        warmup()\n", "path": "python/ray/experimental/sgd/sgd_worker.py"}]}