Separate result setting RPC call

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": DICT_NOT_EMPTY,
+        # "slow_subtasks": DICT_NOT_EMPTY,
     }
 }
 EXPECT_PROFILING_STRUCTURE_NO_SLOW = copy.deepcopy(EXPECT_PROFILING_STRUCTURE)

mars/services/scheduling/supervisor/manager.py

@@ -172,10 +172,13 @@ async def _get_execution_ref(self, band: BandType):
 
         return await mo.actor_ref(SubtaskExecutionActor.default_uid(), address=band[0])
 
-    async def _handle_subtask_result(
-        self, info: SubtaskScheduleInfo, result: SubtaskResult, band: BandType
+    async def set_subtask_result(
+        self, result: SubtaskResult, band: BandType
     ):
+        info = self._subtask_infos[result.subtask_id]
         subtask_id = info.subtask.subtask_id
+        notify_task_service = True
+
         async with redirect_subtask_errors(self, [info.subtask], reraise=False):
             try:
                 info.band_futures[band].set_result(result)
@@ -199,6 +202,7 @@ async def _handle_subtask_result(
                         [info.subtask.priority or tuple()],
                         exclude_bands=set(info.band_futures.keys()),
                     )
+                    notify_task_service = False
                 else:
                     raise ex
             except asyncio.CancelledError:
@@ -236,6 +240,10 @@ async def _handle_subtask_result(
                 if info.num_reschedules > 0:
                     await self._queueing_ref.submit_subtasks.tell()
 
+        if notify_task_service:
+            task_api = await self._get_task_api()
+            await task_api.set_subtask_result(result)
+
     async def finish_subtasks(
         self,
         subtask_results: List[SubtaskResult],
@@ -251,11 +259,6 @@ async def finish_subtasks(
             subtask_info = self._subtask_infos.get(subtask_id, None)
 
             if subtask_info is not None:
-                if subtask_band is not None:
-                    await self._handle_subtask_result(
-                        subtask_info, result, subtask_band
-                    )
-
                 self._finished_subtask_count.record(
                     1,
                     {
@@ -273,7 +276,7 @@ async def finish_subtasks(
                 #  Cancel subtask on other bands.
                 aio_task = subtask_info.band_futures.pop(subtask_band, None)
                 if aio_task:
-                    await aio_task
+                    yield aio_task
                     if schedule_next:
                         band_tasks[subtask_band] += 1
                 if subtask_info.band_futures:

mars/services/scheduling/supervisor/tests/test_manager.py

@@ -23,7 +23,6 @@
 from .....typing import BandType
 from ....cluster import MockClusterAPI
 from ....subtask import Subtask, SubtaskResult, SubtaskStatus
-from ....task import TaskAPI
 from ....task.supervisor.manager import TaskManagerActor
 from ...supervisor import (
     SubtaskQueueingActor,
@@ -91,7 +90,10 @@ async def run_subtask(
         self._run_subtask_events[subtask.subtask_id].set()
 
         async def task_fun():
-            task_api = await TaskAPI.create(subtask.session_id, supervisor_address)
+            manager_ref = await mo.actor_ref(
+                uid=SubtaskManagerActor.gen_uid(subtask.session_id),
+                address=supervisor_address,
+            )
             result = SubtaskResult(
                 subtask_id=subtask.subtask_id,
                 session_id=subtask.session_id,
@@ -107,12 +109,12 @@ async def task_fun():
                 result.status = SubtaskStatus.cancelled
                 result.error = ex
                 result.traceback = ex.__traceback__
-                await task_api.set_subtask_result(result)
+                await manager_ref.set_subtask_result.tell(result, (self.address, band_name))
                 raise
             else:
                 result.status = SubtaskStatus.succeeded
                 result.execution_end_time = time.time()
-                await task_api.set_subtask_result(result)
+                await manager_ref.set_subtask_result.tell(result, (self.address, band_name))
 
         self._subtask_aiotasks[subtask.subtask_id][band_name] = asyncio.create_task(
             task_fun()

mars/services/scheduling/worker/execution.py

@@ -36,7 +36,6 @@
 from ...meta import MetaAPI
 from ...storage import StorageAPI
 from ...subtask import Subtask, SubtaskAPI, SubtaskResult, SubtaskStatus
-from ...task import TaskAPI
 from .quota import QuotaActor
 from .workerslot import BandSlotManagerActor
 
@@ -178,6 +177,17 @@ async def _get_slot_manager_ref(
             BandSlotManagerActor.gen_uid(band), address=self.address
         )
 
+    @classmethod
+    @alru_cache(cache_exceptions=False)
+    async def _get_manager_ref(
+        cls, session_id: str, supervisor_address: str
+    ) -> mo.ActorRefType[BandSlotManagerActor]:
+        from ..supervisor import SubtaskManagerActor
+
+        return await mo.actor_ref(
+            SubtaskManagerActor.gen_uid(session_id), address=supervisor_address
+        )
+
     @alru_cache(cache_exceptions=False)
     async def _get_band_quota_ref(self, band: str) -> mo.ActorRefType[QuotaActor]:
         return await mo.actor_ref(QuotaActor.gen_uid(band), address=self.address)
@@ -415,10 +425,12 @@ async def internal_run_subtask(self, subtask: Subtask, band_name: str):
                 # pop the subtask info at the end is to cancel the job.
                 self._subtask_info.pop(subtask.subtask_id, None)
 
-        task_api = await TaskAPI.create(
+        manager_ref = await self._get_manager_ref(
             subtask.session_id, subtask_info.supervisor_address
         )
-        await task_api.set_subtask_result(subtask_info.result)
+        await manager_ref.set_subtask_result.tell(
+            subtask_info.result, (self.address, subtask_info.band_name)
+        )
         return subtask_info.result
 
     async def _retry_run_subtask(
@@ -557,8 +569,10 @@ async def subtask_caller():
                 )
                 _fill_subtask_result_with_exception(subtask, band_name, res)
 
-                task_api = await TaskAPI.create(subtask.session_id, supervisor_address)
-                await task_api.set_subtask_result(res)
+                manager_ref = await self._get_manager_ref(
+                    subtask.session_id, supervisor_address
+                )
+                await manager_ref.set_subtask_result.tell(res, (self.address, band_name))
             finally:
                 self._subtask_info.pop(subtask_id, None)
                 self._finished_subtask_count.record(1, {"band": self.address})

mars/services/scheduling/worker/tests/test_execution.py

@@ -37,6 +37,7 @@
 from .....resource import Resource
 from .....tensor.fetch import TensorFetch
 from .....tensor.arithmetic import TensorTreeAdd
+from .....typing import BandType
 from .....utils import Timer
 from ....cluster import MockClusterAPI
 from ....lifecycle import MockLifecycleAPI
@@ -47,7 +48,7 @@
 from ....subtask import MockSubtaskAPI, Subtask, SubtaskStatus, SubtaskResult
 from ....task.supervisor.manager import TaskManagerActor
 from ....mutable import MockMutableAPI
-from ...supervisor import GlobalResourceManagerActor
+from ...supervisor import GlobalResourceManagerActor, SubtaskManagerActor
 from ...worker import SubtaskExecutionActor, QuotaActor, BandSlotManagerActor
 
 
@@ -155,6 +156,19 @@ def get_results(self):
         return list(self._results.values())
 
 
+class MockSubtaskManagerActor(mo.Actor):
+    def __init__(self, session_id: str):
+        self._session_id = session_id
+
+    async def __post_create__(self):
+        self._task_manager_ref = await mo.actor_ref(
+            uid=TaskManagerActor.gen_uid(self._session_id), address=self.address
+        )
+
+    async def set_subtask_result(self, result: SubtaskResult, band: BandType):
+        await self._task_manager_ref.set_subtask_result.tell(result)
+
+
 @pytest.fixture
 async def actor_pool(request):
     n_slots, enable_kill = request.param
@@ -221,9 +235,17 @@ async def actor_pool(request):
             address=pool.external_address,
         )
 
+        subtask_manager_ref = await mo.create_actor(
+            MockSubtaskManagerActor,
+            session_id,
+            uid=SubtaskManagerActor.gen_uid(session_id),
+            address=pool.external_address,
+        )
+
         try:
             yield pool, session_id, meta_api, worker_meta_api, storage_api, execution_ref
         finally:
+            await mo.destroy_actor(subtask_manager_ref)
             await mo.destroy_actor(task_manager_ref)
             await mo.destroy_actor(band_slot_ref)
             await mo.destroy_actor(global_resource_ref)