Fix ut

Author: hekaisheng

mars/services/storage/handler.py

@@ -23,7 +23,7 @@
 from ...typing import BandType
 from ...utils import calc_data_size, lazy_import
 from ..cluster import ClusterAPI, StorageInfo
-from ..meta import MetaAPI
+from ..meta import MetaAPI, WorkerMetaAPI
 from .core import (
     StorageQuotaActor,
     DataManagerActor,
@@ -39,7 +39,7 @@
 logger = logging.getLogger(__name__)
 
 
-class StorageHandlerActor(mo.Actor):
+class StorageHandlerActor(mo.StatelessActor):
     """
     Storage handler actor, provide methods like `get`, `put`, etc.
     This actor is stateful and created on worker's sub pools.
@@ -468,29 +468,37 @@ async def _fetch_via_transfer(
         )
         logger.debug("Finish fetching %s from band %s", data_keys, remote_band)
 
-    async def fetch_batch(
-        self,
-        session_id: str,
-        data_keys: List[str],
-        level: StorageLevel,
-        band_name: str,
-        address: str,
-        error: str,
+    async def _get_remote_bands(
+        self, session_id: str, data_keys: List, fetch_keys: List, band_name: str
     ):
-        if error not in ("raise", "ignore"):  # pragma: no cover
-            raise ValueError("error must be raise or ignore")
-
+        # For mapper data, we need to return fetch keys to remove them after
+        # execution(see #2922). However, after introducing prefetch, some
+        # mapper data is fetched in advance and not in missing keys, so here
+        # we need to query worker meta service to check if it is on local.
         meta_api = await self._get_meta_api(session_id)
-        remote_keys = defaultdict(set)
+        mapper_main_keys = defaultdict(list)
         missing_keys = []
-        get_metas = []
         get_info_delays = []
         for data_key in data_keys:
+            if isinstance(data_key, tuple):
+                mapper_main_keys[data_key[0]].append(data_key)
             get_info_delays.append(
                 self._data_manager_ref.get_data_info.delay(
                     session_id, data_key, band_name, error="ignore"
                 )
             )
+        # use worker meta api to check if mapper data is executed locally.
+        worker_meta_api = await WorkerMetaAPI.create(session_id, self.address)
+
+        worker_meta_tasks = [
+            worker_meta_api.get_chunk_meta.delay(key, error="ignore")
+            for key in mapper_main_keys
+        ]
+        mapper_metas = await worker_meta_api.get_chunk_meta.batch(*worker_meta_tasks)
+        for key, meta in zip(mapper_main_keys, mapper_metas):
+            if meta is None:
+                fetch_keys.extend(mapper_main_keys[key])
+
         data_infos = await self._data_manager_ref.get_data_info.batch(*get_info_delays)
         pin_delays = []
         for data_key, info in zip(data_keys, data_infos):
@@ -507,31 +515,50 @@ async def fetch_batch(
             else:
                 # Not exists in local, fetch from remote worker
                 missing_keys.append(data_key)
-        if address is None or band_name is None:
-            # some mapper keys are absent, specify error='ignore'
-            # remember that meta only records those main keys
-            get_metas = [
-                (
-                    meta_api.get_chunk_meta.delay(
-                        data_key[0] if isinstance(data_key, tuple) else data_key,
-                        fields=["bands"],
-                        error="ignore",
-                    )
+        # some mapper keys are absent, specify error='ignore'
+        # remember that meta only records those main keys
+        get_metas = [
+            (
+                meta_api.get_chunk_meta.delay(
+                    data_key[0] if isinstance(data_key, tuple) else data_key,
+                    fields=["bands"],
+                    error="ignore",
                 )
-                for data_key in missing_keys
-            ]
+            )
+            for data_key in missing_keys
+        ]
         await self._data_manager_ref.pin.batch(*pin_delays)
-
-        if get_metas:
-            metas = await meta_api.get_chunk_meta.batch(*get_metas)
-        else:  # pragma: no cover
-            metas = [{"bands": [(address, band_name)]}] * len(missing_keys)
+        metas = await meta_api.get_chunk_meta.batch(*get_metas)
         assert len(metas) == len(missing_keys)
-        for data_key, bands in zip(missing_keys, metas):
+        return missing_keys, metas
+
+    async def fetch_batch(
+        self,
+        session_id: str,
+        data_keys: List[str],
+        level: StorageLevel,
+        band_name: str,
+        address: str,
+        error: str,
+    ):
+        if error not in ("raise", "ignore"):  # pragma: no cover
+            raise ValueError("error must be raise or ignore")
+
+        fetch_keys = []
+        if address and band_name:
+            # if specify address, we don't need to query meta
+            missing_keys = data_keys
+            bands = [{"bands": [(address, band_name)]}] * len(missing_keys)
+        else:
+            missing_keys, bands = await self._get_remote_bands(
+                session_id, data_keys, fetch_keys, band_name
+            )
+
+        remote_keys = defaultdict(set)
+        for data_key, bands in zip(missing_keys, bands):
             if bands is not None:
                 remote_keys[bands["bands"][0]].add(data_key)
         transfer_tasks = []
-        fetch_keys = []
         for band, keys in remote_keys.items():
             if StorageLevel.REMOTE in self._quota_refs:
                 # if storage support remote level, just fetch object id
@@ -550,6 +577,7 @@ async def fetch_batch(
         await asyncio.gather(*transfer_tasks)
 
         append_bands_delays = []
+        meta_api = await self._get_meta_api(session_id)
         for data_key in fetch_keys:
             # skip shuffle keys
             if isinstance(data_key, tuple):

mars/services/subtask/worker/processor.py

@@ -675,17 +675,13 @@ async def run(self, subtask: Subtask, wait_post_run: bool = False):
         try:
             result = yield self._running_aio_task
             logger.info("Finished subtask: %s", subtask.subtask_id)
-            # post run with actor tell which will not block
-            if not wait_post_run:
-                await self.ref().post_run.tell(processor)
-            else:
-                await self.post_run(processor)
+            if wait_post_run:
+                await processor.post_run()
             raise mo.Return(result)
         finally:
             self._processor = self._running_aio_task = None
-
-    async def post_run(self, processor: SubtaskProcessor):
-        await processor.post_run()
+            if not wait_post_run:
+                yield processor.post_run()
 
     async def wait(self):
         return self._processor.is_done.wait()