Add metrics on operand and subtask executions (#2947)

Author: zhongchun

mars/core/operand/core.py

@@ -22,6 +22,7 @@
 except ImportError:  # pragma: no cover
     UFuncTypeError = None
 
+from ...metrics import Metrics
 from ...typing import TileableType, ChunkType, OperandType
 from ...utils import calc_data_size
 from ..context import Context
@@ -39,6 +40,11 @@
 _op_type_to_size_estimator: Dict[Type[OperandType], Callable] = dict()
 
 
+op_executed_number = Metrics.counter(
+    "mars.operand.executed_number", "The number of executed operands.", ("op",)
+)
+
+
 class TileableOperandMixin:
     __slots__ = ()
 
@@ -484,6 +490,7 @@ def execute(results: Dict[str, Any], op: OperandType):
             try:
                 result = executor(results, op)
                 succeeded = True
+                op_executed_number.record(1, {"op": op.__class__.__name__})
                 return result
             except UFuncTypeError as e:  # pragma: no cover
                 raise TypeError(str(e)).with_traceback(sys.exc_info()[2]) from None

mars/metrics/backends/metric.py

@@ -29,9 +29,15 @@ class AbstractMetric(ABC):
     def __init__(
         self, name: str, description: str = "", tag_keys: Optional[Tuple[str]] = None
     ):
+        assert isinstance(name, str), "Argument name should be a str"
+        assert isinstance(description, str), "Argument description should be a str"
+        if tag_keys is not None:
+            assert isinstance(tag_keys, tuple) and all(
+                isinstance(tag, str) for tag in tag_keys
+            ), "Argument tag_keys should be a tuple and its elements should be str"
         self._name = name
         self._description = description
-        self._tag_keys = tuple(tag_keys) if tag_keys else tuple()
+        self._tag_keys = tag_keys or tuple()
         self._init()
 
     @property

mars/metrics/backends/tests/test_metric.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pytest
+
 from ..metric import (
     AbstractMetric,
     AbstractCounter,
@@ -21,6 +23,24 @@
 )
 
 
+def test_illegal_arguments():
+    class DummyMetric(AbstractMetric):
+        pass
+
+    DummyMetric.__abstractmethods__ = set()
+    with pytest.raises(AssertionError):
+        DummyMetric(1)
+
+    with pytest.raises(AssertionError):
+        DummyMetric("dummy_metric", 1)
+
+    with pytest.raises(AssertionError):
+        DummyMetric("dummy_metric", "A test metric", "service")
+
+    with pytest.raises(AssertionError):
+        DummyMetric("dummy_metric", "A test metric", ("service", 1))
+
+
 def test_dummy_metric():
     class DummyMetric(AbstractMetric):
         pass

mars/services/scheduling/supervisor/manager.py

@@ -21,6 +21,7 @@
 
 from .... import oscar as mo
 from ....lib.aio import alru_cache
+from ....metrics import Metrics
 from ....oscar.backends.context import ProfilingContext
 from ....oscar.errors import MarsError
 from ....oscar.profiling import ProfilingData, MARS_ENABLE_PROFILING
@@ -82,6 +83,21 @@ def __init__(
         self._speculation_config = speculation_config or {}
         self._queueing_ref = None
         self._global_resource_ref = None
+        self._submitted_subtask_count = Metrics.counter(
+            "mars.scheduling.submitted_subtask_count",
+            "The count of submitted subtasks to all bands.",
+            ("session_id", "task_id", "stage_id"),
+        )
+        self._finished_subtask_count = Metrics.counter(
+            "mars.scheduling.finished_subtask_count",
+            "The count of finished subtasks of all bands.",
+            ("session_id", "task_id", "stage_id"),
+        )
+        self._canceled_subtask_count = Metrics.counter(
+            "mars.scheduling.canceled_subtask_count",
+            "The count of canceled subtasks of all bands.",
+            ("session_id", "task_id", "stage_id"),
+        )
         logger.info(
             "Created SubtaskManager with subtask_max_reschedules %s, "
             "speculation_config %s",
@@ -167,6 +183,14 @@ async def finish_subtasks(
         for subtask_id, subtask_band in zip(subtask_ids, bands):
             subtask_info = self._subtask_infos.get(subtask_id, None)
             if subtask_info is not None:
+                self._finished_subtask_count.record(
+                    1,
+                    {
+                        "session_id": self._session_id,
+                        "task_id": subtask_info.subtask.task_id,
+                        "stage_id": subtask_info.subtask.stage_id,
+                    },
+                )
                 self._subtask_summaries[subtask_id] = subtask_info.to_summary(
                     is_finished=True
                 )
@@ -236,6 +260,14 @@ async def submit_subtask_to_band(self, subtask_id: str, band: BandType):
                     if enable_profiling
                     else None
                 )
+                self._submitted_subtask_count.record(
+                    1,
+                    {
+                        "session_id": self._session_id,
+                        "task_id": subtask_info.subtask.task_id,
+                        "stage_id": subtask_info.subtask.stage_id,
+                    },
+                )
                 logger.debug("Start run subtask %s in band %s.", subtask_id, band)
                 with Timer() as timer:
                     task = asyncio.create_task(
@@ -388,6 +420,14 @@ async def cancel_single_task(subtask, raw_tasks, cancel_tasks):
                 self._subtask_summaries[subtask_id] = subtask_info.to_summary(
                     is_finished=True, is_cancelled=True
                 )
+                self._canceled_subtask_count.record(
+                    1,
+                    {
+                        "session_id": self._session_id,
+                        "task_id": subtask_info.subtask.task_id,
+                        "stage_id": subtask_info.subtask.stage_id,
+                    },
+                )
         await self._queueing_ref.submit_subtasks.tell()
         logger.info("Subtasks %s canceled.", subtask_ids)
 

mars/services/scheduling/supervisor/queueing.py

@@ -22,6 +22,7 @@
 
 from .... import oscar as mo
 from ....lib.aio import alru_cache
+from ....metrics import Metrics
 from ....resource import ZeroResource
 from ....utils import dataslots
 from ...subtask import Subtask
@@ -70,6 +71,16 @@ def __init__(self, session_id: str, submit_period: Union[float, int] = None):
 
         self._periodical_submit_task = None
         self._submit_period = submit_period or _DEFAULT_SUBMIT_PERIOD
+        self._submitted_subtask_number = Metrics.gauge(
+            "mars.band.submitted_subtask_number",
+            "The number of submitted subtask to a band.",
+            ("session_id", "band"),
+        )
+        self._unsubmitted_subtask_number = Metrics.gauge(
+            "mars.band.unsubmitted_subtask_number",
+            "The number of unsubmitted subtask to a band.",
+            ("session_id", "band"),
+        )
 
     async def __post_create__(self):
         from ...cluster import ClusterAPI
@@ -246,6 +257,12 @@ async def submit_subtasks(self, band: Tuple = None, limit: Optional[int] = None)
                 self, [item.subtask for item in submit_items.values()]
             ):
                 non_submitted_ids = [k for k in submit_items if k not in submitted_ids]
+                tags = {
+                    "session_id": self._session_id,
+                    "band": band[0] if band else "",
+                }
+                self._submitted_subtask_number.record(len(submitted_ids), tags)
+                self._unsubmitted_subtask_number.record(len(non_submitted_ids), tags)
                 if submitted_ids:
                     for stid in subtask_ids:
                         if stid not in submitted_ids:

mars/services/scheduling/worker/execution.py

@@ -26,6 +26,7 @@
 from ....core.graph import DAG
 from ....core.operand import Fetch, FetchShuffle
 from ....lib.aio import alru_cache
+from ....metrics import Metrics
 from ....oscar.errors import MarsError
 from ....storage import StorageLevel
 from ....utils import dataslots, get_chunk_key_to_data_keys, wrap_exception
@@ -145,6 +146,16 @@ def __init__(
         self._data_prepare_timeout = data_prepare_timeout
 
         self._subtask_info = dict()
+        self._submitted_subtask_count = Metrics.counter(
+            "mars.band.submitted_subtask_count",
+            "The count of submitted subtasks to the current band.",
+            ("band",),
+        )
+        self._finished_subtask_count = Metrics.counter(
+            "mars.band.finished_subtask_count",
+            "The count of finished subtasks of the current band.",
+            ("band",),
+        )
 
     async def __post_create__(self):
         self._cluster_api = await ClusterAPI.create(self.address)
@@ -499,6 +510,7 @@ async def run_subtask(
         logger.debug(
             "Start to schedule subtask %s on %s.", subtask.subtask_id, self.address
         )
+        self._submitted_subtask_count.record(1, {"band": self.address})
         with mo.debug.no_message_trace():
             task = asyncio.create_task(
                 self.ref().internal_run_subtask(subtask, band_name)
@@ -519,6 +531,7 @@ async def run_subtask(
         )
         result = await task
         self._subtask_info.pop(subtask.subtask_id, None)
+        self._finished_subtask_count.record(1, {"band": self.address})
         logger.debug("Subtask %s finished with result %s", subtask.subtask_id, result)
         return result