We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Add CLI entry-point to `Worker.distributed_evaluate()`
Part of https://github.com/wangkuiyi/elasticdl/issues/384. Add main CLI entry-point to `Worker.distributed_evaluate()` that will be used in `WorkerManager`.

--- 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/python/elasticdl/worker/worker.py`
Content:
```
1 import logging
2 import traceback
3 import time
4 
5 import tensorflow as tf
6 assert tf.executing_eagerly()
7 
8 import recordio
9 import numpy as np
10 
11 from contextlib import closing
12 from collections import defaultdict
13 from tensorflow.python.ops import math_ops
14 from elasticdl.proto import elasticdl_pb2_grpc
15 from elasticdl.proto import elasticdl_pb2
16 from elasticdl.python.elasticdl.common.ndarray import ndarray_to_tensor, tensor_to_ndarray
17 from elasticdl.python.elasticdl.common.model_helper import load_user_model, build_model
18 from elasticdl.python.data.codec import TFExampleCodec
19 from elasticdl.python.data.codec import BytesCodec
20 
21 # TODO: Move these args to train/evaluate specific functions
22 # the default max number of a minibatch retrain as its gradients are not accepted by master.
23 DEFAULT_MAX_MINIBATCH_RETRAIN_NUM = 64
24 # the default max number of a minibatch retry as its calculated evaluation metrics are not accepted by master.
25 DEFAULT_MAX_MINIBATCH_EVALUATE_RETRY_NUM = 10
26 
27 
28 class Worker(object):
29     """ElasticDL worker"""
30 
31     def __init__(self,
32                  worker_id,
33                  model_file,
34                  channel=None,
35                  max_retrain_num=DEFAULT_MAX_MINIBATCH_RETRAIN_NUM,
36                  max_evaluate_retry_num=DEFAULT_MAX_MINIBATCH_EVALUATE_RETRY_NUM,
37                  codec_type=None):
38         """
39         Arguments:
40             model_module: A module to define the model
41             channel: grpc channel
42             max_retrain_num: max number of a minibatch retrain as its gradients are not accepted by master
43         """
44         self._logger = logging.getLogger(__name__)
45         self._worker_id = worker_id
46         model_module = load_user_model(model_file)
47         self._model = model_module.model
48         self._feature_columns = model_module.feature_columns()
49         build_model(self._model, self._feature_columns)
50         self._input_fn = model_module.input_fn 
51         self._opt_fn = model_module.optimizer
52         self._loss = model_module.loss
53         self._eval_metrics_fn = model_module.eval_metrics_fn
54         all_columns = self._feature_columns + model_module.label_columns()
55         if codec_type == "tf_example":
56             self._codec = TFExampleCodec(all_columns)
57         elif codec_type == "bytes":
58             self._codec = BytesCodec(all_columns)
59         else:
60             raise ValueError("invalid codec_type: " + codec_type)
61 
62         if channel is None:
63             self._stub = None
64         else:
65             self._stub = elasticdl_pb2_grpc.MasterStub(channel)
66         self._max_retrain_num = max_retrain_num
67         self._max_evaluate_retry_num = max_evaluate_retry_num
68         self._model_version = -1
69         self._codec_type = codec_type
70 
71     def get_task(self):
72         """
73         get task from master
74         """
75         req = elasticdl_pb2.GetTaskRequest()
76         req.worker_id = self._worker_id
77         
78         return self._stub.GetTask(req)
79 
80     def get_model(self, min_version):
81         """
82         get model from master, and update model_version
83         """
84         req = elasticdl_pb2.GetModelRequest()
85         req.min_version = min_version
86         model = self._stub.GetModel(req)
87 
88         for var in self._model.trainable_variables:
89             # Assumes all trainable variables exist in model.param.
90             var.assign(
91                 tensor_to_ndarray(model.param[var.name]))
92         self._model_version = model.version
93 
94     def report_task_result(self, task_id, err_msg):
95         """
96         report task result to master
97         """
98         report = elasticdl_pb2.ReportTaskResultRequest()
99         report.task_id = task_id
100         report.err_message = err_msg
101         return self._stub.ReportTaskResult(report)
102 
103     def report_gradient(self, grads):
104         """
105         report gradient to ps, return (accepted, model_version) from rpc call.
106         """
107         req = elasticdl_pb2.ReportGradientRequest()
108         for g, v in zip(grads, self._model.trainable_variables):
109             req.gradient[v.name].CopyFrom(
110                 ndarray_to_tensor(g.numpy()))
111         req.model_version = self._model_version
112         res = self._stub.ReportGradient(req)
113         return res.accepted, res.model_version
114 
115     def report_evaluation_metrics(self, evaluation_metrics):
116         """
117         report evaluation metrics to ps, return (accepted, model_version) from rpc call.
118         """
119         req = elasticdl_pb2.ReportEvaluationMetricsRequest()
120         for k, v in evaluation_metrics.items():
121             req.evaluation_metrics[k].CopyFrom(
122                 ndarray_to_tensor(v))
123         req.model_version = self._model_version
124         res = self._stub.ReportEvaluationMetrics(req)
125         return res.accepted, res.model_version
126 
127     @staticmethod
128     def _get_batch(reader, batch_size, decode):
129         res = []
130         for i in range(batch_size):
131             record = reader.record()
132             if record is None:
133                 break
134             res.append(decode(record))
135         return res
136 
137     def _get_features_and_labels_from_record(self, record_buf):
138         batch_input_data, batch_labels = self._input_fn(record_buf)
139         features = [batch_input_data[f_col.key] for f_col in self._feature_columns]
140         if len(features) == 1:
141             features = features[0]
142         return features, batch_labels
143 
144     def distributed_train(self):
145         """
146         Distributed training.
147         """
148         while True:
149             task = self.get_task()
150             if not task.shard_file_name:
151                 # No more task
152                 break
153             batch_size = task.minibatch_size
154             err_msg = ""
155             try:
156                 with closing(recordio.Scanner(task.shard_file_name, task.start, task.end - task.start)) as reader:
157                     min_model_version = task.model_version
158                     while True:
159                         record_buf = self._get_batch(reader, batch_size, self._codec.decode)
160                         if not record_buf:
161                             break
162 
163                         for _ in range(self._max_retrain_num):
164                             # TODO: optimize the logic to avoid unnecessary get_model call.
165                             self.get_model(
166                                 max(self._model_version, min_model_version))
167 
168                             features, labels = self._get_features_and_labels_from_record(record_buf)
169 
170                             with tf.GradientTape() as tape:
171                                 outputs = self._model.call(features, training=True)
172                                 loss = self._loss(outputs, labels)
173 
174                                 # TODO:  Add regularization loss if any,
175                                 #        which should be divided by the number of contributing workers.
176                             grads = tape.gradient(
177                                 loss, self._model.trainable_variables)
178 
179                             accepted, min_model_version = self.report_gradient(
180                                 grads)
181                             if accepted:
182                                 self._logger.info("Loss is %f" % loss.numpy())
183                                 break
184                         else:
185                             # Worker got stuck, fail the task.
186                             # TODO: stop the worker if it fails to make any progress for some time.
187                             raise RuntimeError("Worker got stuck during training")
188             except Exception as ex:
189                 err_msg = str(ex)
190                 traceback.print_exc()
191             self.report_task_result(task.task_id, err_msg)
192 
193     def local_train(self, file_list, batch_size, epoch=1, kwargs=None):
194         """
195         Local training for local testing. Must in eager mode.
196         Argments:
197             batch_size: batch size in training
198             epoch: the number of epoch in training
199             kwargs: contains a dict of parameters used in training
200         """
201         optimizer = self._opt_fn()
202         for _ in range(epoch):
203             for f in file_list:
204                 with closing(recordio.Scanner(f)) as reader:
205                     while True:
206                         record_buf = self._get_batch(reader, batch_size, self._codec.decode)
207                         if not record_buf:
208                             break
209 
210                         features, labels = self._get_features_and_labels_from_record(record_buf)
211 
212                         with tf.GradientTape() as tape:
213                             outputs = self._model.call(features, training=True)
214                             loss = self._loss(outputs, labels)
215 
216                             # Add regularization loss if any.
217                             # Note: for distributed training, the regularization loss should
218                             #       be divided by the number of contributing workers, which
219                             #       might be difficult for elasticdl.
220                             if self._model.losses:
221                                 loss += math_ops.add_n(self._model.losses)
222                         grads = tape.gradient(
223                             loss, self._model.trainable_variables)
224                         optimizer.apply_gradients(
225                             zip(grads, self._model.trainable_variables))
226                         self._logger.info("Loss is %f" % loss.numpy())
227 
228     def distributed_evaluate(self, steps=None):
229         """
230         Distributed model evaluation.
231 
232         Arguments:
233             steps: Evaluate the model by this many number of steps where the model is evaluated on one batch of samples
234                 for each step. If `None`, evaluation will continue until reaching the end of input.
235         """
236         while True:
237             task = self.get_task()
238             if not task.shard_file_name:
239                 break
240             batch_size = task.minibatch_size
241             err_msg = ""
242             try:
243                 with closing(
244                         recordio.Scanner(task.shard_file_name, task.start, task.end - task.start)
245                 ) as reader:
246                     min_model_version = task.model_version
247                     current_step = 0
248                     evaluation_metrics_collection = defaultdict(list)
249                     while True:
250                         current_step += 1
251                         if steps and current_step > steps:
252                             break
253                         record_buf = self._get_batch(reader, batch_size, self._codec.decode)
254                         if not record_buf:
255                             break
256 
257                         for _ in range(self._max_evaluate_retry_num):
258                             self.get_model(max(self._model_version, min_model_version))
259                             features, labels = self._get_features_and_labels_from_record(record_buf)
260                             outputs = self._model.call(features, training=False)
261                             evaluation_metrics = self._eval_metrics_fn(outputs, labels)
262 
263                             for k, v in evaluation_metrics.items():
264                                 v_np = v.numpy()
265                                 if v_np.size != 1:
266                                     raise Exception("Only metric result of length 1 is supported currently")
267                                 evaluation_metrics_collection[k].append(v_np)
268 
269                             evaluation_metrics = {k: np.mean(v) for k, v in evaluation_metrics_collection.items()}
270                             accepted, min_model_version = self.report_evaluation_metrics(evaluation_metrics)
271 
272                             if accepted:
273                                 self._logger.info("Evaluation metrics: %s" % evaluation_metrics)
274                                 break
275                         else:
276                             raise RuntimeError("Worker got stuck during model evaluation")
277             except Exception as ex:
278                 err_msg = str(ex)
279                 traceback.print_exc()
280             self.report_task_result(task.task_id, err_msg)
281
```

Path: `elasticdl/python/elasticdl/worker/main.py`
Content:
```
1 import argparse
2 import grpc
3 import logging
4 
5 import tensorflow as tf
6 
7 tf.enable_eager_execution()
8 
9 from elasticdl.python.elasticdl.worker.worker import Worker
10 
11 
12 def _parse_args():
13     parser = argparse.ArgumentParser(description="ElasticDL Worker")
14     parser.add_argument(
15         "--worker_id", help="Id unique to the worker", type=int, required=True
16     )
17     parser.add_argument("--master_addr", help="Master ip:port", required=True)
18     parser.add_argument(
19         "--model_file",
20         help="Full file path of user defined neural model",
21         required=True,
22     )
23     parser.add_argument(
24         "--codec_type",
25         default="bytes",
26         choices=["tf_example", "bytes"],
27         help="Type of codec(tf_example or bytes)",
28     )
29     parser.add_argument(
30         "--log_level",
31         choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'],
32         type=str.upper,
33         default='WARNING',
34         help="Set the logging level",
35     )
36 
37     return parser.parse_args()
38 
39 
40 def main():
41     args = _parse_args()
42     channel = grpc.insecure_channel(args.master_addr)
43 
44     # Initialize logger
45     logging.basicConfig(
46         format='%(asctime)s %(name)s %(levelname)-8s '
47         '[%(filename)s:%(lineno)d] %(message)s',
48     )
49     logging.getLogger().setLevel(args.log_level)
50 
51     worker = Worker(
52         args.worker_id,
53         args.model_file,
54         channel=channel,
55         codec_type=args.codec_type,
56     )
57 
58     worker.distributed_train()
59 
60 
61 if __name__ == "__main__":
62     main()
63
```
--- 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/python/elasticdl/worker/main.py b/elasticdl/python/elasticdl/worker/main.py
--- a/elasticdl/python/elasticdl/worker/main.py
+++ b/elasticdl/python/elasticdl/worker/main.py
@@ -33,6 +33,18 @@
         default='WARNING',
         help="Set the logging level",
     )
+    parser.add_argument(
+        "--task_type",
+        default="training",
+        choices=["training", "evaluation"],
+        help="Type of worker task (training or evaluation). Default is training task.",
+    )
+    parser.add_argument(
+        "--evaluate_steps",
+        default=None,
+        help="Evaluate the model by this many number of steps where the model is evaluated on one "
+             "batch of samples for each step. By default, evaluation will continue until reaching the end of input.",
+    )
 
     return parser.parse_args()
 
@@ -55,7 +67,10 @@
         codec_type=args.codec_type,
     )
 
-    worker.distributed_train()
+    if args.task_type == "training":
+        worker.distributed_train()
+    else:
+        worker.distributed_evaluate(steps=args.evaluate_steps)
 
 
 if __name__ == "__main__":
diff --git a/elasticdl/python/elasticdl/worker/worker.py b/elasticdl/python/elasticdl/worker/worker.py
--- a/elasticdl/python/elasticdl/worker/worker.py
+++ b/elasticdl/python/elasticdl/worker/worker.py
@@ -248,7 +248,7 @@
                     evaluation_metrics_collection = defaultdict(list)
                     while True:
                         current_step += 1
-                        if steps and current_step > steps:
+                        if steps and current_step > int(steps):
                             break
                         record_buf = self._get_batch(reader, batch_size, self._codec.decode)
                         if not record_buf:

{"golden_diff": "diff --git a/elasticdl/python/elasticdl/worker/main.py b/elasticdl/python/elasticdl/worker/main.py\n--- a/elasticdl/python/elasticdl/worker/main.py\n+++ b/elasticdl/python/elasticdl/worker/main.py\n@@ -33,6 +33,18 @@\n         default='WARNING',\n         help=\"Set the logging level\",\n     )\n+    parser.add_argument(\n+        \"--task_type\",\n+        default=\"training\",\n+        choices=[\"training\", \"evaluation\"],\n+        help=\"Type of worker task (training or evaluation). Default is training task.\",\n+    )\n+    parser.add_argument(\n+        \"--evaluate_steps\",\n+        default=None,\n+        help=\"Evaluate the model by this many number of steps where the model is evaluated on one \"\n+             \"batch of samples for each step. By default, evaluation will continue until reaching the end of input.\",\n+    )\n \n     return parser.parse_args()\n \n@@ -55,7 +67,10 @@\n         codec_type=args.codec_type,\n     )\n \n-    worker.distributed_train()\n+    if args.task_type == \"training\":\n+        worker.distributed_train()\n+    else:\n+        worker.distributed_evaluate(steps=args.evaluate_steps)\n \n \n if __name__ == \"__main__\":\ndiff --git a/elasticdl/python/elasticdl/worker/worker.py b/elasticdl/python/elasticdl/worker/worker.py\n--- a/elasticdl/python/elasticdl/worker/worker.py\n+++ b/elasticdl/python/elasticdl/worker/worker.py\n@@ -248,7 +248,7 @@\n                     evaluation_metrics_collection = defaultdict(list)\n                     while True:\n                         current_step += 1\n-                        if steps and current_step > steps:\n+                        if steps and current_step > int(steps):\n                             break\n                         record_buf = self._get_batch(reader, batch_size, self._codec.decode)\n                         if not record_buf:\n", "issue": "Add CLI entry-point to `Worker.distributed_evaluate()`\nPart of https://github.com/wangkuiyi/elasticdl/issues/384. Add main CLI entry-point to `Worker.distributed_evaluate()` that will be used in `WorkerManager`.\n", "before_files": [{"content": "import logging\nimport traceback\nimport time\n\nimport tensorflow as tf\nassert tf.executing_eagerly()\n\nimport recordio\nimport numpy as np\n\nfrom contextlib import closing\nfrom collections import defaultdict\nfrom tensorflow.python.ops import math_ops\nfrom elasticdl.proto import elasticdl_pb2_grpc\nfrom elasticdl.proto import elasticdl_pb2\nfrom elasticdl.python.elasticdl.common.ndarray import ndarray_to_tensor, tensor_to_ndarray\nfrom elasticdl.python.elasticdl.common.model_helper import load_user_model, build_model\nfrom elasticdl.python.data.codec import TFExampleCodec\nfrom elasticdl.python.data.codec import BytesCodec\n\n# TODO: Move these args to train/evaluate specific functions\n# the default max number of a minibatch retrain as its gradients are not accepted by master.\nDEFAULT_MAX_MINIBATCH_RETRAIN_NUM = 64\n# the default max number of a minibatch retry as its calculated evaluation metrics are not accepted by master.\nDEFAULT_MAX_MINIBATCH_EVALUATE_RETRY_NUM = 10\n\n\nclass Worker(object):\n    \"\"\"ElasticDL worker\"\"\"\n\n    def __init__(self,\n                 worker_id,\n                 model_file,\n                 channel=None,\n                 max_retrain_num=DEFAULT_MAX_MINIBATCH_RETRAIN_NUM,\n                 max_evaluate_retry_num=DEFAULT_MAX_MINIBATCH_EVALUATE_RETRY_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        self._logger = logging.getLogger(__name__)\n        self._worker_id = worker_id\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        self._eval_metrics_fn = model_module.eval_metrics_fn\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        if channel is None:\n            self._stub = None\n        else:\n            self._stub = elasticdl_pb2_grpc.MasterStub(channel)\n        self._max_retrain_num = max_retrain_num\n        self._max_evaluate_retry_num = max_evaluate_retry_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        req = elasticdl_pb2.GetTaskRequest()\n        req.worker_id = self._worker_id\n        \n        return self._stub.GetTask(req)\n\n    def get_model(self, min_version):\n        \"\"\"\n        get model from master, and update model_version\n        \"\"\"\n        req = elasticdl_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 = elasticdl_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 = elasticdl_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 report_evaluation_metrics(self, evaluation_metrics):\n        \"\"\"\n        report evaluation metrics to ps, return (accepted, model_version) from rpc call.\n        \"\"\"\n        req = elasticdl_pb2.ReportEvaluationMetricsRequest()\n        for k, v in evaluation_metrics.items():\n            req.evaluation_metrics[k].CopyFrom(\n                ndarray_to_tensor(v))\n        req.model_version = self._model_version\n        res = self._stub.ReportEvaluationMetrics(req)\n        return res.accepted, res.model_version\n\n    @staticmethod\n    def _get_batch(reader, batch_size, decode):\n        res = []\n        for i in range(batch_size):\n            record = reader.record()\n            if record is None:\n                break\n            res.append(decode(record))\n        return res\n\n    def _get_features_and_labels_from_record(self, record_buf):\n        batch_input_data, batch_labels = self._input_fn(record_buf)\n        features = [batch_input_data[f_col.key] for f_col in self._feature_columns]\n        if len(features) == 1:\n            features = features[0]\n        return features, batch_labels\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 closing(recordio.Scanner(task.shard_file_name, task.start, task.end - task.start)) as reader:\n                    min_model_version = task.model_version\n                    while True:\n                        record_buf = self._get_batch(reader, batch_size, self._codec.decode)\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                            features, labels = self._get_features_and_labels_from_record(record_buf)\n\n                            with tf.GradientTape() as tape:\n                                outputs = self._model.call(features, training=True)\n                                loss = self._loss(outputs, labels)\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\n                            accepted, min_model_version = self.report_gradient(\n                                grads)\n                            if accepted:\n                                self._logger.info(\"Loss is %f\" % loss.numpy())\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 during training\")\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 closing(recordio.Scanner(f)) as reader:\n                    while True:\n                        record_buf = self._get_batch(reader, batch_size, self._codec.decode)\n                        if not record_buf:\n                            break\n\n                        features, labels = self._get_features_and_labels_from_record(record_buf)\n\n                        with tf.GradientTape() as tape:\n                            outputs = self._model.call(features, 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                        self._logger.info(\"Loss is %f\" % loss.numpy())\n\n    def distributed_evaluate(self, steps=None):\n        \"\"\"\n        Distributed model evaluation.\n\n        Arguments:\n            steps: Evaluate the model by this many number of steps where the model is evaluated on one batch of samples\n                for each step. If `None`, evaluation will continue until reaching the end of input.\n        \"\"\"\n        while True:\n            task = self.get_task()\n            if not task.shard_file_name:\n                break\n            batch_size = task.minibatch_size\n            err_msg = \"\"\n            try:\n                with closing(\n                        recordio.Scanner(task.shard_file_name, task.start, task.end - task.start)\n                ) as reader:\n                    min_model_version = task.model_version\n                    current_step = 0\n                    evaluation_metrics_collection = defaultdict(list)\n                    while True:\n                        current_step += 1\n                        if steps and current_step > steps:\n                            break\n                        record_buf = self._get_batch(reader, batch_size, self._codec.decode)\n                        if not record_buf:\n                            break\n\n                        for _ in range(self._max_evaluate_retry_num):\n                            self.get_model(max(self._model_version, min_model_version))\n                            features, labels = self._get_features_and_labels_from_record(record_buf)\n                            outputs = self._model.call(features, training=False)\n                            evaluation_metrics = self._eval_metrics_fn(outputs, labels)\n\n                            for k, v in evaluation_metrics.items():\n                                v_np = v.numpy()\n                                if v_np.size != 1:\n                                    raise Exception(\"Only metric result of length 1 is supported currently\")\n                                evaluation_metrics_collection[k].append(v_np)\n\n                            evaluation_metrics = {k: np.mean(v) for k, v in evaluation_metrics_collection.items()}\n                            accepted, min_model_version = self.report_evaluation_metrics(evaluation_metrics)\n\n                            if accepted:\n                                self._logger.info(\"Evaluation metrics: %s\" % evaluation_metrics)\n                                break\n                        else:\n                            raise RuntimeError(\"Worker got stuck during model evaluation\")\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", "path": "elasticdl/python/elasticdl/worker/worker.py"}, {"content": "import argparse\nimport grpc\nimport logging\n\nimport tensorflow as tf\n\ntf.enable_eager_execution()\n\nfrom elasticdl.python.elasticdl.worker.worker import Worker\n\n\ndef _parse_args():\n    parser = argparse.ArgumentParser(description=\"ElasticDL Worker\")\n    parser.add_argument(\n        \"--worker_id\", help=\"Id unique to the worker\", type=int, required=True\n    )\n    parser.add_argument(\"--master_addr\", help=\"Master ip:port\", required=True)\n    parser.add_argument(\n        \"--model_file\",\n        help=\"Full file path of user defined neural model\",\n        required=True,\n    )\n    parser.add_argument(\n        \"--codec_type\",\n        default=\"bytes\",\n        choices=[\"tf_example\", \"bytes\"],\n        help=\"Type of codec(tf_example or bytes)\",\n    )\n    parser.add_argument(\n        \"--log_level\",\n        choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'],\n        type=str.upper,\n        default='WARNING',\n        help=\"Set the logging level\",\n    )\n\n    return parser.parse_args()\n\n\ndef main():\n    args = _parse_args()\n    channel = grpc.insecure_channel(args.master_addr)\n\n    # Initialize logger\n    logging.basicConfig(\n        format='%(asctime)s %(name)s %(levelname)-8s '\n        '[%(filename)s:%(lineno)d] %(message)s',\n    )\n    logging.getLogger().setLevel(args.log_level)\n\n    worker = Worker(\n        args.worker_id,\n        args.model_file,\n        channel=channel,\n        codec_type=args.codec_type,\n    )\n\n    worker.distributed_train()\n\n\nif __name__ == \"__main__\":\n    main()\n", "path": "elasticdl/python/elasticdl/worker/main.py"}], "after_files": [{"content": "import logging\nimport traceback\nimport time\n\nimport tensorflow as tf\nassert tf.executing_eagerly()\n\nimport recordio\nimport numpy as np\n\nfrom contextlib import closing\nfrom collections import defaultdict\nfrom tensorflow.python.ops import math_ops\nfrom elasticdl.proto import elasticdl_pb2_grpc\nfrom elasticdl.proto import elasticdl_pb2\nfrom elasticdl.python.elasticdl.common.ndarray import ndarray_to_tensor, tensor_to_ndarray\nfrom elasticdl.python.elasticdl.common.model_helper import load_user_model, build_model\nfrom elasticdl.python.data.codec import TFExampleCodec\nfrom elasticdl.python.data.codec import BytesCodec\n\n# TODO: Move these args to train/evaluate specific functions\n# the default max number of a minibatch retrain as its gradients are not accepted by master.\nDEFAULT_MAX_MINIBATCH_RETRAIN_NUM = 64\n# the default max number of a minibatch retry as its calculated evaluation metrics are not accepted by master.\nDEFAULT_MAX_MINIBATCH_EVALUATE_RETRY_NUM = 10\n\n\nclass Worker(object):\n    \"\"\"ElasticDL worker\"\"\"\n\n    def __init__(self,\n                 worker_id,\n                 model_file,\n                 channel=None,\n                 max_retrain_num=DEFAULT_MAX_MINIBATCH_RETRAIN_NUM,\n                 max_evaluate_retry_num=DEFAULT_MAX_MINIBATCH_EVALUATE_RETRY_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        self._logger = logging.getLogger(__name__)\n        self._worker_id = worker_id\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        self._eval_metrics_fn = model_module.eval_metrics_fn\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        if channel is None:\n            self._stub = None\n        else:\n            self._stub = elasticdl_pb2_grpc.MasterStub(channel)\n        self._max_retrain_num = max_retrain_num\n        self._max_evaluate_retry_num = max_evaluate_retry_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        req = elasticdl_pb2.GetTaskRequest()\n        req.worker_id = self._worker_id\n        \n        return self._stub.GetTask(req)\n\n    def get_model(self, min_version):\n        \"\"\"\n        get model from master, and update model_version\n        \"\"\"\n        req = elasticdl_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 = elasticdl_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 = elasticdl_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 report_evaluation_metrics(self, evaluation_metrics):\n        \"\"\"\n        report evaluation metrics to ps, return (accepted, model_version) from rpc call.\n        \"\"\"\n        req = elasticdl_pb2.ReportEvaluationMetricsRequest()\n        for k, v in evaluation_metrics.items():\n            req.evaluation_metrics[k].CopyFrom(\n                ndarray_to_tensor(v))\n        req.model_version = self._model_version\n        res = self._stub.ReportEvaluationMetrics(req)\n        return res.accepted, res.model_version\n\n    @staticmethod\n    def _get_batch(reader, batch_size, decode):\n        res = []\n        for i in range(batch_size):\n            record = reader.record()\n            if record is None:\n                break\n            res.append(decode(record))\n        return res\n\n    def _get_features_and_labels_from_record(self, record_buf):\n        batch_input_data, batch_labels = self._input_fn(record_buf)\n        features = [batch_input_data[f_col.key] for f_col in self._feature_columns]\n        if len(features) == 1:\n            features = features[0]\n        return features, batch_labels\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 closing(recordio.Scanner(task.shard_file_name, task.start, task.end - task.start)) as reader:\n                    min_model_version = task.model_version\n                    while True:\n                        record_buf = self._get_batch(reader, batch_size, self._codec.decode)\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                            features, labels = self._get_features_and_labels_from_record(record_buf)\n\n                            with tf.GradientTape() as tape:\n                                outputs = self._model.call(features, training=True)\n                                loss = self._loss(outputs, labels)\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\n                            accepted, min_model_version = self.report_gradient(\n                                grads)\n                            if accepted:\n                                self._logger.info(\"Loss is %f\" % loss.numpy())\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 during training\")\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 closing(recordio.Scanner(f)) as reader:\n                    while True:\n                        record_buf = self._get_batch(reader, batch_size, self._codec.decode)\n                        if not record_buf:\n                            break\n\n                        features, labels = self._get_features_and_labels_from_record(record_buf)\n\n                        with tf.GradientTape() as tape:\n                            outputs = self._model.call(features, 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                        self._logger.info(\"Loss is %f\" % loss.numpy())\n\n    def distributed_evaluate(self, steps=None):\n        \"\"\"\n        Distributed model evaluation.\n\n        Arguments:\n            steps: Evaluate the model by this many number of steps where the model is evaluated on one batch of samples\n                for each step. If `None`, evaluation will continue until reaching the end of input.\n        \"\"\"\n        while True:\n            task = self.get_task()\n            if not task.shard_file_name:\n                break\n            batch_size = task.minibatch_size\n            err_msg = \"\"\n            try:\n                with closing(\n                        recordio.Scanner(task.shard_file_name, task.start, task.end - task.start)\n                ) as reader:\n                    min_model_version = task.model_version\n                    current_step = 0\n                    evaluation_metrics_collection = defaultdict(list)\n                    while True:\n                        current_step += 1\n                        if steps and current_step > int(steps):\n                            break\n                        record_buf = self._get_batch(reader, batch_size, self._codec.decode)\n                        if not record_buf:\n                            break\n\n                        for _ in range(self._max_evaluate_retry_num):\n                            self.get_model(max(self._model_version, min_model_version))\n                            features, labels = self._get_features_and_labels_from_record(record_buf)\n                            outputs = self._model.call(features, training=False)\n                            evaluation_metrics = self._eval_metrics_fn(outputs, labels)\n\n                            for k, v in evaluation_metrics.items():\n                                v_np = v.numpy()\n                                if v_np.size != 1:\n                                    raise Exception(\"Only metric result of length 1 is supported currently\")\n                                evaluation_metrics_collection[k].append(v_np)\n\n                            evaluation_metrics = {k: np.mean(v) for k, v in evaluation_metrics_collection.items()}\n                            accepted, min_model_version = self.report_evaluation_metrics(evaluation_metrics)\n\n                            if accepted:\n                                self._logger.info(\"Evaluation metrics: %s\" % evaluation_metrics)\n                                break\n                        else:\n                            raise RuntimeError(\"Worker got stuck during model evaluation\")\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", "path": "elasticdl/python/elasticdl/worker/worker.py"}, {"content": "import argparse\nimport grpc\nimport logging\n\nimport tensorflow as tf\n\ntf.enable_eager_execution()\n\nfrom elasticdl.python.elasticdl.worker.worker import Worker\n\n\ndef _parse_args():\n    parser = argparse.ArgumentParser(description=\"ElasticDL Worker\")\n    parser.add_argument(\n        \"--worker_id\", help=\"Id unique to the worker\", type=int, required=True\n    )\n    parser.add_argument(\"--master_addr\", help=\"Master ip:port\", required=True)\n    parser.add_argument(\n        \"--model_file\",\n        help=\"Full file path of user defined neural model\",\n        required=True,\n    )\n    parser.add_argument(\n        \"--codec_type\",\n        default=\"bytes\",\n        choices=[\"tf_example\", \"bytes\"],\n        help=\"Type of codec(tf_example or bytes)\",\n    )\n    parser.add_argument(\n        \"--log_level\",\n        choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'],\n        type=str.upper,\n        default='WARNING',\n        help=\"Set the logging level\",\n    )\n    parser.add_argument(\n        \"--task_type\",\n        default=\"training\",\n        choices=[\"training\", \"evaluation\"],\n        help=\"Type of worker task (training or evaluation). Default is training task.\",\n    )\n    parser.add_argument(\n        \"--evaluate_steps\",\n        default=None,\n        help=\"Evaluate the model by this many number of steps where the model is evaluated on one \"\n             \"batch of samples for each step. By default, evaluation will continue until reaching the end of input.\",\n    )\n\n    return parser.parse_args()\n\n\ndef main():\n    args = _parse_args()\n    channel = grpc.insecure_channel(args.master_addr)\n\n    # Initialize logger\n    logging.basicConfig(\n        format='%(asctime)s %(name)s %(levelname)-8s '\n        '[%(filename)s:%(lineno)d] %(message)s',\n    )\n    logging.getLogger().setLevel(args.log_level)\n\n    worker = Worker(\n        args.worker_id,\n        args.model_file,\n        channel=channel,\n        codec_type=args.codec_type,\n    )\n\n    if args.task_type == \"training\":\n        worker.distributed_train()\n    else:\n        worker.distributed_evaluate(steps=args.evaluate_steps)\n\n\nif __name__ == \"__main__\":\n    main()\n", "path": "elasticdl/python/elasticdl/worker/main.py"}]}