[Data] Warn when configured memory is below 1.25x max USS

python/ray/data/_internal/execution/callbacks/insert_issue_detectors.py

@@ -21,3 +21,13 @@ def before_execution_starts(self, executor: "StreamingExecutor"):
     def on_execution_step(self, executor: "StreamingExecutor"):
         # Invoke all issue detectors
         executor._issue_detector_manager.invoke_detectors()
+
+    def after_execution_succeeds(self, executor: "StreamingExecutor"):
+        # Force one final issue detection pass.
+        executor._issue_detector_manager.invoke_detectors(force=True)
+
+    def after_execution_fails(
+        self, executor: "StreamingExecutor", error: Exception
+    ) -> None:
+        # Force one final issue detection pass.
+        executor._issue_detector_manager.invoke_detectors(force=True)

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

@@ -802,6 +802,10 @@ def has_completed(self) -> bool:
             and not self.has_next()
         )
 
+    def is_shut_down(self) -> bool:
+        """Return whether shutdown has started for this operator."""
+        return self._shutdown
+
     def get_stats(self) -> StatsDict:
         """Return recorded execution stats for use with DatasetStats."""
         raise NotImplementedError

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

@@ -202,7 +202,7 @@ def detect(self) -> List[Issue]:
         hanging_op_tasks: HangingOpTasks = defaultdict(dict)
 
         for operator in self._operators:
-            if operator.has_execution_finished():
+            if operator.is_shut_down() or operator.has_execution_finished():
                 continue
 
             op_metrics = operator.metrics

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

@@ -71,6 +71,8 @@ def detect(self) -> List[Issue]:
         for op in self._operators:
             if not isinstance(op, HashShuffleOperator):
                 continue
+            if op.is_shut_down():
+                continue
 
             # Skip if operator doesn't have aggregator pool yet
             if op._aggregator_pool is None:

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

@@ -1,6 +1,7 @@
+import math
 import textwrap
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Dict, List
+from typing import TYPE_CHECKING, Dict, List, Optional, Set
 
 from ray.data._internal.execution.operators.map_operator import (
     MapOperator,
@@ -23,16 +24,23 @@
 Operator '{op_name}' uses {memory_per_task} of memory per task on average, but Ray
 only requests {initial_memory_request} per task at the start of the pipeline.
 
-To avoid out-of-memory errors, consider setting `memory={memory_per_task}` in the
-appropriate function or method call. (This might be unnecessary if the number of
-concurrent tasks is low.)
+To avoid out-of-memory errors, consider setting `memory={recommended_memory_bytes}`
+({recommended_memory}) in the appropriate function or method call. (This might be
+unnecessary if the number of concurrent tasks is low.)
 
 To change the frequency of this warning, set
 `DataContext.get_current().issue_detectors_config.high_memory_detector_config.detection_time_interval_s`,
 or disable the warning by setting value to -1. (current value:
 {detection_time_interval_s})
 """  # noqa: E501
 
+HIGH_MEMORY_FINAL_WARNING = """
+Operator '{op_name}' used up to {max_memory} of memory per worker.
+The configured logical memory was {memory_configuration}. To avoid out-of-memory errors, set
+`memory={recommended_memory_bytes}` ({recommended_memory}) in the appropriate
+function or method call.
+"""  # noqa: E501
+
 
 @dataclass
 class HighMemoryIssueDetectorConfig:
@@ -49,13 +57,14 @@ def __init__(
         self._dataset_id = dataset_id
         self._detector_cfg = config
         self._operators = operators
+        self._completion_checked_operators: Set[MapOperator] = set()
 
-        self._initial_memory_requests: Dict[MapOperator, int] = {}
+        self._initial_memory_requests: Dict[MapOperator, Optional[int]] = {}
         for op in operators:
             if isinstance(op, MapOperator):
-                self._initial_memory_requests[op] = (
-                    op._get_dynamic_ray_remote_args().get("memory") or 0
-                )
+                self._initial_memory_requests[
+                    op
+                ] = op._get_dynamic_ray_remote_args().get("memory")
 
     @classmethod
     def from_executor(cls, executor: "StreamingExecutor") -> "HighMemoryIssueDetector":
@@ -81,22 +90,36 @@ def detect(self) -> List[Issue]:
             if not isinstance(op, MapOperator):
                 continue
 
+            if op.is_shut_down() or op.has_completed():
+                issue = self._detect_issue_from_final_metrics(
+                    op, self._initial_memory_requests[op]
+                )
+                if issue is not None:
+                    issues.append(issue)
+                continue
+
             if op.metrics.average_max_uss_per_task is None:
                 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]
+                op.metrics.average_max_uss_per_task
+                > (self._initial_memory_requests[op] or 0)
                 and op.metrics.average_max_uss_per_task >= safe_memory_per_task
             ):
+                recommended_memory = _get_recommended_memory(
+                    op.metrics.average_max_uss_per_task
+                )
                 message = HIGH_MEMORY_PERIODIC_WARNING.format(
                     op_name=op.name,
                     memory_per_task=memory_string(op.metrics.average_max_uss_per_task),
                     initial_memory_request=memory_string(
-                        self._initial_memory_requests[op]
+                        self._initial_memory_requests[op] or 0
                     ),
+                    recommended_memory=memory_string(recommended_memory),
+                    recommended_memory_bytes=recommended_memory,
                     detection_time_interval_s=self.detection_time_interval_s(),
                 )
                 issues.append(
@@ -110,10 +133,46 @@ def detect(self) -> List[Issue]:
 
         return issues
 
+    def _detect_issue_from_final_metrics(
+        self, op: MapOperator, memory_request: Optional[int]
+    ) -> Optional[Issue]:
+        if op in self._completion_checked_operators:
+            return None
+        self._completion_checked_operators.add(op)
+        if memory_request is None:
+            return None
+
+        max_uss_bytes = op.metrics.max_uss_bytes.max
+        if max_uss_bytes is None:
+            return None
+
+        recommended_memory = _get_recommended_memory(max_uss_bytes)
+        if recommended_memory <= memory_request:
+            return None
+
+        message = HIGH_MEMORY_FINAL_WARNING.format(
+            op_name=op.name,
+            max_memory=memory_string(max_uss_bytes),
+            memory_configuration=memory_string(memory_request),
+            recommended_memory=memory_string(recommended_memory),
+            recommended_memory_bytes=recommended_memory,
+        )
+        return Issue(
+            dataset_name=self._dataset_id,
+            operator_id=op.id,
+            issue_type=IssueType.HIGH_MEMORY,
+            message=_format_message(message),
+        )
+
     def detection_time_interval_s(self) -> float:
         return self._detector_cfg.detection_time_interval_s
 
 
+def _get_recommended_memory(memory_usage: float) -> int:
+    # Add 25% headroom to the observed memory usage.
+    return math.ceil(5 * memory_usage / 4)
+
+
 def _format_message(message: str) -> str:
     # Apply some formatting to make the message look nicer when printed.
     formatted_paragraphs = []

python/ray/data/_internal/issue_detection/issue_detector_manager.py

@@ -43,15 +43,16 @@ def __init__(self, executor: "StreamingExecutor"):
         # consumer thread that checks the set of detected issues on shutdown (in the usage callback).
         self._detected_issues_lock = threading.Lock()
 
-    def invoke_detectors(self) -> None:
+    def invoke_detectors(self, force: bool = False) -> None:
         curr_time = time.perf_counter()
         issues = []
         for detector in self._issue_detectors:
             if detector.detection_time_interval_s() == -1:
                 continue
 
             if (
-                curr_time - self._last_detection_times[detector]
+                force
+                or curr_time - self._last_detection_times[detector]
                 > detector.detection_time_interval_s()
             ):
                 issues.extend(detector.detect())

python/ray/data/tests/test_issue_detection.py

@@ -14,6 +14,7 @@
     TaskExecDriverStats,
 )
 from ray.data._internal.execution.interfaces.ref_bundle import BlockEntry
+from ray.data._internal.execution.operators.hash_shuffle import HashShuffleOperator
 from ray.data._internal.execution.operators.input_data_buffer import (
     InputDataBuffer,
 )
@@ -26,6 +27,10 @@
     HangingExecutionIssueDetector,
     HangingExecutionIssueDetectorConfig,
 )
+from ray.data._internal.issue_detection.detectors.hash_shuffle_detector import (
+    HashShuffleAggregatorIssueDetector,
+    HashShuffleAggregatorIssueDetectorConfig,
+)
 from ray.data._internal.issue_detection.detectors.high_memory_detector import (
     HighMemoryIssueDetector,
 )
@@ -201,6 +206,21 @@ def test_hanging_detector_detects_issues(
         assert "has been running or stuck in scheduling for" in issues[0].message
         assert "longer than the average task duration" in issues[0].message
 
+        with patch.object(type(op), "is_shut_down", return_value=True):
+            assert detector.detect() == []
+
+
+def test_hash_shuffle_detector_skips_shutdown_operator():
+    hash_shuffle_operator = MagicMock(spec=HashShuffleOperator)
+    hash_shuffle_operator.is_shut_down.return_value = True
+    hash_shuffle_detector = HashShuffleAggregatorIssueDetector(
+        dataset_id="id",
+        operators=[hash_shuffle_operator],
+        config=HashShuffleAggregatorIssueDetectorConfig(),
+    )
+
+    assert hash_shuffle_detector.detect() == []
+
 
 @pytest.mark.parametrize(
     "configured_memory, actual_memory, should_return_issue",
@@ -226,6 +246,7 @@ def test_high_memory_detection(
         data_context=ctx,
         ray_remote_args={"memory": configured_memory},
     )
+    map_operator.has_completed = MagicMock(return_value=False)
     map_operator._metrics = MagicMock(average_max_uss_per_task=actual_memory)
     topology = {input_data_buffer: MagicMock(), map_operator: MagicMock()}
 
@@ -239,6 +260,68 @@ def test_high_memory_detection(
     issues = detector.detect()
 
     assert should_return_issue == bool(issues)
+    if should_return_issue:
+        normalized_message = " ".join(issues[0].message.split())
+        assert "`memory=10737418240` (10.0GiB)" in normalized_message
+
+
+@pytest.mark.parametrize(
+    "configured_memory, max_memory, expected_memory_configuration, expected_memory, operator_termination",
+    [
+        (10 * GiB, 8 * GiB, None, None, "completed"),
+        (
+            10 * GiB - 1,
+            8 * GiB,
+            "The configured logical memory was 10.0GiB.",
+            10 * GiB,
+            "completed",
+        ),
+        (None, 8 * GiB, None, None, "completed"),
+        (0, 1, "The configured logical memory was 0.0B.", 2, "completed"),
+        (1, None, None, None, "completed"),
+        (1, 1, "The configured logical memory was 1.0B.", 2, "shutdown"),
+    ],
+)
+def test_high_memory_detection_on_operator_termination(
+    configured_memory,
+    max_memory,
+    expected_memory_configuration,
+    expected_memory,
+    operator_termination,
+    restore_data_context,
+):
+    ctx = DataContext.get_current()
+    input_data_buffer = InputDataBuffer(ctx, input_data=[])
+    map_operator = MapOperator.create(
+        map_transformer=MagicMock(),
+        input_op=input_data_buffer,
+        data_context=ctx,
+        ray_remote_args={"memory": configured_memory},
+    )
+    if max_memory is not None:
+        map_operator.metrics.max_uss_bytes.add_sample(max_memory // 2)
+        map_operator.metrics.max_uss_bytes.add_sample(max_memory)
+    if operator_termination == "completed":
+        map_operator.has_completed = MagicMock(return_value=True)
+    else:
+        map_operator.is_shut_down = MagicMock(return_value=True)
+
+    detector = HighMemoryIssueDetector(
+        dataset_id="id",
+        operators=[input_data_buffer, map_operator],
+        config=ctx.issue_detectors_config.high_memory_detector_config,
+    )
+
+    issues = detector.detect()
+
+    assert (expected_memory_configuration is not None) == bool(issues)
+    if expected_memory_configuration is not None:
+        normalized_message = " ".join(issues[0].message.split())
+        assert map_operator.name in normalized_message
+        assert expected_memory_configuration in normalized_message
+        assert f"`memory={expected_memory}`" in normalized_message
+    # Completion checks are one-shot to avoid duplicate warnings.
+    assert detector.detect() == []
 
 
 if __name__ == "__main__":

python/ray/data/tests/test_issue_detection_manager.py

@@ -7,6 +7,9 @@
 
 import ray
 from ray._private import ray_constants
+from ray.data._internal.execution.callbacks.insert_issue_detectors import (
+    IssueDetectionExecutionCallback,
+)
 from ray.data._internal.execution.operators.input_data_buffer import (
     InputDataBuffer,
 )
@@ -113,5 +116,57 @@ def test_report_issues():
     assert detector.get_detected_issues() == expected_issues
 
 
+def test_force_invoke_detectors():
+    ctx = DataContext.get_current()
+    executor = StreamingExecutor(ctx)
+    executor._topology = {}
+    detector = IssueDetectorManager(executor)
+    issue_detector = MagicMock()
+    issue_detector.detection_time_interval_s.return_value = 30
+    issue_detector.detect.return_value = []
+    detector._issue_detectors = [issue_detector]
+    detector._last_detection_times = {
+        issue_detector: float("inf"),
+    }
+
+    detector.invoke_detectors()
+    issue_detector.detect.assert_not_called()
+
+    detector.invoke_detectors(force=True)
+    issue_detector.detect.assert_called_once_with()
+
+
+def test_force_invoke_skips_disabled_detectors():
+    ctx = DataContext.get_current()
+    executor = StreamingExecutor(ctx)
+    executor._topology = {}
+    detector = IssueDetectorManager(executor)
+    issue_detector = MagicMock()
+    issue_detector.detection_time_interval_s.return_value = -1
+    detector._issue_detectors = [issue_detector]
+
+    detector.invoke_detectors(force=True)
+
+    issue_detector.detect.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "callback_name, callback_args",
+    [
+        ("after_execution_succeeds", ()),
+        ("after_execution_fails", (RuntimeError(),)),
+    ],
+)
+def test_issue_detection_callback_forces_final_detection(callback_name, callback_args):
+    callback = IssueDetectionExecutionCallback()
+    executor = MagicMock()
+
+    getattr(callback, callback_name)(executor, *callback_args)
+
+    executor._issue_detector_manager.invoke_detectors.assert_called_once_with(
+        force=True
+    )
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", __file__]))