FEAT: Collect subtask runtime info

Author: UranusSeven

mars/services/scheduling/worker/execution.py

@@ -18,6 +18,7 @@
 import operator
 import pprint
 import sys
+import time
 from collections import defaultdict
 from dataclasses import dataclass, field
 from typing import Dict, Optional
@@ -35,6 +36,7 @@
 from ...meta import MetaAPI
 from ...storage import StorageAPI
 from ...subtask import Subtask, SubtaskAPI, SubtaskResult, SubtaskStatus
+from ...task.task_info_collector import TaskInfoCollector
 from .workerslot import BandSlotManagerActor
 from .quota import QuotaActor
 
@@ -353,6 +355,7 @@ async def internal_run_subtask(self, subtask: Subtask, band_name: str):
         )
         try:
             logger.debug("Preparing data for subtask %s", subtask.subtask_id)
+            fetch_start = time.time()
             prepare_data_task = asyncio.create_task(
                 _retry_run(
                     subtask, subtask_info, self._prepare_input_data, subtask, band_name
@@ -361,6 +364,14 @@ async def internal_run_subtask(self, subtask: Subtask, band_name: str):
             await asyncio.wait_for(
                 prepare_data_task, timeout=self._data_prepare_timeout
             )
+            fetch_end = time.time()
+            collect_task_info = subtask.extra_config and subtask.extra_config.get(
+                "collect_task_info", False
+            )
+            task_info_collector = TaskInfoCollector(self.address, collect_task_info)
+            await task_info_collector.collect_fetch_time(
+                subtask, fetch_start, fetch_end
+            )
 
             input_sizes = await self._collect_input_sizes(
                 subtask, subtask_info.supervisor_address, band_name

mars/services/session/supervisor/core.py

@@ -190,7 +190,7 @@ async def create_services(self):
         )
         if "task" in self._service_config["services"]:
             self._task_api = await TaskAPI.create(
-                session_id=self._session_id, address=self.address
+                session_id=self._session_id, supervisor_address=self.address
             )
 
     async def get_last_idle_time(self):

mars/services/subtask/worker/manager.py

@@ -48,6 +48,7 @@ async def _create_band_runner_actors(self, band_name: str, n_slots: int):
                 SubtaskRunnerActor,
                 band,
                 worker_address=self._worker_address,
+                slot_id=slot_id,
                 subtask_processor_cls=self._subtask_processor_cls,
                 uid=SubtaskRunnerActor.gen_uid(band_name, slot_id),
                 address=self.address,

mars/services/subtask/worker/processor.py

@@ -32,12 +32,13 @@
 from ....optimization.physical import optimize
 from ....serialization import AioSerializer
 from ....typing import BandType, ChunkType
-from ....utils import get_chunk_key_to_data_keys, calc_data_size
+from ....utils import get_chunk_key_to_data_keys, calc_data_size, Timer
 from ...context import ThreadedServiceContext
 from ...meta.api import MetaAPI, WorkerMetaAPI
 from ...session import SessionAPI
 from ...storage import StorageAPI
 from ...task import TaskAPI, task_options
+from ...task.task_info_collector import TaskInfoCollector
 from ..core import Subtask, SubtaskStatus, SubtaskResult
 from ..utils import iter_input_data_keys, iter_output_data, get_mapper_data_keys
 
@@ -77,6 +78,7 @@ def __init__(
         meta_api: MetaAPI,
         worker_meta_api: WorkerMetaAPI,
         band: BandType,
+        slot_id: int,
         supervisor_address: str,
         engines: List[str] = None,
     ):
@@ -91,6 +93,7 @@ def __init__(
             ]
         )
         self._band = band
+        self._slot_id = slot_id
         self._supervisor_address = supervisor_address
         self._engines = engines if engines is not None else task_options.runtime_engines
 
@@ -120,6 +123,10 @@ def __init__(
         self._storage_api = storage_api
         self._meta_api = meta_api
         self._worker_meta_api = worker_meta_api
+        collect_task_info = subtask.extra_config and subtask.extra_config.get(
+            "collect_task_info", False
+        )
+        self._task_info_collector = TaskInfoCollector(self._band[0], collect_task_info)
 
         # add metrics
         self._subtask_execution_time = Metrics.gauge(
@@ -234,35 +241,37 @@ def cb(fut):
                             to_wait.set_result(fut.result())
 
                 future.add_done_callback(cb)
-
-                try:
-                    await to_wait
-                    logger.debug(
-                        "Finish executing operand: %s, chunk: %s, subtask id: %s",
-                        chunk.op,
-                        chunk,
-                        self.subtask.subtask_id,
-                    )
-                except asyncio.CancelledError:
-                    logger.debug(
-                        "Receive cancel instruction for operand: %s,"
-                        "chunk: %s, subtask id: %s",
-                        chunk.op,
-                        chunk,
-                        self.subtask.subtask_id,
-                    )
-                    # wait for this computation to finish
-                    await future
-                    # if cancelled, stop next computation
-                    logger.debug(
-                        "Cancelled operand: %s, chunk: %s, subtask id: %s",
-                        chunk.op,
-                        chunk,
-                        self.subtask.subtask_id,
-                    )
-                    self.result.status = SubtaskStatus.cancelled
-                    raise
-
+                with Timer() as timer:
+                    try:
+                        await to_wait
+                        logger.debug(
+                            "Finish executing operand: %s, chunk: %s, subtask id: %s",
+                            chunk.op,
+                            chunk,
+                            self.subtask.subtask_id,
+                        )
+                    except asyncio.CancelledError:  # pragma: no cover
+                        logger.debug(
+                            "Receive cancel instruction for operand: %s,"
+                            "chunk: %s, subtask id: %s",
+                            chunk.op,
+                            chunk,
+                            self.subtask.subtask_id,
+                        )
+                        # wait for this computation to finish
+                        await future
+                        # if cancelled, stop next computation
+                        logger.debug(
+                            "Cancelled operand: %s, chunk: %s, subtask id: %s",
+                            chunk.op,
+                            chunk,
+                            self.subtask.subtask_id,
+                        )
+                        self.result.status = SubtaskStatus.cancelled
+                        raise
+                await self._task_info_collector.collect_runtime_operand_info(
+                    self.subtask, timer.duration, chunk, self._processor_context
+                )
             self.set_op_progress(chunk.op.key, 1.0)
 
             for inp in chunk_graph.iter_predecessors(chunk):
@@ -541,6 +550,7 @@ async def done(self):
         self.is_done.set()
 
     async def run(self):
+        cost_times = defaultdict(dict)
         self.result.status = SubtaskStatus.running
         input_keys = None
         unpinned = False
@@ -561,29 +571,49 @@ async def run(self):
             }
 
             # load inputs data
+            cost_times["load_data_time"]["start_time"] = time.time()
             input_keys = await self._load_input_data()
+            cost_times["load_data_time"]["end_time"] = time.time()
             try:
                 # execute chunk graph
+                cost_times["execute_time"]["start_time"] = time.time()
                 await self._execute_graph(chunk_graph)
+                cost_times["execute_time"]["end_time"] = time.time()
             finally:
                 # unpin inputs data
                 unpinned = True
+                cost_times["unpin_time"]["start_time"] = time.time()
                 await self._unpin_data(input_keys)
+                cost_times["unpin_time"]["end_time"] = time.time()
             # store results data
+            cost_times["store_result_time"]["start_time"] = time.time()
             (
                 stored_keys,
                 store_sizes,
                 memory_sizes,
                 data_key_to_object_id,
             ) = await self._store_data(chunk_graph)
+            cost_times["store_result_time"]["end_time"] = time.time()
             # store meta
+            cost_times["store_meta_time"]["start_time"] = time.time()
             await self._store_meta(
                 chunk_graph,
                 store_sizes,
                 memory_sizes,
                 data_key_to_object_id,
                 update_meta_chunks,
             )
+            cost_times["store_meta_time"]["end_time"] = time.time()
+            await self._task_info_collector.collect_runtime_subtask_info(
+                self.subtask,
+                self._band,
+                self._slot_id,
+                stored_keys,
+                store_sizes,
+                memory_sizes,
+                cost_times,
+            )
+
         except asyncio.CancelledError:
             self.result.status = SubtaskStatus.cancelled
             self.result.progress = 1.0
@@ -708,7 +738,7 @@ async def _init_context(self, session_id: str) -> ThreadedServiceContext:
         await context.init()
         return context
 
-    async def run(self, subtask: Subtask):
+    async def run(self, subtask: Subtask, slot_id: int):
         logger.info(
             "Start to run subtask: %r on %s. chunk graph contains %s",
             subtask,
@@ -728,6 +758,7 @@ async def run(self, subtask: Subtask):
                 self._meta_api,
                 self._worker_meta_api,
                 self._band,
+                slot_id,
                 self._supervisor_address,
             )
             self._processor = self._last_processor = processor

mars/services/subtask/worker/runner.py

@@ -41,9 +41,14 @@ def gen_uid(cls, band_name: str, slot_id: int):
         return f"slot_{band_name}_{slot_id}_subtask_runner"
 
     def __init__(
-        self, band: BandType, worker_address: str, subtask_processor_cls: Type = None
+        self,
+        band: BandType,
+        worker_address: str,
+        slot_id: int,
+        subtask_processor_cls: Type = None,
     ):
         self._band = band
+        self._slot_id = slot_id
         self._worker_address = worker_address
         self._subtask_processor_cls = self._get_subtask_processor_cls(
             subtask_processor_cls
@@ -122,7 +127,7 @@ async def run_subtask(self, subtask: Subtask):
         processor = self._session_id_to_processors[session_id]
         try:
             self._running_processor = self._last_processor = processor
-            result = yield self._running_processor.run(subtask)
+            result = yield self._running_processor.run(subtask, self._slot_id)
         finally:
             self._running_processor = None
         raise mo.Return(result)

mars/services/task/api/core.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
-from typing import List, Union
+from typing import List, Union, Dict
 
 from ....core import Tileable
 from ..core import TileableGraph, TaskResult
@@ -145,3 +145,8 @@ async def get_last_idle_time(self) -> Union[float, None]:
         last_idle_time: float
             The last idle time if the task manager is idle else None.
         """
+
+    async def save_task_info(self, task_info: Dict, path: str):
+        """
+        Save task information using yaml format.
+        """

mars/services/task/api/oscar.py

@@ -12,46 +12,65 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Union
+from typing import List, Union, Dict
 
 from .... import oscar as mo
 from ....core import Tileable
 from ....lib.aio import alru_cache
 from ...subtask import SubtaskResult
 from ..core import TileableGraph, TaskResult, MapReduceInfo
 from ..supervisor.manager import TaskManagerActor
+from ..task_info_collector import TaskInfoCollectorActor
 from .core import AbstractTaskAPI
 
 
 class TaskAPI(AbstractTaskAPI):
     def __init__(
-        self, session_id: str, task_manager_ref: mo.ActorRefType[TaskManagerActor]
+        self,
+        session_id: str,
+        task_info_collector_ref: mo.ActorRefType[TaskInfoCollectorActor],
+        task_manager_ref: mo.ActorRefType[TaskManagerActor],
     ):
         self._session_id = session_id
+        self._task_info_collector_ref = task_info_collector_ref
         self._task_manager_ref = task_manager_ref
 
     @classmethod
     @alru_cache(cache_exceptions=False)
-    async def create(cls, session_id: str, address: str) -> "TaskAPI":
+    async def create(
+        cls, session_id: str, supervisor_address: str = None, local_address: str = None
+    ) -> "TaskAPI":
         """
         Create Task API.
 
         Parameters
         ----------
         session_id : str
-            Session ID
-        address : str
+            Session ID.
+        supervisor_address : str
             Supervisor address.
+        local_address : str
+            Local address.
 
         Returns
         -------
         task_api
             Task API.
         """
-        task_manager_ref = await mo.actor_ref(
-            address, TaskManagerActor.gen_uid(session_id)
-        )
-        return TaskAPI(session_id, task_manager_ref)
+        try:
+            task_manager_ref = await mo.actor_ref(
+                supervisor_address, TaskManagerActor.gen_uid(session_id)
+            )
+        except (mo.ActorNotExist, ValueError):
+            task_manager_ref = None
+        try:
+            task_info_collector_ref = await mo.actor_ref(
+                local_address, TaskInfoCollectorActor.default_uid()
+            )
+        except (mo.ActorNotExist, ValueError):
+            task_info_collector_ref = None
+
+        return TaskAPI(session_id, task_info_collector_ref, task_manager_ref)
 
     async def get_task_results(self, progress: bool = False) -> List[TaskResult]:
         return await self._task_manager_ref.get_task_results(progress)
@@ -110,3 +129,6 @@ async def get_map_reduce_info(
         self, task_id: str, map_reduce_id: int
     ) -> MapReduceInfo:
         return await self._task_manager_ref.get_map_reduce_info(task_id, map_reduce_id)
+
+    async def save_task_info(self, task_info: Dict, path: str):
+        await self._task_info_collector_ref.save_task_info(task_info, path)

mars/services/task/supervisor/graph_visualizer.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 import itertools
 from io import StringIO
 from typing import Dict, List