[Data] Adding in metrics for number of actors alive, pending and restarting (#51082)

## Why are these changes needed?
Currently we have information about actors that is only displayed in the
progress bar (see
[here](https://github.com/ray-project/ray/blob/678a8d75b47730d6685341a61bc8d14e924a61f8/python/ray/data/_internal/execution/operators/actor_pool_map_operator.py#L783)),
this adds corresponding metrics so they can be included in prometheus /
grafana.

---------

Signed-off-by: Matthew Owen <mowen@anyscale.com>

Author: omatthew98

python/ray/data/_internal/execution/interfaces/op_runtime_metrics.py

@@ -35,6 +35,7 @@ class MetricsGroup(Enum):
     TASKS = "tasks"
     OBJECT_STORE_MEMORY = "object_store_memory"
     MISC = "misc"
+    ACTORS = "actors"
 
 
 @dataclass(frozen=True)
@@ -324,6 +325,23 @@ class OpRuntimeMetrics(metaclass=OpRuntimesMetricsMeta):
         metrics_group=MetricsGroup.TASKS,
     )
 
+    # === Actor-related metrics ===
+    num_alive_actors: int = metric_field(
+        default=0,
+        description="Number of alive actors.",
+        metrics_group=MetricsGroup.ACTORS,
+    )
+    num_restarting_actors: int = metric_field(
+        default=0,
+        description="Number of restarting actors.",
+        metrics_group=MetricsGroup.ACTORS,
+    )
+    num_pending_actors: int = metric_field(
+        default=0,
+        description="Number of pending actors.",
+        metrics_group=MetricsGroup.ACTORS,
+    )
+
     # === Object store memory metrics ===
     obj_store_mem_internal_inqueue_blocks: int = metric_field(
         default=0,

python/ray/data/_internal/execution/interfaces/physical_operator.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Dict, Iterator, List, Optional, Union
+from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
 import uuid
 
 import ray
@@ -569,3 +569,12 @@ def actor_info_progress_str(self) -> str:
         Actors.
         """
         return ""
+
+    def actor_info_counts(self) -> Tuple[int, int, int]:
+        """Returns Actor counts for Alive, Restarting and Pending Actors.
+
+        This method will be called in add_output API in OpState. Subclasses can
+        override it to return counts for Alive, Restarting and Pending
+        Actors.
+        """
+        return 0, 0, 0

python/ray/data/_internal/execution/operators/actor_pool_map_operator.py

@@ -386,6 +386,10 @@ def actor_info_progress_str(self) -> str:
         """Returns Actor progress strings for Alive, Restarting and Pending Actors."""
         return self._actor_pool.actor_info_progress_str()
 
+    def actor_info_counts(self) -> Tuple[int, int, int]:
+        """Returns Actor counts for Alive, Restarting and Pending Actors."""
+        return self._actor_pool.actor_info_counts()
+
 
 class _MapWorker:
     """An actor worker for MapOperator."""
@@ -780,11 +784,16 @@ def _get_location(self, bundle: RefBundle) -> Optional[NodeIdStr]:
         """
         return bundle.get_cached_location()
 
-    def actor_info_progress_str(self) -> str:
-        """Returns Actor progress strings for Alive, Restarting and Pending Actors."""
+    def actor_info_counts(self) -> Tuple[int, int, int]:
+        """Returns Actor counts for Alive, Restarting and Pending Actors."""
         alive = self.num_alive_actors()
         pending = self.num_pending_actors()
         restarting = self.num_restarting_actors()
+        return alive, pending, restarting
+
+    def actor_info_progress_str(self) -> str:
+        """Returns Actor progress strings for Alive, Restarting and Pending Actors."""
+        alive, pending, restarting = self.actor_info_counts()
         total = alive + pending + restarting
         if total == alive:
             return f"; Actors: {total}"

python/ray/data/_internal/execution/streaming_executor_state.py

@@ -253,6 +253,10 @@ def add_output(self, ref: RefBundle) -> None:
                 ref.num_rows() is not None
             ), "RefBundle must have a valid number of rows"
             self.progress_bar.update(ref.num_rows(), self.op.num_output_rows_total())
+        active, restarting, pending = self.op.actor_info_counts()
+        self.op.metrics.num_alive_actors = active
+        self.op.metrics.num_restarting_actors = restarting
+        self.op.metrics.num_pending_actors = pending
 
     def refresh_progress_bar(self, resource_manager: ResourceManager) -> None:
         """Update the console with the latest operator progress."""

python/ray/data/tests/test_stats.py

@@ -74,6 +74,9 @@ def gen_expected_metrics(
                 "'task_submission_backpressure_time': "
                 f"{'N' if task_backpressure else 'Z'}"
             ),
+            "'num_alive_actors': Z",
+            "'num_restarting_actors': Z",
+            "'num_pending_actors': Z",
             "'obj_store_mem_internal_inqueue_blocks': Z",
             "'obj_store_mem_internal_outqueue_blocks': Z",
             "'obj_store_mem_freed': N",
@@ -94,6 +97,9 @@ def gen_expected_metrics(
                 "'task_submission_backpressure_time': "
                 f"{'N' if task_backpressure else 'Z'}"
             ),
+            "'num_alive_actors': Z",
+            "'num_restarting_actors': Z",
+            "'num_pending_actors': Z",
             "'obj_store_mem_internal_inqueue_blocks': Z",
             "'obj_store_mem_internal_outqueue_blocks': Z",
             "'obj_store_mem_used': A",
@@ -566,6 +572,9 @@ def test_dataset__repr__(ray_start_regular_shared, restore_data_context):
         "      num_tasks_failed: Z,\n"
         "      block_generation_time: N,\n"
         "      task_submission_backpressure_time: N,\n"
+        "      num_alive_actors: Z,\n"
+        "      num_restarting_actors: Z,\n"
+        "      num_pending_actors: Z,\n"
         "      obj_store_mem_internal_inqueue_blocks: Z,\n"
         "      obj_store_mem_internal_outqueue_blocks: Z,\n"
         "      obj_store_mem_freed: N,\n"
@@ -681,6 +690,9 @@ def check_stats():
         "      num_tasks_failed: Z,\n"
         "      block_generation_time: N,\n"
         "      task_submission_backpressure_time: N,\n"
+        "      num_alive_actors: Z,\n"
+        "      num_restarting_actors: Z,\n"
+        "      num_pending_actors: Z,\n"
         "      obj_store_mem_internal_inqueue_blocks: Z,\n"
         "      obj_store_mem_internal_outqueue_blocks: Z,\n"
         "      obj_store_mem_freed: N,\n"
@@ -751,6 +763,9 @@ def check_stats():
         "            num_tasks_failed: Z,\n"
         "            block_generation_time: N,\n"
         "            task_submission_backpressure_time: N,\n"
+        "            num_alive_actors: Z,\n"
+        "            num_restarting_actors: Z,\n"
+        "            num_pending_actors: Z,\n"
         "            obj_store_mem_internal_inqueue_blocks: Z,\n"
         "            obj_store_mem_internal_outqueue_blocks: Z,\n"
         "            obj_store_mem_freed: N,\n"
@@ -1370,6 +1385,104 @@ def time_to_seconds(time_str):
         assert isclose(percent, time_s / total_time * 100, rel_tol=0.01)
 
 
+def test_per_node_metrics_basic(ray_start_regular_shared, restore_data_context):
+    """Basic test to ensure per-node metrics are populated."""
+    ctx = DataContext.get_current()
+    ctx.enable_per_node_metrics = True
+
+    def _sum_net_metrics(per_node_metrics: Dict[str, NodeMetrics]) -> Dict[str, float]:
+        sum_metrics = defaultdict(float)
+        for metrics in per_node_metrics.values():
+            for metric, value in metrics.items():
+                sum_metrics[metric] += value
+        return sum_metrics
+
+    with patch("ray.data._internal.stats.StatsManager._stats_actor") as mock_get_actor:
+        mock_actor_handle = MagicMock()
+        mock_get_actor.return_value = mock_actor_handle
+
+        ds = ray.data.range(20).map_batches(lambda batch: batch).materialize()
+        metrics = ds._plan.stats().extra_metrics
+
+        calls = mock_actor_handle.update_execution_metrics.remote.call_args_list
+        assert len(calls) > 0
+
+        last_args, _ = calls[-1]
+        per_node_metrics = last_args[-1]
+
+        assert isinstance(per_node_metrics, dict)
+        assert len(per_node_metrics) >= 1
+
+        for nm in per_node_metrics.values():
+            for f in fields(NodeMetrics):
+                assert f.name in nm
+
+        # basic checks to make sure metrics are populated
+        assert any(nm["num_tasks_finished"] > 0 for nm in per_node_metrics.values())
+        assert any(
+            nm["bytes_outputs_of_finished_tasks"] > 0
+            for nm in per_node_metrics.values()
+        )
+        assert any(
+            nm["blocks_outputs_of_finished_tasks"] > 0
+            for nm in per_node_metrics.values()
+        )
+
+        net_metrics = _sum_net_metrics(per_node_metrics)
+        assert net_metrics["num_tasks_finished"] == metrics["num_tasks_finished"]
+        assert (
+            net_metrics["bytes_outputs_of_finished_tasks"]
+            == metrics["bytes_outputs_of_finished_tasks"]
+        )
+
+
+@pytest.mark.parametrize("enable_metrics", [True, False])
+def test_per_node_metrics_toggle(
+    ray_start_regular_shared, restore_data_context, enable_metrics
+):
+    ctx = DataContext.get_current()
+    ctx.enable_per_node_metrics = enable_metrics
+
+    with patch("ray.data._internal.stats.StatsManager._stats_actor") as mock_get_actor:
+        mock_actor_handle = MagicMock()
+        mock_get_actor.return_value = mock_actor_handle
+
+        ray.data.range(10000).map(lambda x: x).materialize()
+
+        calls = mock_actor_handle.update_execution_metrics.remote.call_args_list
+        assert len(calls) > 0
+
+        last_args, _ = calls[-1]
+        per_node_metrics = last_args[-1]
+
+        if enable_metrics:
+            assert per_node_metrics is not None
+        else:
+            assert per_node_metrics is None
+
+
+def test_task_duration_stats():
+    """Test that OpTaskDurationStats correctly tracks running statistics using Welford's algorithm."""
+    stats = TaskDurationStats()
+
+    # Test initial state
+    assert stats.count() == 0
+    assert stats.mean() == 0.0
+    assert stats.stddev() == 0.0
+
+    # Add some task durations and verify stats
+    durations = [2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0]
+    for d in durations:
+        stats.add_duration(d)
+
+    # Compare with numpy's implementations
+    assert stats.count() == len(durations)
+    assert pytest.approx(stats.mean()) == np.mean(durations)
+    assert pytest.approx(stats.stddev()) == np.std(
+        durations, ddof=1
+    )  # ddof=1 for sample standard deviation
+
+
 # NOTE: All tests above share a Ray cluster, while the tests below do not. These
 # tests should only be carefully reordered to retain this invariant!
 
@@ -1709,104 +1822,6 @@ def update_stats_manager(i):
     wait_for_condition(lambda: not StatsManager._update_thread.is_alive())
 
 
-def test_per_node_metrics_basic(ray_start_regular_shared, restore_data_context):
-    """Basic test to ensure per-node metrics are populated."""
-    ctx = DataContext.get_current()
-    ctx.enable_per_node_metrics = True
-
-    def _sum_net_metrics(per_node_metrics: Dict[str, NodeMetrics]) -> Dict[str, float]:
-        sum_metrics = defaultdict(float)
-        for metrics in per_node_metrics.values():
-            for metric, value in metrics.items():
-                sum_metrics[metric] += value
-        return sum_metrics
-
-    with patch("ray.data._internal.stats.StatsManager._stats_actor") as mock_get_actor:
-        mock_actor_handle = MagicMock()
-        mock_get_actor.return_value = mock_actor_handle
-
-        ds = ray.data.range(20).map_batches(lambda batch: batch).materialize()
-        metrics = ds._plan.stats().extra_metrics
-
-        calls = mock_actor_handle.update_execution_metrics.remote.call_args_list
-        assert len(calls) > 0
-
-        last_args, _ = calls[-1]
-        per_node_metrics = last_args[-1]
-
-        assert isinstance(per_node_metrics, dict)
-        assert len(per_node_metrics) >= 1
-
-        for nm in per_node_metrics.values():
-            for f in fields(NodeMetrics):
-                assert f.name in nm
-
-        # basic checks to make sure metrics are populated
-        assert any(nm["num_tasks_finished"] > 0 for nm in per_node_metrics.values())
-        assert any(
-            nm["bytes_outputs_of_finished_tasks"] > 0
-            for nm in per_node_metrics.values()
-        )
-        assert any(
-            nm["blocks_outputs_of_finished_tasks"] > 0
-            for nm in per_node_metrics.values()
-        )
-
-        net_metrics = _sum_net_metrics(per_node_metrics)
-        assert net_metrics["num_tasks_finished"] == metrics["num_tasks_finished"]
-        assert (
-            net_metrics["bytes_outputs_of_finished_tasks"]
-            == metrics["bytes_outputs_of_finished_tasks"]
-        )
-
-
-@pytest.mark.parametrize("enable_metrics", [True, False])
-def test_per_node_metrics_toggle(
-    ray_start_regular_shared, restore_data_context, enable_metrics
-):
-    ctx = DataContext.get_current()
-    ctx.enable_per_node_metrics = enable_metrics
-
-    with patch("ray.data._internal.stats.StatsManager._stats_actor") as mock_get_actor:
-        mock_actor_handle = MagicMock()
-        mock_get_actor.return_value = mock_actor_handle
-
-        ray.data.range(10000).map(lambda x: x).materialize()
-
-        calls = mock_actor_handle.update_execution_metrics.remote.call_args_list
-        assert len(calls) > 0
-
-        last_args, _ = calls[-1]
-        per_node_metrics = last_args[-1]
-
-        if enable_metrics:
-            assert per_node_metrics is not None
-        else:
-            assert per_node_metrics is None
-
-
-def test_task_duration_stats():
-    """Test that OpTaskDurationStats correctly tracks running statistics using Welford's algorithm."""
-    stats = TaskDurationStats()
-
-    # Test initial state
-    assert stats.count() == 0
-    assert stats.mean() == 0.0
-    assert stats.stddev() == 0.0
-
-    # Add some task durations and verify stats
-    durations = [2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0]
-    for d in durations:
-        stats.add_duration(d)
-
-    # Compare with numpy's implementations
-    assert stats.count() == len(durations)
-    assert pytest.approx(stats.mean()) == np.mean(durations)
-    assert pytest.approx(stats.stddev()) == np.std(
-        durations, ddof=1
-    )  # ddof=1 for sample standard deviation
-
-
 if __name__ == "__main__":
     import sys