Repository: sql-machine-learning/elasticdl

Issue:
Keep track of assigned tasks
Keep a mapping of task id to worker id at master
Add worker_id in GetTask RPC call



Discussion:

Resolving patch:
diff --git a/elasticdl/python/elasticdl/master/servicer.py b/elasticdl/python/elasticdl/master/servicer.py
--- a/elasticdl/python/elasticdl/master/servicer.py
+++ b/elasticdl/python/elasticdl/master/servicer.py
@@ -58,7 +58,7 @@ def GetTask(self, request, context):
         res = master_pb2.Task()
         res.model_version = self._version
         res.minibatch_size = self._minibatch_size
-        task_id, task = self._task_q.get()
+        task_id, task = self._task_q.get(request.worker_id)
         if task:
             res.task_id = task_id
             res.shard_file_name = task.file_name
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
@@ -48,7 +48,7 @@ def _create_tasks(self):
                 )
         random.shuffle(self._todo)
 
-    def get(self):
+    def get(self, worker_id):
         """Return next (task_id, Task) tuple"""
 
         with self._lock:
@@ -65,7 +65,7 @@ def get(self):
             self._task_id += 1
             task = self._todo.pop()
             # TODO: Handle timeout of tasks.
-            self._doing[self._task_id] = task
+            self._doing[self._task_id] = (worker_id, task)
 
             return self._task_id, task
 
@@ -73,7 +73,7 @@ def report(self, task_id, success):
         """Report if the task is successful or not"""
 
         with self._lock:
-            task = self._doing.pop(task_id, None)
+            _, task = self._doing.pop(task_id, (-1, None))
             if not task:
                 self._logger.warning("Unknown task_id: %d" % task_id)
             elif not success:
diff --git a/elasticdl/python/elasticdl/worker/worker.py b/elasticdl/python/elasticdl/worker/worker.py
--- a/elasticdl/python/elasticdl/worker/worker.py
+++ b/elasticdl/python/elasticdl/worker/worker.py
@@ -31,7 +31,7 @@ def __init__(self,
             channel: grpc channel
             max_retrain_num: max number of a minibatch retrain as its gradients are not accepted by master
         """
-        self._worker_id = worker_id,
+        self._worker_id = worker_id
         model_module = load_user_model(model_file)
         self._model = model_module.model
         self._feature_columns = model_module.feature_columns()
@@ -60,7 +60,10 @@ def get_task(self):
         """
         get task from master
         """
-        return self._stub.GetTask(empty_pb2.Empty())
+        req = master_pb2.GetTaskRequest()
+        req.worker_id = self._worker_id
+        
+        return self._stub.GetTask(req)
 
     def get_model(self, min_version):
         """


Test patch:
diff --git a/elasticdl/python/elasticdl/master/servicer_test.py b/elasticdl/python/elasticdl/master/servicer_test.py
--- a/elasticdl/python/elasticdl/master/servicer_test.py
+++ b/elasticdl/python/elasticdl/master/servicer_test.py
@@ -1,4 +1,5 @@
 import logging
+import random
 import unittest
 import numpy as np
 import tensorflow as tf
@@ -45,13 +46,16 @@ def testGetEmptyTask(self):
             _TaskQueue({}, record_per_task=3, num_epoch=2),
         )
 
+        req = master_pb2.GetTaskRequest()
+
         # No task yet, make sure the returned versions are as expected.
-        task = master.GetTask(empty_pb2.Empty(), None)
+        req.worker_id = 1
+        task = master.GetTask(req, None)
         self.assertEqual("", task.shard_file_name)
         self.assertEqual(0, task.model_version)
 
         master._version = 1
-        task = master.GetTask(empty_pb2.Empty(), None)
+        task = master.GetTask(req, None)
         self.assertEqual("", task.shard_file_name)
         self.assertEqual(1, task.model_version)
 
@@ -186,9 +190,12 @@ def testReportTaskResult(self):
         # task to number of runs.
         tasks = defaultdict(int)
         while True:
-            task = master.GetTask(empty_pb2.Empty(), None)
+            req = master_pb2.GetTaskRequest()
+            req.worker_id = random.randint(1, 10)
+            task = master.GetTask(req, None)
             if not task.shard_file_name:
                 break
+            self.assertEqual(task_q._doing[task.task_id][0], req.worker_id)
             task_key = (task.shard_file_name, task.start, task.end)
             tasks[task_key] += 1
             report = master_pb2.ReportTaskResultRequest()
diff --git a/elasticdl/python/elasticdl/worker/worker_test.py b/elasticdl/python/elasticdl/worker/worker_test.py
--- a/elasticdl/python/elasticdl/worker/worker_test.py
+++ b/elasticdl/python/elasticdl/worker/worker_test.py
@@ -84,7 +84,7 @@ def mock_ReportTaskResult(req):
             return master.ReportTaskResult(req, None)
 
         channel = grpc.insecure_channel('localhost:9999')
-        worker = Worker(0, _module_file, channel, codec_type=codec_type)
+        worker = Worker(1, _module_file, channel, codec_type=codec_type)
 
         filename = create_recordio_file(128, codec_type)
         task_q = _TaskQueue(
@@ -111,7 +111,9 @@ def mock_ReportTaskResult(req):
                 res = False
 
         self.assertTrue(res)
-        task = mock_GetTask(empty_pb2.Empty())
+        req = master_pb2.GetTaskRequest()
+        req.worker_id = 1
+        task = mock_GetTask(req)
         # No more task.
         self.assertTrue(not task.shard_file_name)