Repository: sql-machine-learning/elasticdl

Issue:
Recover 'doing' tasks from preempted workers



Discussion:

Resolving patch:
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 @@ def __init__(self):
         #   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 @@ def event_cb(self, event):
         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 @@ def __init__(self, command, args, num_worker=1, cpu_request="1000m", cpu_limit="
         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 @@ def main():
         ]
 
         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 report(self, task_id, success):
     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)


Test patch:
diff --git a/elasticdl/python/elasticdl/master/k8s_worker_manager_test.py b/elasticdl/python/elasticdl/master/k8s_worker_manager_test.py
--- a/elasticdl/python/elasticdl/master/k8s_worker_manager_test.py
+++ b/elasticdl/python/elasticdl/master/k8s_worker_manager_test.py
@@ -6,12 +6,17 @@
 
 tf.enable_eager_execution()
 
+from unittest.mock import MagicMock, call
 from elasticdl.master.k8s_worker_manager import WorkerManager
+from elasticdl.master.task_queue import _TaskQueue
 
 
 class WorkerManagerTest(unittest.TestCase):
     def testCreateDeleteWorkerPod(self):
+        task_q = _TaskQueue({"f": 10}, 1, 1)
+        task_q.recover_tasks = MagicMock()
         worker_servicer = WorkerManager(
+            task_q,
             job_name="test-create-worker-pod",
             #worker_image="elasticdl:dev",
             worker_image="gcr.io/google-samples/hello-app:1.0",
@@ -37,13 +42,19 @@ def testCreateDeleteWorkerPod(self):
             print(counters)
             if not counters:
                 break
+        task_q.recover_tasks.assert_has_calls(
+            [call(0), call(1), call(2)], any_order=True
+        )
 
     def testFailedWorkerPod(self):
         """
         Start a pod running a python program destined to fail with restart_policy="Never"
         to test failed_worker_count
         """
+        task_q = _TaskQueue({"f": 10}, 1, 1)
+        task_q.recover_tasks = MagicMock()
         worker_servicer = WorkerManager(
+            task_q,
             job_name="test-create-worker-pod",
             worker_image="gcr.io/google-samples/hello-app:1.0",
             command=["badcommand"],
@@ -67,6 +78,9 @@ def testFailedWorkerPod(self):
             print(counters)
             if not counters:
                 break
+        task_q.recover_tasks.assert_has_calls(
+            [call(0), call(1), call(2)], any_order=True
+        )
 
 if __name__ == '__main__':
     unittest.main()