Repository: sql-machine-learning/elasticdl

Issue:
Local master
Read local RecordIO file index, split into tasks, assign task to workers, keep track of status of each task.


Discussion:
"keep track of status of each task", does it mean worker need to report the status of the assigned task to the master, or the master has a way to query the status of the task from the worker?

Resolving patch:
diff --git a/python/elasticdl/system/__init__.py b/python/elasticdl/system/__init__.py
new file mode 100644
diff --git a/python/elasticdl/system/master.py b/python/elasticdl/system/master.py
new file mode 100644
--- /dev/null
+++ b/python/elasticdl/system/master.py
@@ -0,0 +1,107 @@
+from binascii import crc32
+from recordio.file import File
+import queue
+import threading
+import time
+
+
+def _id(*parts):
+    # convert each integer part into a 4-byte array, join them and return
+    # crc32
+    return crc32(b''.join(p.to_bytes(4, 'little') for p in parts))
+
+
+class Work(object):
+    '''
+    Represents a unit of work. All data members are read only.
+    '''
+
+    def __init__(self, file_index, offset, epoch, trial=0):
+        self.file_index = file_index
+        self.offset = offset
+        self.epoch = epoch
+        self.trial = trial
+        # Finish earlier epoch first, but randomize the order of chunks within
+        # an epoch.
+        self.id = _id(epoch, file_index, offset)
+        self.priority = (epoch, self.id)
+
+    def next_trial(self):
+        return Work(self.file_index, self.offset, self.epoch, self.trial + 1)
+
+    def next_epoch(self):
+        return Work(self.file_index, self.offset, self.epoch + 1)
+
+
+class WorkQueue(object):
+    def __init__(self, num_epoch, max_trial, files):
+        self._files = list(files)
+        self._num_epoch = num_epoch
+        self._max_trial = max_trial
+        self._q = queue.PriorityQueue()
+        self._in_flight = {}
+        self._lock = threading.Lock()
+
+    def put(self, file_index, offset):
+        with self._lock:
+            work = Work(file_index, offset, 0)
+            self._q.put((work.priority, work))
+
+    def get_work(self, timeout=1.0):
+        with self._lock:
+            work = self._q.get(timeout=1.0)[1]
+            # TODO: support worker timeout
+            self._in_flight[work.id] = work
+            return (work.id, self._files[work.file_index], work.offset)
+
+    def work_done(self, id, succeed):
+        with self._lock:
+            work = self._in_flight.pop(id)
+            next_work = None
+            if not succeed:
+                if work.trial + 1 < self._max_trial:
+                    next_work = work.next_trial()
+                else:
+                    print('work failed', work)
+            elif work.epoch + 1 < self._num_epoch:
+                next_work = work.next_epoch()
+            else:
+                print('work finished', work)
+            if next_work:
+                self._q.put((next_work.priority, next_work))
+            self._q.task_done()
+
+    def join(self):
+        self._q.join()
+
+
+class Master(object):
+    def __init__(self, data_files, num_epoch, max_trial):
+        assert num_epoch > 0
+        assert max_trial > 0
+
+        self._data_files = data_files
+        self._work_queue = WorkQueue(num_epoch, max_trial, files=data_files)
+        self._lock = threading.Lock()
+        self._num_workers = 0
+
+    def register_worker(self):
+        with self._lock:
+            self._num_workers += 1
+            return self._work_queue
+
+    def run(self):
+        # TODO: use a thread pool to build index in parallel.
+        with self._lock:
+            # Technically, we don't need blocking for building indices, but
+            # doing so makes the system cleaner as if there is a corrupted
+            # file, etc., the master will crash and no worker will be started.
+            start = time.time()
+            for i, f in enumerate(self._data_files):
+                with File(f, 'r') as fd:
+                    for chunk in fd.get_index():
+                        self._work_queue.put(i, chunk.offset)
+            print(
+                'Time spent on building index: %s seconds:' %
+                (time.time() - start))
+        self._work_queue.join()


Test patch:
diff --git a/python/elasticdl/system/master_test.py b/python/elasticdl/system/master_test.py
new file mode 100644
--- /dev/null
+++ b/python/elasticdl/system/master_test.py
@@ -0,0 +1,110 @@
+import queue
+import unittest
+import threading
+from unittest.mock import patch
+from collections import Counter
+from recordio.file_index import _ChunkData as C
+from master import Master
+
+# used to keep some shared data across workers.
+
+
+class RunLog(object):
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._failures = set()
+        self._run = Counter()
+
+    # Return False to signal a failure
+    def add_run(self, name, offset):
+        with self._lock:
+            key = (name, offset)
+            # make every work to fail once
+            ret = self._run[key] != 0
+            self._run[key] += 1
+            return ret
+
+    def get(self):
+        return self._run
+
+
+run_log = RunLog()
+
+
+class MockWorkerThread(threading.Thread):
+    def __init__(self, q):
+        self._q = q
+        self._exiting = False
+        threading.Thread.__init__(self)
+
+    def run(self):
+        while not self._exiting:
+            try:
+                work = self._q.get_work(timeout=1.0)
+                success = run_log.add_run(work[1], work[2])
+                self._q.work_done(work[0], success)
+            except queue.Empty:
+                pass
+
+    def exit(self):
+        self._exiting = True
+
+
+class MockRecordIoFile(object):
+    def __init__(self, index):
+        self._index = index
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return
+
+    def get_index(self):
+        return self._index
+
+
+class MockMasterThread(threading.Thread):
+    def __init__(self):
+        self.recordio_data = {'f0': [C(0, 100, 2), C(200, 100, 3)], 'f1': [
+            C(10, 200, 4), C(210, 200, 4)]}
+        self.master = Master(
+            self.recordio_data.keys(),
+            num_epoch=3,
+            max_trial=2)
+        threading.Thread.__init__(self)
+
+    def run(self):
+        # patch Master's recordio calls to inject mock data
+        with patch('master.File', autospec=True) as mock:
+            mock.side_effect = [MockRecordIoFile(
+                index) for index in self.recordio_data.values()]
+            self.master.run()
+
+
+class MasterTest(unittest.TestCase):
+    def test_normal(self):
+        m = MockMasterThread()
+        w1 = MockWorkerThread(m.master.register_worker())
+        w2 = MockWorkerThread(m.master.register_worker())
+
+        m.start()
+        w1.start()
+        w2.start()
+
+        m.join()
+        w1.exit()
+        w2.exit()
+
+        w1.join()
+        w2.join()
+
+        print(run_log.get())
+        # Every work got executed 4 time, including 3 epochs and 1 retry.
+        self.assertCountEqual(
+            run_log.get().elements(), [
+                ('f0', 0), ('f1', 210), ('f1', 10), ('f0', 200)] * 4)
+
+
+if __name__ == '__main__':
+    unittest.main()