Repository: mars-project/mars

Issue:
[BUG] Run DataFrame jobs hangs in distributed mode
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
Run the below code, client hangs when execute the last line.
``` Python
In [41]: df1 = pd.DataFrame({'a': [3, 4, 5, 3, 5, 4, 1, 2, 3], 
    ...:                             'b': [1, 3, 4, 5, 6, 5, 4, 4, 4], 
    ...:                             'c': list('aabaadcce')})                                       

In [42]: mdf = md.DataFrame(df1) 

In [43]: mdf.groupby('c').count().head(2).execute()                                                 
Out[31]: 
   a  b
c      
a  4  4
b  1  1

In [44]: mdf.groupby('c').count().head(4).execute()
```

the log on server side:
```
perand 72674540cbbe1bc51d5edfdcc3763f53(DataFrameShuffleProxy) state from OperandState.FINISHED to OperandState.FREED.
2020-10-29 12:06:44,497 mars.scheduler.chunkmeta 130 DEBUG    Delete chunk meta 908a536690bf65e0df886b0f03d0aa57
2020-10-29 12:06:51,737 mars.scheduler.graph 120 DEBUG    Actor s:0:graph$876b5c913062aed863397b20ed4a8ee5$1dadd914cae606e418a689bbd10f57a9 running in process 120
2020-10-29 12:06:51,740 mars.scheduler.graph 120 DEBUG    Graph 1dadd914cae606e418a689bbd10f57a9 state from GraphState.UNSCHEDULED to GraphState.PREPARING.
2020-10-29 12:06:51,745 mars.scheduler.graph 120 DEBUG    Begin preparing graph 1dadd914cae606e418a689bbd10f57a9 with 3 tileables to chunk graph.
2020-10-29 12:06:51,753 mars.scheduler.graph 120 DEBUG    Terminal chunk keys: {'4a87ad3ec6448362ede5285fe29cd827'}
2020-10-29 12:06:51,754 mars.scheduler.graph 120 DEBUG    Placing initial chunks for graph 1dadd914cae606e418a689bbd10f57a9
2020-10-29 12:06:51,756 mars.scheduler.analyzer 120 DEBUG    Worker assign quotas: {'11.238.144.211:61928': 1, '11.238.144.92:14760': 0}
2020-10-29 12:06:51,756 mars.scheduler.analyzer 120 DEBUG    Scan spread ranges: {}
2020-10-29 12:06:51,756 mars.scheduler.graph 120 DEBUG    Creating operand actors for graph 1dadd914cae606e418a689bbd10f57a9
2020-10-29 12:06:51,768 mars.scheduler.graph 120 DEBUG    Graph 1dadd914cae606e418a689bbd10f57a9 state from GraphState.PREPARING to GraphState.RUNNING.
2020-10-29 12:06:51,771 mars.scheduler.operands.base 132 DEBUG    Operand 72674540cbbe1bc51d5edfdcc3763f53(DataFrameShuffleProxy) state from OperandState.FREED to OperandState.UNSCHEDULED.
2020-10-29 12:06:51,771 mars.scheduler.operands.base 133 DEBUG    Operand fae4a2f925826757d052d69042145b91(DataFrameFuseChunk) state from OperandState.FREED to OperandState.READY.
2020-10-29 12:06:51,771 mars.scheduler.operands.base 132 DEBUG    State of 72674540cbbe1bc51d5edfdcc3763f53(DataFrameShuffleProxy) reset to OperandState.UNSCHEDULED
2020-10-29 12:06:51,771 mars.scheduler.operands.base 133 DEBUG    State of fae4a2f925826757d052d69042145b91(DataFrameFuseChunk) reset to OperandState.READY
2020-10-29 12:06:51,773 mars.scheduler.assigner 124 DEBUG    Operand fae4a2f925826757d052d69042145b91 enqueued
2020-10-29 12:06:51,773 mars.scheduler.assigner 124 DEBUG    1 operands applied for session 876b5c913062aed863397b20ed4a8ee5
2020-10-29 12:06:51,774 mars.scheduler.assigner 120 DEBUG    Operand fae4a2f925826757d052d69042145b91(DataFrameFuseChunk) allocated to run in 11.238.144.211:61928
2020-10-29 12:06:51,777 mars.scheduler.operands.base 133 DEBUG    Operand fae4a2f925826757d052d69042145b91(DataFrameFuseChunk) state from OperandState.READY to OperandState.RUNNING.
2020-10-29 12:06:52,228 mars.scheduler.operands.base 133 DEBUG    Operand fae4a2f925826757d052d69042145b91(DataFrameFuseChunk) state from OperandState.RUNNING to OperandState.FINISHED.
2020-10-29 12:06:52,229 mars.scheduler.operands.shuffle 132 DEBUG    Predecessors of shuffle proxy 72674540cbbe1bc51d5edfdcc3763f53 done, notifying successors
```


Discussion:

Resolving patch:
diff --git a/mars/scheduler/operands/shuffle.py b/mars/scheduler/operands/shuffle.py
--- a/mars/scheduler/operands/shuffle.py
+++ b/mars/scheduler/operands/shuffle.py
@@ -81,6 +81,13 @@ def add_finished_predecessor(self, op_key, worker, output_sizes=None, output_sha
         if all(k in self._finish_preds for k in self._pred_keys):
             self._start_successors()
 
+    def append_graph(self, graph_key, op_info):
+        super().append_graph(graph_key, op_info)
+
+        io_meta = op_info['io_meta']
+        self._shuffle_keys_to_op = dict(zip(io_meta['shuffle_keys'], io_meta['successors']))
+        self._op_to_shuffle_keys = dict(zip(io_meta['successors'], io_meta['shuffle_keys']))
+
     def _start_successors(self):
         self._all_deps_built = True
         futures = []


Test patch:
diff --git a/mars/deploy/local/tests/test_cluster.py b/mars/deploy/local/tests/test_cluster.py
--- a/mars/deploy/local/tests/test_cluster.py
+++ b/mars/deploy/local/tests/test_cluster.py
@@ -601,9 +601,17 @@ def wrapped(*arg, **kwargs):
             r4 = df.groupby('c2').transform(execute_with_session_check(
                 func, session.session_id)).to_pandas()
             expected4 = raw.groupby('c2').transform(func)
-            print(r4, expected4)
             pd.testing.assert_frame_equal(r4.sort_index(), expected4.sort_index())
 
+            # test rerun gropuby
+            df = md.DataFrame(raw.copy(), chunk_size=4)
+            r5 = session.run(df.groupby('c2').count(method='shuffle').max())
+            r6 = session.run(df.groupby('c2').count(method='shuffle').min())
+            expected5 = raw.groupby('c2').count().max()
+            expected6 = raw.groupby('c2').count().min()
+            pd.testing.assert_series_equal(r5, expected5)
+            pd.testing.assert_series_equal(r6, expected6)
+
     def testMultiSessionDecref(self, *_):
         with new_cluster(scheduler_n_process=2, worker_n_process=2,
                          shared_memory='20M', web=True) as cluster: