Reduce RPC cost of oscar by removing unnecessary tasks (#2597)

Author: wjsi

mars/deploy/oscar/session.py

@@ -816,6 +816,28 @@ async def init(
         else:
             return await cls._init(address, session_id, new=new, timeout=timeout)
 
+    async def _update_progress(self, task_id: str, progress: Progress):
+        zero_acc_time = 0
+        delay = 0.5
+        while True:
+            try:
+                last_progress_value = progress.value
+                progress.value = await self._task_api.get_task_progress(task_id)
+                if abs(progress.value - last_progress_value) < 1e-4:
+                    # if percentage does not change, we add delay time by 0.5 seconds every time
+                    zero_acc_time = min(5, zero_acc_time + 0.5)
+                    delay = zero_acc_time
+                else:
+                    # percentage changes, we use percentage speed to calc progress time
+                    zero_acc_time = 0
+                    speed = abs(progress.value - last_progress_value) / delay
+                    # one percent for one second
+                    delay = 0.01 / speed
+                delay = max(0.5, min(delay, 5.0))
+                await asyncio.sleep(delay)
+            except asyncio.CancelledError:
+                break
+
     async def _run_in_background(
         self,
         tileables: list,
@@ -826,55 +848,52 @@ async def _run_in_background(
         with enter_mode(build=True, kernel=True):
             # wait for task to finish
             cancelled = False
+            progress_task = asyncio.create_task(
+                self._update_progress(task_id, progress)
+            )
             start_time = time.time()
-            while True:
-                try:
-                    if not cancelled:
-                        task_result: Union[TaskResult, None] = None
-                        try:
-                            task_result = await self._task_api.wait_task(
-                                task_id, timeout=0.5
-                            )
-                        finally:
-                            if task_result is None:
-                                # not finished, set progress
-                                progress.value = await self._task_api.get_task_progress(
-                                    task_id
-                                )
-                            else:
-                                progress.value = 1.0
-                        if task_result is not None:
-                            break
-                    else:
-                        # wait for task to finish
-                        task_result: TaskResult = await self._task_api.wait_task(
-                            task_id
-                        )
+            task_result: Optional[TaskResult] = None
+            try:
+                if self.timeout is None:
+                    check_interval = 30
+                else:
+                    elapsed = time.time() - start_time
+                    check_interval = min(self.timeout - elapsed, 30)
+
+                while True:
+                    task_result = await self._task_api.wait_task(
+                        task_id, timeout=check_interval
+                    )
+                    if task_result is not None:
                         break
-                except asyncio.CancelledError:
-                    # cancelled
-                    cancelled = True
-                    await self._task_api.cancel_task(task_id)
-                except TimeoutError:  # pragma: no cover
-                    # ignore timeout when waiting for subtask progresses
-                    pass
-                finally:
-                    if (
+                    elif (
                         self.timeout is not None
                         and time.time() - start_time > self.timeout
                     ):
                         raise TimeoutError(
                             f"Task({task_id}) running time > {self.timeout}"
                         )
-            profiling.result = task_result.profiling
-            if task_result.profiling:
-                logger.warning(
-                    "Profile task %s execution result:\n%s",
-                    task_id,
-                    json.dumps(task_result.profiling, indent=4),
-                )
-            if task_result.error:
-                raise task_result.error.with_traceback(task_result.traceback)
+            except asyncio.CancelledError:
+                # cancelled
+                cancelled = True
+                await self._task_api.cancel_task(task_id)
+            finally:
+                progress_task.cancel()
+                if task_result is not None:
+                    progress.value = 1.0
+                else:
+                    # not finished, set progress
+                    progress.value = await self._task_api.get_task_progress(task_id)
+            if task_result is not None:
+                profiling.result = task_result.profiling
+                if task_result.profiling:
+                    logger.warning(
+                        "Profile task %s execution result:\n%s",
+                        task_id,
+                        json.dumps(task_result.profiling, indent=4),
+                    )
+                if task_result.error:
+                    raise task_result.error.with_traceback(task_result.traceback)
             if cancelled:
                 return
             fetch_tileables = await self._task_api.get_fetch_tileables(task_id)

mars/deploy/oscar/tests/test_local.py

@@ -436,7 +436,7 @@ def f1(interval: float, count: int):
             break
         await asyncio.sleep(0.1)
     else:
-        raise Exception("progress test failed.")
+        raise Exception(f"progress test failed, actual value {info.progress()}.")
 
     await info
     assert info.result() is None

mars/oscar/backends/context.py

@@ -71,13 +71,15 @@ def _process_result_message(message: Union[ResultMessage, ErrorMessage]):
 
     async def _wait(self, future: asyncio.Future, address: str, message: _MessageBase):
         try:
-            await asyncio.wait([future])
+            await asyncio.shield(future)
         except asyncio.CancelledError:
             try:
                 await self.cancel(address, message.message_id)
             except CannotCancelTask:
                 # cancel failed, already finished
                 raise asyncio.CancelledError
+        except:  # noqa: E722  # nosec  # pylint: disable=bare-except
+            pass
         return await future
 
     async def create_actor(

mars/oscar/backends/mars/pool.py

@@ -240,7 +240,7 @@ async def kill_sub_pool(
         await asyncio.to_thread(process.join, 5)
 
     async def is_sub_pool_alive(self, process: multiprocessing.Process):
-        return process.is_alive()
+        return await asyncio.to_thread(process.is_alive)
 
     async def recover_sub_pool(self, address: str):
         process_index = self._config.get_process_index(address)

mars/oscar/backends/pool.py

@@ -14,7 +14,6 @@
 
 import asyncio
 import concurrent.futures as futures
-import contextlib
 import itertools
 import logging
 import os
@@ -314,12 +313,10 @@ async def handle_control_command(
 
         return processor.result
 
-    @contextlib.contextmanager
-    def _run_coro(self, message_id: bytes, coro: Coroutine):
-        future = asyncio.create_task(coro)
-        self._process_messages[message_id] = future
+    async def _run_coro(self, message_id: bytes, coro: Coroutine):
+        self._process_messages[message_id] = asyncio.tasks.current_task()
         try:
-            yield future
+            return await coro
         finally:
             self._process_messages.pop(message_id, None)
 
@@ -332,10 +329,9 @@ async def process_message(self, message: _MessageBase, channel: Channel):
                 message,
                 channel,
             ):
-                with self._run_coro(
+                processor.result = await self._run_coro(
                     message.message_id, handler(self, message)
-                ) as future:
-                    processor.result = await future
+                )
         try:
             await channel.send(processor.result)
         except (ChannelClosed, ConnectionResetError):
@@ -464,8 +460,7 @@ async def create_actor(self, message: CreateActorMessage) -> result_message_type
             actor.uid = actor_id
             actor.address = address = self.external_address
             self._actors[actor_id] = actor
-            with self._run_coro(message.message_id, actor.__post_create__()) as future:
-                await future
+            await self._run_coro(message.message_id, actor.__post_create__())
 
             result = ActorRef(address, actor_id)
             # ensemble result message
@@ -491,8 +486,7 @@ async def destroy_actor(self, message: DestroyActorMessage) -> result_message_ty
                 actor = self._actors[actor_id]
             except KeyError:
                 raise ActorNotExist(f"Actor {actor_id} does not exist")
-            with self._run_coro(message.message_id, actor.__pre_destroy__()) as future:
-                await future
+            await self._run_coro(message.message_id, actor.__pre_destroy__())
             del self._actors[actor_id]
 
             processor.result = ResultMessage(
@@ -523,8 +517,7 @@ async def send(self, message: SendMessage) -> result_message_type:
             if actor_id not in self._actors:
                 raise ActorNotExist(f"Actor {actor_id} does not exist")
             coro = self._actors[actor_id].__on_receive__(message.content)
-            with self._run_coro(message.message_id, coro) as future:
-                result = await future
+            result = await self._run_coro(message.message_id, coro)
             processor.result = ResultMessage(
                 message.message_id,
                 result,

mars/oscar/backends/ray/communication.py

@@ -174,7 +174,7 @@ async def send(self, message: Any):
         if self._closed.is_set():  # pragma: no cover
             raise ChannelClosed("Channel already closed, cannot send message")
         # Put ray object ref to queue
-        await self._in_queue.put(
+        self._in_queue.put_nowait(
             (
                 message,
                 self._peer_actor.__on_ray_recv__.remote(
@@ -255,16 +255,15 @@ async def __on_ray_recv__(self, message):
         """This method will be invoked when current process is a ray actor rather than a ray driver"""
         self._msg_recv_counter += 1
         await self._in_queue.put(message)
-        # Avoid hang when channel is closed after `self._out_queue.get()` is awaited.
-        done, _ = await asyncio.wait(
-            [self._out_queue.get(), self._closed.wait()],
-            return_when=asyncio.FIRST_COMPLETED,
-        )
+        result_message = await self._out_queue.get()
         if self._closed.is_set():  # pragma: no cover
             raise ChannelClosed("Channel already closed")
-        if done:
-            result_message = await done.pop()
-            return _ArgWrapper(result_message)
+        return _ArgWrapper(result_message)
+
+    @implements(Channel.close)
+    async def close(self):
+        await super().close()
+        self._out_queue.put_nowait(None)
 
 
 @register_server

mars/oscar/batch.py

@@ -134,7 +134,12 @@ def _gen_args_kwargs_list(delays):
         return args_list, kwargs_list
 
     async def _async_batch(self, *delays):
-        if self.batch_func:
+        # when there is only one call in batch, calling one-pass method
+        # will be more efficient
+        if len(delays) == 1:
+            d = delays[0]
+            return [await self.func(*d.args, **d.kwargs)]
+        elif self.batch_func:
             args_list, kwargs_list = self._gen_args_kwargs_list(delays)
             return await self.batch_func(args_list, kwargs_list)
         else:

mars/oscar/core.pyx

@@ -212,9 +212,9 @@ cdef class _BaseActor:
         return create_actor_ref(self._address, self._uid)
 
     async def _handle_actor_result(self, result):
-        cdef int result_pos
+        cdef int idx
         cdef tuple res_tuple
-        cdef list tasks, values
+        cdef list tasks, coros, coro_poses, values
         cdef object coro
         cdef bint extract_tuple = False
         cdef bint cancelled = False
@@ -228,20 +228,32 @@ cdef class _BaseActor:
 
         if type(result) is tuple:
             res_tuple = result
-            tasks = []
+            coros = []
+            coro_poses = []
             values = []
-            for res_item in res_tuple:
+            for idx, res_item in enumerate(res_tuple):
                 if is_async_generator(res_item):
-                    value = asyncio.create_task(self._run_actor_async_generator(res_item))
-                    tasks.append(value)
+                    value = self._run_actor_async_generator(res_item)
+                    coros.append(value)
+                    coro_poses.append(idx)
                 elif inspect.isawaitable(res_item):
-                    value = asyncio.create_task(res_item)
-                    tasks.append(value)
+                    value = res_item
+                    coros.append(value)
+                    coro_poses.append(idx)
                 else:
                     value = res_item
                 values.append(value)
 
-            if len(tasks) > 0:
+            # when there is only one coroutine, we do not need to use
+            # asyncio.wait as it introduces much overhead
+            if len(coros) == 1:
+                task_result = await coros[0]
+                if extract_tuple:
+                    result = task_result
+                else:
+                    result = tuple(task_result if t is coros[0] else t for t in values)
+            elif len(coros) > 0:
+                tasks = [asyncio.create_task(t) for t in coros]
                 try:
                     dones, pending = await asyncio.wait(tasks)
                 except asyncio.CancelledError:
@@ -254,7 +266,10 @@ cdef class _BaseActor:
                 if extract_tuple:
                     result = list(dones)[0].result()
                 else:
-                    result = tuple(t.result() if t in dones else t for t in values)
+                    for pos in coro_poses:
+                        task = tasks[pos]
+                        values[pos] = task.result()
+                    result = tuple(values)
 
                 if cancelled:
                     # raise in case no CancelledError raised

mars/services/scheduling/supervisor/assigner.py

@@ -92,6 +92,12 @@ def _get_random_band(self, is_gpu: bool):
     async def assign_subtasks(self, subtasks: List[Subtask]):
         inp_keys = set()
         selected_bands = dict()
+
+        if not self._bands:
+            self._update_bands(
+                list(await self._cluster_api.get_all_bands(NodeRole.WORKER))
+            )
+
         for subtask in subtasks:
             is_gpu = any(c.op.gpu for c in subtask.chunk_graph)
             if subtask.expect_bands:
@@ -116,10 +122,6 @@ async def assign_subtasks(self, subtasks: List[Subtask]):
                 if isinstance(indep_chunk.op, Fetch):
                     inp_keys.add(indep_chunk.key)
                 elif isinstance(indep_chunk.op, FetchShuffle):
-                    if not self._bands:
-                        self._update_bands(
-                            list(await self._cluster_api.get_all_bands(NodeRole.WORKER))
-                        )
                     selected_bands[subtask.subtask_id] = [self._get_random_band(is_gpu)]
                     break
 

mars/services/task/supervisor/processor.py

@@ -547,6 +547,8 @@ async def wait(self, timeout: int = None):
         _, pending = yield asyncio.wait(fs, timeout=timeout)
         if not pending:
             raise mo.Return(self.result())
+        else:
+            [fut.cancel() for fut in pending]
 
     async def cancel(self):
         if self._cur_processor: