Repository: mars-project/mars

Issue:
[BUG] Ray shuffle load fetch shuffle raises AssertionError
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
A clear and concise description of what the bug is.

```python
mars/tensor/indexing/tests/test_indexing_execution.py:309 (test_setitem_fancy_index_execution)
setup = <mars.deploy.oscar.session.SyncSession object at 0x1242bbbe0>

    def test_setitem_fancy_index_execution(setup):
        rs = np.random.RandomState(0)
    
        raw = rs.randint(0, 10, size=(11, 12))
    
        # index is a ndarray, value is a scalar
        arr = tensor(raw.copy(), chunk_size=5)
        idx = rs.randint(0, 11, (5,))
        arr[idx] = 20
>       res = arr.execute().fetch()

mars/tensor/indexing/tests/test_indexing_execution.py:319: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
mars/core/entity/tileables.py:462: in execute
    result = self.data.execute(session=session, **kw)
mars/core/entity/executable.py:144: in execute
    return execute(self, session=session, **kw)
mars/deploy/oscar/session.py:1890: in execute
    return session.execute(
mars/deploy/oscar/session.py:1684: in execute
    execution_info: ExecutionInfo = fut.result(
../../.pyenv/versions/3.8.7/lib/python3.8/concurrent/futures/_base.py:439: in result
    return self.__get_result()
../../.pyenv/versions/3.8.7/lib/python3.8/concurrent/futures/_base.py:388: in __get_result
    raise self._exception
mars/deploy/oscar/session.py:1870: in _execute
    await execution_info
../../.pyenv/versions/3.8.7/lib/python3.8/asyncio/tasks.py:695: in _wrap_awaitable
    return (yield from awaitable.__await__())
mars/deploy/oscar/session.py:105: in wait
    return await self._aio_task
mars/deploy/oscar/session.py:953: in _run_in_background
    raise task_result.error.with_traceback(task_result.traceback)
mars/services/task/supervisor/processor.py:369: in run
    await self._process_stage_chunk_graph(*stage_args)
mars/services/task/supervisor/processor.py:247: in _process_stage_chunk_graph
    chunk_to_result = await self._executor.execute_subtask_graph(
mars/services/task/execution/ray/executor.py:485: in execute_subtask_graph
    input_object_refs = await self._load_subtask_inputs(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <mars.services.task.execution.ray.executor.RayTaskExecutor object at 0x13f889ee0>
stage_id = 'xG2oeZ77OGqURoRWnXO9bk0I'
subtask = <Subtask id=SRyg6Deb4AzAjglTkOaDAkBB results=[TensorIndexSetValue(f01fe54c1f4a19cb552c05860d76c114_0)]>
context = {'5ee475fb6e8060867f5a1c97421059b3_0': ObjectRef(16310a0f0a45af5cffffffffffffffffffffffff0100000001000000)}
shuffle_manager = <mars.services.task.execution.ray.shuffle.ShuffleManager object at 0x14295b6a0>

    async def _load_subtask_inputs(
        self,
        stage_id: str,
        subtask: Subtask,
        context: Dict,
        shuffle_manager: ShuffleManager,
    ):
        """
        Load input object refs of subtask from context.
    
        It updates the context if the input object refs are fetched from
        the meta service.
        """
        input_object_refs = []
        key_to_get_meta = {}
        # for non-shuffle chunks, chunk key will be used for indexing object refs.
        # for shuffle chunks, mapper subtasks will have only one mapper chunk, and all outputs for mapper
        # subtask will be shuffle blocks, the downstream reducers will receive inputs in the mappers order.
        start_chunks = _get_start_chunks(subtask.chunk_graph)
        for index, start_chunk in enumerate(start_chunks):
            if isinstance(start_chunk.op, Fetch):
                chunk_key = start_chunk.key
                # pure_depend data is not used, skip it.
                if chunk_key in subtask.pure_depend_keys:
                    input_object_refs.append(None)
                elif chunk_key in context:
                    input_object_refs.append(context[chunk_key])
                else:
                    input_object_refs.append(None)
                    key_to_get_meta[index] = self._meta_api.get_chunk_meta.delay(
                        chunk_key, fields=["object_refs"]
                    )
            elif isinstance(start_chunk.op, FetchShuffle):
>               assert len(start_chunks) == 1, start_chunks
E               AssertionError: [Chunk <op=TensorFetch, key=5ee475fb6e8060867f5a1c97421059b3_0>, Chunk <op=TensorFetchShuffle, key=73d37b69ae018bc2c4f66072dd0bfb7b_0>]

mars/services/task/execution/ray/executor.py:678: AssertionError

```

**To Reproduce**
To help us reproducing this bug, please provide information below:
1. Your Python version
2. The version of Mars you use
3. Versions of crucial packages, such as numpy, scipy and pandas
4. Full stack of the error.
5. Minimized code to reproduce the error.

**Expected behavior**
A clear and concise description of what you expected to happen.

**Additional context**
Add any other context about the problem here.



Discussion:
This is an operand shuffle compatibility issue. In map reduce paradigm, all reducer inputs are just mapper outputs, but in fancy index, reducer depends on not only mapper outputs, but also tensors partitions:
![graphviz (2)](https://user-images.githubusercontent.com/12445254/182757301-07f34d4a-42e3-44ad-8f77-4f81127af697.svg)



Resolving patch:
diff --git a/mars/services/task/execution/ray/executor.py b/mars/services/task/execution/ray/executor.py
--- a/mars/services/task/execution/ray/executor.py
+++ b/mars/services/task/execution/ray/executor.py
@@ -151,27 +151,30 @@ def execute_subtask(
     logger.info("Begin to execute subtask: %s", subtask_id)
     # optimize chunk graph.
     subtask_chunk_graph = _optimize_subtask_graph(subtask_chunk_graph)
-    start_chunks = _get_start_chunks(subtask_chunk_graph)
-    if isinstance(start_chunks[0].op, FetchShuffle):
-        assert len(start_chunks) == 1, start_chunks
+    fetch_chunks, shuffle_chunk = _get_fetch_chunks(subtask_chunk_graph)
+    input_data = {}
+    if shuffle_chunk is not None:
         # the subtask is a reducer subtask
-        n_mappers = len(inputs)
+        n_mappers = shuffle_chunk.op.n_mappers
         # some reducer may have multiple output chunks, see `PSRSshuffle._execute_reduce` and
         # https://user-images.githubusercontent.com/12445254/168569524-f09e42a7-653a-4102-bdf0-cc1631b3168d.png
-        reducer_chunks = subtask_chunk_graph.successors(start_chunks[0])
+        reducer_chunks = subtask_chunk_graph.successors(shuffle_chunk)
         reducer_operands = set(c.op for c in reducer_chunks)
-        if len(reducer_operands) != 1:
+        if len(reducer_operands) != 1:  # pragma: no cover
             raise ValueError(
                 f"Subtask {subtask_id} has more than 1 reduce operands: {subtask_chunk_graph.to_dot()}"
             )
         reducer_operand = reducer_chunks[0].op
         reducer_index = reducer_operand.reducer_index
         # mock input keys, keep this in sync with `MapReducerOperand#_iter_mapper_key_idx_pairs`
-        input_keys = [(i, reducer_index) for i in range(n_mappers)]
-    else:
-        input_keys = [c.key for c in start_chunks if isinstance(c.op, Fetch)]
+        input_data.update(
+            {(i, reducer_index): block for i, block in enumerate(inputs[-n_mappers:])}
+        )
+        inputs = inputs[:-n_mappers]
+    input_keys = [start_chunk.key for start_chunk in fetch_chunks]
+    input_data.update(zip(input_keys, inputs))
     context = RayExecutionWorkerContext(
-        lambda: RayTaskState.get_handle(task_id), zip(input_keys, inputs)
+        lambda: RayTaskState.get_handle(task_id), input_data
     )
 
     for chunk in subtask_chunk_graph.topological_iter():
@@ -227,8 +230,16 @@ def execute_subtask(
     return output_values[0] if len(output_values) == 1 else output_values
 
 
-def _get_start_chunks(chunk_graph):
-    return sorted(chunk_graph.iter_indep(), key=operator.attrgetter("key"))
+def _get_fetch_chunks(chunk_graph):
+    fetch_chunks = []
+    shuffle_chunk = None
+    for start_chunk in chunk_graph.iter_indep():
+        if isinstance(start_chunk.op, FetchShuffle):
+            assert shuffle_chunk is None, shuffle_chunk
+            shuffle_chunk = start_chunk
+        elif isinstance(start_chunk.op, Fetch):
+            fetch_chunks.append(start_chunk)
+    return sorted(fetch_chunks, key=operator.attrgetter("key")), shuffle_chunk
 
 
 def _get_subtask_out_info(
@@ -660,29 +671,28 @@ async def _load_subtask_inputs(
         It updates the context if the input object refs are fetched from
         the meta service.
         """
-        input_object_refs = []
+        normal_object_refs = []
+        shuffle_object_refs = []
         key_to_get_meta = {}
         # for non-shuffle chunks, chunk key will be used for indexing object refs.
         # for shuffle chunks, mapper subtasks will have only one mapper chunk, and all outputs for mapper
         # subtask will be shuffle blocks, the downstream reducers will receive inputs in the mappers order.
-        start_chunks = _get_start_chunks(subtask.chunk_graph)
-        for index, start_chunk in enumerate(start_chunks):
-            if isinstance(start_chunk.op, Fetch):
-                chunk_key = start_chunk.key
-                # pure_depend data is not used, skip it.
-                if chunk_key in subtask.pure_depend_keys:
-                    input_object_refs.append(None)
-                elif chunk_key in context:
-                    input_object_refs.append(context[chunk_key])
-                else:
-                    input_object_refs.append(None)
-                    key_to_get_meta[index] = self._meta_api.get_chunk_meta.delay(
-                        chunk_key, fields=["object_refs"]
-                    )
-            elif isinstance(start_chunk.op, FetchShuffle):
-                assert len(start_chunks) == 1, start_chunks
-                # shuffle meta won't be recorded in meta service, query it from shuffle manager.
-                return shuffle_manager.get_reducer_input_refs(subtask)
+        fetch_chunks, shuffle_chunk = _get_fetch_chunks(subtask.chunk_graph)
+        for index, fetch_chunk in enumerate(fetch_chunks):
+            chunk_key = fetch_chunk.key
+            # pure_depend data is not used, skip it.
+            if chunk_key in subtask.pure_depend_keys:
+                normal_object_refs.append(None)
+            elif chunk_key in context:
+                normal_object_refs.append(context[chunk_key])
+            else:
+                normal_object_refs.append(None)
+                key_to_get_meta[index] = self._meta_api.get_chunk_meta.delay(
+                    chunk_key, fields=["object_refs"]
+                )
+        if shuffle_chunk is not None:
+            # shuffle meta won't be recorded in meta service, query it from shuffle manager.
+            shuffle_object_refs = list(shuffle_manager.get_reducer_input_refs(subtask))
 
         if key_to_get_meta:
             logger.info(
@@ -695,9 +705,9 @@ async def _load_subtask_inputs(
             )
             for index, meta in zip(key_to_get_meta.keys(), meta_list):
                 object_ref = meta["object_refs"][0]
-                input_object_refs[index] = object_ref
-                context[start_chunks[index].key] = object_ref
-        return input_object_refs
+                normal_object_refs[index] = object_ref
+                context[fetch_chunks[index].key] = object_ref
+        return normal_object_refs + shuffle_object_refs
 
     async def _update_progress_and_collect_garbage(
         self,


Test patch:
diff --git a/mars/tensor/indexing/tests/test_indexing_execution.py b/mars/tensor/indexing/tests/test_indexing_execution.py
--- a/mars/tensor/indexing/tests/test_indexing_execution.py
+++ b/mars/tensor/indexing/tests/test_indexing_execution.py
@@ -307,6 +307,7 @@ def test_mixed_indexing_execution(setup):
     np.testing.assert_array_equal(res, expected)
 
 
+@pytest.mark.ray_dag
 def test_setitem_fancy_index_execution(setup):
     rs = np.random.RandomState(0)