Simplify API & fix case

mars/deploy/oscar/tests/fault_injection_config_with_rerun.yml

@@ -7,3 +7,5 @@ scheduling:
 storage:
   # shared-memory38 may lose object if the process crash after put success.
   backends: [plasma]
+  plasma:
+    store_memory: 32M

mars/deploy/oscar/tests/test_local.py

@@ -42,7 +42,7 @@
 from ....storage import StorageLevel
 from ....services.storage import StorageAPI
 from ....tensor.arithmetic.add import TensorAdd
-from ....tests.core import mock, check_dict_structure_same, DICT_NOT_EMPTY
+from ....tests.core import mock, DICT_NOT_EMPTY
 from ..local import new_cluster, _load_config
 from ..session import (
     get_default_async_session,
@@ -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)
@@ -263,10 +263,6 @@ async def test_execute(create_cluster, config):
 
     info = await session.execute(b, extra_config=extra_config)
     await info
-    if extra_config:
-        check_dict_structure_same(info.profiling_result(), expect_profiling_structure)
-    else:
-        assert not info.profiling_result()
     assert info.result() is None
     assert info.exception() is None
     assert info.progress() == 1

mars/deploy/oscar/tests/test_ray.py

@@ -66,8 +66,8 @@
         },
         "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

@@ -16,7 +16,7 @@
 
 from .... import oscar as mo
 from ....lib.aio import alru_cache
-from ...subtask import Subtask, SubtaskResult
+from ...subtask import Subtask
 from ..core import SubtaskScheduleSummary
 from .core import AbstractSchedulingAPI
 
@@ -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,16 @@ 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:
+                    summary_kw = dict(is_finished=True)
+                    if subtask_info.cancel_pending:
+                        summary_kw["is_cancelled"] = True
+                    self._subtask_summaries[subtask_id] = subtask_info.to_summary(
+                        **summary_kw
+                    )
                 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
@@ -321,7 +322,7 @@ async def batch_submit_subtask_to_band(self, args_list, kwargs_list):
                 if info.cancel_pending:
                     res_release_delays.append(
                         self._global_resource_ref.release_subtask_resource.delay(
-                            band, info.subtask.session_id, info.subtask.subtask_id
+                            band, self._session_id, subtask_id
                         )
                     )
                     continue
@@ -330,6 +331,12 @@ async def batch_submit_subtask_to_band(self, args_list, kwargs_list):
                     "Subtask %s is not in added subtasks set, it may be finished or canceled, skip it.",
                     subtask_id,
                 )
+                # in case resource already allocated, do deallocate
+                res_release_delays.append(
+                    self._global_resource_ref.release_subtask_resource.delay(
+                        band, self._session_id, subtask_id
+                    )
+                )
                 continue
             band_to_subtask_ids[band].append(subtask_id)
 
@@ -414,9 +421,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 +441,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

@@ -149,7 +149,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(
@@ -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
@@ -219,7 +218,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:

mars/services/task/tests/test_service.py

@@ -264,6 +264,7 @@ def f1():
     await asyncio.sleep(0.5)
     with Timer() as timer:
         await task_api.cancel_task(task_id)
+        await asyncio.sleep(0.5)
         result = await task_api.get_task_result(task_id)
         assert result.status == TaskStatus.terminated
     assert timer.duration < 20

mars/tests/test_utils.py

@@ -613,7 +613,7 @@ def __call__(self, *args, **kwargs):
     assert get_func_token_values(func) == [func]
 
 
-@pytest.mark.parametrize("id_length", [0, 5, 32, 63])
+@pytest.mark.parametrize("id_length", [0, 5, 32, 63, 254])
 def test_gen_random_id(id_length):
     rnd_id = utils.new_random_id(id_length)
     assert len(rnd_id) == id_length