We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
worker_test failure
```
python -m unittest elasticdl/worker/*_test.py
/usr/local/anaconda3/lib/python3.6/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.
  from ._conv import register_converters as _register_converters
2019-05-21 13:57:07.262725: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
WARNING:tensorflow:From /usr/local/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/resource_variable_ops.py:642: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.
Instructions for updating:
Colocations handled automatically by placer.
must be str, not NoneType
FLoss is  3.090042
Loss is  1.4608976
Loss is  0.913306
Loss is  0.5969497
Loss is  0.66515267
Loss is  0.3935135
Loss is  0.37774342
Loss is  0.289928
.
======================================================================
FAIL: test_distributed_train (elasticdl.worker.worker_test.WorkerTest)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/Users/l.zou/git/elasticdl/elasticdl/worker/worker_test.py", line 96, in test_distributed_train
    self.assertTrue(res)
AssertionError: False is not true

----------------------------------------------------------------------
Ran 2 tests in 0.165s

```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `elasticdl/worker/worker.py`
Content:
```
1 import traceback
2 import tensorflow as tf
3 assert tf.executing_eagerly()
4 
5 from google.protobuf import empty_pb2
6 from tensorflow.python.ops import math_ops
7 from elasticdl.proto import master_pb2_grpc
8 from elasticdl.proto import master_pb2
9 from elasticdl.common.ndarray import ndarray_to_tensor, tensor_to_ndarray
10 from elasticdl.common.model_helper import load_user_model, build_model
11 from edl_data.codec import TFExampleCodec
12 from edl_data.codec import BytesCodec
13 import itertools
14 import recordio
15 
16 # the default max number of a minibatch retrain as its gradients are not accepted by master.
17 DEFAULT_MAX_MINIBATCH_RETRAIN_NUM = 64
18 
19 class Worker(object):
20     """ElasticDL worker"""
21 
22     def __init__(self,
23                  model_file,
24                  channel=None,
25                  max_retrain_num=DEFAULT_MAX_MINIBATCH_RETRAIN_NUM,
26                  codec_type=None):
27         """
28         Arguments:
29             model_module: A module to define the model
30             channel: grpc channel
31             max_retrain_num: max number of a minibatch retrain as its gradients are not accepted by master
32         """
33 
34         model_module = load_user_model(model_file)
35         self._model = model_module.model
36         self._feature_columns = model_module.feature_columns()
37         self._all_columns = self._feature_columns + model_module.label_columns()
38         build_model(self._model, self._feature_columns)
39         self._input_fn = model_module.input_fn 
40         self._opt_fn = model_module.optimizer
41         self._loss = model_module.loss
42 
43         if channel is None:
44             self._stub = None
45         else:
46             self._stub = master_pb2_grpc.MasterStub(channel)
47         self._max_retrain_num = max_retrain_num
48         self._model_version = -1
49         self._codec_type = codec_type
50 
51     def get_task(self):
52         """
53         get task from master
54         """
55         return self._stub.GetTask(empty_pb2.Empty())
56 
57     def get_model(self, min_version):
58         """
59         get model from master, and update model_version
60         """
61         req = master_pb2.GetModelRequest()
62         req.min_version = min_version
63         model = self._stub.GetModel(req)
64 
65         for var in self._model.trainable_variables:
66             # Assumes all trainable variables exist in model.param.
67             var.assign(
68                 tensor_to_ndarray(model.param[var.name]))
69         self._model_version = model.version
70 
71     def report_task_result(self, task_id, err_msg):
72         """
73         report task result to master
74         """
75         report = master_pb2.ReportTaskResultRequest()
76         report.task_id = task_id
77         report.err_message = err_msg
78         return self._stub.ReportTaskResult(report)
79 
80     def report_gradient(self, grads):
81         """
82         report gradient to ps, return (accepted, model_version) from rpc call.
83         """
84         req = master_pb2.ReportGradientRequest()
85         for g, v in zip(grads, self._model.trainable_variables):
86             req.gradient[v.name].CopyFrom(
87                 ndarray_to_tensor(g.numpy()))
88         req.model_version = self._model_version
89         res = self._stub.ReportGradient(req)
90         return res.accepted, res.model_version
91 
92     def distributed_train(self):
93         """
94         Distributed training.
95         """
96         if self._codec_type == "tf_example":
97             codec = TFExampleCodec(self._all_columns)
98         elif self._codec_type == "bytes":
99             codec = BytesCodec(self._all_columns)
100         else:
101             raise ValueError("invalid codec_type: " + self._codec_type)
102         while True:
103             task = self.get_task()
104             if not task.shard_file_name:
105                 # No more task
106                 break
107             batch_size = task.minibatch_size
108             err_msg = ""
109             try:
110                 with recordio.File(task.shard_file_name, "r", decoder=codec.decode) as rdio_r:
111                     reader = rdio_r.get_reader(task.start, task.end)
112                     min_model_version = task.model_version
113                     while True:
114                         record_buf = list(
115                             itertools.islice(reader, 0, batch_size))
116                         if not record_buf:
117                             break
118 
119                         for _ in range(self._max_retrain_num):
120                             # TODO: optimize the logic to avoid unnecessary get_model call.
121                             self.get_model(
122                                 max(self._model_version, min_model_version))
123 
124                             batch_input_data, batch_label = self._input_fn(record_buf)
125 
126                             with tf.GradientTape() as tape:
127                                 inputs = []
128                                 for f_col in self._feature_columns:
129                                     inputs.append(batch_input_data[f_col.key])
130                                 if len(inputs) == 1:
131                                     inputs = inputs[0]
132                                 outputs = self._model.call(inputs, training=True)
133                                 loss = self._loss(outputs, batch_label.flatten())
134 
135                                 # TODO:  Add regularization loss if any,
136                                 #        which should be divided by the number of contributing workers.
137                             grads = tape.gradient(
138                                 loss, self._model.trainable_variables)
139                             print("Loss is ", loss.numpy())
140 
141                             accepted, min_model_version = self.report_gradient(
142                                 grads)
143                             if accepted:
144                                 break
145                         else:
146                             # Worker got stuck, fail the task.
147                             # TODO: stop the worker if it fails to make any progress for some time.
148                             raise RuntimeError("Worker got stuck")
149 
150 
151             except Exception as ex:
152                 err_msg = str(ex)
153                 traceback.print_exc()
154             self.report_task_result(task.task_id, err_msg)
155 
156     def local_train(self, file_list, batch_size, epoch=1, kwargs=None):
157         """
158         Local training for local testing. Must in eager mode.
159         Argments:
160             batch_size: batch size in training
161             epoch: the number of epoch in training
162             kwargs: contains a dict of parameters used in training
163         """
164         optimizer = self._opt_fn()
165         for _ in range(epoch):
166             for f in file_list:
167                 with recordio.File(f, "r") as rdio_r:
168                     reader = rdio_r.get_reader(0, rdio_r.count())
169                     while True:
170                         record_buf = list(
171                             itertools.islice(reader, 0, batch_size))
172                         if not record_buf:
173                             break
174 
175                         data, labels = self._input_fn(record_buf)
176 
177                         with tf.GradientTape() as tape:
178                             inputs = []
179                             for f_col in self._feature_columns:
180                                 inputs.append(data[f_col.key])
181                             if len(inputs) == 1:
182                                 inputs = inputs[0]
183                             outputs = self._model.call(inputs, training=True)
184                             loss = self._loss(outputs, labels)
185 
186                             # Add regularization loss if any.
187                             # Note: for distributed training, the regularization loss should
188                             #       be divided by the number of contributing workers, which
189                             #       might be difficult for elasticdl.
190                             if self._model.losses:
191                                 loss += math_ops.add_n(self._model.losses)
192                         grads = tape.gradient(
193                             loss, self._model.trainable_variables)
194                         optimizer.apply_gradients(
195                             zip(grads, self._model.trainable_variables))
196                         print("Loss is ", loss.numpy())
197
```
--- 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/elasticdl/worker/worker.py b/elasticdl/worker/worker.py
--- a/elasticdl/worker/worker.py
+++ b/elasticdl/worker/worker.py
@@ -30,15 +30,21 @@
             channel: grpc channel
             max_retrain_num: max number of a minibatch retrain as its gradients are not accepted by master
         """
-
         model_module = load_user_model(model_file)
         self._model = model_module.model
         self._feature_columns = model_module.feature_columns()
-        self._all_columns = self._feature_columns + model_module.label_columns()
         build_model(self._model, self._feature_columns)
         self._input_fn = model_module.input_fn 
         self._opt_fn = model_module.optimizer
         self._loss = model_module.loss
+        all_columns = self._feature_columns + model_module.label_columns()
+        if codec_type == "tf_example":
+            self._codec = TFExampleCodec(all_columns)
+        elif codec_type == "bytes":
+            self._codec = BytesCodec(all_columns)
+        else:
+            raise ValueError("invalid codec_type: " + codec_type)
+
 
         if channel is None:
             self._stub = None
@@ -93,12 +99,6 @@
         """
         Distributed training.
         """
-        if self._codec_type == "tf_example":
-            codec = TFExampleCodec(self._all_columns)
-        elif self._codec_type == "bytes":
-            codec = BytesCodec(self._all_columns)
-        else:
-            raise ValueError("invalid codec_type: " + self._codec_type)
         while True:
             task = self.get_task()
             if not task.shard_file_name:
@@ -107,7 +107,7 @@
             batch_size = task.minibatch_size
             err_msg = ""
             try:
-                with recordio.File(task.shard_file_name, "r", decoder=codec.decode) as rdio_r:
+                with recordio.File(task.shard_file_name, "r", decoder=self._codec.decode) as rdio_r:
                     reader = rdio_r.get_reader(task.start, task.end)
                     min_model_version = task.model_version
                     while True:
@@ -164,7 +164,7 @@
         optimizer = self._opt_fn()
         for _ in range(epoch):
             for f in file_list:
-                with recordio.File(f, "r") as rdio_r:
+                with recordio.File(f, "r", decoder=self._codec.decode) as rdio_r:
                     reader = rdio_r.get_reader(0, rdio_r.count())
                     while True:
                         record_buf = list(

{"golden_diff": "diff --git a/elasticdl/worker/worker.py b/elasticdl/worker/worker.py\n--- a/elasticdl/worker/worker.py\n+++ b/elasticdl/worker/worker.py\n@@ -30,15 +30,21 @@\n             channel: grpc channel\n             max_retrain_num: max number of a minibatch retrain as its gradients are not accepted by master\n         \"\"\"\n-\n         model_module = load_user_model(model_file)\n         self._model = model_module.model\n         self._feature_columns = model_module.feature_columns()\n-        self._all_columns = self._feature_columns + model_module.label_columns()\n         build_model(self._model, self._feature_columns)\n         self._input_fn = model_module.input_fn \n         self._opt_fn = model_module.optimizer\n         self._loss = model_module.loss\n+        all_columns = self._feature_columns + model_module.label_columns()\n+        if codec_type == \"tf_example\":\n+            self._codec = TFExampleCodec(all_columns)\n+        elif codec_type == \"bytes\":\n+            self._codec = BytesCodec(all_columns)\n+        else:\n+            raise ValueError(\"invalid codec_type: \" + codec_type)\n+\n \n         if channel is None:\n             self._stub = None\n@@ -93,12 +99,6 @@\n         \"\"\"\n         Distributed training.\n         \"\"\"\n-        if self._codec_type == \"tf_example\":\n-            codec = TFExampleCodec(self._all_columns)\n-        elif self._codec_type == \"bytes\":\n-            codec = BytesCodec(self._all_columns)\n-        else:\n-            raise ValueError(\"invalid codec_type: \" + self._codec_type)\n         while True:\n             task = self.get_task()\n             if not task.shard_file_name:\n@@ -107,7 +107,7 @@\n             batch_size = task.minibatch_size\n             err_msg = \"\"\n             try:\n-                with recordio.File(task.shard_file_name, \"r\", decoder=codec.decode) as rdio_r:\n+                with recordio.File(task.shard_file_name, \"r\", decoder=self._codec.decode) as rdio_r:\n                     reader = rdio_r.get_reader(task.start, task.end)\n                     min_model_version = task.model_version\n                     while True:\n@@ -164,7 +164,7 @@\n         optimizer = self._opt_fn()\n         for _ in range(epoch):\n             for f in file_list:\n-                with recordio.File(f, \"r\") as rdio_r:\n+                with recordio.File(f, \"r\", decoder=self._codec.decode) as rdio_r:\n                     reader = rdio_r.get_reader(0, rdio_r.count())\n                     while True:\n                         record_buf = list(\n", "issue": "worker_test failure\n```\r\npython -m unittest elasticdl/worker/*_test.py\r\n/usr/local/anaconda3/lib/python3.6/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\r\n  from ._conv import register_converters as _register_converters\r\n2019-05-21 13:57:07.262725: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA\r\nWARNING:tensorflow:From /usr/local/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/resource_variable_ops.py:642: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\r\nInstructions for updating:\r\nColocations handled automatically by placer.\r\nmust be str, not NoneType\r\nFLoss is  3.090042\r\nLoss is  1.4608976\r\nLoss is  0.913306\r\nLoss is  0.5969497\r\nLoss is  0.66515267\r\nLoss is  0.3935135\r\nLoss is  0.37774342\r\nLoss is  0.289928\r\n.\r\n======================================================================\r\nFAIL: test_distributed_train (elasticdl.worker.worker_test.WorkerTest)\r\n----------------------------------------------------------------------\r\nTraceback (most recent call last):\r\n  File \"/Users/l.zou/git/elasticdl/elasticdl/worker/worker_test.py\", line 96, in test_distributed_train\r\n    self.assertTrue(res)\r\nAssertionError: False is not true\r\n\r\n----------------------------------------------------------------------\r\nRan 2 tests in 0.165s\r\n\r\n```\n", "before_files": [{"content": "import traceback\nimport tensorflow as tf\nassert tf.executing_eagerly()\n\nfrom google.protobuf import empty_pb2\nfrom tensorflow.python.ops import math_ops\nfrom elasticdl.proto import master_pb2_grpc\nfrom elasticdl.proto import master_pb2\nfrom elasticdl.common.ndarray import ndarray_to_tensor, tensor_to_ndarray\nfrom elasticdl.common.model_helper import load_user_model, build_model\nfrom edl_data.codec import TFExampleCodec\nfrom edl_data.codec import BytesCodec\nimport itertools\nimport recordio\n\n# the default max number of a minibatch retrain as its gradients are not accepted by master.\nDEFAULT_MAX_MINIBATCH_RETRAIN_NUM = 64\n\nclass Worker(object):\n    \"\"\"ElasticDL worker\"\"\"\n\n    def __init__(self,\n                 model_file,\n                 channel=None,\n                 max_retrain_num=DEFAULT_MAX_MINIBATCH_RETRAIN_NUM,\n                 codec_type=None):\n        \"\"\"\n        Arguments:\n            model_module: A module to define the model\n            channel: grpc channel\n            max_retrain_num: max number of a minibatch retrain as its gradients are not accepted by master\n        \"\"\"\n\n        model_module = load_user_model(model_file)\n        self._model = model_module.model\n        self._feature_columns = model_module.feature_columns()\n        self._all_columns = self._feature_columns + model_module.label_columns()\n        build_model(self._model, self._feature_columns)\n        self._input_fn = model_module.input_fn \n        self._opt_fn = model_module.optimizer\n        self._loss = model_module.loss\n\n        if channel is None:\n            self._stub = None\n        else:\n            self._stub = master_pb2_grpc.MasterStub(channel)\n        self._max_retrain_num = max_retrain_num\n        self._model_version = -1\n        self._codec_type = codec_type\n\n    def get_task(self):\n        \"\"\"\n        get task from master\n        \"\"\"\n        return self._stub.GetTask(empty_pb2.Empty())\n\n    def get_model(self, min_version):\n        \"\"\"\n        get model from master, and update model_version\n        \"\"\"\n        req = master_pb2.GetModelRequest()\n        req.min_version = min_version\n        model = self._stub.GetModel(req)\n\n        for var in self._model.trainable_variables:\n            # Assumes all trainable variables exist in model.param.\n            var.assign(\n                tensor_to_ndarray(model.param[var.name]))\n        self._model_version = model.version\n\n    def report_task_result(self, task_id, err_msg):\n        \"\"\"\n        report task result to master\n        \"\"\"\n        report = master_pb2.ReportTaskResultRequest()\n        report.task_id = task_id\n        report.err_message = err_msg\n        return self._stub.ReportTaskResult(report)\n\n    def report_gradient(self, grads):\n        \"\"\"\n        report gradient to ps, return (accepted, model_version) from rpc call.\n        \"\"\"\n        req = master_pb2.ReportGradientRequest()\n        for g, v in zip(grads, self._model.trainable_variables):\n            req.gradient[v.name].CopyFrom(\n                ndarray_to_tensor(g.numpy()))\n        req.model_version = self._model_version\n        res = self._stub.ReportGradient(req)\n        return res.accepted, res.model_version\n\n    def distributed_train(self):\n        \"\"\"\n        Distributed training.\n        \"\"\"\n        if self._codec_type == \"tf_example\":\n            codec = TFExampleCodec(self._all_columns)\n        elif self._codec_type == \"bytes\":\n            codec = BytesCodec(self._all_columns)\n        else:\n            raise ValueError(\"invalid codec_type: \" + self._codec_type)\n        while True:\n            task = self.get_task()\n            if not task.shard_file_name:\n                # No more task\n                break\n            batch_size = task.minibatch_size\n            err_msg = \"\"\n            try:\n                with recordio.File(task.shard_file_name, \"r\", decoder=codec.decode) as rdio_r:\n                    reader = rdio_r.get_reader(task.start, task.end)\n                    min_model_version = task.model_version\n                    while True:\n                        record_buf = list(\n                            itertools.islice(reader, 0, batch_size))\n                        if not record_buf:\n                            break\n\n                        for _ in range(self._max_retrain_num):\n                            # TODO: optimize the logic to avoid unnecessary get_model call.\n                            self.get_model(\n                                max(self._model_version, min_model_version))\n\n                            batch_input_data, batch_label = self._input_fn(record_buf)\n\n                            with tf.GradientTape() as tape:\n                                inputs = []\n                                for f_col in self._feature_columns:\n                                    inputs.append(batch_input_data[f_col.key])\n                                if len(inputs) == 1:\n                                    inputs = inputs[0]\n                                outputs = self._model.call(inputs, training=True)\n                                loss = self._loss(outputs, batch_label.flatten())\n\n                                # TODO:  Add regularization loss if any,\n                                #        which should be divided by the number of contributing workers.\n                            grads = tape.gradient(\n                                loss, self._model.trainable_variables)\n                            print(\"Loss is \", loss.numpy())\n\n                            accepted, min_model_version = self.report_gradient(\n                                grads)\n                            if accepted:\n                                break\n                        else:\n                            # Worker got stuck, fail the task.\n                            # TODO: stop the worker if it fails to make any progress for some time.\n                            raise RuntimeError(\"Worker got stuck\")\n\n\n            except Exception as ex:\n                err_msg = str(ex)\n                traceback.print_exc()\n            self.report_task_result(task.task_id, err_msg)\n\n    def local_train(self, file_list, batch_size, epoch=1, kwargs=None):\n        \"\"\"\n        Local training for local testing. Must in eager mode.\n        Argments:\n            batch_size: batch size in training\n            epoch: the number of epoch in training\n            kwargs: contains a dict of parameters used in training\n        \"\"\"\n        optimizer = self._opt_fn()\n        for _ in range(epoch):\n            for f in file_list:\n                with recordio.File(f, \"r\") as rdio_r:\n                    reader = rdio_r.get_reader(0, rdio_r.count())\n                    while True:\n                        record_buf = list(\n                            itertools.islice(reader, 0, batch_size))\n                        if not record_buf:\n                            break\n\n                        data, labels = self._input_fn(record_buf)\n\n                        with tf.GradientTape() as tape:\n                            inputs = []\n                            for f_col in self._feature_columns:\n                                inputs.append(data[f_col.key])\n                            if len(inputs) == 1:\n                                inputs = inputs[0]\n                            outputs = self._model.call(inputs, training=True)\n                            loss = self._loss(outputs, labels)\n\n                            # Add regularization loss if any.\n                            # Note: for distributed training, the regularization loss should\n                            #       be divided by the number of contributing workers, which\n                            #       might be difficult for elasticdl.\n                            if self._model.losses:\n                                loss += math_ops.add_n(self._model.losses)\n                        grads = tape.gradient(\n                            loss, self._model.trainable_variables)\n                        optimizer.apply_gradients(\n                            zip(grads, self._model.trainable_variables))\n                        print(\"Loss is \", loss.numpy())\n", "path": "elasticdl/worker/worker.py"}], "after_files": [{"content": "import traceback\nimport tensorflow as tf\nassert tf.executing_eagerly()\n\nfrom google.protobuf import empty_pb2\nfrom tensorflow.python.ops import math_ops\nfrom elasticdl.proto import master_pb2_grpc\nfrom elasticdl.proto import master_pb2\nfrom elasticdl.common.ndarray import ndarray_to_tensor, tensor_to_ndarray\nfrom elasticdl.common.model_helper import load_user_model, build_model\nfrom edl_data.codec import TFExampleCodec\nfrom edl_data.codec import BytesCodec\nimport itertools\nimport recordio\n\n# the default max number of a minibatch retrain as its gradients are not accepted by master.\nDEFAULT_MAX_MINIBATCH_RETRAIN_NUM = 64\n\nclass Worker(object):\n    \"\"\"ElasticDL worker\"\"\"\n\n    def __init__(self,\n                 model_file,\n                 channel=None,\n                 max_retrain_num=DEFAULT_MAX_MINIBATCH_RETRAIN_NUM,\n                 codec_type=None):\n        \"\"\"\n        Arguments:\n            model_module: A module to define the model\n            channel: grpc channel\n            max_retrain_num: max number of a minibatch retrain as its gradients are not accepted by master\n        \"\"\"\n        model_module = load_user_model(model_file)\n        self._model = model_module.model\n        self._feature_columns = model_module.feature_columns()\n        build_model(self._model, self._feature_columns)\n        self._input_fn = model_module.input_fn \n        self._opt_fn = model_module.optimizer\n        self._loss = model_module.loss\n        all_columns = self._feature_columns + model_module.label_columns()\n        if codec_type == \"tf_example\":\n            self._codec = TFExampleCodec(all_columns)\n        elif codec_type == \"bytes\":\n            self._codec = BytesCodec(all_columns)\n        else:\n            raise ValueError(\"invalid codec_type: \" + codec_type)\n\n\n        if channel is None:\n            self._stub = None\n        else:\n            self._stub = master_pb2_grpc.MasterStub(channel)\n        self._max_retrain_num = max_retrain_num\n        self._model_version = -1\n        self._codec_type = codec_type\n\n    def get_task(self):\n        \"\"\"\n        get task from master\n        \"\"\"\n        return self._stub.GetTask(empty_pb2.Empty())\n\n    def get_model(self, min_version):\n        \"\"\"\n        get model from master, and update model_version\n        \"\"\"\n        req = master_pb2.GetModelRequest()\n        req.min_version = min_version\n        model = self._stub.GetModel(req)\n\n        for var in self._model.trainable_variables:\n            # Assumes all trainable variables exist in model.param.\n            var.assign(\n                tensor_to_ndarray(model.param[var.name]))\n        self._model_version = model.version\n\n    def report_task_result(self, task_id, err_msg):\n        \"\"\"\n        report task result to master\n        \"\"\"\n        report = master_pb2.ReportTaskResultRequest()\n        report.task_id = task_id\n        report.err_message = err_msg\n        return self._stub.ReportTaskResult(report)\n\n    def report_gradient(self, grads):\n        \"\"\"\n        report gradient to ps, return (accepted, model_version) from rpc call.\n        \"\"\"\n        req = master_pb2.ReportGradientRequest()\n        for g, v in zip(grads, self._model.trainable_variables):\n            req.gradient[v.name].CopyFrom(\n                ndarray_to_tensor(g.numpy()))\n        req.model_version = self._model_version\n        res = self._stub.ReportGradient(req)\n        return res.accepted, res.model_version\n\n    def distributed_train(self):\n        \"\"\"\n        Distributed training.\n        \"\"\"\n        while True:\n            task = self.get_task()\n            if not task.shard_file_name:\n                # No more task\n                break\n            batch_size = task.minibatch_size\n            err_msg = \"\"\n            try:\n                with recordio.File(task.shard_file_name, \"r\", decoder=self._codec.decode) as rdio_r:\n                    reader = rdio_r.get_reader(task.start, task.end)\n                    min_model_version = task.model_version\n                    while True:\n                        record_buf = list(\n                            itertools.islice(reader, 0, batch_size))\n                        if not record_buf:\n                            break\n\n                        for _ in range(self._max_retrain_num):\n                            # TODO: optimize the logic to avoid unnecessary get_model call.\n                            self.get_model(\n                                max(self._model_version, min_model_version))\n\n                            batch_input_data, batch_label = self._input_fn(record_buf)\n\n                            with tf.GradientTape() as tape:\n                                inputs = []\n                                for f_col in self._feature_columns:\n                                    inputs.append(batch_input_data[f_col.key])\n                                if len(inputs) == 1:\n                                    inputs = inputs[0]\n                                outputs = self._model.call(inputs, training=True)\n                                loss = self._loss(outputs, batch_label.flatten())\n\n                                # TODO:  Add regularization loss if any,\n                                #        which should be divided by the number of contributing workers.\n                            grads = tape.gradient(\n                                loss, self._model.trainable_variables)\n                            print(\"Loss is \", loss.numpy())\n\n                            accepted, min_model_version = self.report_gradient(\n                                grads)\n                            if accepted:\n                                break\n                        else:\n                            # Worker got stuck, fail the task.\n                            # TODO: stop the worker if it fails to make any progress for some time.\n                            raise RuntimeError(\"Worker got stuck\")\n\n\n            except Exception as ex:\n                err_msg = str(ex)\n                traceback.print_exc()\n            self.report_task_result(task.task_id, err_msg)\n\n    def local_train(self, file_list, batch_size, epoch=1, kwargs=None):\n        \"\"\"\n        Local training for local testing. Must in eager mode.\n        Argments:\n            batch_size: batch size in training\n            epoch: the number of epoch in training\n            kwargs: contains a dict of parameters used in training\n        \"\"\"\n        optimizer = self._opt_fn()\n        for _ in range(epoch):\n            for f in file_list:\n                with recordio.File(f, \"r\", decoder=self._codec.decode) as rdio_r:\n                    reader = rdio_r.get_reader(0, rdio_r.count())\n                    while True:\n                        record_buf = list(\n                            itertools.islice(reader, 0, batch_size))\n                        if not record_buf:\n                            break\n\n                        data, labels = self._input_fn(record_buf)\n\n                        with tf.GradientTape() as tape:\n                            inputs = []\n                            for f_col in self._feature_columns:\n                                inputs.append(data[f_col.key])\n                            if len(inputs) == 1:\n                                inputs = inputs[0]\n                            outputs = self._model.call(inputs, training=True)\n                            loss = self._loss(outputs, labels)\n\n                            # Add regularization loss if any.\n                            # Note: for distributed training, the regularization loss should\n                            #       be divided by the number of contributing workers, which\n                            #       might be difficult for elasticdl.\n                            if self._model.losses:\n                                loss += math_ops.add_n(self._model.losses)\n                        grads = tape.gradient(\n                            loss, self._model.trainable_variables)\n                        optimizer.apply_gradients(\n                            zip(grads, self._model.trainable_variables))\n                        print(\"Loss is \", loss.numpy())\n", "path": "elasticdl/worker/worker.py"}]}