We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Recover 'doing' tasks from preempted workers


--- 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/master/k8s_worker_manager.py`
Content:
```
1 import logging
2 import os
3 import time
4 
5 from collections import Counter
6 from elasticdl.master import k8s
7 
8 
9 class WorkerTracker(object):
10     def __init__(self):
11         # pod name to phase mapping
12         # phase: Pending/Running/Succeeded/Failed/Unknown
13         #   Pending: worker pod not started yet
14         #   Running: worker pod is running
15         #   Succeeded: worker pod finishes all jobs and terminates with no issue.
16         #   Failed: worker pod is killed for some reason
17         #   Unknown: unkown
18         self._pods_phase = {}
19 
20     def get_counters(self):
21         return Counter(self._pods_phase.values())
22 
23     def event_cb(self, event):
24         pod_name = event["object"].metadata.name
25         self._pods_phase[pod_name] = event["object"].status.phase
26         if event["type"] == "DELETED":
27             del self._pods_phase[pod_name]
28 
29 
30 class WorkerManager(object):
31     def __init__(self, command, args, num_worker=1, cpu_request="1000m", cpu_limit="1000m",
32             memory_request="4096Mi", memory_limit="4096Mi", pod_priority=None, **kwargs):
33         self._logger = logging.getLogger("WorkerManager")
34         self._command = command
35         self._args = args
36         self._num_worker = num_worker
37         self._cpu_request = cpu_request 
38         self._cpu_limit = cpu_limit
39         self._memory_request = memory_request
40         self._memory_limit = memory_limit
41         self._pod_priority = pod_priority
42         self._worker_tracker = WorkerTracker()
43         self._k8s_client = k8s.Client(
44             event_callback=self._worker_tracker.event_cb, **kwargs
45         )
46 
47     def start_workers(self, restart_policy="OnFailure"):
48         for i in range(self._num_worker):
49             self._logger.warning("Starting worker: %d" % i)
50             self._add_worker(i, restart_policy=restart_policy)
51 
52     def remove_workers(self):
53         for i in range(self._num_worker):
54             pod_name = self._k8s_client.get_pod_name(i)
55             if pod_name in self._worker_tracker._pods_phase:
56                 self._logger.warning("Deleting worker: %d", i)
57                 self._delete_worker(i)
58 
59     def _add_worker(self, worker_id, restart_policy):
60         self._k8s_client.create_worker(
61             worker_id,
62             self._cpu_request,
63             self._cpu_limit,
64             self._memory_request,
65             self._memory_limit,
66             self._pod_priority,
67             command=self._command,
68             args=self._args + ["--worker_id", str(worker_id)],
69             restart_policy=restart_policy,
70         )
71 
72     def _delete_worker(self, worker_id):
73         self._k8s_client.delete_worker(worker_id)
74 
75     def get_counters(self):
76         return self._worker_tracker.get_counters()
77
```

Path: `elasticdl/python/elasticdl/master/task_queue.py`
Content:
```
1 """TaskQueue Implementation"""
2 
3 import logging
4 import random
5 import threading
6 
7 
8 class _Task(object):
9     """Internal representation of a task"""
10 
11     def __init__(self, *, file_name, start, end):
12         self.file_name = file_name
13         self.start = start
14         self.end = end
15 
16 
17 class _TaskQueue(object):
18     """Creates and dispatches Tasks. Keep track of a Task's lifecycle."""
19 
20     def __init__(self, shards, record_per_task, num_epoch):
21         """
22         shards: a dictionary from RecordIO file name to number of records
23         """
24         self._logger = logging.getLogger("TaskQueue")
25         self._lock = threading.Lock()
26 
27         self._num_epoch = num_epoch
28         self._epoch = 0
29         self._shards = shards
30         self._record_per_task = record_per_task
31 
32         self._todo = []
33         # dictionary from task id to Task.
34         self._doing = {}
35         self._task_id = 0
36 
37         self._create_tasks()
38 
39     def _create_tasks(self):
40         for name, num_records in self._shards.items():
41             for start in range(0, num_records, self._record_per_task):
42                 self._todo.append(
43                     _Task(
44                         file_name=name,
45                         start=start,
46                         end=min(start + self._record_per_task, num_records),
47                     )
48                 )
49         random.shuffle(self._todo)
50 
51     def get(self, worker_id):
52         """Return next (task_id, Task) tuple"""
53 
54         with self._lock:
55             if not self._todo and self._epoch < self._num_epoch - 1:
56                 # Start a new epoch
57                 self._create_tasks()
58                 self._epoch += 1
59                 self._logger.warning("Starting epoch %d" % self._epoch)
60 
61             if not self._todo:
62                 # No more tasks
63                 return -1, None
64 
65             self._task_id += 1
66             task = self._todo.pop()
67             # TODO: Handle timeout of tasks.
68             self._doing[self._task_id] = (worker_id, task)
69 
70             return self._task_id, task
71 
72     def report(self, task_id, success):
73         """Report if the task is successful or not"""
74 
75         with self._lock:
76             _, task = self._doing.pop(task_id, (-1, None))
77             if not task:
78                 self._logger.warning("Unknown task_id: %d" % task_id)
79             elif not success:
80                 # TODO: keep count of retries.
81                 self._todo.append(task)
82 
83     def finished(self):
84         """Return if all tasks are done"""
85         return not self._todo and not self._doing
86
```

Path: `elasticdl/python/elasticdl/master/main.py`
Content:
```
1 import logging
2 import time
3 import argparse
4 import os
5 
6 import grpc
7 import tensorflow as tf
8 
9 tf.enable_eager_execution()
10 
11 from concurrent import futures
12 from recordio import File
13 from elasticdl.proto import master_pb2_grpc
14 from elasticdl.master.servicer import MasterServicer
15 from elasticdl.master.task_queue import _TaskQueue
16 from elasticdl.master.k8s_worker_manager import WorkerManager
17 from elasticdl.common.model_helper import load_user_model, build_model
18 
19 
20 def _make_task_queue(data_dir, record_per_task, num_epoch):
21     f_records = {}
22     for f in os.listdir(data_dir):
23         p = os.path.join(data_dir, f)
24         with File(p, "r") as rio:
25             f_records[p] = rio.count()
26     return _TaskQueue(f_records, record_per_task, num_epoch)
27 
28 
29 def _parse_args():
30     parser = argparse.ArgumentParser(description="ElasticDL Master")
31     parser.add_argument(
32         "--model_file",
33         help="Full file path of user defined neural model",
34         required=True,
35     )
36     parser.add_argument(
37         "--train_data_dir",
38         help="Training data directory. Files should be in RecordIO format",
39         required=True,
40     )
41     parser.add_argument("--record_per_task", type=int, required=True)
42     parser.add_argument("--num_epoch", type=int, required=True)
43     parser.add_argument(
44         "--grads_to_wait",
45         type=int,
46         help="Number of gradients to wait before updating model",
47         required=True,
48     )
49     parser.add_argument(
50         "--minibatch_size",
51         type=int,
52         help="Minibatch size used by workers to compute gradients",
53         required=True,
54     )
55     parser.add_argument(
56         "--num_worker",
57         type=int,
58         help="the number of workers used in training",
59         default=0,
60     )
61     parser.add_argument(
62         "--worker_cpu_request",
63         help="the minimal cpu required by worker in training",
64         default="1000m",
65     )
66     parser.add_argument(
67         "--worker_cpu_limit",
68         help="the maximal cpu used by worker in training",
69         default="1000m",
70     )
71     parser.add_argument(
72         "--worker_memory_request",
73         help="the minimal memory required by worker in training",
74         default="4096Mi",
75     )
76     parser.add_argument(
77         "--worker_memory_limit",
78         help="the maximal memory used by worker in training",
79         default="4096Mi",
80     )
81     parser.add_argument(
82         "--worker_pod_priority",
83         help="the requested priority of worker pod")
84     parser.add_argument(
85         "--worker_image", help="docker image for worker", default=None
86     )
87     parser.add_argument("--job_name", help="job name", required=True)
88     parser.add_argument(
89         "--codec_type",
90         default="bytes",
91         choices=["tf_example", "bytes"],
92         help="Type of codec(tf_example or bytes)",
93     )
94     return parser.parse_args()
95 
96 
97 def main():
98     # TODO: pass port via flags.
99     PORT = 50001
100     logger = logging.getLogger("master")
101     args = _parse_args()
102     task_q = _make_task_queue(
103         args.train_data_dir, args.record_per_task, args.num_epoch
104     )
105     model_module = load_user_model(args.model_file)
106     model_inst = model_module.model
107     build_model(model_inst, model_module.feature_columns())
108     optimizer = model_module.optimizer()
109 
110     server = grpc.server(futures.ThreadPoolExecutor(max_workers=64))
111     master_pb2_grpc.add_MasterServicer_to_server(
112         MasterServicer(
113             logger,
114             args.grads_to_wait,
115             args.minibatch_size,
116             optimizer,
117             task_q,
118             init_var=model_inst.trainable_variables,
119         ),
120         server,
121     )
122     server.add_insecure_port("[::]:{}".format(PORT))
123     server.start()
124     logger.warning("Server started at port: %d", PORT)
125 
126     if args.num_worker:
127         master_addr = "%s:%d" % (os.getenv("MY_POD_IP", "localhost"), PORT)
128         worker_command = ["python"]
129         worker_args = [
130             "-m",
131             "elasticdl.worker.main",
132             "--model_file",
133             args.model_file,
134             "--master_addr",
135             master_addr,
136             "--codec_type",
137             args.codec_type
138         ]
139 
140         worker_manager = WorkerManager(
141             job_name=args.job_name,
142             worker_image=args.worker_image,
143             command=worker_command,
144             args=worker_args,
145             namespace="default",
146             num_worker=args.num_worker,
147             cpu_request=args.worker_cpu_request,
148             cpu_limit=args.worker_cpu_limit,
149             memory_request=args.worker_memory_request,
150             memory_limit=args.worker_memory_limit,
151             pod_priority=args.worker_pod_priority,
152         )
153         worker_manager.start_workers(restart_policy="Never")
154 
155     try:
156         while True:
157             if task_q.finished():
158                 break
159             time.sleep(30)
160     except KeyboardInterrupt:
161         logger.warning("Server stopping")
162 
163     if args.num_worker:
164         # TODO: worker_manager.remove_workers supports synchronized call
165         worker_manager.remove_workers()
166 
167     server.stop(0)
168 
169 
170 if __name__ == "__main__":
171     logging.basicConfig()
172     main()
173
```
--- 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/master/k8s_worker_manager.py b/elasticdl/python/elasticdl/master/k8s_worker_manager.py
--- a/elasticdl/python/elasticdl/master/k8s_worker_manager.py
+++ b/elasticdl/python/elasticdl/master/k8s_worker_manager.py
@@ -6,8 +6,8 @@
 from elasticdl.master import k8s
 
 
-class WorkerTracker(object):
-    def __init__(self):
+class _WorkerTracker(object):
+    def __init__(self, task_q):
         # pod name to phase mapping
         # phase: Pending/Running/Succeeded/Failed/Unknown
         #   Pending: worker pod not started yet
@@ -16,6 +16,7 @@
         #   Failed: worker pod is killed for some reason
         #   Unknown: unkown
         self._pods_phase = {}
+        self._task_q = task_q
 
     def get_counters(self):
         return Counter(self._pods_phase.values())
@@ -25,10 +26,15 @@
         self._pods_phase[pod_name] = event["object"].status.phase
         if event["type"] == "DELETED":
             del self._pods_phase[pod_name]
+            self._task_q.recover_tasks(
+                # TODO: move worker_id and pod name mapping to a separate class 
+                int(pod_name.rsplit("-", 1)[1])
+            )
+
 
 
 class WorkerManager(object):
-    def __init__(self, command, args, num_worker=1, cpu_request="1000m", cpu_limit="1000m",
+    def __init__(self, task_q, command, args, num_worker=1, cpu_request="1000m", cpu_limit="1000m",
             memory_request="4096Mi", memory_limit="4096Mi", pod_priority=None, **kwargs):
         self._logger = logging.getLogger("WorkerManager")
         self._command = command
@@ -39,7 +45,7 @@
         self._memory_request = memory_request
         self._memory_limit = memory_limit
         self._pod_priority = pod_priority
-        self._worker_tracker = WorkerTracker()
+        self._worker_tracker = _WorkerTracker(task_q)
         self._k8s_client = k8s.Client(
             event_callback=self._worker_tracker.event_cb, **kwargs
         )
diff --git a/elasticdl/python/elasticdl/master/main.py b/elasticdl/python/elasticdl/master/main.py
--- a/elasticdl/python/elasticdl/master/main.py
+++ b/elasticdl/python/elasticdl/master/main.py
@@ -138,6 +138,7 @@
         ]
 
         worker_manager = WorkerManager(
+            task_q,
             job_name=args.job_name,
             worker_image=args.worker_image,
             command=worker_command,
diff --git a/elasticdl/python/elasticdl/master/task_queue.py b/elasticdl/python/elasticdl/master/task_queue.py
--- a/elasticdl/python/elasticdl/master/task_queue.py
+++ b/elasticdl/python/elasticdl/master/task_queue.py
@@ -83,3 +83,10 @@
     def finished(self):
         """Return if all tasks are done"""
         return not self._todo and not self._doing
+
+    def recover_tasks(self, worker_id):
+        """Recover doing tasks for a dead worker"""
+
+        with self._lock:
+            tasks = [task for wid, task in self._doing if wid == worker_id]
+            self._todo.extend(tasks)

{"golden_diff": "diff --git a/elasticdl/python/elasticdl/master/k8s_worker_manager.py b/elasticdl/python/elasticdl/master/k8s_worker_manager.py\n--- a/elasticdl/python/elasticdl/master/k8s_worker_manager.py\n+++ b/elasticdl/python/elasticdl/master/k8s_worker_manager.py\n@@ -6,8 +6,8 @@\n from elasticdl.master import k8s\n \n \n-class WorkerTracker(object):\n-    def __init__(self):\n+class _WorkerTracker(object):\n+    def __init__(self, task_q):\n         # pod name to phase mapping\n         # phase: Pending/Running/Succeeded/Failed/Unknown\n         #   Pending: worker pod not started yet\n@@ -16,6 +16,7 @@\n         #   Failed: worker pod is killed for some reason\n         #   Unknown: unkown\n         self._pods_phase = {}\n+        self._task_q = task_q\n \n     def get_counters(self):\n         return Counter(self._pods_phase.values())\n@@ -25,10 +26,15 @@\n         self._pods_phase[pod_name] = event[\"object\"].status.phase\n         if event[\"type\"] == \"DELETED\":\n             del self._pods_phase[pod_name]\n+            self._task_q.recover_tasks(\n+                # TODO: move worker_id and pod name mapping to a separate class \n+                int(pod_name.rsplit(\"-\", 1)[1])\n+            )\n+\n \n \n class WorkerManager(object):\n-    def __init__(self, command, args, num_worker=1, cpu_request=\"1000m\", cpu_limit=\"1000m\",\n+    def __init__(self, task_q, command, args, num_worker=1, cpu_request=\"1000m\", cpu_limit=\"1000m\",\n             memory_request=\"4096Mi\", memory_limit=\"4096Mi\", pod_priority=None, **kwargs):\n         self._logger = logging.getLogger(\"WorkerManager\")\n         self._command = command\n@@ -39,7 +45,7 @@\n         self._memory_request = memory_request\n         self._memory_limit = memory_limit\n         self._pod_priority = pod_priority\n-        self._worker_tracker = WorkerTracker()\n+        self._worker_tracker = _WorkerTracker(task_q)\n         self._k8s_client = k8s.Client(\n             event_callback=self._worker_tracker.event_cb, **kwargs\n         )\ndiff --git a/elasticdl/python/elasticdl/master/main.py b/elasticdl/python/elasticdl/master/main.py\n--- a/elasticdl/python/elasticdl/master/main.py\n+++ b/elasticdl/python/elasticdl/master/main.py\n@@ -138,6 +138,7 @@\n         ]\n \n         worker_manager = WorkerManager(\n+            task_q,\n             job_name=args.job_name,\n             worker_image=args.worker_image,\n             command=worker_command,\ndiff --git a/elasticdl/python/elasticdl/master/task_queue.py b/elasticdl/python/elasticdl/master/task_queue.py\n--- a/elasticdl/python/elasticdl/master/task_queue.py\n+++ b/elasticdl/python/elasticdl/master/task_queue.py\n@@ -83,3 +83,10 @@\n     def finished(self):\n         \"\"\"Return if all tasks are done\"\"\"\n         return not self._todo and not self._doing\n+\n+    def recover_tasks(self, worker_id):\n+        \"\"\"Recover doing tasks for a dead worker\"\"\"\n+\n+        with self._lock:\n+            tasks = [task for wid, task in self._doing if wid == worker_id]\n+            self._todo.extend(tasks)\n", "issue": "Recover 'doing' tasks from preempted workers\n\n", "before_files": [{"content": "import logging\nimport os\nimport time\n\nfrom collections import Counter\nfrom elasticdl.master import k8s\n\n\nclass WorkerTracker(object):\n    def __init__(self):\n        # pod name to phase mapping\n        # phase: Pending/Running/Succeeded/Failed/Unknown\n        #   Pending: worker pod not started yet\n        #   Running: worker pod is running\n        #   Succeeded: worker pod finishes all jobs and terminates with no issue.\n        #   Failed: worker pod is killed for some reason\n        #   Unknown: unkown\n        self._pods_phase = {}\n\n    def get_counters(self):\n        return Counter(self._pods_phase.values())\n\n    def event_cb(self, event):\n        pod_name = event[\"object\"].metadata.name\n        self._pods_phase[pod_name] = event[\"object\"].status.phase\n        if event[\"type\"] == \"DELETED\":\n            del self._pods_phase[pod_name]\n\n\nclass WorkerManager(object):\n    def __init__(self, command, args, num_worker=1, cpu_request=\"1000m\", cpu_limit=\"1000m\",\n            memory_request=\"4096Mi\", memory_limit=\"4096Mi\", pod_priority=None, **kwargs):\n        self._logger = logging.getLogger(\"WorkerManager\")\n        self._command = command\n        self._args = args\n        self._num_worker = num_worker\n        self._cpu_request = cpu_request \n        self._cpu_limit = cpu_limit\n        self._memory_request = memory_request\n        self._memory_limit = memory_limit\n        self._pod_priority = pod_priority\n        self._worker_tracker = WorkerTracker()\n        self._k8s_client = k8s.Client(\n            event_callback=self._worker_tracker.event_cb, **kwargs\n        )\n\n    def start_workers(self, restart_policy=\"OnFailure\"):\n        for i in range(self._num_worker):\n            self._logger.warning(\"Starting worker: %d\" % i)\n            self._add_worker(i, restart_policy=restart_policy)\n\n    def remove_workers(self):\n        for i in range(self._num_worker):\n            pod_name = self._k8s_client.get_pod_name(i)\n            if pod_name in self._worker_tracker._pods_phase:\n                self._logger.warning(\"Deleting worker: %d\", i)\n                self._delete_worker(i)\n\n    def _add_worker(self, worker_id, restart_policy):\n        self._k8s_client.create_worker(\n            worker_id,\n            self._cpu_request,\n            self._cpu_limit,\n            self._memory_request,\n            self._memory_limit,\n            self._pod_priority,\n            command=self._command,\n            args=self._args + [\"--worker_id\", str(worker_id)],\n            restart_policy=restart_policy,\n        )\n\n    def _delete_worker(self, worker_id):\n        self._k8s_client.delete_worker(worker_id)\n\n    def get_counters(self):\n        return self._worker_tracker.get_counters()\n", "path": "elasticdl/python/elasticdl/master/k8s_worker_manager.py"}, {"content": "\"\"\"TaskQueue Implementation\"\"\"\n\nimport logging\nimport random\nimport threading\n\n\nclass _Task(object):\n    \"\"\"Internal representation of a task\"\"\"\n\n    def __init__(self, *, file_name, start, end):\n        self.file_name = file_name\n        self.start = start\n        self.end = end\n\n\nclass _TaskQueue(object):\n    \"\"\"Creates and dispatches Tasks. Keep track of a Task's lifecycle.\"\"\"\n\n    def __init__(self, shards, record_per_task, num_epoch):\n        \"\"\"\n        shards: a dictionary from RecordIO file name to number of records\n        \"\"\"\n        self._logger = logging.getLogger(\"TaskQueue\")\n        self._lock = threading.Lock()\n\n        self._num_epoch = num_epoch\n        self._epoch = 0\n        self._shards = shards\n        self._record_per_task = record_per_task\n\n        self._todo = []\n        # dictionary from task id to Task.\n        self._doing = {}\n        self._task_id = 0\n\n        self._create_tasks()\n\n    def _create_tasks(self):\n        for name, num_records in self._shards.items():\n            for start in range(0, num_records, self._record_per_task):\n                self._todo.append(\n                    _Task(\n                        file_name=name,\n                        start=start,\n                        end=min(start + self._record_per_task, num_records),\n                    )\n                )\n        random.shuffle(self._todo)\n\n    def get(self, worker_id):\n        \"\"\"Return next (task_id, Task) tuple\"\"\"\n\n        with self._lock:\n            if not self._todo and self._epoch < self._num_epoch - 1:\n                # Start a new epoch\n                self._create_tasks()\n                self._epoch += 1\n                self._logger.warning(\"Starting epoch %d\" % self._epoch)\n\n            if not self._todo:\n                # No more tasks\n                return -1, None\n\n            self._task_id += 1\n            task = self._todo.pop()\n            # TODO: Handle timeout of tasks.\n            self._doing[self._task_id] = (worker_id, task)\n\n            return self._task_id, task\n\n    def report(self, task_id, success):\n        \"\"\"Report if the task is successful or not\"\"\"\n\n        with self._lock:\n            _, task = self._doing.pop(task_id, (-1, None))\n            if not task:\n                self._logger.warning(\"Unknown task_id: %d\" % task_id)\n            elif not success:\n                # TODO: keep count of retries.\n                self._todo.append(task)\n\n    def finished(self):\n        \"\"\"Return if all tasks are done\"\"\"\n        return not self._todo and not self._doing\n", "path": "elasticdl/python/elasticdl/master/task_queue.py"}, {"content": "import logging\nimport time\nimport argparse\nimport os\n\nimport grpc\nimport tensorflow as tf\n\ntf.enable_eager_execution()\n\nfrom concurrent import futures\nfrom recordio import File\nfrom elasticdl.proto import master_pb2_grpc\nfrom elasticdl.master.servicer import MasterServicer\nfrom elasticdl.master.task_queue import _TaskQueue\nfrom elasticdl.master.k8s_worker_manager import WorkerManager\nfrom elasticdl.common.model_helper import load_user_model, build_model\n\n\ndef _make_task_queue(data_dir, record_per_task, num_epoch):\n    f_records = {}\n    for f in os.listdir(data_dir):\n        p = os.path.join(data_dir, f)\n        with File(p, \"r\") as rio:\n            f_records[p] = rio.count()\n    return _TaskQueue(f_records, record_per_task, num_epoch)\n\n\ndef _parse_args():\n    parser = argparse.ArgumentParser(description=\"ElasticDL Master\")\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        \"--train_data_dir\",\n        help=\"Training data directory. Files should be in RecordIO format\",\n        required=True,\n    )\n    parser.add_argument(\"--record_per_task\", type=int, required=True)\n    parser.add_argument(\"--num_epoch\", type=int, required=True)\n    parser.add_argument(\n        \"--grads_to_wait\",\n        type=int,\n        help=\"Number of gradients to wait before updating model\",\n        required=True,\n    )\n    parser.add_argument(\n        \"--minibatch_size\",\n        type=int,\n        help=\"Minibatch size used by workers to compute gradients\",\n        required=True,\n    )\n    parser.add_argument(\n        \"--num_worker\",\n        type=int,\n        help=\"the number of workers used in training\",\n        default=0,\n    )\n    parser.add_argument(\n        \"--worker_cpu_request\",\n        help=\"the minimal cpu required by worker in training\",\n        default=\"1000m\",\n    )\n    parser.add_argument(\n        \"--worker_cpu_limit\",\n        help=\"the maximal cpu used by worker in training\",\n        default=\"1000m\",\n    )\n    parser.add_argument(\n        \"--worker_memory_request\",\n        help=\"the minimal memory required by worker in training\",\n        default=\"4096Mi\",\n    )\n    parser.add_argument(\n        \"--worker_memory_limit\",\n        help=\"the maximal memory used by worker in training\",\n        default=\"4096Mi\",\n    )\n    parser.add_argument(\n        \"--worker_pod_priority\",\n        help=\"the requested priority of worker pod\")\n    parser.add_argument(\n        \"--worker_image\", help=\"docker image for worker\", default=None\n    )\n    parser.add_argument(\"--job_name\", help=\"job name\", required=True)\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    return parser.parse_args()\n\n\ndef main():\n    # TODO: pass port via flags.\n    PORT = 50001\n    logger = logging.getLogger(\"master\")\n    args = _parse_args()\n    task_q = _make_task_queue(\n        args.train_data_dir, args.record_per_task, args.num_epoch\n    )\n    model_module = load_user_model(args.model_file)\n    model_inst = model_module.model\n    build_model(model_inst, model_module.feature_columns())\n    optimizer = model_module.optimizer()\n\n    server = grpc.server(futures.ThreadPoolExecutor(max_workers=64))\n    master_pb2_grpc.add_MasterServicer_to_server(\n        MasterServicer(\n            logger,\n            args.grads_to_wait,\n            args.minibatch_size,\n            optimizer,\n            task_q,\n            init_var=model_inst.trainable_variables,\n        ),\n        server,\n    )\n    server.add_insecure_port(\"[::]:{}\".format(PORT))\n    server.start()\n    logger.warning(\"Server started at port: %d\", PORT)\n\n    if args.num_worker:\n        master_addr = \"%s:%d\" % (os.getenv(\"MY_POD_IP\", \"localhost\"), PORT)\n        worker_command = [\"python\"]\n        worker_args = [\n            \"-m\",\n            \"elasticdl.worker.main\",\n            \"--model_file\",\n            args.model_file,\n            \"--master_addr\",\n            master_addr,\n            \"--codec_type\",\n            args.codec_type\n        ]\n\n        worker_manager = WorkerManager(\n            job_name=args.job_name,\n            worker_image=args.worker_image,\n            command=worker_command,\n            args=worker_args,\n            namespace=\"default\",\n            num_worker=args.num_worker,\n            cpu_request=args.worker_cpu_request,\n            cpu_limit=args.worker_cpu_limit,\n            memory_request=args.worker_memory_request,\n            memory_limit=args.worker_memory_limit,\n            pod_priority=args.worker_pod_priority,\n        )\n        worker_manager.start_workers(restart_policy=\"Never\")\n\n    try:\n        while True:\n            if task_q.finished():\n                break\n            time.sleep(30)\n    except KeyboardInterrupt:\n        logger.warning(\"Server stopping\")\n\n    if args.num_worker:\n        # TODO: worker_manager.remove_workers supports synchronized call\n        worker_manager.remove_workers()\n\n    server.stop(0)\n\n\nif __name__ == \"__main__\":\n    logging.basicConfig()\n    main()\n", "path": "elasticdl/python/elasticdl/master/main.py"}], "after_files": [{"content": "import logging\nimport os\nimport time\n\nfrom collections import Counter\nfrom elasticdl.master import k8s\n\n\nclass _WorkerTracker(object):\n    def __init__(self, task_q):\n        # pod name to phase mapping\n        # phase: Pending/Running/Succeeded/Failed/Unknown\n        #   Pending: worker pod not started yet\n        #   Running: worker pod is running\n        #   Succeeded: worker pod finishes all jobs and terminates with no issue.\n        #   Failed: worker pod is killed for some reason\n        #   Unknown: unkown\n        self._pods_phase = {}\n        self._task_q = task_q\n\n    def get_counters(self):\n        return Counter(self._pods_phase.values())\n\n    def event_cb(self, event):\n        pod_name = event[\"object\"].metadata.name\n        self._pods_phase[pod_name] = event[\"object\"].status.phase\n        if event[\"type\"] == \"DELETED\":\n            del self._pods_phase[pod_name]\n            self._task_q.recover_tasks(\n                # TODO: move worker_id and pod name mapping to a separate class \n                int(pod_name.rsplit(\"-\", 1)[1])\n            )\n\n\n\nclass WorkerManager(object):\n    def __init__(self, task_q, command, args, num_worker=1, cpu_request=\"1000m\", cpu_limit=\"1000m\",\n            memory_request=\"4096Mi\", memory_limit=\"4096Mi\", pod_priority=None, **kwargs):\n        self._logger = logging.getLogger(\"WorkerManager\")\n        self._command = command\n        self._args = args\n        self._num_worker = num_worker\n        self._cpu_request = cpu_request \n        self._cpu_limit = cpu_limit\n        self._memory_request = memory_request\n        self._memory_limit = memory_limit\n        self._pod_priority = pod_priority\n        self._worker_tracker = _WorkerTracker(task_q)\n        self._k8s_client = k8s.Client(\n            event_callback=self._worker_tracker.event_cb, **kwargs\n        )\n\n    def start_workers(self, restart_policy=\"OnFailure\"):\n        for i in range(self._num_worker):\n            self._logger.warning(\"Starting worker: %d\" % i)\n            self._add_worker(i, restart_policy=restart_policy)\n\n    def remove_workers(self):\n        for i in range(self._num_worker):\n            pod_name = self._k8s_client.get_pod_name(i)\n            if pod_name in self._worker_tracker._pods_phase:\n                self._logger.warning(\"Deleting worker: %d\", i)\n                self._delete_worker(i)\n\n    def _add_worker(self, worker_id, restart_policy):\n        self._k8s_client.create_worker(\n            worker_id,\n            self._cpu_request,\n            self._cpu_limit,\n            self._memory_request,\n            self._memory_limit,\n            self._pod_priority,\n            command=self._command,\n            args=self._args + [\"--worker_id\", str(worker_id)],\n            restart_policy=restart_policy,\n        )\n\n    def _delete_worker(self, worker_id):\n        self._k8s_client.delete_worker(worker_id)\n\n    def get_counters(self):\n        return self._worker_tracker.get_counters()\n", "path": "elasticdl/python/elasticdl/master/k8s_worker_manager.py"}, {"content": "\"\"\"TaskQueue Implementation\"\"\"\n\nimport logging\nimport random\nimport threading\n\n\nclass _Task(object):\n    \"\"\"Internal representation of a task\"\"\"\n\n    def __init__(self, *, file_name, start, end):\n        self.file_name = file_name\n        self.start = start\n        self.end = end\n\n\nclass _TaskQueue(object):\n    \"\"\"Creates and dispatches Tasks. Keep track of a Task's lifecycle.\"\"\"\n\n    def __init__(self, shards, record_per_task, num_epoch):\n        \"\"\"\n        shards: a dictionary from RecordIO file name to number of records\n        \"\"\"\n        self._logger = logging.getLogger(\"TaskQueue\")\n        self._lock = threading.Lock()\n\n        self._num_epoch = num_epoch\n        self._epoch = 0\n        self._shards = shards\n        self._record_per_task = record_per_task\n\n        self._todo = []\n        # dictionary from task id to Task.\n        self._doing = {}\n        self._task_id = 0\n\n        self._create_tasks()\n\n    def _create_tasks(self):\n        for name, num_records in self._shards.items():\n            for start in range(0, num_records, self._record_per_task):\n                self._todo.append(\n                    _Task(\n                        file_name=name,\n                        start=start,\n                        end=min(start + self._record_per_task, num_records),\n                    )\n                )\n        random.shuffle(self._todo)\n\n    def get(self, worker_id):\n        \"\"\"Return next (task_id, Task) tuple\"\"\"\n\n        with self._lock:\n            if not self._todo and self._epoch < self._num_epoch - 1:\n                # Start a new epoch\n                self._create_tasks()\n                self._epoch += 1\n                self._logger.warning(\"Starting epoch %d\" % self._epoch)\n\n            if not self._todo:\n                # No more tasks\n                return -1, None\n\n            self._task_id += 1\n            task = self._todo.pop()\n            # TODO: Handle timeout of tasks.\n            self._doing[self._task_id] = (worker_id, task)\n\n            return self._task_id, task\n\n    def report(self, task_id, success):\n        \"\"\"Report if the task is successful or not\"\"\"\n\n        with self._lock:\n            _, task = self._doing.pop(task_id, (-1, None))\n            if not task:\n                self._logger.warning(\"Unknown task_id: %d\" % task_id)\n            elif not success:\n                # TODO: keep count of retries.\n                self._todo.append(task)\n\n    def finished(self):\n        \"\"\"Return if all tasks are done\"\"\"\n        return not self._todo and not self._doing\n\n    def recover_tasks(self, worker_id):\n        \"\"\"Recover doing tasks for a dead worker\"\"\"\n\n        with self._lock:\n            tasks = [task for wid, task in self._doing if wid == worker_id]\n            self._todo.extend(tasks)\n", "path": "elasticdl/python/elasticdl/master/task_queue.py"}, {"content": "import logging\nimport time\nimport argparse\nimport os\n\nimport grpc\nimport tensorflow as tf\n\ntf.enable_eager_execution()\n\nfrom concurrent import futures\nfrom recordio import File\nfrom elasticdl.proto import master_pb2_grpc\nfrom elasticdl.master.servicer import MasterServicer\nfrom elasticdl.master.task_queue import _TaskQueue\nfrom elasticdl.master.k8s_worker_manager import WorkerManager\nfrom elasticdl.common.model_helper import load_user_model, build_model\n\n\ndef _make_task_queue(data_dir, record_per_task, num_epoch):\n    f_records = {}\n    for f in os.listdir(data_dir):\n        p = os.path.join(data_dir, f)\n        with File(p, \"r\") as rio:\n            f_records[p] = rio.count()\n    return _TaskQueue(f_records, record_per_task, num_epoch)\n\n\ndef _parse_args():\n    parser = argparse.ArgumentParser(description=\"ElasticDL Master\")\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        \"--train_data_dir\",\n        help=\"Training data directory. Files should be in RecordIO format\",\n        required=True,\n    )\n    parser.add_argument(\"--record_per_task\", type=int, required=True)\n    parser.add_argument(\"--num_epoch\", type=int, required=True)\n    parser.add_argument(\n        \"--grads_to_wait\",\n        type=int,\n        help=\"Number of gradients to wait before updating model\",\n        required=True,\n    )\n    parser.add_argument(\n        \"--minibatch_size\",\n        type=int,\n        help=\"Minibatch size used by workers to compute gradients\",\n        required=True,\n    )\n    parser.add_argument(\n        \"--num_worker\",\n        type=int,\n        help=\"the number of workers used in training\",\n        default=0,\n    )\n    parser.add_argument(\n        \"--worker_cpu_request\",\n        help=\"the minimal cpu required by worker in training\",\n        default=\"1000m\",\n    )\n    parser.add_argument(\n        \"--worker_cpu_limit\",\n        help=\"the maximal cpu used by worker in training\",\n        default=\"1000m\",\n    )\n    parser.add_argument(\n        \"--worker_memory_request\",\n        help=\"the minimal memory required by worker in training\",\n        default=\"4096Mi\",\n    )\n    parser.add_argument(\n        \"--worker_memory_limit\",\n        help=\"the maximal memory used by worker in training\",\n        default=\"4096Mi\",\n    )\n    parser.add_argument(\n        \"--worker_pod_priority\",\n        help=\"the requested priority of worker pod\")\n    parser.add_argument(\n        \"--worker_image\", help=\"docker image for worker\", default=None\n    )\n    parser.add_argument(\"--job_name\", help=\"job name\", required=True)\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    return parser.parse_args()\n\n\ndef main():\n    # TODO: pass port via flags.\n    PORT = 50001\n    logger = logging.getLogger(\"master\")\n    args = _parse_args()\n    task_q = _make_task_queue(\n        args.train_data_dir, args.record_per_task, args.num_epoch\n    )\n    model_module = load_user_model(args.model_file)\n    model_inst = model_module.model\n    build_model(model_inst, model_module.feature_columns())\n    optimizer = model_module.optimizer()\n\n    server = grpc.server(futures.ThreadPoolExecutor(max_workers=64))\n    master_pb2_grpc.add_MasterServicer_to_server(\n        MasterServicer(\n            logger,\n            args.grads_to_wait,\n            args.minibatch_size,\n            optimizer,\n            task_q,\n            init_var=model_inst.trainable_variables,\n        ),\n        server,\n    )\n    server.add_insecure_port(\"[::]:{}\".format(PORT))\n    server.start()\n    logger.warning(\"Server started at port: %d\", PORT)\n\n    if args.num_worker:\n        master_addr = \"%s:%d\" % (os.getenv(\"MY_POD_IP\", \"localhost\"), PORT)\n        worker_command = [\"python\"]\n        worker_args = [\n            \"-m\",\n            \"elasticdl.worker.main\",\n            \"--model_file\",\n            args.model_file,\n            \"--master_addr\",\n            master_addr,\n            \"--codec_type\",\n            args.codec_type\n        ]\n\n        worker_manager = WorkerManager(\n            task_q,\n            job_name=args.job_name,\n            worker_image=args.worker_image,\n            command=worker_command,\n            args=worker_args,\n            namespace=\"default\",\n            num_worker=args.num_worker,\n            cpu_request=args.worker_cpu_request,\n            cpu_limit=args.worker_cpu_limit,\n            memory_request=args.worker_memory_request,\n            memory_limit=args.worker_memory_limit,\n            pod_priority=args.worker_pod_priority,\n        )\n        worker_manager.start_workers(restart_policy=\"Never\")\n\n    try:\n        while True:\n            if task_q.finished():\n                break\n            time.sleep(30)\n    except KeyboardInterrupt:\n        logger.warning(\"Server stopping\")\n\n    if args.num_worker:\n        # TODO: worker_manager.remove_workers supports synchronized call\n        worker_manager.remove_workers()\n\n    server.stop(0)\n\n\nif __name__ == \"__main__\":\n    logging.basicConfig()\n    main()\n", "path": "elasticdl/python/elasticdl/master/main.py"}]}