[Data] Export task USS distribution metrics

Signed-off-by: viiccwen <vicwen@apache.org>

Author: viiccwen

python/ray/dashboard/modules/metrics/dashboards/data_dashboard_panels.py

@@ -431,6 +431,21 @@
     stack=False,
 )
 
+MAX_USS_PER_TASK_PANEL = Panel(
+    id=91,
+    title="Max Task USS Memory per Operator",
+    description="Maximum unique set size (USS) memory usage in bytes among completed tasks for each operator.",
+    unit="bytes",
+    targets=[
+        Target(
+            expr='max(ray_data_max_uss_bytes_max{{{global_filters}, operator=~"$Operator"}}) by (dataset, operator)',
+            legend="Max Task USS: {{dataset}}, {{operator}}",
+        )
+    ],
+    fill=0,
+    stack=False,
+)
+
 TASK_THROUGHPUT_BY_NODE_PANEL = Panel(
     id=46,
     title="Task Throughput (by Node)",
@@ -1479,6 +1494,7 @@
             TASK_COMPLETION_TIME_WITHOUT_BACKPRESSURE_PANEL,
             TASK_OUTPUT_BACKPRESSURE_TIME_PANEL,
             TASK_SUBMISSION_BACKPRESSURE_PANEL,
+            MAX_USS_PER_TASK_PANEL,
             TASK_THROUGHPUT_BY_NODE_PANEL,
             TASKS_WITH_OUTPUT_PANEL,
             SUBMITTED_TASKS_PANEL,

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

@@ -54,6 +54,7 @@ class MetricsType(Enum):
     Gauge = 1
     Histogram = 2
     Unsupported = 3
+    Distribution = 4
 
 
 @dataclass(frozen=True)
@@ -899,21 +900,12 @@ def op_task_duration_stats(self) -> DistributionTracker:
     @metric_property(
         description="Distribution of max USS bytes across tasks.",
         metrics_group=MetricsGroup.TASKS,
-        metrics_type=MetricsType.Unsupported,
+        metrics_type=MetricsType.Distribution,
+        metrics_args={"statistics": ("mean", "max")},
     )
     def max_uss_bytes(self) -> DistributionTracker:
         return self._max_uss_bytes
 
-    @metric_property(
-        description="Average USS usage of tasks.",
-        metrics_group=MetricsGroup.TASKS,
-    )
-    def average_max_uss_per_task(self) -> Optional[float]:
-        """Average max USS usage of tasks."""
-        if self.max_uss_bytes.num_samples == 0:
-            return None
-        return self.max_uss_bytes.mean
-
     @metric_property(
         description="Indicates if the operator is hanging.",
         metrics_group=MetricsGroup.MISC,

python/ray/data/_internal/issue_detection/detectors/high_memory_detector.py

@@ -81,19 +81,20 @@ def detect(self) -> List[Issue]:
             if not isinstance(op, MapOperator):
                 continue
 
-            if op.metrics.average_max_uss_per_task is None:
+            max_uss_bytes = op.metrics.max_uss_bytes
+            if max_uss_bytes.num_samples == 0:
                 continue
 
             remote_args = op._get_dynamic_ray_remote_args()
             safe_memory_per_task = get_safe_default_logical_memory(remote_args)
 
             if (
-                op.metrics.average_max_uss_per_task > self._initial_memory_requests[op]
-                and op.metrics.average_max_uss_per_task >= safe_memory_per_task
+                max_uss_bytes.mean > self._initial_memory_requests[op]
+                and max_uss_bytes.mean >= safe_memory_per_task
             ):
                 message = HIGH_MEMORY_PERIODIC_WARNING.format(
                     op_name=op.name,
-                    memory_per_task=memory_string(op.metrics.average_max_uss_per_task),
+                    memory_per_task=memory_string(max_uss_bytes.mean),
                     initial_memory_request=memory_string(
                         self._initial_memory_requests[op]
                     ),

python/ray/data/_internal/stats.py

@@ -681,7 +681,7 @@ def __init__(self, max_stats=1000):
 
     def _create_prometheus_metrics_for_execution_metrics(
         self, metrics_group: MetricsGroup, tag_keys: Tuple[str, ...]
-    ) -> Dict[str, Metric]:
+    ) -> Dict[str, Union[Metric, Dict[str, Gauge]]]:
         metrics = {}
         for metric in OpRuntimeMetrics.get_metrics():
             if not metric.metrics_group == metrics_group:
@@ -709,6 +709,15 @@ def _create_prometheus_metrics_for_execution_metrics(
                     description=metric_description,
                     tag_keys=tag_keys,
                 )
+            elif metric.metrics_type == MetricsType.Distribution:
+                metrics[metric.name] = {
+                    statistic: Gauge(
+                        f"{metric_name}_{statistic}",
+                        description=f"{metric_description} ({statistic})",
+                        tag_keys=tag_keys,
+                    )
+                    for statistic in metric.metrics_args["statistics"]
+                }
         return metrics
 
     def _create_prometheus_metrics_for_per_node_metrics(self) -> Dict[str, Gauge]:
@@ -731,14 +740,14 @@ def gen_dataset_id(self) -> str:
     def update_execution_metrics(
         self,
         dataset_tag: str,
-        op_metrics: List[Dict[str, int | float]],
+        op_metrics: List[Dict[str, Any]],
         operator_tags: List[str],
         state: Dict[str, Any],
         per_node_metrics: Optional[Dict[str, Dict[str, int | float]]] = None,
     ):
         def _record(
-            prom_metric: Metric,
-            value: Union[int, float, List[int]],
+            prom_metric: Union[Metric, Dict[str, Gauge]],
+            value: Any,
             tags: Dict[str, str] = None,
         ):
             if isinstance(prom_metric, Gauge):
@@ -748,6 +757,13 @@ def _record(
             elif isinstance(prom_metric, Histogram):
                 if isinstance(value, RuntimeMetricsHistogram):
                     value.export_to(prom_metric, tags)
+            elif isinstance(prom_metric, dict) and isinstance(value, dict):
+                if value.get("num_samples", 0) == 0:
+                    return
+                for statistic, gauge in prom_metric.items():
+                    statistic_value = value.get(statistic)
+                    if statistic_value is not None:
+                        gauge.set(statistic_value, tags)
 
         for stats, operator_tag in zip(op_metrics, operator_tags):
             tags = self._create_tags(dataset_tag, operator_tag)

python/ray/data/tests/test_issue_detection.py

@@ -226,7 +226,9 @@ def test_high_memory_detection(
         data_context=ctx,
         ray_remote_args={"memory": configured_memory},
     )
-    map_operator._metrics = MagicMock(average_max_uss_per_task=actual_memory)
+    map_operator._metrics = MagicMock()
+    map_operator._metrics.max_uss_bytes.num_samples = 1
+    map_operator._metrics.max_uss_bytes.mean = actual_memory
     topology = {input_data_buffer: MagicMock(), map_operator: MagicMock()}
 
     operators = list(topology.keys())

python/ray/data/tests/test_op_runtime_metrics.py

@@ -17,11 +17,13 @@
 from ray.data.context import DataContext
 
 
-def test_average_max_uss_per_task():
+def test_max_uss_bytes_distribution():
     op = MagicMock()
     op.data_context.enable_get_object_locations_for_metrics = False
     metrics = OpRuntimeMetrics(op)
-    assert metrics.average_max_uss_per_task is None
+    assert metrics.max_uss_bytes.num_samples == 0
+    assert metrics.max_uss_bytes.mean == 0
+    assert metrics.max_uss_bytes.max is None
 
     input_bundle = RefBundle([], owns_blocks=False, schema=None)
 
@@ -33,7 +35,9 @@ def test_average_max_uss_per_task():
         TaskExecWorkerStats(task_wall_time_s=1.0, max_uss_bytes=100),
         TaskExecDriverStats(task_output_backpressure_s=0),
     )
-    assert metrics.average_max_uss_per_task == 100
+    assert metrics.max_uss_bytes.num_samples == 1
+    assert metrics.max_uss_bytes.mean == 100
+    assert metrics.max_uss_bytes.max == 100
 
     # Submit and finish second task with USS of 300 bytes.
     metrics.on_task_submitted(1, input_bundle)
@@ -43,7 +47,9 @@ def test_average_max_uss_per_task():
         TaskExecWorkerStats(task_wall_time_s=1.0, max_uss_bytes=300),
         TaskExecDriverStats(task_output_backpressure_s=0),
     )
-    assert metrics.average_max_uss_per_task == 200  # (100 + 300) / 2
+    assert metrics.max_uss_bytes.num_samples == 2
+    assert metrics.max_uss_bytes.mean == 200  # (100 + 300) / 2
+    assert metrics.max_uss_bytes.max == 300
 
 
 def test_task_completion_time_histogram():

python/ray/data/tests/test_stats.py

@@ -56,6 +56,7 @@
 from ray.data.context import DataContext
 from ray.data.tests.util import column_udf
 from ray.tests.conftest import *  # noqa
+from ray.util.metrics import Gauge
 
 
 @dataclass(frozen=True)
@@ -353,7 +354,6 @@ def gen_expected_metrics(
             "'average_rows_outputs_per_task': N",
             "'op_task_duration_stats': {'num_samples': N, 'mean': N, 'variance': N, 'min': N, 'max': N, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P}",
             "'max_uss_bytes': H",
-            "'average_max_uss_per_task': H",
             "'num_inputs_received': N",
             "'num_row_inputs_received': N",
             "'bytes_inputs_received': N",
@@ -443,7 +443,6 @@ def gen_expected_metrics(
             "'average_rows_outputs_per_task': None",
             "'op_task_duration_stats': {'num_samples': Z, 'mean': Z, 'variance': Z, 'min': None, 'max': None, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P}",
             "'max_uss_bytes': H",
-            "'average_max_uss_per_task': H",
             "'num_inputs_received': N",
             "'num_row_inputs_received': N",
             "'bytes_inputs_received': N",
@@ -635,11 +634,6 @@ def canonicalize(
     # Replace tabs with spaces.
     canonicalized_stats = re.sub("\t", "    ", canonicalized_stats)
 
-    canonicalized_stats = re.sub(
-        r"(average_max_uss_per_task:|'average_max_uss_per_task':) (?:N|Z|None)\b",
-        r"\g<1> H",
-        canonicalized_stats,
-    )
     # Percentile values in DistributionTracker dicts can be None (when datasketches
     # is not installed) or a number (canonicalized to N). Normalize to P.
     canonicalized_stats = re.sub(
@@ -932,7 +926,6 @@ def test_dataset__repr__(ray_start_regular_shared, restore_data_context):
         "      average_rows_outputs_per_task: N,\n"
         "      op_task_duration_stats: {'num_samples': N, 'mean': N, 'variance': N, 'min': N, 'max': N, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P},\n"
         "      max_uss_bytes: H,\n"
-        "      average_max_uss_per_task: H,\n"
         "      num_inputs_received: N,\n"
         "      num_row_inputs_received: N,\n"
         "      bytes_inputs_received: N,\n"
@@ -1097,7 +1090,6 @@ def check_stats():
         "      average_rows_outputs_per_task: N,\n"
         "      op_task_duration_stats: {'num_samples': N, 'mean': N, 'variance': N, 'min': N, 'max': N, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P},\n"
         "      max_uss_bytes: H,\n"
-        "      average_max_uss_per_task: H,\n"
         "      num_inputs_received: N,\n"
         "      num_row_inputs_received: N,\n"
         "      bytes_inputs_received: N,\n"
@@ -1215,7 +1207,6 @@ def check_stats():
         "            average_rows_outputs_per_task: N,\n"
         "            op_task_duration_stats: {'num_samples': N, 'mean': N, 'variance': N, 'min': N, 'max': N, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P},\n"
         "            max_uss_bytes: H,\n"
-        "            average_max_uss_per_task: H,\n"
         "            num_inputs_received: N,\n"
         "            num_row_inputs_received: N,\n"
         "            bytes_inputs_received: N,\n"
@@ -2017,6 +2008,46 @@ def test_stats_actor_iter_metrics():
     assert update_fn.call_args_list[-1].args[2] is None
 
 
+def test_stats_actor_exports_distribution_metrics():
+    actor = _StatsActor.__ray_metadata__.modified_class()
+
+    metrics = actor.execution_metrics_tasks["max_uss_bytes"]
+    assert set(metrics) == {"mean", "max"}
+    for statistic, metric in metrics.items():
+        assert isinstance(metric, Gauge)
+        assert metric.info["name"] == f"data_max_uss_bytes_{statistic}"
+        assert metric.info["tag_keys"] == ("dataset", "operator")
+
+    actor.update_dataset = MagicMock()
+
+    with (
+        patch.object(metrics["mean"], "set") as set_mean,
+        patch.object(metrics["max"], "set") as set_max,
+    ):
+        actor.update_execution_metrics(
+            "dataset_1",
+            [{"max_uss_bytes": {"num_samples": 0, "mean": 0, "max": None}}],
+            ["MapBatches_1"],
+            {},
+        )
+        set_mean.assert_not_called()
+        set_max.assert_not_called()
+
+        actor.update_execution_metrics(
+            "dataset_1",
+            [{"max_uss_bytes": {"num_samples": 2, "mean": 200, "max": 300}}],
+            ["MapBatches_1"],
+            {},
+        )
+
+    set_mean.assert_called_once_with(
+        200, {"dataset": "dataset_1", "operator": "MapBatches_1"}
+    )
+    set_max.assert_called_once_with(
+        300, {"dataset": "dataset_1", "operator": "MapBatches_1"}
+    )
+
+
 @pytest.mark.parametrize(
     "split_index_arg, expected_split_label",
     [