wehub-resource-sync
118 lines
4.4 KiB
Python
118 lines
4.4 KiB
Python
import logging
|
|
import os
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from typing import Any, Dict
|
|
|
|
import ray
|
|
from ray.train.torch.config import (
|
|
TorchConfig,
|
|
_TorchBackend,
|
|
)
|
|
from ray.train.v2._internal.constants import TORCHFT_LIGHTHOUSE_ADDR_ENV_VAR
|
|
from ray.train.v2._internal.execution.worker_group import ReplicaGroup, WorkerGroup
|
|
from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class TorchftConfig(TorchConfig):
|
|
"""Configuration for torchft-based fault tolerant training.
|
|
|
|
See https://github.com/meta-pytorch/torchft for more info.
|
|
|
|
Args:
|
|
lighthouse_kwargs: Keyword arguments to pass to the torchft.Lighthouse constructor.
|
|
**kwargs: Additional keyword arguments to pass to the TorchConfig constructor.
|
|
"""
|
|
|
|
def __init__(self, lighthouse_kwargs: Dict[str, Any], **kwargs):
|
|
self.lighthouse_kwargs = lighthouse_kwargs
|
|
super().__init__(**kwargs)
|
|
|
|
@property
|
|
def backend_cls(self):
|
|
return _TorchftBackend
|
|
|
|
|
|
@ray.remote
|
|
class LighthouseServerActor:
|
|
"""Actor that runs the torchft.Lighthouse server.
|
|
|
|
ray.remote(LighthouseServer) does not work because it is a PyO3 type.
|
|
"""
|
|
|
|
def __init__(self, lighthouse_kwargs: Dict[str, Any]):
|
|
from torchft.coordination import LighthouseServer
|
|
|
|
self.lighthouse = LighthouseServer(**lighthouse_kwargs)
|
|
|
|
def address(self) -> str:
|
|
return self.lighthouse.address()
|
|
|
|
|
|
class _TorchftBackend(_TorchBackend):
|
|
"""Backend for torchft-based fault-tolerant training with replica groups.
|
|
|
|
Creates a separate process group per replica group by calling
|
|
the parent _TorchBackend.on_start() once per replica group.
|
|
"""
|
|
|
|
has_replica_groups: bool = True
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.lighthouse_actor = None
|
|
|
|
def _maybe_create_lighthouse_actor(self, backend_config: TorchftConfig) -> str:
|
|
"""Create lighthouse actor if it doesn't exist and return its address."""
|
|
if self.lighthouse_actor is not None:
|
|
# Intentionally read address from actor in case it was restarted.
|
|
return ray.get(self.lighthouse_actor.address.remote())
|
|
|
|
# Let the OS pick a free port by default
|
|
if "bind" in backend_config.lighthouse_kwargs:
|
|
lighthouse_kwargs = backend_config.lighthouse_kwargs
|
|
else:
|
|
lighthouse_kwargs = {"bind": "[::]:0"} | backend_config.lighthouse_kwargs
|
|
|
|
# Store reference so the actor lives as long as the backend/controller.
|
|
self.lighthouse_actor = LighthouseServerActor.options(
|
|
# Schedule lightweight lighthouse actor on head node
|
|
scheduling_strategy=NodeAffinitySchedulingStrategy(
|
|
node_id=ray.get_runtime_context().get_node_id(),
|
|
soft=False,
|
|
)
|
|
).remote(lighthouse_kwargs=lighthouse_kwargs)
|
|
lighthouse_address = ray.get(self.lighthouse_actor.address.remote())
|
|
logger.info(f"Created torchft lighthouse at {lighthouse_address}")
|
|
return lighthouse_address
|
|
|
|
def on_start(self, worker_group, backend_config: TorchftConfig):
|
|
lighthouse_address = self._maybe_create_lighthouse_actor(backend_config)
|
|
|
|
# Push the lighthouse address to all workers in this group.
|
|
# Necessary because workers were already started before on_start runs.
|
|
def _set_lighthouse_address(addr: str):
|
|
os.environ[TORCHFT_LIGHTHOUSE_ADDR_ENV_VAR] = addr
|
|
|
|
worker_group.execute(_set_lighthouse_address, lighthouse_address)
|
|
|
|
# Bind super() eagerly — the zero-arg form relies on a __class__ cell
|
|
# that doesn't transfer correctly into ThreadPoolExecutor submissions.
|
|
parent_on_start = super().on_start
|
|
|
|
if isinstance(worker_group, ReplicaGroup):
|
|
# Single replica group replacement — just start this one.
|
|
parent_on_start(worker_group, backend_config)
|
|
else:
|
|
# Full worker group startup — start all replica groups in parallel.
|
|
assert isinstance(worker_group, WorkerGroup)
|
|
replica_groups = worker_group.get_replica_groups()
|
|
with ThreadPoolExecutor() as executor:
|
|
futures = [
|
|
executor.submit(parent_on_start, rg, backend_config)
|
|
for rg in replica_groups
|
|
]
|
|
for f in futures:
|
|
f.result()
|