chore: import upstream snapshot with attribution

python/ray/train/v2/_internal/callbacks/preemption_callback.py

@@ -0,0 +1,93 @@
+import logging
+import os
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+import ray
+from ray.actor import ActorHandle
+from ray.train.v2._internal.constants import (
+    DEFAULT_PREEMPTION_POLL_INTERVAL_S,
+    PREEMPTION_POLL_INTERVAL_S_ENV_VAR,
+)
+from ray.train.v2._internal.execution.callback import WorkerGroupCallback
+from ray.train.v2._internal.execution.preemption import PreemptionWatcher
+
+if TYPE_CHECKING:
+    from ray.train.v2._internal.execution.worker_group import (
+        WorkerGroup,
+        WorkerGroupContext,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+class PreemptionCallback(WorkerGroupCallback):
+    """Manages a :class:`PreemptionWatcher` across worker-group lifecycles.
+
+    Spawns a fresh watcher in :meth:`after_worker_group_start` and stops it on
+    every teardown path (shutdown and abort). Each worker group gets its own
+    watcher and failure-domain map, so elastic resizes and restarts never
+    leak stale state.
+    """
+
+    def __init__(self) -> None:
+        self._poll_interval_s: float = float(
+            os.getenv(
+                PREEMPTION_POLL_INTERVAL_S_ENV_VAR,
+                str(DEFAULT_PREEMPTION_POLL_INTERVAL_S),
+            )
+        )
+        self._watcher: Optional[ActorHandle] = None
+
+    def after_worker_group_start(self, worker_group: "WorkerGroup") -> None:
+        # Tear down any watcher from a previous worker group first. Worker-group
+        # startup can fail after this hook without running the shutdown hook, so
+        # this also prevents leaking an orphaned watcher across a reschedule.
+        self._stop_watcher()
+
+        # These handles are captured once per worker-group start. With the
+        # standard backend, any worker replacement goes through a full worker
+        # group restart (this hook runs again with fresh handles), so they
+        # never go stale.
+        # TODO(lehui): refresh worker handles on in-place replica replacement
+        # when adding preemption support for replica groups (TorchFT).
+        node_to_ranks: Dict[str, List[int]] = {}
+        worker_actors_by_rank: Dict[int, ActorHandle] = {}
+        for w in worker_group.get_workers():
+            rank = w.distributed_context.world_rank
+            node_to_ranks.setdefault(w.metadata.node_id, []).append(rank)
+            worker_actors_by_rank[rank] = w.actor
+
+        watcher_cls = ray.remote(num_cpus=0, max_restarts=-1)(PreemptionWatcher)
+        self._watcher = watcher_cls.remote(
+            node_to_ranks=node_to_ranks,
+            poll_interval_s=self._poll_interval_s,
+            worker_actors_by_rank=worker_actors_by_rank,
+        )
+
+        logger.debug(
+            "PreemptionCallback: started watcher for %d node(s).",
+            len(node_to_ranks),
+        )
+
+    def before_worker_group_shutdown(self, worker_group: "WorkerGroup") -> None:
+        self._stop_watcher()
+
+    def after_worker_group_abort(
+        self, worker_group_context: "WorkerGroupContext"
+    ) -> None:
+        # abort() doesn't run the shutdown hook, so tear the watcher down here
+        # too — otherwise it keeps polling GCS until the cluster reaps it.
+        self._stop_watcher()
+
+    def _stop_watcher(self) -> None:
+        if self._watcher is None:
+            return
+        watcher = self._watcher
+        self._watcher = None
+        # Force-kill (non-blocking) rather than a synchronous graceful stop, so
+        # we never block the controller's event loop. The watcher's daemon poll
+        # thread dies with the actor process and holds no external resources.
+        try:
+            ray.kill(watcher)
+        except Exception:
+            logger.warning("Failed to kill PreemptionWatcher actor.", exc_info=True)