Repository: mars-project/mars

Issue:
[BUG] Shuffle task raised `KeyError`
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
Shuffle task raised `KeyError` when cluster has more than one workers.

``` Python
s = np.random.RandomState(0)
raw = pd.DataFrame(s.rand(100, 4), columns=list('abcd'))
df = md.DataFrame(raw, chunk_size=20)
df.sort_values(by='a').execute()
```

error message:

```
../session.py:151: in _run_in_background
    raise task_result.error.with_traceback(task_result.traceback)
../../../services/scheduling/worker/execution.py:223: in internal_run_subtask
    yield asyncio.shield(run_aiotask)
mars/oscar/core.pyx:252: in mars.oscar.core._Actor._run_actor_async_generator
    ???
mars/oscar/core.pyx:194: in _handle_actor_result
    ???
../../../services/subtask/api.py:52: in run_subtask_in_slot
    return await ref.run_subtask(subtask)
../../../oscar/backends/context.py:149: in send
    return self._process_result_message(result)
../../../oscar/backends/context.py:58: in _process_result_message
    raise message.error.with_traceback(message.traceback)
../../../oscar/backends/pool.py:481: in send
    result = await future
../../../oscar/api.py:111: in __on_receive__
    return await super().__on_receive__(message)
mars/oscar/core.pyx:302: in __on_receive__
    ???
mars/oscar/core.pyx:194: in _handle_actor_result
    ???
../../../services/subtask/worker/subtask.py:440: in _run_subtask
    return await processor.run()
../../../services/subtask/worker/subtask.py:227: in run
    await future
/Users/hekaisheng/miniconda3/envs/py3.8/lib/python3.8/concurrent/futures/thread.py:57: in run
    result = self.fn(*self.args, **self.kwargs)
../../../services/subtask/worker/subtask.py:194: in _execute_operand
    return execute(ctx, op)
../../../core/operand/core.py:347: in execute
    return executor(results, op)
../../../dataframe/sort/psrs.py:645: in execute
    cls._execute_reduce(ctx, op)
../../../dataframe/sort/psrs.py:610: in _execute_reduce
    raw_inputs = list(op.iter_mapper_data(ctx, pop=False))
../../../core/operand/shuffle.py:90: in iter_mapper_data
    for _idx, data in self.iter_mapper_data_with_index(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

>   yield idx, ctx[key, self.reducer_index]
E   KeyError: ('d357be89b8d14061e80d86e321d793a1', (2,))

../../../core/operand/shuffle.py:81: KeyError
```




Discussion:

Resolving patch:
diff --git a/mars/services/storage/api/oscar.py b/mars/services/storage/api/oscar.py
--- a/mars/services/storage/api/oscar.py
+++ b/mars/services/storage/api/oscar.py
@@ -233,7 +233,7 @@ async def open_writer(self,
 
     async def list(self, level: StorageLevel) -> List:
         """
-        List all stored objects in storage.
+        List all stored data_keys in storage.
 
         Parameters
         ----------
@@ -244,7 +244,7 @@ async def list(self, level: StorageLevel) -> List:
         -------
             list of data keys
         """
-        return await self._storage_handler_ref.list(level=level)
+        return await self._storage_manager_ref.list(level=level)
 
 
 class MockStorageAPI(StorageAPI):
diff --git a/mars/services/storage/core.py b/mars/services/storage/core.py
--- a/mars/services/storage/core.py
+++ b/mars/services/storage/core.py
@@ -138,6 +138,9 @@ def __init__(self):
         # mapping value is list of InternalDataInfo
         self._data_key_to_info: Dict[tuple, List[InternalDataInfo]] = defaultdict(list)
         self._data_info_list = dict()
+        # data key may be a tuple in some cases,
+        # we record main key to manage their lifecycle
+        self._main_key_to_sub_keys = defaultdict(set)
         for level in StorageLevel.__members__.values():
             self._data_info_list[level] = dict()
 
@@ -149,11 +152,19 @@ def put(self,
         info = InternalDataInfo(data_info, object_info)
         self._data_key_to_info[(session_id, data_key)].append(info)
         self._data_info_list[data_info.level][(session_id, data_key)] = object_info
+        if isinstance(data_key, tuple):
+            self._main_key_to_sub_keys[(session_id, data_key[0])].update([data_key])
 
     def get_infos(self,
                   session_id: str,
                   data_key: str) -> List[DataInfo]:
-        if (session_id, data_key) not in self._data_key_to_info:  # pragma: no cover
+        if (session_id, data_key) not in self._data_key_to_info:
+            if (session_id, data_key) in self._main_key_to_sub_keys:
+                infos = []
+                for sub_key in self._main_key_to_sub_keys[(session_id, data_key)]:
+                    infos.extend([info.data_info for info in
+                                  self._data_key_to_info.get((session_id, sub_key))])
+                return infos
             raise DataNotExist(f'Data key {session_id, data_key} not exists.')
         return [info.data_info for info in
                 self._data_key_to_info.get((session_id, data_key))]
@@ -177,14 +188,24 @@ def delete(self,
                session_id: str,
                data_key: str,
                level: StorageLevel):
-        if (session_id, data_key) in self._data_key_to_info:
-            self._data_info_list[level].pop((session_id, data_key))
-            infos = self._data_key_to_info[(session_id, data_key)]
-            rest = [info for info in infos if info.data_info.level != level]
-            if len(rest) == 0:
-                del self._data_key_to_info[(session_id, data_key)]
-            else:  # pragma: no cover
-                self._data_key_to_info[(session_id, data_key)] = rest
+        if (session_id, data_key) in self._main_key_to_sub_keys:
+            to_delete_keys = self._main_key_to_sub_keys[(session_id, data_key)]
+        else:
+            to_delete_keys = [data_key]
+        logger.debug(f'Begin to delete data keys in data manager: {to_delete_keys}')
+        for key in to_delete_keys:
+            if (session_id, key) in self._data_key_to_info:
+                self._data_info_list[level].pop((session_id, key))
+                infos = self._data_key_to_info[(session_id, key)]
+                rest = [info for info in infos if info.data_info.level != level]
+                if len(rest) == 0:
+                    del self._data_key_to_info[(session_id, key)]
+                else:  # pragma: no cover
+                    self._data_key_to_info[(session_id, key)] = rest
+        logger.debug(f'Finish deleting data keys in data manager: {to_delete_keys}')
+
+    def list(self, level: StorageLevel):
+        return list(self._data_info_list[level].keys())
 
 
 class StorageHandlerActor(mo.Actor):
@@ -311,19 +332,22 @@ async def batch_put(self, args_list, kwargs_list):
         data_infos = []
         put_infos = []
         for size, data_key, obj in zip(sizes, data_keys, objs):
+            logger.debug(f'Begin to put data key {data_key}')
             object_info = await self._clients[level].put(obj)
             data_info = _build_data_info(object_info, level, size)
             data_infos.append(data_info)
             put_infos.append(
                 self._storage_manager_ref.put_data_info.delay(
                     session_id, data_key, data_info, object_info))
+            logger.debug(f'Finish putting data key {data_key}, size is {size}, '
+                         f'object_id is {data_info.object_id}')
         await self._storage_manager_ref.put_data_info.batch(*put_infos)
         return data_infos
 
     def _get_data_infos_arg(self,
-                           session_id: str,
-                           data_key: str,
-                           error: str):
+                            session_id: str,
+                            data_key: str,
+                            error: str):
         infos = self._storage_manager_ref.get_data_infos.delay(
             session_id, data_key, error)
         return infos, session_id, data_key
@@ -379,8 +403,10 @@ async def batch_delete(self, args_list, kwargs_list):
             return
 
         await self._storage_manager_ref.delete_data_info.batch(*delete_infos)
+        logger.debug(f'Begin to delete batch data {to_removes}')
         for level, object_id in to_removes:
             yield self._clients[level].delete(object_id)
+        logger.debug(f'Finish deleting batch data {to_removes}')
         releases = []
         for level, size in level_sizes.items():
             releases.append(self._storage_manager_ref.release_quota.delay(size, level))
@@ -421,9 +447,12 @@ async def fetch(self,
                     data_key: str):
         if StorageLevel.REMOTE not in self._clients:
             raise NotImplementedError
-        else:  # pragma: no cover
+        else:
+            logger.debug(f'Begin to fetch data {data_key} info from remote worker')
             data_info = await self._storage_manager_ref.fetch_data_info(
                 session_id, data_key)
+            logger.debug(f'Finish fetching data {data_key} info from remote worker,'
+                         f'object_id is {data_info.object_id}')
             await self._clients[StorageLevel.REMOTE].fetch(data_info.object_id)
 
 
@@ -575,10 +604,14 @@ async def fetch(self,
                         main_key = data_key[0] if isinstance(data_key, tuple) else data_key
                         address = (await meta_api.get_chunk_meta(
                             main_key, fields=['bands']))['bands'][0][0]
+                    if address == self.address:
+                        return
+                    logger.debug(f'Begin to fetch data {data_key} from {address}')
                     sender_ref = await mo.actor_ref(
                         address=address, uid=SenderManagerActor.default_uid())
                     yield sender_ref.send_data(session_id, data_key,
                                                self.address, level)
+                    logger.debug(f'finish fetching data {data_key} from {address}')
                     if not isinstance(data_key, tuple):
                         # no need to update meta for shuffle data
                         await meta_api.add_chunk_bands(
@@ -676,6 +709,9 @@ def batch_delete_data_info(self, args_list, kwargs_list):
         for args, kwargs in zip(args_list, kwargs_list):
             self._data_manager.delete(*args, *kwargs)
 
+    def list(self, level: StorageLevel) -> List:
+        return self._data_manager.list(level)
+
     def pin(self, object_id):
         self._pinned_keys.append(object_id)
 
diff --git a/mars/services/task/supervisor/task_manager.py b/mars/services/task/supervisor/task_manager.py
--- a/mars/services/task/supervisor/task_manager.py
+++ b/mars/services/task/supervisor/task_manager.py
@@ -26,7 +26,7 @@
 from ....core import TileableGraph, ChunkGraph, ChunkGraphBuilder, \
     TileableType, ChunkType, enter_mode
 from ....core.context import set_context
-from ....core.operand import Fetch, Fuse
+from ....core.operand import Fetch, Fuse, MapReduceOperand, ShuffleProxy
 from ....optimization.logical.core import OptimizationRecords
 from ....optimization.logical.chunk import optimize as optimize_chunk_graph
 from ....optimization.logical.tileable import optimize as optimize_tileable_graph
@@ -632,8 +632,14 @@ async def _incref_after_analyze(self,
         for subtask in subtask_graph:
             # for subtask has successors, incref number of successors
             n = subtask_graph.count_successors(subtask)
-            incref_chunk_keys.extend(
-                [c.key for c in subtask.chunk_graph.results] * n)
+            for c in subtask.chunk_graph.results:
+                incref_chunk_keys.extend([c.key] * n)
+                # for shuffle reducer, incref its mapper chunk
+                for pre_graph in subtask_graph.predecessors(subtask):
+                    for chk in pre_graph.chunk_graph.results:
+                        if isinstance(chk.op, ShuffleProxy):
+                            incref_chunk_keys.extend(
+                                [map_chunk.key for map_chunk in chk.inputs])
         incref_chunk_keys.extend([c.key for c in chunk_results])
         await self._lifecycle_api.incref_chunks(incref_chunk_keys)
 
@@ -764,7 +770,13 @@ async def _decref_input_subtasks(self,
                                      subtask_graph: SubtaskGraph):
         decref_chunks = []
         for in_subtask in subtask_graph.iter_predecessors(subtask):
-            decref_chunks.extend(in_subtask.chunk_graph.results)
+            for result_chunk in in_subtask.chunk_graph.results:
+                # for reducer chunk, decref mapper chunks
+                if isinstance(result_chunk.op, ShuffleProxy):
+                    outputs_num = len([r for r in subtask.chunk_graph.results
+                                       if isinstance(r.op, MapReduceOperand)])
+                    decref_chunks.extend(result_chunk.inputs * outputs_num)
+                decref_chunks.append(result_chunk)
         await self._lifecycle_api.decref_chunks([c.key for c in decref_chunks])
 
     def get_task_progress(self, task_id: str) -> float:


Test patch:
diff --git a/mars/deploy/oscar/tests/test_local.py b/mars/deploy/oscar/tests/test_local.py
--- a/mars/deploy/oscar/tests/test_local.py
+++ b/mars/deploy/oscar/tests/test_local.py
@@ -92,6 +92,23 @@ async def test_iterative_tiling(create_cluster):
     assert df2.index_value.max_val <= 30
 
 
+@pytest.mark.asyncio
+async def test_execute_describe(create_cluster):
+    s = np.random.RandomState(0)
+    raw = pd.DataFrame(s.rand(100, 4), columns=list('abcd'))
+    df = md.DataFrame(raw, chunk_size=30)
+
+    session = get_default_session()
+    r = df.describe()
+    info = await session.execute(r)
+    await info
+    assert info.result() is None
+    assert info.exception() is None
+    assert info.progress() == 1
+    res = await session.fetch(r)
+    pd.testing.assert_frame_equal(res[0], raw.describe())
+
+
 async def _run_web_session_test(web_address):
     session_id = str(uuid.uuid4())
     session = await Session.init(web_address, session_id)
diff --git a/mars/deploy/oscar/tests/test_ray.py b/mars/deploy/oscar/tests/test_ray.py
--- a/mars/deploy/oscar/tests/test_ray.py
+++ b/mars/deploy/oscar/tests/test_ray.py
@@ -65,6 +65,12 @@ async def test_iterative_tiling(ray_cluster, create_cluster):
     await test_local.test_iterative_tiling(create_cluster)
 
 
+@require_ray
+@pytest.mark.asyncio
+async def test_execute_describe(ray_cluster, create_cluster):
+    await test_local.test_execute_describe(create_cluster)
+
+
 @require_ray
 @pytest.mark.asyncio
 def test_sync_execute(ray_cluster, create_cluster):
diff --git a/mars/services/task/supervisor/tests/test_task_manager.py b/mars/services/task/supervisor/tests/test_task_manager.py
--- a/mars/services/task/supervisor/tests/test_task_manager.py
+++ b/mars/services/task/supervisor/tests/test_task_manager.py
@@ -218,8 +218,9 @@ async def test_iterative_tiling(actor_pool):
 async def test_shuffle(actor_pool):
     pool, session_id, meta_api, lifecycle_api, storage_api, manager = actor_pool
 
-    raw = np.random.rand(10, 10)
-    raw2 = np.random.randint(10, size=(10,))
+    rs = np.random.RandomState(0)
+    raw = rs.rand(10, 10)
+    raw2 = rs.randint(10, size=(10,))
     a = mt.tensor(raw, chunk_size=5)
     b = mt.tensor(raw2, chunk_size=5)
     c = a[b]
@@ -247,6 +248,13 @@ async def test_shuffle(actor_pool):
     assert (await lifecycle_api.get_tileable_ref_counts([c.key]))[0] == 1
     assert (await lifecycle_api.get_chunk_ref_counts(
         [c.key for c in result_tileable.chunks])) == [1] * len(result_tileable.chunks)
+    await lifecycle_api.decref_tileables([c.key])
+    ref_counts = await lifecycle_api.get_all_chunk_ref_counts()
+    assert len(ref_counts) == 0
+
+    # test if exists in storage
+    from mars.storage import StorageLevel
+    assert len(await storage_api.list(level=StorageLevel.MEMORY)) == 0
 
 
 @pytest.mark.asyncio