chore: import upstream snapshot with attribution

python/ray/train/v2/torch/torchft_config.py

@@ -0,0 +1,117 @@
+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()