Fix duplicate decref of subtask input chunk (#2611)

Author: Catch-Bull

mars/services/task/supervisor/processor.py

@@ -215,8 +215,9 @@ def _get_decref_stage_chunk_keys(
             if error_or_cancelled:
                 # error or cancel, rollback incref for subtask results
                 for subtask in subtask_graph:
-                    if stage_processor.subtask_results.get(subtask):
+                    if subtask.subtask_id in stage_processor.decref_subtask:
                         continue
+                    stage_processor.decref_subtask.add(subtask.subtask_id)
                     # if subtask not executed, rollback incref of predecessors
                     for inp_subtask in subtask_graph.predecessors(subtask):
                         for result_chunk in inp_subtask.chunk_graph.results:
@@ -821,9 +822,12 @@ async def set_subtask_result(self, subtask_result: SubtaskResult):
                 # Since every worker will call supervisor to set subtask result,
                 # we need to release actor lock to make `decref_chunks` parallel to avoid blocking
                 # other `set_subtask_result` calls.
-                yield self._decref_input_subtasks(
-                    subtask, stage_processor.subtask_graph
-                )
+                if subtask.subtask_id not in stage_processor.decref_subtask:
+                    stage_processor.decref_subtask.add(subtask.subtask_id)
+                    yield self._decref_input_subtasks(
+                        subtask, stage_processor.subtask_graph
+                    )
+
             except:  # noqa: E722  # nosec  # pylint: disable=bare-except  # pragma: no cover
                 _, err, tb = sys.exc_info()
                 if subtask_result.status not in (

mars/services/task/supervisor/stage.py

@@ -67,6 +67,9 @@ def __init__(
         self.subtask_results: Dict[Subtask, SubtaskResult] = dict()
         self._submitted_subtask_ids = set()
 
+        # All subtask IDs whose input chunk reference count is reduced.
+        self.decref_subtask = set()
+
         self._band_manager: Dict[BandType, mo.ActorRef] = dict()
 
         # result