[core][train][data] Keep strong references to fire-and-forget asyncio tasks (ray-project#63291)

## Why are these changes needed?

Python's asyncio documentation explicitly warns that the event loop only
holds **weak** references to tasks created via `asyncio.create_task` /
`loop.create_task`. A task without any other strong reference can be
garbage collected at **any** time — even before it finishes:

> Important: Save a reference to the result of this function, to avoid a
task disappearing mid-execution. The event loop only keeps weak
references to tasks. A task that isn't referenced elsewhere may be
garbage collected at any time, even before it's done.
>
> —
https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task

Three call sites in the repository were violating this contract (caught
by ruff's `RUF006` rule). All three create background tasks whose return
value is discarded, so the resulting `Task` object is only referenced by
the event loop's weak set and may be GC'd at any moment:

- **`python/ray/_private/async_utils.py`** — `enable_monitor_loop_lag()`
The event-loop lag monitor task could be silently collected, stopping
lag monitoring without warning.

-
**`python/ray/train/v2/_internal/execution/checkpoint/checkpoint_manager.py`**
— `CheckpointManager._notify()`
The notify task wakes up coroutines waiting on `self._condition`. If
it's collected before `notify_all()` runs, listeners may wait forever.

- **`python/ray/data/_internal/planner/plan_udf_map_op.py`** —
`_generate_transform_fn_for_async_map().._execute_transform()`
The `_reorder` background task is responsible for forwarding completed
results from `completed_tasks_queue` into the output queue in
deterministic order. Losing this task to GC would silently break the
entire async map operation.

## What was changed

Each of the three sites now retains a strong reference to the created
task using the pattern recommended by the asyncio docs:

| File | Pattern |
|---|---|
| `async_utils.py` | Module-level `_BACKGROUND_TASKS: Set[asyncio.Task]`
+ `task.add_done_callback(_BACKGROUND_TASKS.discard)` |
| `checkpoint_manager.py` | Per-instance `self._background_tasks: set` +
`add_done_callback(self._background_tasks.discard)` |
| `plan_udf_map_op.py` | Local variable `reorder_task =
asyncio.create_task(_reorder())`, awaited in the `finally` block of
`_execute_transform` (also propagates unexpected exceptions) |

No public API change. No behavior change in the happy path — only
correctness under GC pressure.

## Related issues

N/A (silences the `RUF006` lint warnings for these three files).

## Checks

- [x] I've signed off every commit (DCO).
- [x] `ruff check --select RUF006` passes on the three modified files.
- [x] `python -m py_compile` passes for all three files.
- [ ] I've made sure the tests are passing.

---------

Signed-off-by: forwardxu <forwardxu@apache.org>
Signed-off-by: phattruong <23120318@student.hcmus.edu.vn>

Author: XuQianJin-Stars

python/ray/_private/async_utils.py

@@ -21,7 +21,13 @@
 
 import asyncio
 import asyncio.events
-from typing import Callable, Optional
+from typing import Callable, Optional, Set
+
+# Strong references to background tasks to prevent them from being garbage
+# collected mid-execution. The asyncio event loop only keeps weak references
+# to tasks, so a task with no other strong references can be collected at
+# any time. See https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task
+_BACKGROUND_TASKS: Set[asyncio.Task] = set()
 
 
 def enable_monitor_loop_lag(
@@ -49,4 +55,7 @@ async def monitor():
             lag = loop.time() - t0 - interval_s  # Should be close to zero.
             callback(lag)
 
-    loop.create_task(monitor(), name="async_utils.monitor_loop_lag")
+    task = loop.create_task(monitor(), name="async_utils.monitor_loop_lag")
+    # Keep a strong reference so the task isn't garbage collected mid-execution.
+    _BACKGROUND_TASKS.add(task)
+    task.add_done_callback(_BACKGROUND_TASKS.discard)

python/ray/data/_internal/planner/plan_udf_map_op.py

@@ -876,8 +876,12 @@ async def _reorder() -> None:
             finally:
                 output_queue.put(sentinel)
 
-        # NOTE: Reordering is an async process
-        asyncio.create_task(_reorder())
+        # NOTE: Reordering is an async process. Keep a strong reference to
+        # the created task: ``loop.create_task`` only registers a weak
+        # reference with the event loop, so without a strong reference the
+        # task could be garbage collected mid-execution and the reordering
+        # would silently stop.
+        reorder_task = loop.create_task(_reorder())
 
         cur_task_map: Dict[asyncio.Task, int] = dict()
         consumed = False
@@ -926,6 +930,11 @@ async def _reorder() -> None:
         finally:
             assert len(cur_task_map) == 0, f"{cur_task_map}"
             await completed_tasks_queue.put((sentinel, None))
+            # Wait for the reorder task to finish draining ``completed_tasks_queue``
+            # and pushing remaining results to the output queue. This both keeps a
+            # strong reference to the task alive until completion (preventing GC)
+            # and surfaces any unexpected exception raised inside ``_reorder``.
+            await reorder_task
 
     def _transform(batch_iter: Iterable[T], task_context: TaskContext) -> Iterable[U]:
         output_queue = queue.Queue(maxsize=max_concurrency)

python/ray/train/v2/_internal/execution/checkpoint/checkpoint_manager.py

@@ -124,6 +124,11 @@ def __init__(
 
         self._condition = asyncio.Condition()
 
+        # Strong references to background tasks created via
+        # ``asyncio.create_task`` to prevent them from being garbage
+        # collected mid-execution. The event loop only keeps weak refs.
+        self._background_tasks: set = set()
+
         self._collective_warn_interval_s = env_float(
             COLLECTIVE_WARN_INTERVAL_S_ENV_VAR,
             DEFAULT_COLLECTIVE_WARN_INTERVAL_S,
@@ -248,7 +253,12 @@ async def async_notify():
             async with self._condition:
                 self._condition.notify_all()
 
-        asyncio.create_task(async_notify())
+        # Keep a strong reference to the task so it isn't garbage
+        # collected before completing, which would silently drop
+        # the notification and could leave listeners waiting forever.
+        task = asyncio.create_task(async_notify())
+        self._background_tasks.add(task)
+        task.add_done_callback(self._background_tasks.discard)
 
     def _save_state_and_delete_old_checkpoints(self):
         """Delete the old checkpoints."""