Repository: sql-machine-learning/elasticdl

Issue:
Choose better image/cmd for some unittest jobs
Currently, pods launched by k8s_client_test.py will execute a command `/bin/ls --worker-id 1`, which would cause errors and make pod crash. Although it doesn't affect the test results, it would be better to make the pod a meaningful one.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/elasticdl/master/k8s.py b/elasticdl/python/elasticdl/master/k8s.py
--- a/elasticdl/python/elasticdl/master/k8s.py
+++ b/elasticdl/python/elasticdl/master/k8s.py
@@ -109,7 +109,7 @@ def create_worker(self, worker_id, resource_requests, resource_limits, priority=
             worker_id, resource_requests, resource_limits, priority,
             mount_path, volume_name, command=command,
             args=args, restart_policy=restart_policy)
-        self._v1.create_namespaced_pod(self._ns, pod)
+        return self._v1.create_namespaced_pod(self._ns, pod)
 
     def delete_worker(self, worker_id):
         self._logger.info("Deleting worker: " + str(worker_id))
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
@@ -1,3 +1,4 @@
+import itertools
 import logging
 import threading
 
@@ -19,6 +20,7 @@ def __init__(
             pod_priority=None,
             mount_path=None,
             volume_name=None,
+            restart_policy="OnFailure",
             **kwargs):
         self._logger = logging.getLogger("WorkerManager")
         self._command = command
@@ -32,52 +34,77 @@ def __init__(
             "cpu": cpu_limit,
             "memory": memory_limit
         }
+        self._restart_policy = restart_policy
         self._pod_priority = pod_priority
         self._mount_path = mount_path
         self._volume_name = volume_name
         self._task_q = task_q
+        self._next_worker_id = itertools.count().__next__
 
         # protects followed variables, which are accessed from event_cb.
         self._lock = threading.Lock()
-        # pod name to phase mapping
-        # phase: Pending/Running/Succeeded/Failed/Unknown
+        # worker id to (pod name, phase) mapping
+        # phase: None/Pending/Running/Succeeded/Failed/Unknown
+        #   None: worker was just launched, haven't received event yet.
         #   Pending: worker pod not started yet
         #   Running: worker pod is running
-        #   Succeeded: worker pod finishes all jobs and terminates with no issue.
+        #   Succeeded: worker pod finishes all tasks and terminates with
+        #       no issue.
         #   Failed: worker pod is killed for some reason
         #   Unknown: unknown
         self._pods_phase = {}
+        # pod name to worker id mapping
+        self._pod_name_to_id = {}
+
+        self._relaunch_deleted_live_worker = True
 
         self._k8s_client = k8s.Client(
             event_callback=self._event_cb, **kwargs
         )
 
-    def start_workers(self, restart_policy="OnFailure"):
-        for i in range(self._num_worker):
-            self._logger.info("Starting worker: %d" % i)
-            self._k8s_client.create_worker(
-                i,
+    def set_relaunch_deleted_live_worker(self, val):
+        self._relaunch_deleted_live_worker = bool(val)
+
+    def _start_worker(self, worker_id):
+        self._logger.info("Starting worker: %d" % worker_id)
+        with self._lock:
+            pod = self._k8s_client.create_worker(
+                worker_id,
                 self._resource_requests,
                 self._resource_limits,
                 self._pod_priority,
                 self._mount_path,
                 self._volume_name,
                 command=self._command,
-                args=self._args + ["--worker_id", str(i)],
-                restart_policy=restart_policy,
+                args=self._args + ["--worker_id", str(worker_id)],
+                restart_policy=self._restart_policy,
             )
+            name = pod.metadata.name
+            self._pod_name_to_id[name] = worker_id
+            self._pods_phase[worker_id] = (name, None)
 
-    def remove_workers(self):
+    def start_workers(self):
         for i in range(self._num_worker):
-            pod_name = self._k8s_client.get_pod_name(i)
-            with self._lock:
-                if pod_name in self._pods_phase:
-                    self._logger.info("Deleting worker: %d", i)
-                    self._k8s_client.delete_worker(i)
+            self._start_worker(self._next_worker_id())
+
+    def _remove_worker(self, worker_id):
+        with self._lock:
+            if worker_id not in self._pods_phase:
+                self._logger.error("Unknown worker id: %s" % worker_id)
+                return
+
+        # TODO: change _k8s_client to accept pod name instead of worker id.
+        self._k8s_client.delete_worker(worker_id)
+
+    def stop_relaunch_and_remove_workers(self):
+        with self._lock:
+            self._relaunch_deleted_live_worker = False
+            for worker_id in self._pods_phase:
+                self._k8s_client.delete_worker(worker_id)
 
     def get_counters(self):
         with self._lock:
-            return Counter(self._pods_phase.values())
+            return Counter([v for _, v in self._pods_phase.values()])
 
     def _event_cb(self, event):
         evt_obj = event.get("object")
@@ -85,13 +112,25 @@ def _event_cb(self, event):
         if not evt_obj or not evt_type:
             self._logger.error("Event doesn't have object or type: %s" % event)
             return
+
         pod_name = evt_obj.metadata.name
+        relaunch = False
         with self._lock:
-            self._pods_phase[pod_name] = evt_obj.status.phase
+            worker_id = self._pod_name_to_id.get(pod_name)
+            if worker_id is None:
+                self._logger.error("Unknown pod name: %s" % pod_name)
+                return
+
+            phase = evt_obj.status.phase
+            self._pods_phase[worker_id] = (pod_name, phase)
             if evt_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])
-                )
+                del self._pods_phase[worker_id]
+                del self._pod_name_to_id[pod_name]
+                self._task_q.recover_tasks(worker_id)
+
+                # If the pod being deleted was 'Running' or 'Pending'. Launch
+                # a worker.
+                relaunch = self._relaunch_deleted_live_worker and phase in [
+                        "Running", "Pending"]
+        if relaunch:
+            self._start_worker(self._next_worker_id())
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
@@ -173,8 +173,9 @@ def main():
             pod_priority=args.worker_pod_priority,
             mount_path=args.mount_path,
             volume_name=args.volume_name,
+            restart_policy="Never",
         )
-        worker_manager.start_workers(restart_policy="Never")
+        worker_manager.start_workers()
 
     try:
         while True:


Test patch:
diff --git a/elasticdl/python/tests/k8s_client_test.py b/elasticdl/python/tests/k8s_client_test.py
--- a/elasticdl/python/tests/k8s_client_test.py
+++ b/elasticdl/python/tests/k8s_client_test.py
@@ -12,7 +12,7 @@ def __init__(self):
         self._count = 0
 
     def event_cb(self, event):
-        print("----- %s -----\n" % event["type"])
+        print("----- %s %s -----\n" % (event["type"], event["object"].status))
         if event["type"] == "ADDED":
             self._count += 1
         elif event["type"] == "DELETED":
@@ -34,7 +34,7 @@ def test_client(self):
         # Start 3 workers
         resource = {"cpu": "100m", "memory": "64M"}
         for i in range(3):
-            c.create_worker("worker-%d" % i, resource, resource)
+            _ = c.create_worker("worker-%d" % i, resource, resource, command=["echo"])
             time.sleep(5)
 
         # wait for workers to be added
diff --git a/elasticdl/python/tests/k8s_worker_manager_test.py b/elasticdl/python/tests/k8s_worker_manager_test.py
--- a/elasticdl/python/tests/k8s_worker_manager_test.py
+++ b/elasticdl/python/tests/k8s_worker_manager_test.py
@@ -20,11 +20,12 @@ def testCreateDeleteWorkerPod(self):
             task_q,
             job_name="test-create-worker-pod",
             worker_image="gcr.io/google-samples/hello-app:1.0",
-            command=["/bin/ls"],
+            command=["echo"],
             args=[],
             namespace="default",
             num_worker=3
         )
+
         worker_servicer.start_workers()
         max_check_num = 20
         for _ in range(max_check_num):
@@ -34,7 +35,7 @@ def testCreateDeleteWorkerPod(self):
             if counters["Succeeded"] == 3:
                 break
 
-        worker_servicer.remove_workers()
+        worker_servicer.stop_relaunch_and_remove_workers()
         for _ in range(max_check_num):
             time.sleep(3)
             counters = worker_servicer.get_counters()
@@ -60,9 +61,10 @@ def testFailedWorkerPod(self):
             command=["badcommand"],
             args=["badargs"],
             namespace="default",
-            num_worker=3
+            num_worker=3,
+            restart_policy="Never"
         )
-        worker_servicer.start_workers(restart_policy="Never")
+        worker_servicer.start_workers()
         max_check_num = 20
         for _ in range(max_check_num):
             time.sleep(3)
@@ -71,7 +73,7 @@ def testFailedWorkerPod(self):
             if counters["Failed"] == 3:
                 break
 
-        worker_servicer.remove_workers()
+        worker_servicer.stop_relaunch_and_remove_workers()
         for _ in range(max_check_num):
             time.sleep(3)
             counters = worker_servicer.get_counters()
@@ -82,6 +84,57 @@ def testFailedWorkerPod(self):
             [call(0), call(1), call(2)], any_order=True
         )
 
+    @unittest.skipIf(os.environ.get('K8S_TESTS', 'True') == 'False', 'No Kubernetes cluster available')
+    def testRelaunchWorkerPod(self):
+        task_q = _TaskQueue({"f": 10}, 1, 1)
+        worker_servicer = WorkerManager(
+            task_q,
+            job_name="test-relaunch-worker-pod",
+            worker_image="gcr.io/google-samples/hello-app:1.0",
+            command=["sleep 10"],
+            args=[],
+            namespace="default",
+            num_worker=3
+        )
+
+        worker_servicer.start_workers()
+        
+        max_check_num = 60
+        for _ in range(max_check_num):
+            time.sleep(1)
+            counters = worker_servicer.get_counters()
+            print(counters)
+            if counters["Runing"] + counters["Pending"] > 0:
+                break
+        # Note: There is a slight chance of race condition.
+        # Hack to find a worker to remove
+        current_workers = set()
+        live_workers = set()
+        with worker_servicer._lock:
+            for k, (_, phase) in worker_servicer._pods_phase.items():
+                current_workers.add(k)
+                if phase in ["Running", "Pending"]:
+                    live_workers.add(k)
+        self.assertTrue(live_workers)
+
+        worker_servicer._remove_worker(live_workers.pop())
+        # verify a new worker get launched
+        found = False
+        print(current_workers)
+        for _ in range(max_check_num):
+            if found:
+                break
+            time.sleep(1)
+            counters = worker_servicer.get_counters()
+            print(counters)
+            with worker_servicer._lock:
+                for k in worker_servicer._pods_phase:
+                    if k not in current_workers:
+                        found = True
+        else:
+            self.fail("Failed to find newly launched worker.")
+        
+        worker_servicer.stop_relaunch_and_remove_workers()
 
 if __name__ == '__main__':
     unittest.main()