Repository: mars-project/mars

Issue:
[BUG] release_free_slot got wrong slot_id
<!--
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.

**To Reproduce**
To help us reproducing this bug, please provide information below:
1. Your Python version
2. The version of Mars you use: 8f322b496df58022f6b7dd3839f1ce6b2d119c73
3. Versions of crucial packages, such as numpy, scipy and pandas
4. Full stack of the error.
```
2022-03-14 11:57:44,895 ERROR api.py:115 -- Got unhandled error when handling message ('release_free_slot', 0, (0, ('t7rQ0udN4qHUWr9x3qg1baTI', 'XPRUxfx1ld4Dzz797R5KZeGW')), {}) in actor b'numa-0_band_slot_manager' at 127.0.0.1:14887
Traceback (most recent call last):
  File "mars/oscar/core.pyx", line 478, in mars.oscar.core._BaseActor.__on_receive__
    with debug_async_timeout('actor_lock_timeout',
  File "mars/oscar/core.pyx", line 481, in mars.oscar.core._BaseActor.__on_receive__
    async with self._lock:
  File "mars/oscar/core.pyx", line 482, in mars.oscar.core._BaseActor.__on_receive__
    result = func(*args, **kwargs)
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/services/scheduling/worker/workerslot.py", line 169, in release_free_slot
    assert acquired_slot_id == slot_id, f"acquired_slot_id {acquired_slot_id} != slot_id {slot_id}"
AssertionError: acquired_slot_id 1 != slot_id 0
2022-03-14 11:57:44,897 ERROR execution.py:120 -- Failed to run subtask XPRUxfx1ld4Dzz797R5KZeGW on band numa-0
Traceback (most recent call last):
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/services/scheduling/worker/execution.py", line 331, in internal_run_subtask
    subtask_info.result = await self._retry_run_subtask(
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/services/scheduling/worker/execution.py", line 420, in _retry_run_subtask
    return await _retry_run(subtask, subtask_info, _run_subtask_once)
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/services/scheduling/worker/execution.py", line 107, in _retry_run
    raise ex
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/services/scheduling/worker/execution.py", line 67, in _retry_run
    return await target_async_func(*args)
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/services/scheduling/worker/execution.py", line 412, in _run_subtask_once
    await slot_manager_ref.release_free_slot(
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/oscar/backends/context.py", line 189, in send
    return self._process_result_message(result)
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/oscar/backends/context.py", line 70, in _process_result_message
    raise message.as_instanceof_cause()
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/oscar/backends/pool.py", line 542, in send
    result = await self._run_coro(message.message_id, coro)
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/oscar/backends/pool.py", line 333, in _run_coro
    return await coro
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/oscar/api.py", line 115, in __on_receive__
    return await super().__on_receive__(message)
  File "mars/oscar/core.pyx", line 506, in __on_receive__
    raise ex
  File "mars/oscar/core.pyx", line 478, in mars.oscar.core._BaseActor.__on_receive__
    with debug_async_timeout('actor_lock_timeout',
  File "mars/oscar/core.pyx", line 481, in mars.oscar.core._BaseActor.__on_receive__
    async with self._lock:
  File "mars/oscar/core.pyx", line 482, in mars.oscar.core._BaseActor.__on_receive__
    result = func(*args, **kwargs)
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/services/scheduling/worker/workerslot.py", line 169, in release_free_slot
    assert acquired_slot_id == slot_id, f"acquired_slot_id {acquired_slot_id} != slot_id {slot_id}"
AssertionError: [address=127.0.0.1:14887, pid=42643] acquired_slot_id 1 != slot_id 0
2022-03-14 11:57:44,900 INFO processor.py:508 -- Time consuming to execute a subtask is 0.023977041244506836s with session_id t7rQ0udN4qHUWr9x3qg1baTI, subtask_id XPRUxfx1ld4Dzz797R5KZeGW
2022-03-14 11:57:44,903 ERROR api.py:115 -- Got unhandled error when handling message ('release_free_slot', 0, (1, ('t7rQ0udN4qHUWr9x3qg1baTI', 'XPRUxfx1ld4Dzz797R5KZeGW')), {}) in actor b'numa-0_band_slot_manager' at 127.0.0.1:14887
Traceback (most recent call last):
  File "mars/oscar/core.pyx", line 478, in mars.oscar.core._BaseActor.__on_receive__
    with debug_async_timeout('actor_lock_timeout',
  File "mars/oscar/core.pyx", line 481, in mars.oscar.core._BaseActor.__on_receive__
    async with self._lock:
  File "mars/oscar/core.pyx", line 482, in mars.oscar.core._BaseActor.__on_receive__
    result = func(*args, **kwargs)
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/services/scheduling/worker/workerslot.py", line 168, in release_free_slot
    acquired_slot_id = self._session_stid_to_slot.pop(acquired_session_stid)
KeyError: ('t7rQ0udN4qHUWr9x3qg1baTI', 'XPRUxfx1ld4Dzz797R5KZeGW')
2022-03-14 11:57:44,904 ERROR execution.py:120 -- Failed to run subtask XPRUxfx1ld4Dzz797R5KZeGW on band numa-0
Traceback (most recent call last):
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/services/scheduling/worker/execution.py", line 331, in internal_run_subtask
    subtask_info.result = await self._retry_run_subtask(
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/services/scheduling/worker/execution.py", line 420, in _retry_run_subtask
    return await _retry_run(subtask, subtask_info, _run_subtask_once)
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/services/scheduling/worker/execution.py", line 107, in _retry_run
    raise ex
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/services/scheduling/worker/execution.py", line 67, in _retry_run
    return await target_async_func(*args)
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/services/scheduling/worker/execution.py", line 412, in _run_subtask_once
    await slot_manager_ref.release_free_slot(
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/oscar/backends/context.py", line 189, in send
    return self._process_result_message(result)
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/oscar/backends/context.py", line 70, in _process_result_message
    raise message.as_instanceof_cause()
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/oscar/backends/pool.py", line 542, in send
    result = await self._run_coro(message.message_id, coro)
  File "/Users/chaokunyang/ant/Development/DevProjects/python/mars/mars/oscar/backends/pool.py", line 333, in _run_coro
    return await coro

```
6. Minimized code to reproduce the error.
` pytest --log-level=INFO -v -s mars/dataframe/indexing/tests/test_indexing_execution.py::test_series_getitem`


Discussion:
Turns out same subtask has been executed on same band multiple times:
![image](https://user-images.githubusercontent.com/12445254/158124779-6305be02-36f6-4d70-8c6c-d459c110cfbe.png)

It's a bug in `stage_processor._schedule_subtasks` which may schedule subtasks multiple times.

Resolving patch:
diff --git a/mars/services/scheduling/worker/workerslot.py b/mars/services/scheduling/worker/workerslot.py
--- a/mars/services/scheduling/worker/workerslot.py
+++ b/mars/services/scheduling/worker/workerslot.py
@@ -166,7 +166,9 @@ def release_free_slot(self, slot_id: int, session_stid: Tuple[str, str]):
                 f"the releasing session_stid: {session_stid}"
             )
         acquired_slot_id = self._session_stid_to_slot.pop(acquired_session_stid)
-        assert acquired_slot_id == slot_id
+        assert (
+            acquired_slot_id == slot_id
+        ), f"{acquired_session_stid}: acquired_slot_id {acquired_slot_id} != slot_id {slot_id}"
 
         logger.debug("Slot %d released", slot_id)
 
diff --git a/mars/services/task/supervisor/processor.py b/mars/services/task/supervisor/processor.py
--- a/mars/services/task/supervisor/processor.py
+++ b/mars/services/task/supervisor/processor.py
@@ -887,6 +887,11 @@ async def set_subtask_result(self, subtask_result: SubtaskResult):
             stage_processor.subtask_snapshots.get(subtask)
         )
         if subtask_result.status.is_done:
+            # update stage_processor.subtask_results to avoid concurrent set_subtask_result
+            # since we release lock when `_decref_input_subtasks`.
+            stage_processor.subtask_results[subtask] = subtask_result.update(
+                stage_processor.subtask_results.get(subtask)
+            )
             try:
                 # Since every worker will call supervisor to set subtask result,
                 # we need to release actor lock to make `decref_chunks` parallel to avoid blocking
diff --git a/mars/services/task/supervisor/stage.py b/mars/services/task/supervisor/stage.py
--- a/mars/services/task/supervisor/stage.py
+++ b/mars/services/task/supervisor/stage.py
@@ -96,6 +96,11 @@ def is_cancelled(self):
         return self._cancelled.is_set()
 
     async def _schedule_subtasks(self, subtasks: List[Subtask]):
+        subtasks = [
+            subtask
+            for subtask in subtasks
+            if subtask.subtask_id not in self._submitted_subtask_ids
+        ]
         if not subtasks:
             return
         self._submitted_subtask_ids.update(subtask.subtask_id for subtask in subtasks)
@@ -125,7 +130,7 @@ def _schedule_done(self):
 
     async def set_subtask_result(self, result: SubtaskResult):
         subtask = self.subtask_id_to_subtask[result.subtask_id]
-        self.subtask_results[subtask] = result.update(self.subtask_results.get(subtask))
+        #  update subtask_results in `TaskProcessorActor.set_subtask_result`
         self._submitted_subtask_ids.difference_update([result.subtask_id])
 
         all_done = len(self.subtask_results) == len(self.subtask_graph)


Test patch: