Fix timeout for wait_task (mars-project#2883)

Author: hekaisheng

mars/services/task/api/web.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import base64
 import json
 from typing import Callable, List, Optional, Union
@@ -152,8 +153,18 @@ async def wait_task(self, session_id: str, task_id: str):
         timeout = self.get_argument("timeout", None) or None
         timeout = float(timeout) if timeout is not None else None
         oscar_api = await self._get_oscar_task_api(session_id)
-        res = await oscar_api.wait_task(task_id, timeout)
-        self.write(json.dumps(_json_serial_task_result(res)))
+        if timeout:
+            try:
+                res = await asyncio.wait_for(
+                    asyncio.shield(oscar_api.wait_task(task_id, timeout)),
+                    timeout=timeout,
+                )
+                self.write(json.dumps(_json_serial_task_result(res)))
+            except asyncio.TimeoutError:
+                self.write(json.dumps({}))
+        else:
+            res = await oscar_api.wait_task(task_id, timeout)
+            self.write(json.dumps(_json_serial_task_result(res)))
 
     @web_api("(?P<task_id>[^/]+)", method="delete")
     async def cancel_task(self, session_id: str, task_id: str):

mars/services/task/supervisor/manager.py

@@ -34,7 +34,7 @@
 from ..core import Task, new_task_id, TaskStatus
 from ..errors import TaskNotExist
 from .preprocessor import TaskPreprocessor
-from .processor import TaskProcessorActor
+from .processor import TaskProcessorActor, TaskProcessor
 
 logger = logging.getLogger(__name__)
 
@@ -43,15 +43,18 @@ class TaskConfigurationActor(mo.Actor):
     def __init__(
         self,
         task_conf: Dict[str, Any],
+        task_processor_cls: Type[TaskProcessor] = None,
         task_preprocessor_cls: Type[TaskPreprocessor] = None,
     ):
         for name, value in task_conf.items():
             setattr(task_options, name, value)
+        self._task_processor_cls = task_processor_cls
         self._task_preprocessor_cls = task_preprocessor_cls
 
     def get_config(self):
         return {
             "task_options": task_options,
+            "task_processor_cls": self._task_processor_cls,
             "task_preprocessor_cls": self._task_preprocessor_cls,
         }
 
@@ -77,6 +80,7 @@ def __init__(self, session_id: str):
         self._session_id = session_id
 
         self._config = None
+        self._task_processor_cls = None
         self._task_preprocessor_cls = None
         self._last_idle_time = None
 
@@ -104,9 +108,10 @@ async def __post_create__(self):
             TaskConfigurationActor.default_uid(), address=self.address
         )
         task_conf = await configuration_ref.get_config()
-        self._config, self._task_preprocessor_cls = (
+        self._config, self._task_preprocessor_cls, self._task_processor_cls = (
             task_conf["task_options"],
             task_conf["task_preprocessor_cls"],
+            task_conf["task_processor_cls"],
         )
         self._task_preprocessor_cls = self._get_task_preprocessor_cls()
 
@@ -159,6 +164,7 @@ async def submit_tileable_graph(
                 self._session_id,
                 parent_task_id,
                 task_name=task_name,
+                task_processor_cls=self._task_processor_cls,
                 address=self.address,
                 uid=uid,
             )

mars/services/task/supervisor/processor.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+import importlib
 import itertools
 import logging
 import operator
@@ -525,11 +526,18 @@ class TaskProcessorActor(mo.Actor):
     _processed_task_ids: Set[str]
     _cur_processor: Optional[TaskProcessor]
 
-    def __init__(self, session_id: str, task_id: str, task_name: str = None):
+    def __init__(
+        self,
+        session_id: str,
+        task_id: str,
+        task_name: str = None,
+        task_processor_cls: Type[TaskPreprocessor] = None,
+    ):
         self.session_id = session_id
         self.task_id = task_id
         self.task_name = task_name
 
+        self._task_processor_cls = self._get_task_processor_cls(task_processor_cls)
         self._task_id_to_processor = dict()
         self._cur_processor = None
         self._subtask_decref_events = dict()
@@ -559,7 +567,7 @@ async def add_task(
         task_preprocessor = task_preprocessor_cls(
             task, tiled_context=tiled_context, config=config
         )
-        processor = TaskProcessor(
+        processor = self._task_processor_cls(
             task,
             task_preprocessor,
             self._cluster_api,
@@ -572,6 +580,15 @@ async def add_task(
         # tell self to start running
         await self.ref().start.tell()
 
+    @classmethod
+    def _get_task_processor_cls(cls, task_processor_cls):
+        if task_processor_cls is not None:
+            assert isinstance(task_processor_cls, str)
+            module, name = task_processor_cls.rsplit(".", 1)
+            return getattr(importlib.import_module(module), name)
+        else:
+            return TaskProcessor
+
     def _get_unprocessed_task_processor(self):
         for processor in self._task_id_to_processor.values():
             if processor.result.status == TaskStatus.pending:

mars/services/task/supervisor/service.py

@@ -37,10 +37,12 @@ class TaskSupervisorService(AbstractService):
     async def start(self):
         task_config = self._config.get("task", dict())
         options = task_config.get("default_config", dict())
+        task_processor_cls = task_config.get("task_processor_cls")
         task_preprocessor_cls = task_config.get("task_preprocessor_cls")
         await mo.create_actor(
             TaskConfigurationActor,
             options,
+            task_processor_cls=task_processor_cls,
             task_preprocessor_cls=task_preprocessor_cls,
             address=self._address,
             uid=TaskConfigurationActor.default_uid(),

mars/services/task/tests/test_service.py

@@ -32,6 +32,7 @@
 from ...web import WebActor
 from ...meta import MetaAPI
 from .. import TaskAPI, TaskStatus, WebTaskAPI
+from ..supervisor.processor import TaskProcessor
 from ..errors import TaskNotExist
 
 
@@ -45,7 +46,7 @@ async def start_pool(is_worker: bool):
                 subprocess_start_method="spawn",
             )
         else:
-            kw = dict(n_process=0, subprocess_start_method="spawn")
+            kw = dict(n_process=1, subprocess_start_method="spawn")
         pool = await mo.create_actor_pool("127.0.0.1", **kw)
         await pool.start()
         return pool
@@ -57,11 +58,9 @@ async def start_pool(is_worker: bool):
         await asyncio.gather(sv_pool.stop(), worker_pool.stop())
 
 
-@pytest.mark.parametrize(indirect=True)
-@pytest.fixture(params=[False, True])
-async def start_test_service(actor_pools, request):
-    sv_pool, worker_pool = actor_pools
-
+async def _start_services(
+    supervisor_pool, worker_pool, request, task_processor_cls=None
+):
     config = {
         "services": [
             "cluster",
@@ -75,37 +74,49 @@ async def start_test_service(actor_pools, request):
         ],
         "cluster": {
             "backend": "fixed",
-            "lookup_address": sv_pool.external_address,
+            "lookup_address": supervisor_pool.external_address,
             "resource": {"numa-0": Resource(num_cpus=2)},
         },
         "meta": {"store": "dict"},
         "scheduling": {},
         "task": {},
     }
+    if task_processor_cls:
+        config["task"]["task_processor_cls"] = task_processor_cls
     if request:
         config["services"].append("web")
-
-    await start_services(NodeRole.SUPERVISOR, config, address=sv_pool.external_address)
+    await start_services(
+        NodeRole.SUPERVISOR, config, address=supervisor_pool.external_address
+    )
     await start_services(NodeRole.WORKER, config, address=worker_pool.external_address)
 
     session_id = "test_session"
-    session_api = await SessionAPI.create(sv_pool.external_address)
+    session_api = await SessionAPI.create(supervisor_pool.external_address)
     await session_api.create_session(session_id)
 
     if not request.param:
-        task_api = await TaskAPI.create(session_id, sv_pool.external_address)
+        task_api = await TaskAPI.create(session_id, supervisor_pool.external_address)
     else:
         web_actor = await mo.actor_ref(
-            WebActor.default_uid(), address=sv_pool.external_address
+            WebActor.default_uid(), address=supervisor_pool.external_address
         )
         web_address = await web_actor.get_web_address()
         task_api = WebTaskAPI(session_id, web_address)
 
     assert await task_api.get_task_results() == []
 
     # create mock meta and storage APIs
-    _ = await MetaAPI.create(session_id, sv_pool.external_address)
+    _ = await MetaAPI.create(session_id, supervisor_pool.external_address)
     storage_api = await MockStorageAPI.create(session_id, worker_pool.external_address)
+    return task_api, storage_api, config
+
+
+@pytest.mark.parametrize(indirect=True)
+@pytest.fixture(params=[False, True])
+async def start_test_service(actor_pools, request):
+    sv_pool, worker_pool = actor_pools
+
+    task_api, storage_api, config = await _start_services(sv_pool, worker_pool, request)
 
     try:
         yield sv_pool.external_address, task_api, storage_api
@@ -115,6 +126,66 @@ async def start_test_service(actor_pools, request):
         await stop_services(NodeRole.SUPERVISOR, config, sv_pool.external_address)
 
 
+class MockTaskProcessor(TaskProcessor):
+    @classmethod
+    def _get_decref_stage_chunk_keys(cls, stage_processor):
+        import time
+
+        # time.sleep to block async thread
+        time.sleep(5)
+        return super()._get_decref_stage_chunk_keys(stage_processor)
+
+
+@pytest.mark.parametrize(indirect=True)
+@pytest.fixture(params=[True])
+async def start_test_service_with_mock(actor_pools, request):
+    sv_pool, worker_pool = actor_pools
+
+    task_api, storage_api, config = await _start_services(
+        sv_pool,
+        worker_pool,
+        request,
+        task_processor_cls="mars.services.task.tests.test_service.MockTaskProcessor",
+    )
+
+    try:
+        yield sv_pool.external_address, task_api, storage_api
+    finally:
+        await MockStorageAPI.cleanup(worker_pool.external_address)
+        await stop_services(NodeRole.WORKER, config, worker_pool.external_address)
+        await stop_services(NodeRole.SUPERVISOR, config, sv_pool.external_address)
+
+
+@pytest.mark.asyncio
+async def test_task_timeout_execution(start_test_service_with_mock):
+    _sv_pool_address, task_api, storage_api = start_test_service_with_mock
+
+    def f1():
+        return np.arange(5)
+
+    def f2():
+        return np.arange(5, 10)
+
+    def f3(f1r, f2r):
+        return np.concatenate([f1r, f2r]).sum()
+
+    r1 = mr.spawn(f1)
+    r2 = mr.spawn(f2)
+    r3 = mr.spawn(f3, args=(r1, r2))
+
+    graph = TileableGraph([r3.data])
+    next(TileableGraphBuilder(graph).build())
+
+    task_id = await task_api.submit_tileable_graph(graph, fuse_enabled=False)
+    assert await task_api.get_last_idle_time() is None
+    assert isinstance(task_id, str)
+
+    await task_api.wait_task(task_id, timeout=2)
+    task_result = await task_api.get_task_result(task_id)
+
+    assert task_result.status == TaskStatus.terminated
+
+
 @pytest.mark.asyncio
 async def test_task_execution(start_test_service):
     _sv_pool_address, task_api, storage_api = start_test_service