Fix cases

Author: wjsi

mars/deploy/oscar/tests/test_local.py

@@ -93,8 +93,8 @@
         "serialization": {},
         "most_calls": DICT_NOT_EMPTY,
         "slow_calls": DICT_NOT_EMPTY,
-        # "band_subtasks": DICT_NOT_EMPTY,
-        # "slow_subtasks": DICT_NOT_EMPTY,
+        "band_subtasks": {},
+        "slow_subtasks": {},
     }
 }
 EXPECT_PROFILING_STRUCTURE_NO_SLOW = copy.deepcopy(EXPECT_PROFILING_STRUCTURE)

mars/services/scheduling/api/oscar.py

@@ -99,7 +99,6 @@ async def cancel_subtasks(
         self,
         subtask_ids: List[str],
         kill_timeout: Union[float, int] = None,
-        wait: bool = False,
     ):
         """
         Cancel pending and running subtasks.
@@ -111,14 +110,7 @@ async def cancel_subtasks(
         kill_timeout
             timeout seconds to kill actor process forcibly
         """
-        if wait:
-            await self._manager_ref.cancel_subtasks(
-                subtask_ids, kill_timeout=kill_timeout
-            )
-        else:
-            await self._manager_ref.cancel_subtasks.tell(
-                subtask_ids, kill_timeout=kill_timeout
-            )
+        await self._manager_ref.cancel_subtasks(subtask_ids, kill_timeout=kill_timeout)
 
     async def finish_subtasks(
         self,
@@ -139,7 +131,9 @@ async def finish_subtasks(
         schedule_next
             whether to schedule succeeding subtasks
         """
-        await self._manager_ref.finish_subtasks(subtask_results, bands, schedule_next)
+        await self._manager_ref.finish_subtasks.tell(
+            subtask_results, bands, schedule_next
+        )
 
 
 class MockSchedulingAPI(SchedulingAPI):

mars/services/scheduling/supervisor/manager.py

@@ -435,13 +435,13 @@ async def cancel_task_in_band(band):
                     )
                     band_to_futures[band].append(future)
 
-        for band in band_to_futures:
-            cancel_tasks.append(asyncio.create_task(cancel_task_in_band(band)))
-
+        # Dequeue first as it is possible to leak subtasks from queues
         if queued_subtask_ids:
-            # Don't use `finish_subtasks` because it may remove queued
             await self._queueing_ref.remove_queued_subtasks(queued_subtask_ids)
 
+        for band in band_to_futures:
+            cancel_tasks.append(asyncio.create_task(cancel_task_in_band(band)))
+
         if cancel_tasks:
             yield asyncio.gather(*cancel_tasks)
 

mars/services/task/execution/mars/stage.py

@@ -184,8 +184,7 @@ async def set_subtask_result(self, result: SubtaskResult, band: BandType = None)
                     )
                     # if error or cancel, cancel all submitted subtasks
                     await self._scheduling_api.cancel_subtasks(
-                        list(self._submitted_subtask_ids),
-                        wait=False,
+                        list(self._submitted_subtask_ids)
                     )
                 self._schedule_done()
                 cost_time_secs = self.result.end_time - self.result.start_time