Repository: mars-project/mars

Issue:
Assign workers given data distribution among workers
**Is your feature request related to a problem? Please describe.**
Currently for eager mode (#28), initial assignments for new graphs are done without information of distribution of existing data among workers, which may lead to unnecessary data transfer. A similar scenario exists for shuffle (#134).

**Describe the solution you'd like**
Given that we are going to have a separate GraphAnalyzer(#193), some heuristic rules can be added to help assigning workers given existing input data.



Discussion:

Resolving patch:
diff --git a/mars/scheduler/analyzer.py b/mars/scheduler/analyzer.py
--- a/mars/scheduler/analyzer.py
+++ b/mars/scheduler/analyzer.py
@@ -85,6 +85,44 @@ def calc_descendant_sizes(self):
                 sizes[ni.op.key] += sizes[n.op.key]
         return sizes
 
+    def collect_external_input_chunks(self, initial=True):
+        """
+        Collect keys of input chunks not in current graph, for instance,
+        chunks as inputs of initial operands in eager mode.
+        :param initial: collect initial chunks only
+        :return: dict mapping operand key to its input keys
+        """
+        graph = self._graph
+        chunk_keys = set(n.key for n in graph)
+        visited = set()
+        results = dict()
+        for n in graph:
+            op_key = n.op.key
+            if op_key in visited or not n.inputs:
+                continue
+            if initial and graph.count_predecessors(n):
+                continue
+            ext_keys = [c.key for c in n.inputs if c.key not in chunk_keys]
+            if not ext_keys:
+                continue
+            visited.add(op_key)
+            results[op_key] = ext_keys
+        return results
+
+    @staticmethod
+    def _get_workers_with_max_size(worker_to_size):
+        """Get workers with maximal size"""
+        max_workers = set()
+        max_size = 0
+        for w, size in worker_to_size.items():
+            if size > max_size:
+                max_size = size
+                max_workers = {w}
+            elif size == max_size:
+                max_workers.add(w)
+        max_workers.difference_update([None])
+        return max_size, list(max_workers)
+
     def _iter_successor_assigns(self, existing_assigns):
         """Iterate over all successors to get allocations of successor nodes"""
         graph = self._graph
@@ -107,19 +145,19 @@ def _iter_successor_assigns(self, existing_assigns):
                 else:
                     worker_involved.add(None)
             # get the worker occupying most of the data
-            max_worker = None
-            max_size = 0
-            for w, size in pred_sizes.items():
-                if size > max_size:
-                    max_size = size
-                    max_worker = w
+            max_size, max_workers = self._get_workers_with_max_size(pred_sizes)
             # if there is a dominant worker, return it
-            if max_size > total_size / max(2, len(worker_involved)) and \
-                    max_worker is not None:
+            if max_size > total_size / max(2, len(worker_involved)) and max_workers:
+                max_worker = random.choice(max_workers)
                 worker_assigns[n.op.key] = max_worker
                 yield n.op.key, max_worker
 
     def _calc_worker_initial_limits(self, initial_count, occupied=None):
+        """
+        Calculate limitation of number of initial operands for workers
+        :param initial_count: num of nodes in READY state
+        :param occupied: worker -> num of initials already assigned
+        """
         occupied = occupied or dict()
         actual_count = initial_count - sum(occupied.values())
 
@@ -133,6 +171,11 @@ def _calc_worker_initial_limits(self, initial_count, occupied=None):
         counts = initial_count * endpoint_cores / endpoint_cores.sum()
         for idx, ep in enumerate(endpoints):
             counts[idx] = max(0, counts[idx] - occupied.get(ep, 0))
+
+        # all assigned, nothing to do
+        if counts.sum() == 0:
+            return dict((ep, 0) for ep in endpoints)
+
         counts = (actual_count * counts / counts.sum()).astype(np.int32)
 
         # assign remaining nodes
@@ -171,11 +214,39 @@ def _collect_zero_degrees(self):
         # note that different orders can contribute to different efficiency
         return sorted(zero_degrees, key=lambda n: descendant_sizes[n.op.key])
 
+    def _iter_assignments_by_transfer_sizes(self, worker_quotas, input_chunk_metas):
+        """
+        Assign chunks by input sizes
+        :type input_chunk_metas: dict[str, dict[str, mars.scheduler.chunkmeta.WorkerMeta]]
+        """
+        total_transfers = dict((k, sum(v.chunk_size for v in chunk_to_meta.values()))
+                               for k, chunk_to_meta in input_chunk_metas.items())
+        # operands with largest amount of data will be allocated first
+        sorted_chunks = sorted(total_transfers.keys(), reverse=True,
+                               key=lambda k: total_transfers[k])
+        for op_key in sorted_chunks:
+            # compute data amounts held in workers
+            worker_stores = defaultdict(lambda: 0)
+            for meta in input_chunk_metas[op_key].values():
+                for w in meta.workers:
+                    worker_stores[w] += meta.chunk_size
+
+            max_size, max_workers = self._get_workers_with_max_size(worker_stores)
+            if max_workers and max_size > 0.5 * total_transfers[op_key]:
+                max_worker = random.choice(max_workers)
+                if worker_quotas.get(max_worker, 0) <= 0:
+                    continue
+                worker_quotas[max_worker] -= 1
+                yield op_key, max_worker
+
     def _assign_by_bfs(self, start, worker, initial_sizes, spread_limits, assigned_record):
         """
         Assign initial nodes using Breadth-first Search given initial sizes and
         limitations of spread range.
         """
+        if initial_sizes[worker] <= 0:
+            return
+
         graph = self._graph
         if self._undigraph is None:
             undigraph = self._undigraph = graph.build_undirected()
@@ -203,10 +274,13 @@ def _successor_fun(n):
                 break
         initial_sizes[worker] -= assigned
 
-    def calc_initial_assignments(self):
+    def calc_initial_assignments(self, input_chunk_metas=None):
         """
         Decide target worker for initial chunks. This function works when
         initializing a new graph, or recovering from worker losses.
+
+        :param input_chunk_metas: chunk metas for graph-level inputs, grouped by initial chunks
+        :type input_chunk_metas: dict[str, dict[str, mars.scheduler.chunkmeta.WorkerMeta]]
         :return: dict mapping operand keys into worker endpoints
         """
         graph = self._graph
@@ -215,25 +289,44 @@ def calc_initial_assignments(self):
 
         # collect chunks with no inputs or all inputs from shuffle, and nodes
         zero_degrees = self._collect_zero_degrees()
+        zero_degree_op_keys = set(n.op.key for n in zero_degrees)
 
         descendant_readies = set()
 
         assigned_initial_counts = defaultdict(lambda: 0)
-        worker_op_keys = defaultdict(list)
+        worker_op_keys = defaultdict(set)
         if cur_assigns:
-            # calculate ranges of nodes already assigned
-            for op_key, worker in self._iter_successor_assigns(cur_assigns):
-                if op_states.get(op_key) == OperandState.READY:
+            for op_key, state in op_states.items():
+                if op_key not in zero_degree_op_keys and state == OperandState.READY:
                     descendant_readies.add(op_key)
-                    assigned_initial_counts[worker] += 1
-                cur_assigns[op_key] = worker
-                worker_op_keys[worker].append(op_key)
+                    assigned_initial_counts[cur_assigns[op_key]] += 1
 
         # calculate the number of initial nodes to be assigned to every worker
         # given number of workers and existing assignments
-        worker_quotas = self._calc_worker_initial_limits(
+        pre_worker_quotas = self._calc_worker_initial_limits(
             len(zero_degrees) + len(descendant_readies), assigned_initial_counts)
 
+        # pre-assign nodes given pre-determined transfer sizes
+        if not input_chunk_metas:
+            worker_quotas = pre_worker_quotas
+        else:
+            for op_key, worker in self._iter_assignments_by_transfer_sizes(
+                    pre_worker_quotas, input_chunk_metas):
+                if op_key in cur_assigns:
+                    continue
+                assigned_initial_counts[worker] += 1
+                cur_assigns[op_key] = worker
+                worker_op_keys[worker].add(op_key)
+
+            worker_quotas = self._calc_worker_initial_limits(
+                len(zero_degrees) + len(descendant_readies), assigned_initial_counts)
+
+        if cur_assigns:
+            # calculate ranges of nodes already assigned
+            for op_key, worker in self._iter_successor_assigns(cur_assigns):
+                cur_assigns[op_key] = worker
+                worker_op_keys[worker].add(op_key)
+
         logger.debug('Worker initial quotas: %r', worker_quotas)
 
         # calculate expected descendant count (spread range) of
@@ -255,7 +348,7 @@ def calc_initial_assignments(self):
             start_chunks = reduce(operator.add, (key_to_chunks[op_key] for op_key in worker_op_keys[worker]))
             self._assign_by_bfs(start_chunks, worker, worker_quotas, spread_ranges, cur_assigns)
 
-        # assign pass 2: assign from initial nodes
+        # assign pass 2: assign from other initial nodes
         sorted_initials = [v for v in zero_degrees]
         while max(worker_quotas.values()):
             worker = max(worker_quotas, key=lambda k: worker_quotas[k])
diff --git a/mars/scheduler/graph.py b/mars/scheduler/graph.py
--- a/mars/scheduler/graph.py
+++ b/mars/scheduler/graph.py
@@ -14,9 +14,10 @@
 
 import itertools
 import logging
+import operator
+import os
 import random
 import time
-import os
 from collections import deque, defaultdict
 
 from .analyzer import GraphAnalyzer
@@ -26,7 +27,7 @@
 from .kvstore import KVStoreActor
 from .session import SessionActor
 from .utils import SchedulerActor, GraphState, OperandState
-from ..compat import six, functools32, OrderedDict
+from ..compat import six, functools32, reduce, OrderedDict
 from ..errors import ExecutionInterrupted, GraphNotExists
 from ..graph import DAG
 from ..tiles import handler, DataNotReady
@@ -487,7 +488,20 @@ def analyze_graph(self, **kwargs):
 
         if kwargs.get('do_placement', True):
             logger.debug('Placing initial chunks for graph %s', self._graph_key)
-            for k, v in analyzer.calc_initial_assignments().items():
+
+            # collect external inputs for eager mode
+            ext_chunks_to_inputs = analyzer.collect_external_input_chunks(initial=True)
+            ext_chunk_keys = reduce(operator.add, ext_chunks_to_inputs.values(), [])
+            metas = dict(zip(ext_chunk_keys,
+                             self._chunk_meta_ref.batch_get_chunk_meta(self._session_id, ext_chunk_keys)))
+            input_chunk_metas = defaultdict(dict)
+            for chunk_key, input_chunk_keys in ext_chunks_to_inputs.items():
+                chunk_metas = input_chunk_metas[chunk_key]
+                for k in input_chunk_keys:
+                    chunk_metas[k] = metas[k]
+
+            # do placements
+            for k, v in analyzer.calc_initial_assignments(input_chunk_metas=input_chunk_metas).items():
                 operand_infos[k]['target_worker'] = v
 
     @log_unhandled


Test patch:
diff --git a/mars/scheduler/tests/test_analyzer.py b/mars/scheduler/tests/test_analyzer.py
--- a/mars/scheduler/tests/test_analyzer.py
+++ b/mars/scheduler/tests/test_analyzer.py
@@ -56,6 +56,31 @@ def testDescendantSize(self):
             self.assertGreaterEqual(descendants[nodes[idx].op.key],
                                     descendants[nodes[idx + 1].op.key])
 
+    def testInitialWithInputs(self):
+        import numpy as np
+        from mars.tensor.expressions.random import TensorRandint
+        from mars.tensor.expressions.arithmetic import TensorTreeAdd
+
+        n1 = TensorRandint(dtype=np.float32()).new_chunk(None, (10, 10))
+        n2 = TensorRandint(dtype=np.float32()).new_chunk(None, (10, 10))
+
+        n3 = TensorTreeAdd(dtype=np.float32()).new_chunk(None, (10, 10))
+        n3.op._inputs = [n1, n2]
+        n4 = TensorTreeAdd(dtype=np.float32()).new_chunk(None, (10, 10))
+        n4.op._inputs = [n3]
+
+        graph = DAG()
+        graph.add_node(n1)
+        graph.add_node(n3)
+        graph.add_node(n4)
+        graph.add_edge(n1, n3)
+        graph.add_edge(n3, n4)
+
+        analyzer = GraphAnalyzer(graph, {})
+        ext_chunks = analyzer.collect_external_input_chunks(initial=False)
+        self.assertListEqual(ext_chunks[n3.op.key], [n2.key])
+        self.assertEqual(len(analyzer.collect_external_input_chunks(initial=True)), 0)
+
     def testFullInitialAssign(self):
         import numpy as np
         from mars.tensor.expressions.random import TensorRandint
@@ -121,6 +146,83 @@ def testSameKeyAssign(self):
         assignments = analyzer.calc_initial_assignments()
         self.assertEqual(len(assignments), 6)
 
+    def testAssignWithPreviousData(self):
+        import numpy as np
+        from mars.scheduler.chunkmeta import WorkerMeta
+        from mars.tensor.expressions.random import TensorRandint
+        from mars.tensor.expressions.arithmetic import TensorTreeAdd
+
+        graph = DAG()
+
+        r"""
+        Proper initial allocation should divide the graph like
+        
+         U   U  |  U   U  |  U   U
+          \ /   |   \ /   |   \ /
+           U    |    U    |    U
+        """
+
+        inputs = [
+            tuple(TensorRandint(_key=str(i * 2 + j), dtype=np.float32()).new_chunk(None, (10, 10)) for j in range(2))
+            for i in range(3)
+        ]
+        results = [TensorTreeAdd(dtype=np.float32()).new_chunk(None, (10, 10)) for _ in range(3)]
+        for inp, r in zip(inputs, results):
+            r.op._inputs = list(inp)
+
+            graph.add_node(r)
+            for n in inp:
+                graph.add_node(n)
+                graph.add_edge(n, r)
+
+        # assign with partial mismatch
+        data_dist = {
+            '0': dict(c00=WorkerMeta(chunk_size=5, workers=('w1',)),
+                      c01=WorkerMeta(chunk_size=5, workers=('w2',))),
+            '1': dict(c10=WorkerMeta(chunk_size=10, workers=('w1',))),
+            '2': dict(c20=WorkerMeta(chunk_size=10, workers=('w3',))),
+            '3': dict(c30=WorkerMeta(chunk_size=10, workers=('w3',))),
+            '4': dict(c40=WorkerMeta(chunk_size=7, workers=('w3',))),
+        }
+        analyzer = GraphAnalyzer(graph, dict(w1=24, w2=24, w3=24))
+        assignments = analyzer.calc_initial_assignments(input_chunk_metas=data_dist)
+
+        self.assertEqual(len(assignments), 6)
+
+        # explanation of the result:
+        # for '1', all data are in w1, hence assigned to w1
+        # '0' assigned to w1 according to connectivity
+        # '2' and '3' assigned to w3 according to connectivity
+        # '4' assigned to w2 because it has fewer data, and the slots of w3 is used up
+
+        self.assertEqual(assignments['0'], 'w1')
+        self.assertEqual(assignments['1'], 'w1')
+        self.assertEqual(assignments['2'], 'w3')
+        self.assertEqual(assignments['3'], 'w3')
+        self.assertEqual(assignments['4'], 'w2')
+        self.assertEqual(assignments['5'], 'w2')
+
+        # assign with full mismatch
+        data_dist = {
+            '0': dict(c00=WorkerMeta(chunk_size=5, workers=('w1',)),
+                      c01=WorkerMeta(chunk_size=5, workers=('w1', 'w2',))),
+            '1': dict(c10=WorkerMeta(chunk_size=10, workers=('w1',))),
+            '2': dict(c20=WorkerMeta(chunk_size=10, workers=('w3',))),
+            '3': dict(c30=WorkerMeta(chunk_size=10, workers=('w3',))),
+            '4': dict(c40=WorkerMeta(chunk_size=7, workers=('w2',))),
+            '5': dict(c50=WorkerMeta(chunk_size=7, workers=('w2',))),
+        }
+        analyzer = GraphAnalyzer(graph, dict(w1=24, w2=24, w3=24))
+        assignments = analyzer.calc_initial_assignments(input_chunk_metas=data_dist)
+
+        self.assertEqual(len(assignments), 6)
+        self.assertEqual(assignments['0'], 'w1')
+        self.assertEqual(assignments['1'], 'w1')
+        self.assertEqual(assignments['2'], 'w3')
+        self.assertEqual(assignments['3'], 'w3')
+        self.assertEqual(assignments['4'], 'w2')
+        self.assertEqual(assignments['5'], 'w2')
+
     def testAssignOnWorkerAdd(self):
         import numpy as np
         from mars.scheduler import OperandState
@@ -164,6 +266,7 @@ def testAssignOnWorkerAdd(self):
             for i in range(2):
                 fixed_assigns[inputs[idx][i].op.key] = 'w%d' % (idx + 1)
                 op_states[results[idx][i].op.key] = OperandState.READY
+                fixed_assigns[results[idx][i].op.key] = 'w%d' % (idx + 1)
 
         for inp in inputs:
             for n in inp:
@@ -229,6 +332,7 @@ def testAssignOnWorkerLost(self):
             for i in range(2):
                 fixed_assigns[inputs[idx][i].op.key] = 'w%d' % (idx % 2 + 1)
                 op_states[inputs[idx][i].op.key] = OperandState.FINISHED
+                fixed_assigns[results[idx][i].op.key] = 'w%d' % (idx % 2 + 1)
                 op_states[results[idx][i].op.key] = OperandState.READY
 
         for inp in inputs:
diff --git a/mars/scheduler/tests/test_graph.py b/mars/scheduler/tests/test_graph.py
--- a/mars/scheduler/tests/test_graph.py
+++ b/mars/scheduler/tests/test_graph.py
@@ -193,7 +193,7 @@ def testErrorOnPrepare(self, *_):
             graph_ref = pool.create_actor(GraphActor, session_id, graph_key, serialized_graph,
                                           uid=GraphActor.gen_name(session_id, graph_key))
 
-            def _mock_raises(*_):
+            def _mock_raises(*_, **__):
                 raise RuntimeError
 
             with patch_method(GraphActor.create_operand_actors, new=_mock_raises):
@@ -207,7 +207,7 @@ def _mock_raises(*_):
             graph_ref = pool.create_actor(GraphActor, session_id, graph_key, serialized_graph,
                                           uid=GraphActor.gen_name(session_id, graph_key))
 
-            def _mock_cancels(*_):
+            def _mock_cancels(*_, **__):
                 graph_meta_ref = pool.actor_ref(GraphMetaActor.gen_name(session_id, graph_key))
                 graph_meta_ref.set_state(GraphState.CANCELLING)
 
@@ -220,7 +220,7 @@ def _mock_cancels(*_):
             graph_ref = pool.create_actor(GraphActor, session_id, graph_key, serialized_graph,
                                           uid=GraphActor.gen_name(session_id, graph_key))
 
-            def _mock_cancels(*_):
+            def _mock_cancels(*_, **__):
                 graph_meta_ref = pool.actor_ref(GraphMetaActor.gen_name(session_id, graph_key))
                 graph_meta_ref.set_state(GraphState.CANCELLING)
                 return dict()
diff --git a/mars/tensor/execution/tests/test_datasource_execute.py b/mars/tensor/execution/tests/test_datasource_execute.py
--- a/mars/tensor/execution/tests/test_datasource_execute.py
+++ b/mars/tensor/execution/tests/test_datasource_execute.py
@@ -89,12 +89,12 @@ def testEmptyExecution(self):
         res = self.executor.execute_tensor(t, concat=True)
         self.assertEqual(res[0].shape, (20, 30))
         self.assertEqual(res[0].dtype, np.int64)
-        self.assertFalse(np.array_equal(res, np.zeros((20, 30))))
 
         t = empty((20, 30), chunk_size=5)
 
         res = self.executor.execute_tensor(t, concat=True)
-        self.assertFalse(np.allclose(res, np.zeros((20, 30))))
+        self.assertEqual(res[0].shape, (20, 30))
+        self.assertEqual(res[0].dtype, np.float64)
 
         t2 = empty_like(t)
         res = self.executor.execute_tensor(t2, concat=True)
diff --git a/mars/tests/test_graph.py b/mars/tests/test_graph.py
--- a/mars/tests/test_graph.py
+++ b/mars/tests/test_graph.py
@@ -43,7 +43,7 @@ def testDAG(self):
         with self.assertRaises(KeyError):
             dag.add_edge(10, 1)
 
-        self.assertEqual(set(dag[2]), set([5, 6]))
+        self.assertEqual(set(dag[2]), {5, 6})
         self.assertEqual(list(dag.topological_iter()), [3, 2, 5, 6, 1, 4])
 
         self.assertEqual(list(dag.dfs()), [3, 2, 5, 6, 1, 4])
@@ -81,6 +81,12 @@ def testDAG(self):
             self.assertTrue(all(undigraph.has_successor(n, pred)
                                 for pred in dag.predecessors(n)))
 
+        dag_copy = dag.copy()
+        for n in dag:
+            self.assertIn(n, dag_copy)
+            self.assertTrue(all(dag_copy.has_successor(pred, n)
+                                for pred in dag_copy.predecessors(n)))
+
     def testToDot(self):
         import mars.tensor as mt