We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Master needs to relaunch a failed ps/worker pod for OOM workaround
TensorFlow 2.0 has memory leak issues in the eager mode. A ps/worker pod may fail due to OOM after running for a while.

To work around it, the master can delete the failed pod so it can be relaunched.

We can remove this workaround after all memory leak issues are resolved.


--- 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/master/k8s_instance_manager.py`
Content:
```
1 import copy
2 import itertools
3 import threading
4 from collections import Counter
5 
6 from elasticdl.python.common import k8s_client as k8s
7 from elasticdl.python.common.log_utils import default_logger as logger
8 
9 _SERVICE_ADDR_SEP = ","
10 
11 
12 class InstanceManager(object):
13     def __init__(
14         self,
15         task_d,
16         num_workers=1,
17         worker_command=None,
18         worker_args=None,
19         worker_resource_request="cpu=1,memory=4096Mi",
20         worker_resource_limit="cpu=1,memory=4096Mi",
21         worker_pod_priority=None,
22         num_ps=0,
23         ps_command=None,
24         ps_args=None,
25         ps_resource_request="cpu=1,memory=4096Mi",
26         ps_resource_limit="cpu=1,memory=4096Mi",
27         ps_pod_priority=None,
28         volume=None,
29         image_pull_policy=None,
30         restart_policy="Never",
31         envs=None,
32         **kwargs
33     ):
34         self._num_workers = num_workers
35         self._worker_command = worker_command
36         self._worker_args = worker_args
37         self._worker_resource_request = worker_resource_request
38         self._worker_resource_limit = worker_resource_limit
39         self._worker_pod_priority = worker_pod_priority
40 
41         self._num_ps = num_ps
42         self._ps_command = ps_command
43         self._ps_args = ps_args
44         self._ps_resource_request = ps_resource_request
45         self._ps_resource_limit = ps_resource_limit
46         self._ps_pod_priority = ps_pod_priority
47 
48         self._restart_policy = restart_policy
49         self._volume = volume
50         self._image_pull_policy = image_pull_policy
51         self._envs = envs
52         self._task_d = task_d
53         self._next_worker_id = itertools.count().__next__
54 
55         # Protects followed variables, which are accessed from event_cb.
56         self._lock = threading.Lock()
57         # worker id to (pod name, phase) mapping
58         # phase: None/Pending/Running/Succeeded/Failed/Unknown
59         #   None: worker was just launched, haven't received event yet.
60         #   Pending: worker pod not started yet
61         #   Running: worker pod is running
62         #   Succeeded: worker pod finishes all tasks and terminates with
63         #       no issue.
64         #   Failed: worker pod is killed for some reason
65         #   Unknown: unknown
66         self._worker_pods_phase = {}
67         # pod name to worker id mapping
68         self._worker_pod_name_to_id = {}
69 
70         self._relaunch_deleted_live_worker = True
71 
72         self._ps_pods_phase = {}
73         self._ps_pod_name_to_id = {}
74         self._relaunch_deleted_live_ps = True
75 
76         self._k8s_client = k8s.Client(event_callback=self._event_cb, **kwargs)
77         self._ps_addrs = self._get_addrs(
78             self._num_ps, self._k8s_client.get_ps_service_address
79         )
80         # TODO: Select a worker address to be used for broadcasting model
81         # parameters under allreduce-strategy.
82         self._worker_addrs = self._get_addrs(
83             self._num_workers, self._k8s_client.get_worker_service_address
84         )
85 
86     def _start_worker(self, worker_id):
87         logger.info("Starting worker: %d" % worker_id)
88         with self._lock:
89             pod = self._k8s_client.create_worker(
90                 worker_id=worker_id,
91                 resource_requests=self._worker_resource_request,
92                 resource_limits=self._worker_resource_limit,
93                 pod_priority=self._worker_pod_priority,
94                 volume=self._volume,
95                 image_pull_policy=self._image_pull_policy,
96                 command=self._worker_command,
97                 args=self._worker_args
98                 + ["--worker_id", str(worker_id)]
99                 + ["--ps_addrs", self._ps_addrs],
100                 restart_policy=self._restart_policy,
101                 ps_addrs=self._ps_addrs,
102                 envs=copy.deepcopy(self._envs),
103             )
104             name = pod.metadata.name
105             self._worker_pod_name_to_id[name] = worker_id
106             self._worker_pods_phase[worker_id] = (name, None)
107             self._k8s_client.create_worker_service(worker_id)
108 
109     def _start_ps(self, ps_id):
110         logger.info("Starting PS: %d" % ps_id)
111         with self._lock:
112             pod = self._k8s_client.create_ps(
113                 ps_id=ps_id,
114                 resource_requests=self._ps_resource_request,
115                 resource_limits=self._ps_resource_limit,
116                 pod_priority=self._ps_pod_priority,
117                 volume=self._volume,
118                 image_pull_policy=self._image_pull_policy,
119                 command=self._ps_command,
120                 args=self._ps_args + ["--ps_id", str(ps_id)],
121                 restart_policy=self._restart_policy,
122                 envs=copy.deepcopy(self._envs),
123             )
124             name = pod.metadata.name
125             self._ps_pod_name_to_id[name] = ps_id
126             self._ps_pods_phase[ps_id] = (name, None)
127             self._k8s_client.create_ps_service(ps_id)
128 
129     def _get_addrs(self, num_addrs, addr_get_fn):
130         addrs = []
131         for addr_id in range(num_addrs):
132             addrs.append(addr_get_fn(addr_id))
133         return _SERVICE_ADDR_SEP.join(addrs)
134 
135     @staticmethod
136     def _update_addr(old_addr, new_addr, addrs, addr_get_fn):
137         addrs_list = addrs.split(_SERVICE_ADDR_SEP)
138         addrs_list[addrs_list.index(addr_get_fn(old_addr))] = addr_get_fn(
139             new_addr
140         )
141         return _SERVICE_ADDR_SEP.join(addrs_list)
142 
143     def update_status(self, status):
144         master_name = self._k8s_client.get_master_pod_name()
145         self._k8s_client.patch_labels_to_pod(
146             master_name, labels_dict={"status": status}
147         )
148 
149     def start_workers(self):
150         for _ in range(self._num_workers):
151             self._start_worker(self._next_worker_id())
152 
153     def start_parameter_servers(self):
154         for i in range(self._num_ps):
155             self._start_ps(i)
156 
157     def _remove_worker(self, worker_id):
158         logger.info("Removing worker: %d", worker_id)
159         with self._lock:
160             if worker_id not in self._worker_pods_phase:
161                 logger.error("Unknown worker id: %s" % worker_id)
162                 return
163 
164         # TODO: change _k8s_client to accept pod name instead of worker id.
165         self._k8s_client.delete_worker(worker_id)
166 
167     def _remove_ps(self, ps_id):
168         logger.info("Removing PS: %d", ps_id)
169         with self._lock:
170             if ps_id not in self._ps_pods_phase:
171                 logger.error("Unknown PS id: %s" % ps_id)
172                 return
173 
174         self._k8s_client.delete_ps(ps_id)
175 
176     def stop_relaunch_and_remove_workers(self):
177         with self._lock:
178             self._relaunch_deleted_live_worker = False
179             for worker_id in self._worker_pods_phase:
180                 self._k8s_client.delete_worker(worker_id)
181 
182     def stop_relaunch_and_remove_all_ps(self):
183         with self._lock:
184             self._relaunch_deleted_live_ps = False
185             for ps_id in self._ps_pods_phase:
186                 self._k8s_client.delete_ps(ps_id)
187 
188     def get_worker_counter(self):
189         with self._lock:
190             return Counter([v for _, v in self._worker_pods_phase.values()])
191 
192     def get_ps_counter(self):
193         with self._lock:
194             return Counter([v for _, v in self._ps_pods_phase.values()])
195 
196     def _event_cb(self, event):
197         evt_obj = event.get("object")
198         evt_type = event.get("type")
199         if not evt_obj or not evt_type:
200             logger.error("Event doesn't have object or type: %s" % event)
201             return
202 
203         if evt_obj.kind != "Pod":
204             # We only care about pod related events
205             return
206 
207         pod_name = evt_obj.metadata.name
208         phase = evt_obj.status.phase
209         logger.info(
210             "Got event %s, phase %s for pod: %s" % (evt_type, phase, pod_name)
211         )
212         if pod_name == self._k8s_client.get_master_pod_name():
213             # No need to care about master pod
214             return
215 
216         relaunch_worker = False
217         relaunch_ps = False
218         worker_id = -1
219         ps_id = -1
220         with self._lock:
221             if pod_name in self._worker_pod_name_to_id:
222                 worker_id = self._worker_pod_name_to_id.get(pod_name)
223                 self._worker_pods_phase[worker_id] = (pod_name, phase)
224                 if evt_type == "DELETED":
225                     del self._worker_pods_phase[worker_id]
226                     del self._worker_pod_name_to_id[pod_name]
227                     self._task_d.recover_tasks(worker_id)
228 
229                     # If a deleted pod was not "Succeeded", relaunch a worker.
230                     relaunch_worker = (
231                         self._relaunch_deleted_live_worker
232                         and phase != "Succeeded"
233                     )
234 
235             elif pod_name in self._ps_pod_name_to_id:
236                 ps_id = self._ps_pod_name_to_id.get(pod_name)
237                 self._ps_pods_phase[ps_id] = (pod_name, phase)
238                 if evt_type == "DELETED":
239                     del self._ps_pods_phase[ps_id]
240                     del self._ps_pod_name_to_id[pod_name]
241                     relaunch_ps = self._relaunch_deleted_live_ps
242             else:
243                 logger.error("Unknown pod name: %s" % pod_name)
244                 return
245 
246         if relaunch_worker and worker_id >= 0:
247             logger.info("Relaunching worker.")
248             new_worker_id = self._next_worker_id()
249             self._start_worker(new_worker_id)
250             self._update_addr(
251                 worker_id,
252                 new_worker_id,
253                 self._worker_addrs,
254                 addr_get_fn=self._k8s_client.get_worker_service_address,
255             )
256         elif relaunch_ps:
257             logger.info("Relaunching ps.")
258             # Note: the ID and service address for relaunched parameter
259             # server are intentionally left unchanged to support fault
260             # tolerance.
261             self._start_ps(ps_id)
262 
263     @property
264     def ps_addrs(self):
265         return self._ps_addrs
266
```
--- 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/master/k8s_instance_manager.py b/elasticdl/python/master/k8s_instance_manager.py
--- a/elasticdl/python/master/k8s_instance_manager.py
+++ b/elasticdl/python/master/k8s_instance_manager.py
@@ -73,6 +73,8 @@
         self._ps_pod_name_to_id = {}
         self._relaunch_deleted_live_ps = True
 
+        self._failed_pods = []
+
         self._k8s_client = k8s.Client(event_callback=self._event_cb, **kwargs)
         self._ps_addrs = self._get_addrs(
             self._num_ps, self._k8s_client.get_ps_service_address
@@ -218,10 +220,18 @@
         worker_id = -1
         ps_id = -1
         with self._lock:
+            if pod_name in self._failed_pods:
+                return
             if pod_name in self._worker_pod_name_to_id:
                 worker_id = self._worker_pod_name_to_id.get(pod_name)
                 self._worker_pods_phase[worker_id] = (pod_name, phase)
-                if evt_type == "DELETED":
+                # Workaround for memory leak issues in tf eager mode.
+                # A pod may fail due to OOM from tf eager mode memory leak.
+                failed_pod = False
+                if evt_type == "MODIFIED" and phase == "Failed":
+                    self._failed_pods.append(pod_name)
+                    failed_pod = True
+                if evt_type == "DELETED" or failed_pod:
                     del self._worker_pods_phase[worker_id]
                     del self._worker_pod_name_to_id[pod_name]
                     self._task_d.recover_tasks(worker_id)
@@ -235,7 +245,13 @@
             elif pod_name in self._ps_pod_name_to_id:
                 ps_id = self._ps_pod_name_to_id.get(pod_name)
                 self._ps_pods_phase[ps_id] = (pod_name, phase)
-                if evt_type == "DELETED":
+                # Workaround for memory leak issues in tf eager mode.
+                # A pod may fail due to OOM from tf eager mode memory leak.
+                failed_pod = False
+                if evt_type == "MODIFIED" and phase == "Failed":
+                    self._failed_pods.append(pod_name)
+                    failed_pod = True
+                if evt_type == "DELETED" or failed_pod:
                     del self._ps_pods_phase[ps_id]
                     del self._ps_pod_name_to_id[pod_name]
                     relaunch_ps = self._relaunch_deleted_live_ps

{"golden_diff": "diff --git a/elasticdl/python/master/k8s_instance_manager.py b/elasticdl/python/master/k8s_instance_manager.py\n--- a/elasticdl/python/master/k8s_instance_manager.py\n+++ b/elasticdl/python/master/k8s_instance_manager.py\n@@ -73,6 +73,8 @@\n         self._ps_pod_name_to_id = {}\n         self._relaunch_deleted_live_ps = True\n \n+        self._failed_pods = []\n+\n         self._k8s_client = k8s.Client(event_callback=self._event_cb, **kwargs)\n         self._ps_addrs = self._get_addrs(\n             self._num_ps, self._k8s_client.get_ps_service_address\n@@ -218,10 +220,18 @@\n         worker_id = -1\n         ps_id = -1\n         with self._lock:\n+            if pod_name in self._failed_pods:\n+                return\n             if pod_name in self._worker_pod_name_to_id:\n                 worker_id = self._worker_pod_name_to_id.get(pod_name)\n                 self._worker_pods_phase[worker_id] = (pod_name, phase)\n-                if evt_type == \"DELETED\":\n+                # Workaround for memory leak issues in tf eager mode.\n+                # A pod may fail due to OOM from tf eager mode memory leak.\n+                failed_pod = False\n+                if evt_type == \"MODIFIED\" and phase == \"Failed\":\n+                    self._failed_pods.append(pod_name)\n+                    failed_pod = True\n+                if evt_type == \"DELETED\" or failed_pod:\n                     del self._worker_pods_phase[worker_id]\n                     del self._worker_pod_name_to_id[pod_name]\n                     self._task_d.recover_tasks(worker_id)\n@@ -235,7 +245,13 @@\n             elif pod_name in self._ps_pod_name_to_id:\n                 ps_id = self._ps_pod_name_to_id.get(pod_name)\n                 self._ps_pods_phase[ps_id] = (pod_name, phase)\n-                if evt_type == \"DELETED\":\n+                # Workaround for memory leak issues in tf eager mode.\n+                # A pod may fail due to OOM from tf eager mode memory leak.\n+                failed_pod = False\n+                if evt_type == \"MODIFIED\" and phase == \"Failed\":\n+                    self._failed_pods.append(pod_name)\n+                    failed_pod = True\n+                if evt_type == \"DELETED\" or failed_pod:\n                     del self._ps_pods_phase[ps_id]\n                     del self._ps_pod_name_to_id[pod_name]\n                     relaunch_ps = self._relaunch_deleted_live_ps\n", "issue": "Master needs to relaunch a failed ps/worker pod for OOM workaround\nTensorFlow 2.0 has memory leak issues in the eager mode. A ps/worker pod may fail due to OOM after running for a while.\r\n\r\nTo work around it, the master can delete the failed pod so it can be relaunched.\r\n\r\nWe can remove this workaround after all memory leak issues are resolved.\r\n\n", "before_files": [{"content": "import copy\nimport itertools\nimport threading\nfrom collections import Counter\n\nfrom elasticdl.python.common import k8s_client as k8s\nfrom elasticdl.python.common.log_utils import default_logger as logger\n\n_SERVICE_ADDR_SEP = \",\"\n\n\nclass InstanceManager(object):\n    def __init__(\n        self,\n        task_d,\n        num_workers=1,\n        worker_command=None,\n        worker_args=None,\n        worker_resource_request=\"cpu=1,memory=4096Mi\",\n        worker_resource_limit=\"cpu=1,memory=4096Mi\",\n        worker_pod_priority=None,\n        num_ps=0,\n        ps_command=None,\n        ps_args=None,\n        ps_resource_request=\"cpu=1,memory=4096Mi\",\n        ps_resource_limit=\"cpu=1,memory=4096Mi\",\n        ps_pod_priority=None,\n        volume=None,\n        image_pull_policy=None,\n        restart_policy=\"Never\",\n        envs=None,\n        **kwargs\n    ):\n        self._num_workers = num_workers\n        self._worker_command = worker_command\n        self._worker_args = worker_args\n        self._worker_resource_request = worker_resource_request\n        self._worker_resource_limit = worker_resource_limit\n        self._worker_pod_priority = worker_pod_priority\n\n        self._num_ps = num_ps\n        self._ps_command = ps_command\n        self._ps_args = ps_args\n        self._ps_resource_request = ps_resource_request\n        self._ps_resource_limit = ps_resource_limit\n        self._ps_pod_priority = ps_pod_priority\n\n        self._restart_policy = restart_policy\n        self._volume = volume\n        self._image_pull_policy = image_pull_policy\n        self._envs = envs\n        self._task_d = task_d\n        self._next_worker_id = itertools.count().__next__\n\n        # Protects followed variables, which are accessed from event_cb.\n        self._lock = threading.Lock()\n        # worker id to (pod name, phase) mapping\n        # phase: None/Pending/Running/Succeeded/Failed/Unknown\n        #   None: worker was just launched, haven't received event yet.\n        #   Pending: worker pod not started yet\n        #   Running: worker pod is running\n        #   Succeeded: worker pod finishes all tasks and terminates with\n        #       no issue.\n        #   Failed: worker pod is killed for some reason\n        #   Unknown: unknown\n        self._worker_pods_phase = {}\n        # pod name to worker id mapping\n        self._worker_pod_name_to_id = {}\n\n        self._relaunch_deleted_live_worker = True\n\n        self._ps_pods_phase = {}\n        self._ps_pod_name_to_id = {}\n        self._relaunch_deleted_live_ps = True\n\n        self._k8s_client = k8s.Client(event_callback=self._event_cb, **kwargs)\n        self._ps_addrs = self._get_addrs(\n            self._num_ps, self._k8s_client.get_ps_service_address\n        )\n        # TODO: Select a worker address to be used for broadcasting model\n        # parameters under allreduce-strategy.\n        self._worker_addrs = self._get_addrs(\n            self._num_workers, self._k8s_client.get_worker_service_address\n        )\n\n    def _start_worker(self, worker_id):\n        logger.info(\"Starting worker: %d\" % worker_id)\n        with self._lock:\n            pod = self._k8s_client.create_worker(\n                worker_id=worker_id,\n                resource_requests=self._worker_resource_request,\n                resource_limits=self._worker_resource_limit,\n                pod_priority=self._worker_pod_priority,\n                volume=self._volume,\n                image_pull_policy=self._image_pull_policy,\n                command=self._worker_command,\n                args=self._worker_args\n                + [\"--worker_id\", str(worker_id)]\n                + [\"--ps_addrs\", self._ps_addrs],\n                restart_policy=self._restart_policy,\n                ps_addrs=self._ps_addrs,\n                envs=copy.deepcopy(self._envs),\n            )\n            name = pod.metadata.name\n            self._worker_pod_name_to_id[name] = worker_id\n            self._worker_pods_phase[worker_id] = (name, None)\n            self._k8s_client.create_worker_service(worker_id)\n\n    def _start_ps(self, ps_id):\n        logger.info(\"Starting PS: %d\" % ps_id)\n        with self._lock:\n            pod = self._k8s_client.create_ps(\n                ps_id=ps_id,\n                resource_requests=self._ps_resource_request,\n                resource_limits=self._ps_resource_limit,\n                pod_priority=self._ps_pod_priority,\n                volume=self._volume,\n                image_pull_policy=self._image_pull_policy,\n                command=self._ps_command,\n                args=self._ps_args + [\"--ps_id\", str(ps_id)],\n                restart_policy=self._restart_policy,\n                envs=copy.deepcopy(self._envs),\n            )\n            name = pod.metadata.name\n            self._ps_pod_name_to_id[name] = ps_id\n            self._ps_pods_phase[ps_id] = (name, None)\n            self._k8s_client.create_ps_service(ps_id)\n\n    def _get_addrs(self, num_addrs, addr_get_fn):\n        addrs = []\n        for addr_id in range(num_addrs):\n            addrs.append(addr_get_fn(addr_id))\n        return _SERVICE_ADDR_SEP.join(addrs)\n\n    @staticmethod\n    def _update_addr(old_addr, new_addr, addrs, addr_get_fn):\n        addrs_list = addrs.split(_SERVICE_ADDR_SEP)\n        addrs_list[addrs_list.index(addr_get_fn(old_addr))] = addr_get_fn(\n            new_addr\n        )\n        return _SERVICE_ADDR_SEP.join(addrs_list)\n\n    def update_status(self, status):\n        master_name = self._k8s_client.get_master_pod_name()\n        self._k8s_client.patch_labels_to_pod(\n            master_name, labels_dict={\"status\": status}\n        )\n\n    def start_workers(self):\n        for _ in range(self._num_workers):\n            self._start_worker(self._next_worker_id())\n\n    def start_parameter_servers(self):\n        for i in range(self._num_ps):\n            self._start_ps(i)\n\n    def _remove_worker(self, worker_id):\n        logger.info(\"Removing worker: %d\", worker_id)\n        with self._lock:\n            if worker_id not in self._worker_pods_phase:\n                logger.error(\"Unknown worker id: %s\" % worker_id)\n                return\n\n        # TODO: change _k8s_client to accept pod name instead of worker id.\n        self._k8s_client.delete_worker(worker_id)\n\n    def _remove_ps(self, ps_id):\n        logger.info(\"Removing PS: %d\", ps_id)\n        with self._lock:\n            if ps_id not in self._ps_pods_phase:\n                logger.error(\"Unknown PS id: %s\" % ps_id)\n                return\n\n        self._k8s_client.delete_ps(ps_id)\n\n    def stop_relaunch_and_remove_workers(self):\n        with self._lock:\n            self._relaunch_deleted_live_worker = False\n            for worker_id in self._worker_pods_phase:\n                self._k8s_client.delete_worker(worker_id)\n\n    def stop_relaunch_and_remove_all_ps(self):\n        with self._lock:\n            self._relaunch_deleted_live_ps = False\n            for ps_id in self._ps_pods_phase:\n                self._k8s_client.delete_ps(ps_id)\n\n    def get_worker_counter(self):\n        with self._lock:\n            return Counter([v for _, v in self._worker_pods_phase.values()])\n\n    def get_ps_counter(self):\n        with self._lock:\n            return Counter([v for _, v in self._ps_pods_phase.values()])\n\n    def _event_cb(self, event):\n        evt_obj = event.get(\"object\")\n        evt_type = event.get(\"type\")\n        if not evt_obj or not evt_type:\n            logger.error(\"Event doesn't have object or type: %s\" % event)\n            return\n\n        if evt_obj.kind != \"Pod\":\n            # We only care about pod related events\n            return\n\n        pod_name = evt_obj.metadata.name\n        phase = evt_obj.status.phase\n        logger.info(\n            \"Got event %s, phase %s for pod: %s\" % (evt_type, phase, pod_name)\n        )\n        if pod_name == self._k8s_client.get_master_pod_name():\n            # No need to care about master pod\n            return\n\n        relaunch_worker = False\n        relaunch_ps = False\n        worker_id = -1\n        ps_id = -1\n        with self._lock:\n            if pod_name in self._worker_pod_name_to_id:\n                worker_id = self._worker_pod_name_to_id.get(pod_name)\n                self._worker_pods_phase[worker_id] = (pod_name, phase)\n                if evt_type == \"DELETED\":\n                    del self._worker_pods_phase[worker_id]\n                    del self._worker_pod_name_to_id[pod_name]\n                    self._task_d.recover_tasks(worker_id)\n\n                    # If a deleted pod was not \"Succeeded\", relaunch a worker.\n                    relaunch_worker = (\n                        self._relaunch_deleted_live_worker\n                        and phase != \"Succeeded\"\n                    )\n\n            elif pod_name in self._ps_pod_name_to_id:\n                ps_id = self._ps_pod_name_to_id.get(pod_name)\n                self._ps_pods_phase[ps_id] = (pod_name, phase)\n                if evt_type == \"DELETED\":\n                    del self._ps_pods_phase[ps_id]\n                    del self._ps_pod_name_to_id[pod_name]\n                    relaunch_ps = self._relaunch_deleted_live_ps\n            else:\n                logger.error(\"Unknown pod name: %s\" % pod_name)\n                return\n\n        if relaunch_worker and worker_id >= 0:\n            logger.info(\"Relaunching worker.\")\n            new_worker_id = self._next_worker_id()\n            self._start_worker(new_worker_id)\n            self._update_addr(\n                worker_id,\n                new_worker_id,\n                self._worker_addrs,\n                addr_get_fn=self._k8s_client.get_worker_service_address,\n            )\n        elif relaunch_ps:\n            logger.info(\"Relaunching ps.\")\n            # Note: the ID and service address for relaunched parameter\n            # server are intentionally left unchanged to support fault\n            # tolerance.\n            self._start_ps(ps_id)\n\n    @property\n    def ps_addrs(self):\n        return self._ps_addrs\n", "path": "elasticdl/python/master/k8s_instance_manager.py"}], "after_files": [{"content": "import copy\nimport itertools\nimport threading\nfrom collections import Counter\n\nfrom elasticdl.python.common import k8s_client as k8s\nfrom elasticdl.python.common.log_utils import default_logger as logger\n\n_SERVICE_ADDR_SEP = \",\"\n\n\nclass InstanceManager(object):\n    def __init__(\n        self,\n        task_d,\n        num_workers=1,\n        worker_command=None,\n        worker_args=None,\n        worker_resource_request=\"cpu=1,memory=4096Mi\",\n        worker_resource_limit=\"cpu=1,memory=4096Mi\",\n        worker_pod_priority=None,\n        num_ps=0,\n        ps_command=None,\n        ps_args=None,\n        ps_resource_request=\"cpu=1,memory=4096Mi\",\n        ps_resource_limit=\"cpu=1,memory=4096Mi\",\n        ps_pod_priority=None,\n        volume=None,\n        image_pull_policy=None,\n        restart_policy=\"Never\",\n        envs=None,\n        **kwargs\n    ):\n        self._num_workers = num_workers\n        self._worker_command = worker_command\n        self._worker_args = worker_args\n        self._worker_resource_request = worker_resource_request\n        self._worker_resource_limit = worker_resource_limit\n        self._worker_pod_priority = worker_pod_priority\n\n        self._num_ps = num_ps\n        self._ps_command = ps_command\n        self._ps_args = ps_args\n        self._ps_resource_request = ps_resource_request\n        self._ps_resource_limit = ps_resource_limit\n        self._ps_pod_priority = ps_pod_priority\n\n        self._restart_policy = restart_policy\n        self._volume = volume\n        self._image_pull_policy = image_pull_policy\n        self._envs = envs\n        self._task_d = task_d\n        self._next_worker_id = itertools.count().__next__\n\n        # Protects followed variables, which are accessed from event_cb.\n        self._lock = threading.Lock()\n        # worker id to (pod name, phase) mapping\n        # phase: None/Pending/Running/Succeeded/Failed/Unknown\n        #   None: worker was just launched, haven't received event yet.\n        #   Pending: worker pod not started yet\n        #   Running: worker pod is running\n        #   Succeeded: worker pod finishes all tasks and terminates with\n        #       no issue.\n        #   Failed: worker pod is killed for some reason\n        #   Unknown: unknown\n        self._worker_pods_phase = {}\n        # pod name to worker id mapping\n        self._worker_pod_name_to_id = {}\n\n        self._relaunch_deleted_live_worker = True\n\n        self._ps_pods_phase = {}\n        self._ps_pod_name_to_id = {}\n        self._relaunch_deleted_live_ps = True\n\n        self._failed_pods = []\n\n        self._k8s_client = k8s.Client(event_callback=self._event_cb, **kwargs)\n        self._ps_addrs = self._get_addrs(\n            self._num_ps, self._k8s_client.get_ps_service_address\n        )\n        # TODO: Select a worker address to be used for broadcasting model\n        # parameters under allreduce-strategy.\n        self._worker_addrs = self._get_addrs(\n            self._num_workers, self._k8s_client.get_worker_service_address\n        )\n\n    def _start_worker(self, worker_id):\n        logger.info(\"Starting worker: %d\" % worker_id)\n        with self._lock:\n            pod = self._k8s_client.create_worker(\n                worker_id=worker_id,\n                resource_requests=self._worker_resource_request,\n                resource_limits=self._worker_resource_limit,\n                pod_priority=self._worker_pod_priority,\n                volume=self._volume,\n                image_pull_policy=self._image_pull_policy,\n                command=self._worker_command,\n                args=self._worker_args\n                + [\"--worker_id\", str(worker_id)]\n                + [\"--ps_addrs\", self._ps_addrs],\n                restart_policy=self._restart_policy,\n                ps_addrs=self._ps_addrs,\n                envs=copy.deepcopy(self._envs),\n            )\n            name = pod.metadata.name\n            self._worker_pod_name_to_id[name] = worker_id\n            self._worker_pods_phase[worker_id] = (name, None)\n            self._k8s_client.create_worker_service(worker_id)\n\n    def _start_ps(self, ps_id):\n        logger.info(\"Starting PS: %d\" % ps_id)\n        with self._lock:\n            pod = self._k8s_client.create_ps(\n                ps_id=ps_id,\n                resource_requests=self._ps_resource_request,\n                resource_limits=self._ps_resource_limit,\n                pod_priority=self._ps_pod_priority,\n                volume=self._volume,\n                image_pull_policy=self._image_pull_policy,\n                command=self._ps_command,\n                args=self._ps_args + [\"--ps_id\", str(ps_id)],\n                restart_policy=self._restart_policy,\n                envs=copy.deepcopy(self._envs),\n            )\n            name = pod.metadata.name\n            self._ps_pod_name_to_id[name] = ps_id\n            self._ps_pods_phase[ps_id] = (name, None)\n            self._k8s_client.create_ps_service(ps_id)\n\n    def _get_addrs(self, num_addrs, addr_get_fn):\n        addrs = []\n        for addr_id in range(num_addrs):\n            addrs.append(addr_get_fn(addr_id))\n        return _SERVICE_ADDR_SEP.join(addrs)\n\n    @staticmethod\n    def _update_addr(old_addr, new_addr, addrs, addr_get_fn):\n        addrs_list = addrs.split(_SERVICE_ADDR_SEP)\n        addrs_list[addrs_list.index(addr_get_fn(old_addr))] = addr_get_fn(\n            new_addr\n        )\n        return _SERVICE_ADDR_SEP.join(addrs_list)\n\n    def update_status(self, status):\n        master_name = self._k8s_client.get_master_pod_name()\n        self._k8s_client.patch_labels_to_pod(\n            master_name, labels_dict={\"status\": status}\n        )\n\n    def start_workers(self):\n        for _ in range(self._num_workers):\n            self._start_worker(self._next_worker_id())\n\n    def start_parameter_servers(self):\n        for i in range(self._num_ps):\n            self._start_ps(i)\n\n    def _remove_worker(self, worker_id):\n        logger.info(\"Removing worker: %d\", worker_id)\n        with self._lock:\n            if worker_id not in self._worker_pods_phase:\n                logger.error(\"Unknown worker id: %s\" % worker_id)\n                return\n\n        # TODO: change _k8s_client to accept pod name instead of worker id.\n        self._k8s_client.delete_worker(worker_id)\n\n    def _remove_ps(self, ps_id):\n        logger.info(\"Removing PS: %d\", ps_id)\n        with self._lock:\n            if ps_id not in self._ps_pods_phase:\n                logger.error(\"Unknown PS id: %s\" % ps_id)\n                return\n\n        self._k8s_client.delete_ps(ps_id)\n\n    def stop_relaunch_and_remove_workers(self):\n        with self._lock:\n            self._relaunch_deleted_live_worker = False\n            for worker_id in self._worker_pods_phase:\n                self._k8s_client.delete_worker(worker_id)\n\n    def stop_relaunch_and_remove_all_ps(self):\n        with self._lock:\n            self._relaunch_deleted_live_ps = False\n            for ps_id in self._ps_pods_phase:\n                self._k8s_client.delete_ps(ps_id)\n\n    def get_worker_counter(self):\n        with self._lock:\n            return Counter([v for _, v in self._worker_pods_phase.values()])\n\n    def get_ps_counter(self):\n        with self._lock:\n            return Counter([v for _, v in self._ps_pods_phase.values()])\n\n    def _event_cb(self, event):\n        evt_obj = event.get(\"object\")\n        evt_type = event.get(\"type\")\n        if not evt_obj or not evt_type:\n            logger.error(\"Event doesn't have object or type: %s\" % event)\n            return\n\n        if evt_obj.kind != \"Pod\":\n            # We only care about pod related events\n            return\n\n        pod_name = evt_obj.metadata.name\n        phase = evt_obj.status.phase\n        logger.info(\n            \"Got event %s, phase %s for pod: %s\" % (evt_type, phase, pod_name)\n        )\n        if pod_name == self._k8s_client.get_master_pod_name():\n            # No need to care about master pod\n            return\n\n        relaunch_worker = False\n        relaunch_ps = False\n        worker_id = -1\n        ps_id = -1\n        with self._lock:\n            if pod_name in self._failed_pods:\n                return\n            if pod_name in self._worker_pod_name_to_id:\n                worker_id = self._worker_pod_name_to_id.get(pod_name)\n                self._worker_pods_phase[worker_id] = (pod_name, phase)\n                # Workaround for memory leak issues in tf eager mode.\n                # A pod may fail due to OOM from tf eager mode memory leak.\n                failed_pod = False\n                if evt_type == \"MODIFIED\" and phase == \"Failed\":\n                    self._failed_pods.append(pod_name)\n                    failed_pod = True\n                if evt_type == \"DELETED\" or failed_pod:\n                    del self._worker_pods_phase[worker_id]\n                    del self._worker_pod_name_to_id[pod_name]\n                    self._task_d.recover_tasks(worker_id)\n\n                    # If a deleted pod was not \"Succeeded\", relaunch a worker.\n                    relaunch_worker = (\n                        self._relaunch_deleted_live_worker\n                        and phase != \"Succeeded\"\n                    )\n\n            elif pod_name in self._ps_pod_name_to_id:\n                ps_id = self._ps_pod_name_to_id.get(pod_name)\n                self._ps_pods_phase[ps_id] = (pod_name, phase)\n                # Workaround for memory leak issues in tf eager mode.\n                # A pod may fail due to OOM from tf eager mode memory leak.\n                failed_pod = False\n                if evt_type == \"MODIFIED\" and phase == \"Failed\":\n                    self._failed_pods.append(pod_name)\n                    failed_pod = True\n                if evt_type == \"DELETED\" or failed_pod:\n                    del self._ps_pods_phase[ps_id]\n                    del self._ps_pod_name_to_id[pod_name]\n                    relaunch_ps = self._relaunch_deleted_live_ps\n            else:\n                logger.error(\"Unknown pod name: %s\" % pod_name)\n                return\n\n        if relaunch_worker and worker_id >= 0:\n            logger.info(\"Relaunching worker.\")\n            new_worker_id = self._next_worker_id()\n            self._start_worker(new_worker_id)\n            self._update_addr(\n                worker_id,\n                new_worker_id,\n                self._worker_addrs,\n                addr_get_fn=self._k8s_client.get_worker_service_address,\n            )\n        elif relaunch_ps:\n            logger.info(\"Relaunching ps.\")\n            # Note: the ID and service address for relaunched parameter\n            # server are intentionally left unchanged to support fault\n            # tolerance.\n            self._start_ps(ps_id)\n\n    @property\n    def ps_addrs(self):\n        return self._ps_addrs\n", "path": "elasticdl/python/master/k8s_instance_manager.py"}]}