Revert "[data] continue grabbing task state until response is not Non…e (#60592)"

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

@@ -2,7 +2,7 @@
 import time
 from collections import defaultdict
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Dict, List, Optional, Set
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union
 
 import ray
 from ray.data._internal.issue_detection.issue_detector import (
@@ -35,7 +35,6 @@
 class HangingExecutionState:
     operator_id: str
     task_idx: int
-    task_id: ray.TaskID
     task_state: Optional[TaskState]
     bytes_output: int
     start_time_hanging: float
@@ -113,7 +112,7 @@ def _create_issues(
                     attempt_number = state.task_state.attempt_number
 
                 message = (
-                    f"A task (task_id={state.task_id}) of operator {op_name} (pid={pid}, node_id={node_id}, attempt={attempt_number}) has been running for {duration:.2f}s, which is longer"
+                    f"A task of operator {op_name} (pid={pid}, node_id={node_id}, attempt={attempt_number}) has been running for {duration:.2f}s, which is longer"
                     f" than the average task duration of this operator ({avg_duration:.2f}s)."
                     f" If this message persists, please check the stack trace of the "
                     "task for potential hanging issues."
@@ -154,11 +153,29 @@ def detect(self) -> List[Issue]:
                         prev_state_value is None
                         or bytes_output != prev_state_value.bytes_output
                     ):
+                        task_state = None
+                        try:
+                            task_state: Union[
+                                TaskState, List[TaskState]
+                            ] = ray.util.state.get_task(
+                                task_info.task_id.hex(),
+                                timeout=1.0,
+                                _explain=True,
+                            )
+                            if isinstance(task_state, list):
+                                # get the latest task
+                                task_state = max(
+                                    task_state, key=lambda ts: ts.attempt_number
+                                )
+                        except Exception as e:
+                            logger.debug(
+                                f"Failed to grab task state with task_index={task_idx}, task_id={task_info.task_id}: {e}"
+                            )
+                            pass
                         self._state_map[operator.id][task_idx] = HangingExecutionState(
                             operator_id=operator.id,
                             task_idx=task_idx,
-                            task_id=task_info.task_id,
-                            task_state=None,
+                            task_state=task_state,
                             bytes_output=bytes_output,
                             start_time_hanging=time.perf_counter(),
                         )
@@ -177,10 +194,6 @@ def detect(self) -> List[Issue]:
             for task_idx, state_value in op_state_values.items():
                 curr_time = time.perf_counter() - state_value.start_time_hanging
                 if op_task_stats.count() >= self._op_task_stats_min_count:
-                    if state_value.task_state is None:
-                        state_value.task_state = get_latest_state_for_task(
-                            state_value.task_id
-                        )
                     mean = op_task_stats.mean()
                     stddev = op_task_stats.stddev()
                     threshold = mean + self._op_task_stats_std_factor_threshold * stddev
@@ -198,20 +211,3 @@ def detect(self) -> List[Issue]:
 
     def detection_time_interval_s(self) -> float:
         return self._detector_cfg.detection_time_interval_s
-
-
-def get_latest_state_for_task(task_id: ray.TaskID) -> TaskState | None:
-    try:
-        task_state: TaskState | List[TaskState] | None = ray.util.state.get_task(
-            task_id.hex(),
-            timeout=1.0,
-            _explain=True,
-        )
-        if isinstance(task_state, list):
-            # get the latest task
-            task_state = max(task_state, key=lambda ts: ts.attempt_number)
-        return task_state
-    except Exception as e:
-        logger.debug(f"Failed to grab task state with task_id={task_id}: {e}")
-        pass
-    return None

python/ray/data/tests/test_issue_detection.py

@@ -123,7 +123,7 @@ def f1(x):
         _ = ray.data.range(1).map(f1).materialize()
 
         log_output = log_capture.getvalue()
-        warn_msg = r"A task \(task_id=.+\) .+ \(pid=.+, node_id=.+, attempt=.+\) has been running for [\d\.]+s"
+        warn_msg = r"A task of operator .+ \(pid=.+, node_id=.+, attempt=.+\) has been running for [\d\.]+s"
         assert re.search(warn_msg, log_output) is None, log_output
 
         # # test hanging does log hanging warning