chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,243 @@
|
||||
import logging
|
||||
import threading
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional, Set
|
||||
|
||||
import ray
|
||||
from ray.actor import ActorHandle
|
||||
from ray.train.v2._internal.constants import DEFAULT_PREEMPTION_POLL_INTERVAL_S
|
||||
from ray.util.tpu import get_tpu_slice_name_from_node
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.v2._internal.worker import RayTrainWorker
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PreemptionInfo:
|
||||
"""Information about an imminent preemption event.
|
||||
|
||||
Attributes:
|
||||
deadline_ms: Earliest preemption deadline (UNIX time in milliseconds)
|
||||
across all preempted nodes. ``None`` if no deadline was reported.
|
||||
preempted_node_to_ranks: Map of preempted ``node_id`` to the worker ``world_rank``s affected when that node
|
||||
is preempted.
|
||||
"""
|
||||
|
||||
deadline_ms: Optional[int]
|
||||
preempted_node_to_ranks: Dict[str, List[int]]
|
||||
|
||||
@property
|
||||
def preempted_node_ids(self) -> List[str]:
|
||||
"""Preempted node IDs, sorted lexicographically."""
|
||||
return sorted(self.preempted_node_to_ranks)
|
||||
|
||||
@property
|
||||
def preempted_ranks(self) -> List[int]:
|
||||
"""All affected ranks across the preempted nodes, sorted ascending."""
|
||||
return sorted(
|
||||
{r for ranks in self.preempted_node_to_ranks.values() for r in ranks}
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class PreemptionContext:
|
||||
"""Thread-shared preemption signal for one worker actor."""
|
||||
|
||||
_preemption_info: Optional[PreemptionInfo] = field(default=None, init=False)
|
||||
_lock: threading.Lock = field(default_factory=threading.Lock, init=False)
|
||||
|
||||
def set(self, info: PreemptionInfo) -> None:
|
||||
with self._lock:
|
||||
self._preemption_info = info
|
||||
|
||||
def get(self) -> Optional[PreemptionInfo]:
|
||||
"""Return the current preemption signal, or ``None`` if none received."""
|
||||
with self._lock:
|
||||
return self._preemption_info
|
||||
|
||||
|
||||
def _get_draining_nodes() -> Dict[str, int]:
|
||||
"""Ray Core's draining nodes as ``{node_id_hex: deadline_ms}`` (0 = no deadline)."""
|
||||
return ray._private.state.state.get_draining_nodes()
|
||||
|
||||
|
||||
class PreemptionWatcher:
|
||||
"""Polls Ray Core for node drains and logs detected preemption events.
|
||||
|
||||
One watcher per worker group, spawned as a ``num_cpus=0`` actor by
|
||||
``PreemptionCallback``. The poll loop runs in a background thread. The
|
||||
failure-domain map is built once on construction and is immutable for the
|
||||
watcher's lifetime.
|
||||
|
||||
The failure-domain map records which of our ranks are affected if a node is
|
||||
preempted: for a GPU node, the ranks on that node; for a TPU node, every
|
||||
rank in the node's slice, since a TPU slice is preempted atomically.
|
||||
|
||||
Args:
|
||||
node_to_ranks: Map ``node_id_hex -> [ranks on that node]``. Used both
|
||||
as the set of nodes we care about (drains elsewhere are ignored)
|
||||
and as the seed for failure-domain expansion.
|
||||
poll_interval_s: Seconds between drain-state polls.
|
||||
worker_actors_by_rank: Map ``world_rank -> worker actor handle``. On a
|
||||
detected preemption, ``mark_preempt`` is called on every worker.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
node_to_ranks: Dict[str, List[int]],
|
||||
poll_interval_s: float = DEFAULT_PREEMPTION_POLL_INTERVAL_S,
|
||||
worker_actors_by_rank: Optional[
|
||||
Dict[int, ActorHandle["RayTrainWorker"]]
|
||||
] = None,
|
||||
):
|
||||
self._node_to_ranks: Dict[str, List[int]] = {
|
||||
nid: sorted(ranks) for nid, ranks in node_to_ranks.items()
|
||||
}
|
||||
self._poll_interval_s = poll_interval_s
|
||||
self._worker_actors_by_rank: Dict[int, ActorHandle["RayTrainWorker"]] = (
|
||||
worker_actors_by_rank or {}
|
||||
)
|
||||
self._failure_domain_map: Dict[str, List[int]] = self._build_failure_domain_map(
|
||||
self._node_to_ranks
|
||||
)
|
||||
|
||||
self._stop_event = threading.Event()
|
||||
self._last_drained: Dict[str, int] = {}
|
||||
self._latest_info: Optional[PreemptionInfo] = None
|
||||
|
||||
self._monitor_thread = threading.Thread(
|
||||
target=self._watch_loop,
|
||||
name="PreemptionWatcher",
|
||||
daemon=True,
|
||||
)
|
||||
self._monitor_thread.start()
|
||||
|
||||
@staticmethod
|
||||
def _build_failure_domain_map(
|
||||
node_to_ranks: Dict[str, List[int]],
|
||||
) -> Dict[str, List[int]]:
|
||||
"""Map each node we host to all ranks in its failure domain.
|
||||
|
||||
- Non-TPU (e.g. GPU) clusters: the failure domain is the node itself,
|
||||
so a drain on a node flags only the ranks this job runs there.
|
||||
- TPU multislice: every host in a slice is reclaimed atomically, so a
|
||||
drain on any host is fate-shared with the rest.
|
||||
"""
|
||||
per_node = {nid: sorted(set(ranks)) for nid, ranks in node_to_ranks.items()}
|
||||
|
||||
try:
|
||||
all_nodes = ray.nodes()
|
||||
|
||||
# Slice label for each node we host (None for non-TPU nodes).
|
||||
node_to_slice: Dict[str, Optional[str]] = {
|
||||
node["NodeID"]: get_tpu_slice_name_from_node(node)
|
||||
for node in all_nodes
|
||||
if node["NodeID"] in node_to_ranks
|
||||
}
|
||||
|
||||
# Union our ranks per slice.
|
||||
slice_to_ranks: Dict[str, Set[int]] = {}
|
||||
for node_id, ranks in node_to_ranks.items():
|
||||
slice_label = node_to_slice.get(node_id)
|
||||
if slice_label:
|
||||
slice_to_ranks.setdefault(slice_label, set()).update(ranks)
|
||||
|
||||
# Non-TPU cluster (or none of our nodes are on a slice): per-node.
|
||||
if not slice_to_ranks:
|
||||
return per_node
|
||||
|
||||
result: Dict[str, List[int]] = {}
|
||||
for node_id, ranks in node_to_ranks.items():
|
||||
slice_label = node_to_slice.get(node_id)
|
||||
if slice_label:
|
||||
result[node_id] = sorted(slice_to_ranks[slice_label])
|
||||
else:
|
||||
result[node_id] = sorted(set(ranks))
|
||||
return result
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"Could not build failure-domain map; falling back to per-node "
|
||||
"domains (no TPU-slice expansion).",
|
||||
exc_info=True,
|
||||
)
|
||||
return per_node
|
||||
|
||||
def get_latest_preemption_info(self) -> Optional[PreemptionInfo]:
|
||||
"""Most recent :class:`PreemptionInfo` observed, or ``None``."""
|
||||
return self._latest_info
|
||||
|
||||
def _watch_loop(self) -> None:
|
||||
logger.debug(
|
||||
"PreemptionWatcher polling %d node(s) every %.1fs.",
|
||||
len(self._node_to_ranks),
|
||||
self._poll_interval_s,
|
||||
)
|
||||
while not self._stop_event.is_set():
|
||||
self._poll_once()
|
||||
self._stop_event.wait(timeout=self._poll_interval_s)
|
||||
logger.debug("PreemptionWatcher stopped.")
|
||||
|
||||
def _poll_once(self) -> None:
|
||||
"""Poll the drain source once and dispatch on change.
|
||||
|
||||
Per-poll exceptions are caught and logged so a transient GCS hiccup
|
||||
doesn't kill the watcher loop.
|
||||
"""
|
||||
try:
|
||||
drained = _get_draining_nodes() or {}
|
||||
# Keep only drains on this job's own nodes (others are ignored).
|
||||
# That's complete for TPU — an SPMD job fully occupies its slice, so
|
||||
# every fate-shared host is one of our nodes and a drain on any slice
|
||||
# host appears here. For GPU, a drain on a host we don't run on is
|
||||
# correctly irrelevant.
|
||||
relevant = {
|
||||
n: d for n, d in drained.items() if n in self._failure_domain_map
|
||||
}
|
||||
if relevant != self._last_drained:
|
||||
self._on_drain_change(relevant)
|
||||
self._last_drained = relevant
|
||||
except Exception:
|
||||
# TODO(lehui): consider exponential backoff when the drain API keeps
|
||||
# failing, instead of retrying at the fixed poll interval.
|
||||
logger.warning("PreemptionWatcher poll failed", exc_info=True)
|
||||
|
||||
def _on_drain_change(self, drained: Dict[str, int]) -> None:
|
||||
"""Handle a change in the drained-node set.
|
||||
|
||||
``drained`` has already been narrowed to this job's nodes by the
|
||||
caller (``_poll_once``).
|
||||
"""
|
||||
if not drained:
|
||||
return
|
||||
|
||||
affected_node_ids = sorted(drained.keys())
|
||||
preempted_node_to_ranks = {
|
||||
node_id: self._failure_domain_map[node_id] for node_id in affected_node_ids
|
||||
}
|
||||
|
||||
# Earliest deadline across the preempted nodes; None if none reported one
|
||||
# (Ray Core uses 0 for "no deadline", which is falsy and filtered out).
|
||||
reported_deadlines = [drained[n] for n in affected_node_ids if drained[n]]
|
||||
deadline_ms = min(reported_deadlines) if reported_deadlines else None
|
||||
|
||||
info = PreemptionInfo(
|
||||
deadline_ms=deadline_ms,
|
||||
preempted_node_to_ranks=preempted_node_to_ranks,
|
||||
)
|
||||
self._latest_info = info
|
||||
|
||||
logger.warning(
|
||||
"PreemptionWatcher: preemption detected — "
|
||||
"preempted_node_ids=%s, preempted_ranks=%s, deadline_ms=%s",
|
||||
info.preempted_node_ids,
|
||||
info.preempted_ranks,
|
||||
deadline_ms,
|
||||
)
|
||||
|
||||
for rank, actor in self._worker_actors_by_rank.items():
|
||||
actor.mark_preempt.remote(info)
|
||||
# TODO(lehui): coalesce preemptions seen within one window into a single
|
||||
# worker-group restart, so a staggered drain (node A at t, node B at
|
||||
# t+60s) doesn't cause back-to-back restarts.
|
||||
Reference in New Issue
Block a user