Use tell when remove mapper data after execution (#3027)

Author: hekaisheng

mars/lib/tbcode.py

@@ -24,9 +24,7 @@
 from collections import defaultdict
 
 
-def dump_traceback_code(
-    tb: types.TracebackType, number_of_lines_of_context: int = 5
-):
+def dump_traceback_code(tb: types.TracebackType, number_of_lines_of_context: int = 5):
     """
     Dump codes before and after lines of tracebacks.
 
@@ -56,9 +54,7 @@ def dump_traceback_code(
                 dict(left=left_range, right=right_range, code=fragment)
             )
             results[file_name].update(
-                dict(
-                    size=cache_data[0], lines=len(cache_data[2])
-                )
+                dict(size=cache_data[0], lines=len(cache_data[2]))
             )
         tb = tb.tb_next
     return dict(results)
@@ -90,7 +86,10 @@ def load_traceback_code(code_frags: dict, cache: dict = None):
         if file_name not in cache:
             # keep field 1 (mtime) as None to ensure lazy cache
             cache[file_name] = (
-                profile["size"], None, [""] * profile["lines"], file_name
+                profile["size"],
+                None,
+                [""] * profile["lines"],
+                file_name,
             )
         for fragment in profile["fragments"]:
             left_range, right_range = fragment["left"], fragment["right"]

mars/services/scheduling/worker/execution.py

@@ -20,7 +20,7 @@
 import sys
 from collections import defaultdict
 from dataclasses import dataclass, field
-from typing import Dict, Optional, Union
+from typing import Dict, List, Optional, Union
 
 from .... import oscar as mo
 from ....core import ExecutionError
@@ -344,6 +344,20 @@ def _check_cancelling(cls, subtask_info: SubtaskExecutionInfo):
         if subtask_info.cancelling:
             raise asyncio.CancelledError
 
+    async def _remove_mapper_data(
+        self, session_id: str, band_name: str, remote_mapper_keys: List
+    ):
+        storage_api = await StorageAPI.create(
+            session_id, address=self.address, band_name=band_name
+        )
+        logger.debug("Delete mapper data %s", remote_mapper_keys)
+        await storage_api.delete.batch(
+            *[
+                storage_api.delete.delay(key, error="ignore")
+                for key in remote_mapper_keys
+            ]
+        )
+
     async def internal_run_subtask(self, subtask: Subtask, band_name: str):
         subtask_api = SubtaskAPI(self.address)
         subtask_info = self._subtask_info[subtask.subtask_id]
@@ -379,15 +393,8 @@ async def internal_run_subtask(self, subtask: Subtask, band_name: str):
                 subtask, band_name, subtask_api, batch_quota_req
             )
             if remote_mapper_keys:
-                storage_api = await StorageAPI.create(
-                    subtask.session_id, address=self.address, band_name=band_name
-                )
-                logger.debug("Delete mapper data %s", remote_mapper_keys)
-                await storage_api.delete.batch(
-                    *[
-                        storage_api.delete.delay(key, error="ignore")
-                        for key in remote_mapper_keys
-                    ]
+                await self.ref()._remove_mapper_data.tell(
+                    subtask.session_id, band_name, remote_mapper_keys
                 )
         except:  # noqa: E722  # pylint: disable=bare-except
             _fill_subtask_result_with_exception(subtask, subtask_info)

mars/services/storage/transfer.py

@@ -318,7 +318,7 @@ async def open_writers(
                 raise
 
     async def do_write(self, message: TransferMessage):
-        # close may be a high cost operation, use create_task
+        # close may be a high-cost operation, use create_task
         close_tasks = []
         finished_keys = []
         session_id = message.session_id

mars/services/task/api/web.py

@@ -39,8 +39,12 @@ def _json_serial_task_result(result: Optional[TaskResult]):
         "profiling": result.profiling,
     }
     if result.error is not None:
-        res_json["error"] = base64.b64encode(serialize_serializable(result.error)).decode()
-        res_json["traceback"] = base64.b64encode(serialize_serializable(result.traceback)).decode()
+        res_json["error"] = base64.b64encode(
+            serialize_serializable(result.error)
+        ).decode()
+        res_json["traceback"] = base64.b64encode(
+            serialize_serializable(result.traceback)
+        ).decode()
         res_json["traceback_code"] = dump_traceback_code(result.traceback)
     return res_json
 

pyproject.toml

@@ -9,7 +9,7 @@ requires = [
 build-backend = "setuptools.build_meta"
 
 [tool.black]
-extend-exclude = [
-    "^/mars/_version.py",
-    "^/mars/lib/tblib/.*",
-]
+include = '\.pyi?$'
+extend-exclude = '''
+^/mars/(_version.py|lib/tblib/.*)
+'''