Repository: mars-project/mars

Issue:
[core] Optimize mars graph building performance
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Is your feature request related to a problem? Please describe.**
```python
num_rows = 100_0000_0000
df1 = md.DataFrame(
    mt.random.rand(num_rows, 4, chunk_size=100_0000),
    columns=list('abcd'))
df2 = md.DataFrame(
    mt.random.rand(num_rows, 4, chunk_size=100_0000),
    columns=list('abcd'))
df1.merge(df2, left_on='a', right_on='a').execute()
```
20000subtasks and 50 16C32Gworker
* without coloring, graph build time:
  * tile: 128s
  * assign: 40min
![image](https://user-images.githubusercontent.com/12445254/159617020-a37d596f-8476-4b45-a352-57734373ff60.png)
![image](https://user-images.githubusercontent.com/12445254/159617058-f0818955-1eff-4ff6-a4c1-6cbeb081e9fb.png)


**Describe the solution you'd like**
A clear and concise description of what you want to happen.

**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.

**Additional context**
Add any other context or screenshots about the feature request here.



Discussion:
The process can be reproduced by:
```
from ..... import dataframe as md
from ..... import tensor as mt
from .....core.graph import TileableGraph, TileableGraphBuilder, ChunkGraphBuilder
from ..analyzer import GraphAnalyzer
from ..assigner import GraphAssigner


def test_assigner():
    num_rows = 100_0000_0000
    df1 = md.DataFrame(
        mt.random.rand(num_rows, 4, chunk_size=100_0000),
        columns=list('abcd'))
    df2 = md.DataFrame(
        mt.random.rand(num_rows, 4, chunk_size=100_0000),
        columns=list('abcd'))
    merged_df = df1.merge(df2, left_on='a', right_on='a')
    graph = TileableGraph([merged_df.data])
    next(TileableGraphBuilder(graph).build())
    chunk_graph = next(ChunkGraphBuilder(graph, fuse_enabled=False).build())
    start_ops = list(GraphAnalyzer._iter_start_ops(chunk_graph))
    band_slots = {
        (f"worker-{i}", "numa-0"): 16 for i in range(50)
    }
    current_assign = {}
    assigner = GraphAssigner(chunk_graph, start_ops, band_slots)
    assigned_result = assigner.assign(current_assign)
    print(len(assigned_result))
```

Resolving patch:
diff --git a/asv_bench/benchmarks/graph_assigner.py b/asv_bench/benchmarks/graph_assigner.py
new file mode 100644
--- /dev/null
+++ b/asv_bench/benchmarks/graph_assigner.py
@@ -0,0 +1,46 @@
+# Copyright 1999-2022 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import mars.tensor as mt
+import mars.dataframe as md
+from mars.core.graph import TileableGraph, TileableGraphBuilder, ChunkGraphBuilder
+from mars.services.task.analyzer import GraphAnalyzer
+from mars.services.task.analyzer.assigner import GraphAssigner
+
+
+class ChunkGraphAssignerSuite:
+    """
+    Benchmark that times performance of chunk graph assigner
+    """
+
+    def setup(self):
+        num_rows = 10000
+        df1 = md.DataFrame(
+            mt.random.rand(num_rows, 4, chunk_size=10), columns=list("abcd")
+        )
+        df2 = md.DataFrame(
+            mt.random.rand(num_rows, 4, chunk_size=10), columns=list("abcd")
+        )
+        merged_df = df1.merge(df2, left_on="a", right_on="a")
+        graph = TileableGraph([merged_df.data])
+        next(TileableGraphBuilder(graph).build())
+        self.chunk_graph = next(ChunkGraphBuilder(graph, fuse_enabled=False).build())
+
+    def time_assigner(self):
+        start_ops = list(GraphAnalyzer._iter_start_ops(self.chunk_graph))
+        band_slots = {(f"worker-{i}", "numa-0"): 16 for i in range(50)}
+        current_assign = {}
+        assigner = GraphAssigner(self.chunk_graph, start_ops, band_slots)
+        assigned_result = assigner.assign(current_assign)
+        assert len(assigned_result) == len(start_ops)
diff --git a/mars/services/task/analyzer/assigner.py b/mars/services/task/analyzer/assigner.py
--- a/mars/services/task/analyzer/assigner.py
+++ b/mars/services/task/analyzer/assigner.py
@@ -15,7 +15,7 @@
 from abc import ABC, abstractmethod
 from collections import defaultdict
 from operator import itemgetter
-from typing import List, Dict, Set
+from typing import List, Dict, Set, Union
 
 import numpy as np
 
@@ -130,7 +130,7 @@ def _assign_by_bfs(
         initial_sizes: Dict[BandType, int],
         spread_limits: Dict[BandType, float],
         key_to_assign: Set[str],
-        assigned_record: Dict[str, int],
+        assigned_record: Dict[str, Union[str, BandType]],
     ):
         """
         Assign initial nodes using breath-first search given initial sizes and
@@ -151,9 +151,10 @@ def _assign_by_bfs(
             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.
+            assigned_record[op_key] = band
             if op_key not in key_to_assign:
                 continue
-            assigned_record[op_key] = band
             assigned += 1
             if spread_range >= spread_limits[band] or assigned >= initial_sizes[band]:
                 break


Test patch: