wehub-resource-sync
94 lines
3.6 KiB
Python
94 lines
3.6 KiB
Python
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)
|