Simplify API & fix case

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,18 +110,11 @@ 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,
-        subtask_results: List[SubtaskResult],
+        subtask_ids: List[str],
         bands: List[Tuple] = None,
         schedule_next: bool = True,
     ):
@@ -132,14 +124,14 @@ async def finish_subtasks(
 
         Parameters
         ----------
-        subtask_results
+        subtask_ids
             results of subtasks, must in finished states
         bands
             bands of subtasks to mark as finished
         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_ids, bands, schedule_next)
 
 
 class MockSchedulingAPI(SchedulingAPI):

mars/services/scheduling/supervisor/manager.py

@@ -244,16 +244,14 @@ async def set_subtask_result(self, result: SubtaskResult, band: BandType):
 
     async def finish_subtasks(
         self,
-        subtask_results: List[SubtaskResult],
+        subtask_ids: List[str],
         bands: List[BandType] = None,
         schedule_next: bool = True,
     ):
-        subtask_ids = [result.subtask_id for result in subtask_results]
         logger.debug("Finished subtasks %s.", subtask_ids)
         band_tasks = defaultdict(lambda: 0)
         bands = bands or [None] * len(subtask_ids)
-        for result, subtask_band in zip(subtask_results, bands):
-            subtask_id = result.subtask_id
+        for subtask_id, subtask_band in zip(subtask_ids, bands):
             subtask_info = self._subtask_infos.get(subtask_id, None)
 
             if subtask_info is not None:
@@ -265,13 +263,13 @@ async def finish_subtasks(
                         "stage_id": subtask_info.subtask.stage_id,
                     },
                 )
-                self._subtask_summaries[subtask_id] = subtask_info.to_summary(
-                    is_finished=True,
-                    is_cancelled=result.status == SubtaskStatus.cancelled,
-                )
+                if subtask_id not in self._subtask_summaries:
+                    self._subtask_summaries[subtask_id] = subtask_info.to_summary(
+                        is_finished=True,
+                    )
                 subtask_info.end_time = time.time()
                 self._speculation_execution_scheduler.finish_subtask(subtask_info)
-                #  Cancel subtask on other bands.
+                # Cancel subtask on other bands.
                 aio_task = subtask_info.band_futures.pop(subtask_band, None)
                 if aio_task:
                     yield aio_task
@@ -414,9 +412,8 @@ async def cancel_task_in_band(band):
 
             info = self._subtask_infos[subtask_id]
             info.cancel_pending = True
-            raw_tasks_to_cancel = list(info.band_futures.values())
 
-            if not raw_tasks_to_cancel:
+            if not info.band_futures:
                 # not submitted yet: mark subtasks as cancelled
                 result = SubtaskResult(
                     subtask_id=info.subtask.subtask_id,
@@ -435,13 +432,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/scheduling/tests/test_service.py

@@ -50,7 +50,9 @@ async def set_subtask_result(self, subtask_result: SubtaskResult):
         for event in self._events[subtask_result.subtask_id]:
             event.set()
         self._events.pop(subtask_result.subtask_id, None)
-        await scheduling_api.finish_subtasks([subtask_result], subtask_result.bands)
+        await scheduling_api.finish_subtasks(
+            [subtask_result.subtask_id], subtask_result.bands
+        )
 
     def _return_result(self, subtask_id: str):
         result = self._results[subtask_id]

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

@@ -82,6 +82,7 @@ def __init__(
         # status
         self._done = asyncio.Event()
         self._cancelled = asyncio.Event()
+        self._terminated = asyncio.Event()
 
         # add metrics
         self._stage_execution_time = Metrics.gauge(
@@ -149,7 +150,7 @@ async def set_subtask_result(self, result: SubtaskResult, band: BandType = None)
         if all_done or error_or_cancelled:
             # tell scheduling to finish subtasks
             await self._scheduling_api.finish_subtasks(
-                [result], bands=[band], schedule_next=not error_or_cancelled
+                [result.subtask_id], bands=[band], schedule_next=not error_or_cancelled
             )
             if self.result.status != TaskStatus.terminated:
                 self.result = TaskResult(
@@ -162,6 +163,7 @@ async def set_subtask_result(self, result: SubtaskResult, band: BandType = None)
                     error=result.error,
                     traceback=result.traceback,
                 )
+                self._terminated.set()
                 if not all_done and error_or_cancelled:
                     if result.status == SubtaskStatus.errored:
                         logger.exception(
@@ -184,8 +186,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
@@ -219,7 +220,9 @@ async def set_subtask_result(self, result: SubtaskResult, band: BandType = None)
                     # all predecessors finished
                     to_schedule_subtasks.append(succ_subtask)
             await self._schedule_subtasks(to_schedule_subtasks)
-            await self._scheduling_api.finish_subtasks([result], bands=[band])
+            await self._scheduling_api.finish_subtasks(
+                [result.subtask_id], bands=[band]
+            )
 
     async def run(self):
         if len(self.subtask_graph) == 0:
@@ -234,6 +237,7 @@ async def run(self):
 
         # wait for completion
         await self._done.wait()
+        await self._terminated.wait()
         if self.error_or_cancelled():
             if self.result.error is not None:
                 raise self.result.error.with_traceback(self.result.traceback)