Make assigner more balanced when encountering reassign_workers

Author: 继盛

mars/services/task/analyzer/assigner.py

@@ -15,12 +15,13 @@
 from abc import ABC, abstractmethod
 from collections import defaultdict
 from operator import itemgetter
-from typing import List, Dict, Set, Union
+from typing import List, Dict, Union
 
 import numpy as np
 
 from ....core import ChunkGraph, ChunkData
 from ....core.operand import Operand
+from ....lib.ordered_set import OrderedSet
 from ....resource import Resource
 from ....typing import BandType
 from ....utils import implements
@@ -77,8 +78,9 @@ def __init__(
         band_resource: Dict[BandType, Resource],
     ):
         super().__init__(chunk_graph, start_ops, band_resource)
-        self._undirected_chunk_graph = None
-        self._op_keys: Set[str] = {start_op.key for start_op in start_ops}
+        self._op_keys: OrderedSet[str] = OrderedSet(
+            [start_op.key for start_op in start_ops]
+        )
 
     def _calc_band_assign_limits(
         self, initial_count: int, occupied: Dict[BandType, int]
@@ -124,13 +126,15 @@ def _calc_band_assign_limits(
             pos = (pos + 1) % len(counts)
         return dict(zip(bands, counts))
 
+    @classmethod
     def _assign_by_bfs(
-        self,
+        cls,
+        undirected_chunk_graph: ChunkGraph,
         start: ChunkData,
         band: BandType,
         initial_sizes: Dict[BandType, int],
         spread_limits: Dict[BandType, float],
-        key_to_assign: Set[str],
+        key_to_assign: OrderedSet[str],
         assigned_record: Dict[str, Union[str, BandType]],
     ):
         """
@@ -140,19 +144,15 @@ def _assign_by_bfs(
         if initial_sizes[band] <= 0:
             return
 
-        graph = self._chunk_graph
-        if self._undirected_chunk_graph is None:
-            self._undirected_chunk_graph = graph.build_undirected()
-        undirected_chunk_graph = self._undirected_chunk_graph
-
         assigned = 0
         spread_range = 0
         for chunk in undirected_chunk_graph.bfs(start=start, visit_predicate="all"):
             op_key = chunk.op.key
             if op_key in assigned_record:
                 continue
             spread_range += 1
-            # `op_key` may not be in `key_to_assign`, but we need to record it to avoid iterate the node repeatedly.
+            # `op_key` may not be in `key_to_assign`,
+            # but we need to record it to avoid iterate the node repeatedly.
             assigned_record[op_key] = band
             if op_key not in key_to_assign:
                 continue
@@ -161,8 +161,23 @@ def _assign_by_bfs(
                 break
         initial_sizes[band] -= assigned
 
+    def _build_undirected_chunk_graph(
+        self, chunk_to_assign: List[ChunkData]
+    ) -> ChunkGraph:
+        chunk_graph = ChunkGraph()
+        self._chunk_graph.copyto(chunk_graph)
+        # remove edges for all chunk_to_assign which may contain chunks
+        # that need be reassigned
+        for chunk in chunk_to_assign:
+            if chunk_graph.count_predecessors(chunk) > 0:
+                for pred in list(chunk_graph.predecessors(chunk)):
+                    chunk_graph.remove_edge(pred, chunk)
+        return chunk_graph.build_undirected()
+
     @implements(AbstractGraphAssigner.assign)
-    def assign(self, cur_assigns: Dict[str, str] = None) -> Dict[ChunkData, BandType]:
+    def assign(
+        self, cur_assigns: Dict[str, BandType] = None
+    ) -> Dict[ChunkData, BandType]:
         graph = self._chunk_graph
         assign_result = dict()
         cur_assigns = cur_assigns or dict()
@@ -173,7 +188,7 @@ def assign(self, cur_assigns: Dict[str, str] = None) -> Dict[ChunkData, BandType
         for chunk in graph:
             op_key_to_chunks[chunk.op.key].append(chunk)
 
-        op_keys = set(self._op_keys)
+        op_keys = OrderedSet(self._op_keys)
         chunk_to_assign = [
             op_key_to_chunks[op_key][0]
             for op_key in op_keys
@@ -183,6 +198,9 @@ def assign(self, cur_assigns: Dict[str, str] = None) -> Dict[ChunkData, BandType
         for band in cur_assigns.values():
             assigned_counts[band] += 1
 
+        # build undirected graph
+        undirected_chunk_graph = self._build_undirected_chunk_graph(chunk_to_assign)
+
         # calculate the number of chunks to be assigned to each band
         # given number of bands and existing assignments
         band_quotas = self._calc_band_assign_limits(
@@ -195,14 +213,20 @@ def assign(self, cur_assigns: Dict[str, str] = None) -> Dict[ChunkData, BandType
         spread_ranges = defaultdict(lambda: average_spread_range)
         # assign from other chunks to be assigned
         # TODO: sort by what?
-        sorted_candidates = [v for v in chunk_to_assign]
+        sorted_candidates = chunk_to_assign.copy()
         while max(band_quotas.values()):
             band = max(band_quotas, key=lambda k: band_quotas[k])
             cur = sorted_candidates.pop()
             while cur.op.key in cur_assigns:
                 cur = sorted_candidates.pop()
             self._assign_by_bfs(
-                cur, band, band_quotas, spread_ranges, op_keys, cur_assigns
+                undirected_chunk_graph,
+                cur,
+                band,
+                band_quotas,
+                spread_ranges,
+                op_keys,
+                cur_assigns,
             )
 
         key_to_assign = {n.op.key for n in chunk_to_assign} | initial_assigned_op_keys

mars/services/task/analyzer/tests/test_assigner.py

@@ -13,9 +13,13 @@
 # limitations under the License.
 
 import numpy as np
+import pandas as pd
 
+from ..... import dataframe as md
 from .....config import Config
 from .....core import ChunkGraph
+from .....core.graph.builder.utils import build_graph
+from .....core.operand import OperandStage
 from .....tensor.random import TensorRand
 from .....tensor.arithmetic import TensorAdd
 from .....tensor.fetch import TensorFetch
@@ -71,3 +75,34 @@ def test_assigner_with_fetch_inputs():
             for inp in input_chunks:
                 if not isinstance(inp.op, TensorFetch):
                     assert subtask.expect_band == key_to_assign[inp.key]
+
+
+def test_shuffle_assign():
+    band_num = 8
+    all_bands = [(f"address_{i}", "numa-0") for i in range(band_num)]
+
+    pdf = pd.DataFrame(np.random.rand(32, 4))
+    df = md.DataFrame(pdf, chunk_size=4)
+    r = df.groupby(0).sum(method="shuffle")
+    chunk_graph = build_graph([r], tile=True)
+
+    band_resource = dict((band, Resource(num_cpus=1)) for band in all_bands)
+
+    reassign_worker_ops = [
+        chunk.op for chunk in chunk_graph if chunk.op.reassign_worker
+    ]
+    start_ops = list(GraphAnalyzer._iter_start_ops(chunk_graph))
+    to_assign_ops = start_ops + reassign_worker_ops
+
+    assigner = GraphAssigner(chunk_graph, to_assign_ops, band_resource)
+    assigns = assigner.assign()
+    assert len(assigns) == 16
+    init_assigns = set()
+    reducer_assigns = set()
+    for chunk, assign in assigns.items():
+        if chunk.op.stage == OperandStage.reduce:
+            reducer_assigns.add(assign)
+        else:
+            init_assigns.add(assign)
+    # init and reducers are assigned on all bands
+    assert len(init_assigns) == len(reducer_assigns) == 8