We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Use tf.function-decorated functions to accelerate model compute


--- 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/worker/worker.py`
Content:
```
1 import logging
2 import time
3 import traceback
4 from contextlib import closing
5 
6 import recordio
7 import tensorflow as tf
8 from tensorflow.python.ops import math_ops
9 
10 from elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc
11 from elasticdl.python.common.model_helper import (
12     load_model_from_module,
13     load_module,
14 )
15 from elasticdl.python.common.ndarray import (
16     ndarray_to_tensor,
17     tensor_to_ndarray,
18 )
19 
20 # The default maximum number of a minibatch retry as its results
21 # (e.g. gradients) are not accepted by master.
22 DEFAULT_MAX_MINIBATCH_RETRY_NUM = 64
23 
24 
25 class Worker(object):
26     """ElasticDL worker"""
27 
28     def __init__(
29         self,
30         worker_id,
31         model_file,
32         input_fn="input_fn",
33         loss="loss",
34         optimizer="optimizer",
35         eval_metrics_fn="eval_metrics_fn",
36         channel=None,
37         model_def=None,
38         model_params="",
39         max_minibatch_retry_num=DEFAULT_MAX_MINIBATCH_RETRY_NUM,
40     ):
41         """
42         Arguments:
43             model_file: A module to define the model
44             channel: grpc channel
45             max_minibatch_retry_num: The maximum number of a minibatch retry
46                 as its results (e.g. gradients) are not accepted by master.
47         """
48         self._logger = logging.getLogger(__name__)
49         self._worker_id = worker_id
50         model_module = load_module(model_file).__dict__
51         self._model = load_model_from_module(
52             model_def, model_module, model_params
53         )
54         self._var_created = self._model.built
55         self._input_fn = model_module[input_fn]
56         self._opt_fn = model_module[optimizer]
57         self._loss = model_module[loss]
58         self._eval_metrics_fn = model_module[eval_metrics_fn]
59 
60         if channel is None:
61             self._stub = None
62         else:
63             self._stub = elasticdl_pb2_grpc.MasterStub(channel)
64         self._max_minibatch_retry_num = max_minibatch_retry_num
65         self._model_version = -1
66 
67     def get_task(self):
68         """
69         get task from master
70         """
71         req = elasticdl_pb2.GetTaskRequest()
72         req.worker_id = self._worker_id
73 
74         return self._stub.GetTask(req)
75 
76     def get_model(self, version, method):
77         """
78         get model from master, and update model_version
79         """
80         req = elasticdl_pb2.GetModelRequest()
81         req.version = version
82         req.method = method
83         model = self._stub.GetModel(req)
84 
85         for var in self._model.trainable_variables:
86             # Assumes all trainable variables exist in model.param.
87             var.assign(tensor_to_ndarray(model.param[var.name]))
88         self._model_version = model.version
89 
90     def report_task_result(self, task_id, err_msg):
91         """
92         report task result to master
93         """
94         report = elasticdl_pb2.ReportTaskResultRequest()
95         report.task_id = task_id
96         report.err_message = err_msg
97         return self._stub.ReportTaskResult(report)
98 
99     def report_variable(self):
100         """
101         report variable to ps.
102         """
103         req = elasticdl_pb2.ReportVariableRequest()
104         for v in self._model.trainable_variables:
105             req.variable[v.name].CopyFrom(ndarray_to_tensor(v.numpy()))
106         self._stub.ReportVariable(req)
107 
108     def report_gradient(self, grads):
109         """
110         report gradient to ps, return (accepted, model_version) from rpc call.
111         """
112         req = elasticdl_pb2.ReportGradientRequest()
113         for g, v in zip(grads, self._model.trainable_variables):
114             req.gradient[v.name].CopyFrom(ndarray_to_tensor(g.numpy()))
115         req.model_version = self._model_version
116         res = self._stub.ReportGradient(req)
117         return res.accepted, res.model_version
118 
119     def report_evaluation_metrics(self, evaluation_metrics):
120         """
121         report evaluation metrics to ps, return (accepted, model_version)
122         from rpc call.
123         """
124         req = elasticdl_pb2.ReportEvaluationMetricsRequest()
125         for k, v in evaluation_metrics.items():
126             v_np = v.numpy()
127             # If scalar, convert to numpy 1D array with size 1
128             if not v_np.shape:
129                 v_np = v_np.reshape(1)
130             req.evaluation_metrics[k].CopyFrom(ndarray_to_tensor(v_np))
131         req.model_version = self._model_version
132         res = self._stub.ReportEvaluationMetrics(req)
133         return res.accepted, res.model_version
134 
135     def report_prediction_outputs(self, predictions):
136         self._logger.info("Predicted: %f" % predictions.numpy())
137         # TODO: Decide whether we need to send results to master first
138         # or write results to destination directly from workers.
139         # Also, need to think about how users configure where to
140         # write results.
141         return True
142 
143     def _get_batch(self, reader, batch_size):
144         res = []
145         for i in range(batch_size):
146             record = reader.record()
147             if record is None:
148                 break
149             res.append(record)
150         return res
151 
152     def _create_variable_and_report(self, features):
153         # Use model.call to create variables, then report to ps
154         _ = self._model.call(features)
155         self.report_variable()
156         self._var_created = True
157 
158     def _run_training_task(self, features, labels):
159         with tf.GradientTape() as tape:
160             outputs = self._model.call(features, training=True)
161             loss = self._loss(outputs, labels)
162             # Add regularization loss if any
163             if self._model.losses:
164                 loss += math_ops.add_n(self._model.losses)
165         grads = tape.gradient(loss, self._model.trainable_variables)
166         accepted, min_model_version = self.report_gradient(grads)
167         return accepted, min_model_version, loss
168 
169     def _run_evaluation_task(self, features, labels):
170         outputs = self._model.call(features, training=False)
171         evaluation_metrics = self._eval_metrics_fn(outputs, labels)
172         return self.report_evaluation_metrics(evaluation_metrics)
173 
174     def _run_prediction_task(self, features):
175         predictions = self._model.call(features, training=False)
176         return self.report_prediction_outputs(predictions)
177 
178     def _handle_task(self, task):
179         min_model_version = task.model_version
180         with closing(
181             recordio.Scanner(
182                 task.shard_file_name, task.start, task.end - task.start
183             )
184         ) as reader:
185             while True:
186                 record_buf = self._get_batch(reader, task.minibatch_size)
187                 if not record_buf:
188                     break
189                 min_model_version = self._process_minibatch(
190                     task, record_buf, min_model_version
191                 )
192 
193     def _process_minibatch(self, task, record_buf, min_model_version):
194         # TODO: Discuss how we separate input_fn for different tasks
195         features, labels = self._input_fn(record_buf)
196         if not self._var_created:
197             self._create_variable_and_report(features)
198         for _ in range(self._max_minibatch_retry_num):
199             if task.type == elasticdl_pb2.EVALUATION:
200                 if min_model_version == -1:
201                     if self._model_version < 0:
202                         self.get_model(0, elasticdl_pb2.MINIMUM)
203                 elif self._model_version != min_model_version:
204                     self.get_model(min_model_version, elasticdl_pb2.FIXED)
205                 accepted, _ = self._run_evaluation_task(features, labels)
206                 if accepted:
207                     break
208             elif task.type == elasticdl_pb2.TRAINING:
209                 # TODO: optimize the logic to avoid unnecessary
210                 #       get_model call.
211                 self.get_model(
212                     max(self._model_version, min_model_version),
213                     elasticdl_pb2.MINIMUM,
214                 )
215                 accepted, min_model_version, loss = self._run_training_task(
216                     features, labels
217                 )
218                 if accepted:
219                     self._logger.info("Loss is %f" % loss.numpy())
220                     break
221             elif task.type == elasticdl_pb2.PREDICTION:
222                 if self._model_version != min_model_version:
223                     self.get_model(min_model_version, elasticdl_pb2.FIXED)
224                 accepted = self._run_prediction_task(features)
225                 if accepted:
226                     break
227             else:
228                 raise RuntimeError("Unrecognized task type, %s" % task.type)
229         else:
230             # Worker got stuck, fail the task.
231             # TODO: stop the worker if it fails to make any
232             #       progress for some time.
233             raise RuntimeError("Worker got stuck")
234         return min_model_version
235 
236     def run(self):
237         """
238         Fetches task from master and performs training or evaluation.
239         """
240         while True:
241             task = self.get_task()
242             if not task.shard_file_name:
243                 if task.type == elasticdl_pb2.WAIT:
244                     # Wait a few seconds then try to get_task again
245                     time.sleep(5)
246                     continue
247                 else:
248                     # No more task
249                     self._logger.info("No more task, stopping")
250                     break
251             self._logger.info("Receive a new task: %d", task.task_id)
252             err_msg = ""
253             try:
254                 self._handle_task(task)
255             except RuntimeError as err:
256                 err_msg = str(err)
257                 traceback.print_exc()
258             except Exception as ex:
259                 err_msg = str(ex)
260                 traceback.print_exc()
261                 raise ex
262             self.report_task_result(task.task_id, err_msg)
263
```
--- 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/worker/worker.py b/elasticdl/python/worker/worker.py
--- a/elasticdl/python/worker/worker.py
+++ b/elasticdl/python/worker/worker.py
@@ -155,7 +155,8 @@
         self.report_variable()
         self._var_created = True
 
-    def _run_training_task(self, features, labels):
+    @tf.function
+    def training_process(self, features, labels):
         with tf.GradientTape() as tape:
             outputs = self._model.call(features, training=True)
             loss = self._loss(outputs, labels)
@@ -163,16 +164,30 @@
             if self._model.losses:
                 loss += math_ops.add_n(self._model.losses)
         grads = tape.gradient(loss, self._model.trainable_variables)
+        return loss, grads
+
+    @tf.function
+    def evaluation_process(self, features, labels):
+        outputs = self._model.call(features, training=False)
+        evaluation_metrics = self._eval_metrics_fn(outputs, labels)
+        return evaluation_metrics
+
+    @tf.function
+    def predict_process(self, features):
+        outputs = self._model.call(features, training=False)
+        return outputs
+
+    def _run_training_task(self, features, labels):
+        loss, grads = self.training_process(features, labels)
         accepted, min_model_version = self.report_gradient(grads)
         return accepted, min_model_version, loss
 
     def _run_evaluation_task(self, features, labels):
-        outputs = self._model.call(features, training=False)
-        evaluation_metrics = self._eval_metrics_fn(outputs, labels)
+        evaluation_metrics = self.evaluation_process(features, labels)
         return self.report_evaluation_metrics(evaluation_metrics)
 
     def _run_prediction_task(self, features):
-        predictions = self._model.call(features, training=False)
+        predictions = self.predict_process(features)
         return self.report_prediction_outputs(predictions)
 
     def _handle_task(self, task):

{"golden_diff": "diff --git a/elasticdl/python/worker/worker.py b/elasticdl/python/worker/worker.py\n--- a/elasticdl/python/worker/worker.py\n+++ b/elasticdl/python/worker/worker.py\n@@ -155,7 +155,8 @@\n         self.report_variable()\n         self._var_created = True\n \n-    def _run_training_task(self, features, labels):\n+    @tf.function\n+    def training_process(self, features, labels):\n         with tf.GradientTape() as tape:\n             outputs = self._model.call(features, training=True)\n             loss = self._loss(outputs, labels)\n@@ -163,16 +164,30 @@\n             if self._model.losses:\n                 loss += math_ops.add_n(self._model.losses)\n         grads = tape.gradient(loss, self._model.trainable_variables)\n+        return loss, grads\n+\n+    @tf.function\n+    def evaluation_process(self, features, labels):\n+        outputs = self._model.call(features, training=False)\n+        evaluation_metrics = self._eval_metrics_fn(outputs, labels)\n+        return evaluation_metrics\n+\n+    @tf.function\n+    def predict_process(self, features):\n+        outputs = self._model.call(features, training=False)\n+        return outputs\n+\n+    def _run_training_task(self, features, labels):\n+        loss, grads = self.training_process(features, labels)\n         accepted, min_model_version = self.report_gradient(grads)\n         return accepted, min_model_version, loss\n \n     def _run_evaluation_task(self, features, labels):\n-        outputs = self._model.call(features, training=False)\n-        evaluation_metrics = self._eval_metrics_fn(outputs, labels)\n+        evaluation_metrics = self.evaluation_process(features, labels)\n         return self.report_evaluation_metrics(evaluation_metrics)\n \n     def _run_prediction_task(self, features):\n-        predictions = self._model.call(features, training=False)\n+        predictions = self.predict_process(features)\n         return self.report_prediction_outputs(predictions)\n \n     def _handle_task(self, task):\n", "issue": "Use tf.function-decorated functions to accelerate model compute\n\n", "before_files": [{"content": "import logging\nimport time\nimport traceback\nfrom contextlib import closing\n\nimport recordio\nimport tensorflow as tf\nfrom tensorflow.python.ops import math_ops\n\nfrom elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc\nfrom elasticdl.python.common.model_helper import (\n    load_model_from_module,\n    load_module,\n)\nfrom elasticdl.python.common.ndarray import (\n    ndarray_to_tensor,\n    tensor_to_ndarray,\n)\n\n# The default maximum number of a minibatch retry as its results\n# (e.g. gradients) are not accepted by master.\nDEFAULT_MAX_MINIBATCH_RETRY_NUM = 64\n\n\nclass Worker(object):\n    \"\"\"ElasticDL worker\"\"\"\n\n    def __init__(\n        self,\n        worker_id,\n        model_file,\n        input_fn=\"input_fn\",\n        loss=\"loss\",\n        optimizer=\"optimizer\",\n        eval_metrics_fn=\"eval_metrics_fn\",\n        channel=None,\n        model_def=None,\n        model_params=\"\",\n        max_minibatch_retry_num=DEFAULT_MAX_MINIBATCH_RETRY_NUM,\n    ):\n        \"\"\"\n        Arguments:\n            model_file: A module to define the model\n            channel: grpc channel\n            max_minibatch_retry_num: The maximum number of a minibatch retry\n                as its results (e.g. gradients) are not accepted by master.\n        \"\"\"\n        self._logger = logging.getLogger(__name__)\n        self._worker_id = worker_id\n        model_module = load_module(model_file).__dict__\n        self._model = load_model_from_module(\n            model_def, model_module, model_params\n        )\n        self._var_created = self._model.built\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\n        if channel is None:\n            self._stub = None\n        else:\n            self._stub = elasticdl_pb2_grpc.MasterStub(channel)\n        self._max_minibatch_retry_num = max_minibatch_retry_num\n        self._model_version = -1\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, version, method):\n        \"\"\"\n        get model from master, and update model_version\n        \"\"\"\n        req = elasticdl_pb2.GetModelRequest()\n        req.version = version\n        req.method = method\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(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_variable(self):\n        \"\"\"\n        report variable to ps.\n        \"\"\"\n        req = elasticdl_pb2.ReportVariableRequest()\n        for v in self._model.trainable_variables:\n            req.variable[v.name].CopyFrom(ndarray_to_tensor(v.numpy()))\n        self._stub.ReportVariable(req)\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(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)\n        from rpc call.\n        \"\"\"\n        req = elasticdl_pb2.ReportEvaluationMetricsRequest()\n        for k, v in evaluation_metrics.items():\n            v_np = v.numpy()\n            # If scalar, convert to numpy 1D array with size 1\n            if not v_np.shape:\n                v_np = v_np.reshape(1)\n            req.evaluation_metrics[k].CopyFrom(ndarray_to_tensor(v_np))\n        req.model_version = self._model_version\n        res = self._stub.ReportEvaluationMetrics(req)\n        return res.accepted, res.model_version\n\n    def report_prediction_outputs(self, predictions):\n        self._logger.info(\"Predicted: %f\" % predictions.numpy())\n        # TODO: Decide whether we need to send results to master first\n        # or write results to destination directly from workers.\n        # Also, need to think about how users configure where to\n        # write results.\n        return True\n\n    def _get_batch(self, reader, batch_size):\n        res = []\n        for i in range(batch_size):\n            record = reader.record()\n            if record is None:\n                break\n            res.append(record)\n        return res\n\n    def _create_variable_and_report(self, features):\n        # Use model.call to create variables, then report to ps\n        _ = self._model.call(features)\n        self.report_variable()\n        self._var_created = True\n\n    def _run_training_task(self, features, labels):\n        with tf.GradientTape() as tape:\n            outputs = self._model.call(features, training=True)\n            loss = self._loss(outputs, labels)\n            # Add regularization loss if any\n            if self._model.losses:\n                loss += math_ops.add_n(self._model.losses)\n        grads = tape.gradient(loss, self._model.trainable_variables)\n        accepted, min_model_version = self.report_gradient(grads)\n        return accepted, min_model_version, loss\n\n    def _run_evaluation_task(self, features, labels):\n        outputs = self._model.call(features, training=False)\n        evaluation_metrics = self._eval_metrics_fn(outputs, labels)\n        return self.report_evaluation_metrics(evaluation_metrics)\n\n    def _run_prediction_task(self, features):\n        predictions = self._model.call(features, training=False)\n        return self.report_prediction_outputs(predictions)\n\n    def _handle_task(self, task):\n        min_model_version = task.model_version\n        with closing(\n            recordio.Scanner(\n                task.shard_file_name, task.start, task.end - task.start\n            )\n        ) as reader:\n            while True:\n                record_buf = self._get_batch(reader, task.minibatch_size)\n                if not record_buf:\n                    break\n                min_model_version = self._process_minibatch(\n                    task, record_buf, min_model_version\n                )\n\n    def _process_minibatch(self, task, record_buf, min_model_version):\n        # TODO: Discuss how we separate input_fn for different tasks\n        features, labels = self._input_fn(record_buf)\n        if not self._var_created:\n            self._create_variable_and_report(features)\n        for _ in range(self._max_minibatch_retry_num):\n            if task.type == elasticdl_pb2.EVALUATION:\n                if min_model_version == -1:\n                    if self._model_version < 0:\n                        self.get_model(0, elasticdl_pb2.MINIMUM)\n                elif self._model_version != min_model_version:\n                    self.get_model(min_model_version, elasticdl_pb2.FIXED)\n                accepted, _ = self._run_evaluation_task(features, labels)\n                if accepted:\n                    break\n            elif task.type == elasticdl_pb2.TRAINING:\n                # TODO: optimize the logic to avoid unnecessary\n                #       get_model call.\n                self.get_model(\n                    max(self._model_version, min_model_version),\n                    elasticdl_pb2.MINIMUM,\n                )\n                accepted, min_model_version, loss = self._run_training_task(\n                    features, labels\n                )\n                if accepted:\n                    self._logger.info(\"Loss is %f\" % loss.numpy())\n                    break\n            elif task.type == elasticdl_pb2.PREDICTION:\n                if self._model_version != min_model_version:\n                    self.get_model(min_model_version, elasticdl_pb2.FIXED)\n                accepted = self._run_prediction_task(features)\n                if accepted:\n                    break\n            else:\n                raise RuntimeError(\"Unrecognized task type, %s\" % task.type)\n        else:\n            # Worker got stuck, fail the task.\n            # TODO: stop the worker if it fails to make any\n            #       progress for some time.\n            raise RuntimeError(\"Worker got stuck\")\n        return min_model_version\n\n    def run(self):\n        \"\"\"\n        Fetches task from master and performs training or evaluation.\n        \"\"\"\n        while True:\n            task = self.get_task()\n            if not task.shard_file_name:\n                if task.type == elasticdl_pb2.WAIT:\n                    # Wait a few seconds then try to get_task again\n                    time.sleep(5)\n                    continue\n                else:\n                    # No more task\n                    self._logger.info(\"No more task, stopping\")\n                    break\n            self._logger.info(\"Receive a new task: %d\", task.task_id)\n            err_msg = \"\"\n            try:\n                self._handle_task(task)\n            except RuntimeError as err:\n                err_msg = str(err)\n                traceback.print_exc()\n            except Exception as ex:\n                err_msg = str(ex)\n                traceback.print_exc()\n                raise ex\n            self.report_task_result(task.task_id, err_msg)\n", "path": "elasticdl/python/worker/worker.py"}], "after_files": [{"content": "import logging\nimport time\nimport traceback\nfrom contextlib import closing\n\nimport recordio\nimport tensorflow as tf\nfrom tensorflow.python.ops import math_ops\n\nfrom elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc\nfrom elasticdl.python.common.model_helper import (\n    load_model_from_module,\n    load_module,\n)\nfrom elasticdl.python.common.ndarray import (\n    ndarray_to_tensor,\n    tensor_to_ndarray,\n)\n\n# The default maximum number of a minibatch retry as its results\n# (e.g. gradients) are not accepted by master.\nDEFAULT_MAX_MINIBATCH_RETRY_NUM = 64\n\n\nclass Worker(object):\n    \"\"\"ElasticDL worker\"\"\"\n\n    def __init__(\n        self,\n        worker_id,\n        model_file,\n        input_fn=\"input_fn\",\n        loss=\"loss\",\n        optimizer=\"optimizer\",\n        eval_metrics_fn=\"eval_metrics_fn\",\n        channel=None,\n        model_def=None,\n        model_params=\"\",\n        max_minibatch_retry_num=DEFAULT_MAX_MINIBATCH_RETRY_NUM,\n    ):\n        \"\"\"\n        Arguments:\n            model_file: A module to define the model\n            channel: grpc channel\n            max_minibatch_retry_num: The maximum number of a minibatch retry\n                as its results (e.g. gradients) are not accepted by master.\n        \"\"\"\n        self._logger = logging.getLogger(__name__)\n        self._worker_id = worker_id\n        model_module = load_module(model_file).__dict__\n        self._model = load_model_from_module(\n            model_def, model_module, model_params\n        )\n        self._var_created = self._model.built\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\n        if channel is None:\n            self._stub = None\n        else:\n            self._stub = elasticdl_pb2_grpc.MasterStub(channel)\n        self._max_minibatch_retry_num = max_minibatch_retry_num\n        self._model_version = -1\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, version, method):\n        \"\"\"\n        get model from master, and update model_version\n        \"\"\"\n        req = elasticdl_pb2.GetModelRequest()\n        req.version = version\n        req.method = method\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(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_variable(self):\n        \"\"\"\n        report variable to ps.\n        \"\"\"\n        req = elasticdl_pb2.ReportVariableRequest()\n        for v in self._model.trainable_variables:\n            req.variable[v.name].CopyFrom(ndarray_to_tensor(v.numpy()))\n        self._stub.ReportVariable(req)\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(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)\n        from rpc call.\n        \"\"\"\n        req = elasticdl_pb2.ReportEvaluationMetricsRequest()\n        for k, v in evaluation_metrics.items():\n            v_np = v.numpy()\n            # If scalar, convert to numpy 1D array with size 1\n            if not v_np.shape:\n                v_np = v_np.reshape(1)\n            req.evaluation_metrics[k].CopyFrom(ndarray_to_tensor(v_np))\n        req.model_version = self._model_version\n        res = self._stub.ReportEvaluationMetrics(req)\n        return res.accepted, res.model_version\n\n    def report_prediction_outputs(self, predictions):\n        self._logger.info(\"Predicted: %f\" % predictions.numpy())\n        # TODO: Decide whether we need to send results to master first\n        # or write results to destination directly from workers.\n        # Also, need to think about how users configure where to\n        # write results.\n        return True\n\n    def _get_batch(self, reader, batch_size):\n        res = []\n        for i in range(batch_size):\n            record = reader.record()\n            if record is None:\n                break\n            res.append(record)\n        return res\n\n    def _create_variable_and_report(self, features):\n        # Use model.call to create variables, then report to ps\n        _ = self._model.call(features)\n        self.report_variable()\n        self._var_created = True\n\n    @tf.function\n    def training_process(self, features, labels):\n        with tf.GradientTape() as tape:\n            outputs = self._model.call(features, training=True)\n            loss = self._loss(outputs, labels)\n            # Add regularization loss if any\n            if self._model.losses:\n                loss += math_ops.add_n(self._model.losses)\n        grads = tape.gradient(loss, self._model.trainable_variables)\n        return loss, grads\n\n    @tf.function\n    def evaluation_process(self, features, labels):\n        outputs = self._model.call(features, training=False)\n        evaluation_metrics = self._eval_metrics_fn(outputs, labels)\n        return evaluation_metrics\n\n    @tf.function\n    def predict_process(self, features):\n        outputs = self._model.call(features, training=False)\n        return outputs\n\n    def _run_training_task(self, features, labels):\n        loss, grads = self.training_process(features, labels)\n        accepted, min_model_version = self.report_gradient(grads)\n        return accepted, min_model_version, loss\n\n    def _run_evaluation_task(self, features, labels):\n        evaluation_metrics = self.evaluation_process(features, labels)\n        return self.report_evaluation_metrics(evaluation_metrics)\n\n    def _run_prediction_task(self, features):\n        predictions = self.predict_process(features)\n        return self.report_prediction_outputs(predictions)\n\n    def _handle_task(self, task):\n        min_model_version = task.model_version\n        with closing(\n            recordio.Scanner(\n                task.shard_file_name, task.start, task.end - task.start\n            )\n        ) as reader:\n            while True:\n                record_buf = self._get_batch(reader, task.minibatch_size)\n                if not record_buf:\n                    break\n                min_model_version = self._process_minibatch(\n                    task, record_buf, min_model_version\n                )\n\n    def _process_minibatch(self, task, record_buf, min_model_version):\n        # TODO: Discuss how we separate input_fn for different tasks\n        features, labels = self._input_fn(record_buf)\n        if not self._var_created:\n            self._create_variable_and_report(features)\n        for _ in range(self._max_minibatch_retry_num):\n            if task.type == elasticdl_pb2.EVALUATION:\n                if min_model_version == -1:\n                    if self._model_version < 0:\n                        self.get_model(0, elasticdl_pb2.MINIMUM)\n                elif self._model_version != min_model_version:\n                    self.get_model(min_model_version, elasticdl_pb2.FIXED)\n                accepted, _ = self._run_evaluation_task(features, labels)\n                if accepted:\n                    break\n            elif task.type == elasticdl_pb2.TRAINING:\n                # TODO: optimize the logic to avoid unnecessary\n                #       get_model call.\n                self.get_model(\n                    max(self._model_version, min_model_version),\n                    elasticdl_pb2.MINIMUM,\n                )\n                accepted, min_model_version, loss = self._run_training_task(\n                    features, labels\n                )\n                if accepted:\n                    self._logger.info(\"Loss is %f\" % loss.numpy())\n                    break\n            elif task.type == elasticdl_pb2.PREDICTION:\n                if self._model_version != min_model_version:\n                    self.get_model(min_model_version, elasticdl_pb2.FIXED)\n                accepted = self._run_prediction_task(features)\n                if accepted:\n                    break\n            else:\n                raise RuntimeError(\"Unrecognized task type, %s\" % task.type)\n        else:\n            # Worker got stuck, fail the task.\n            # TODO: stop the worker if it fails to make any\n            #       progress for some time.\n            raise RuntimeError(\"Worker got stuck\")\n        return min_model_version\n\n    def run(self):\n        \"\"\"\n        Fetches task from master and performs training or evaluation.\n        \"\"\"\n        while True:\n            task = self.get_task()\n            if not task.shard_file_name:\n                if task.type == elasticdl_pb2.WAIT:\n                    # Wait a few seconds then try to get_task again\n                    time.sleep(5)\n                    continue\n                else:\n                    # No more task\n                    self._logger.info(\"No more task, stopping\")\n                    break\n            self._logger.info(\"Receive a new task: %d\", task.task_id)\n            err_msg = \"\"\n            try:\n                self._handle_task(task)\n            except RuntimeError as err:\n                err_msg = str(err)\n                traceback.print_exc()\n            except Exception as ex:\n                err_msg = str(ex)\n                traceback.print_exc()\n                raise ex\n            self.report_task_result(task.task_id, err_msg)\n", "path": "elasticdl/python/worker/worker.py"}]}