Support master-mode ports in RayWorkerManager

miles/utils/workers/ray_worker_manager.py

@@ -1,6 +1,6 @@
 import asyncio
 from collections.abc import Callable
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any, Generic, TypeVar
 
 import ray
@@ -77,9 +77,8 @@ def initial(cls, spec: BaseWorkerSpec, manager_ref: RayWorkerManager) -> "_PoolM
                     actors=[
                         # TODO support Serve mode
                         _CommandActorManager(
-                            cell_index=cell_index,
                             worker_in_cell_index=worker_in_cell_index,
-                            manager_ref=manager_ref,
+                            manager=manager_ref,
                             spec=spec,
                             actor_handle=None,
                             generation=1,
@@ -97,14 +96,18 @@ class _CellManager:
     cell_index: int
     actors: list["_BaseActorManager"]
 
+    def __post_init__(self):
+        for actor in self.actors:
+            actor.parent = self
+
 
 SpecT = TypeVar("SpecT", bound=BaseWorkerSpec)
 
 
 @dataclass(kw_only=True)
 class _BaseActorManager(Generic[SpecT]):
-    manager_ref: RayWorkerManager
-    cell_index: int
+    manager: RayWorkerManager
+    parent: _CellManager = field(init=False)
     worker_in_cell_index: int
     spec: SpecT
     actor_handle: ray.actor.ActorHandle | None
@@ -124,14 +127,18 @@ async def post_setup(self) -> None:
     def name(self) -> str:
         return compute_worker_name(
             pool_id=self.spec.name,
-            cell_index=self.cell_index,
+            cell_index=self.parent.cell_index,
             worker_in_cell_index=self.worker_in_cell_index,
         )
 
     @property
     def primary_addr(self) -> HostAndPort:
         return self.self_addrs["primary"]
 
+    @property
+    def master_mode_addrs(self) -> NamedHostAndPorts:
+        return {info.name: self.self_addrs[info.name] for info in self.spec.port_infos if info.mode == "master"}
+
 
 @dataclass
 class _CommandActorManager(_BaseActorManager[CommandWorkerSpec]):
@@ -153,12 +160,14 @@ async def alloc_ports(self) -> None:
 
         node_ip = await self.actor_handle._get_node_ip.remote()
         for port_info in self.spec.port_infos:
+            if self.worker_in_cell_index != 0 and port_info.mode == "master":
+                continue
             if port_info.allow_dynamic:
-                port = self.manager_ref.port_allocator.alloc(
+                port = self.manager.port_allocator.alloc(
                     self.actor_handle, node_ip=node_ip, consecutive=port_info.num_consecutive
                 )
             else:
-                port = port_info.static_port + (self.cell_index if port_info.offset_by_cell else 0)
+                port = port_info.static_port + (self.parent.cell_index if port_info.offset_by_cell else 0)
                 await self._assert_static_port_is_free(port=port, port_name=port_info.name, node_ip=node_ip)
             self.self_addrs[port_info.name] = HostAndPort(host=_wrap_ipv6(node_ip), port=port)
 
@@ -171,10 +180,13 @@ async def _assert_static_port_is_free(self, *, port: int, port_name: str, node_i
 
     async def post_setup(self) -> None:
         ctx = LaunchCommandContext(
-            cell_index=self.cell_index,
+            cell_index=self.parent.cell_index,
             worker_in_cell_index=self.worker_in_cell_index,
-            self_addrs=self.self_addrs,
-            pool_addrs=self.manager_ref.get_addrs(),
+            self_addrs={
+                **self.self_addrs,
+                **self.parent.actors[0].master_mode_addrs,
+            },
+            pool_addrs=self.manager.get_addrs(),
             gpu_ids=[],  # TODO
         )
         launch_cmd = self.spec.launch_command(ctx)

tests/fast/utils/workers/real_ray/test_ray_worker_manager.py

@@ -221,3 +221,46 @@ def test_each_spec_only_gets_its_own_env(self, manager_factory, worker_probe_fac
 
         assert router_record["env"] == {"ROUTER_ONLY": "r", "ENGINE_ONLY": None}
         assert engine_record["env"] == {"ROUTER_ONLY": None, "ENGINE_ONLY": "e"}
+
+
+class TestMasterPortsOnRealRay:
+    def test_all_ranks_of_a_cell_are_launched_with_their_cells_master_endpoint(
+        self, manager_factory, worker_probe_factory
+    ):
+        """Every worker of a cell receives the same master endpoint, allocated once by rank 0."""
+        probe = worker_probe_factory()
+        handle = manager_factory(
+            [
+                make_command_spec(
+                    "engine",
+                    num_cells=2,
+                    num_workers_per_cell=2,
+                    launch_command=probe.launch_command,
+                    port_infos=[
+                        PortInfo(name="primary", static_port=8000, allow_dynamic=True),
+                        PortInfo(
+                            name="dist_init", static_port=9000, mode="master", allow_dynamic=True, num_consecutive=4
+                        ),
+                    ],
+                )
+            ]
+        )
+
+        records = probe.wait_for_records(4)
+        masters = {name: record["context"]["self_addrs"]["dist_init"] for name, record in records.items()}
+
+        assert masters["0-0"] == masters["0-1"]
+        assert masters["1-0"] == masters["1-1"]
+        assert masters["0-0"] != masters["1-0"]
+        addrs = ray.get(handle.get_addrs.remote())["engine"]
+        assert [sorted(addr) for addr in addrs] == [
+            ["dist_init", "primary"],
+            ["primary"],
+            ["dist_init", "primary"],
+            ["primary"],
+        ]
+        primary_ports = [addr["primary"].port for addr in addrs]
+        assert len(set(primary_ports)) == 4
+        for master in [masters["0-0"], masters["1-0"]]:
+            reserved = range(master["port"], master["port"] + 4)
+            assert not set(reserved) & set(primary_ports)

tests/fast/utils/workers/test_ray_worker_manager.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 from collections.abc import Callable
 from dataclasses import dataclass, field
 
@@ -572,3 +573,133 @@ async def test_a_disabled_group_is_listed_as_empty(self, fake_ray_cluster: FakeR
 
         assert manager.get_addrs()["session-server"] == []
         assert len(manager.get_addrs()["router"]) == 1
+
+
+class TestMasterPorts:
+    async def test_a_master_port_is_allocated_once_per_cell(self, fake_ray_cluster: FakeRayCluster):
+        """Only worker 0 reserves the cell's master port, so peers cannot each take their own."""
+        spec = _make_spec(
+            "engine",
+            num_cells=2,
+            num_workers_per_cell=3,
+            port_infos=[
+                PortInfo(name="primary", static_port=8000, allow_dynamic=True),
+                PortInfo(name="dist_init", static_port=9000, mode="master", allow_dynamic=True, num_consecutive=5),
+            ],
+        )
+        manager = await _launch([spec])
+
+        assert [call.kwargs["count"] for call in fake_ray_cluster.calls_of("_get_free_port_block")].count(5) == 2
+        addrs = manager.get_addrs()["engine"]
+        assert [sorted(addr) for addr in addrs] == [
+            ["dist_init", "primary"],
+            ["primary"],
+            ["primary"],
+            ["dist_init", "primary"],
+            ["primary"],
+            ["primary"],
+        ]
+
+    async def test_a_static_master_port_is_recorded_only_on_worker_zero(self, fake_ray_cluster: FakeRayCluster):
+        """A pinned master port is not allocated, and still belongs to worker 0 alone."""
+        spec = _make_spec(
+            "engine",
+            num_workers_per_cell=2,
+            port_infos=[
+                PortInfo(name="primary", static_port=8000, allow_dynamic=True),
+                PortInfo(name="dist_init", static_port=9123, mode="master", allow_dynamic=False),
+            ],
+        )
+        manager = await _launch([spec])
+
+        addrs = manager.get_addrs()["engine"]
+        assert addrs[0]["dist_init"].port == 9123
+        assert "dist_init" not in addrs[1]
+        assert len(fake_ray_cluster.calls_of("_get_free_port_block")) == 2
+
+    async def test_every_worker_of_a_cell_launches_with_that_cells_master_addr(self, fake_ray_cluster: FakeRayCluster):
+        """All ranks of a cell must be told the same master endpoint, and never another cell's."""
+        recorder = _LaunchRecorder()
+        spec = _make_spec(
+            "engine",
+            num_cells=2,
+            num_workers_per_cell=2,
+            port_infos=[
+                PortInfo(name="primary", static_port=8000, allow_dynamic=True),
+                PortInfo(name="dist_init", static_port=9000, mode="master", allow_dynamic=True),
+            ],
+            launch_command=recorder.command,
+        )
+        await _launch([spec])
+
+        masters = {
+            (cell_index, worker_in_cell_index): recorder.context_of(
+                cell_index=cell_index, worker_in_cell_index=worker_in_cell_index
+            ).self_addrs["dist_init"]
+            for cell_index in range(2)
+            for worker_in_cell_index in range(2)
+        }
+        assert masters[(0, 0)] == masters[(0, 1)]
+        assert masters[(1, 0)] == masters[(1, 1)]
+        assert masters[(0, 0)] != masters[(1, 0)]
+
+    async def test_a_master_addr_keeps_rank_zeros_host_on_another_node(self, fake_ray_cluster: FakeRayCluster):
+        """A peer on a second node must dial rank 0's host, not its own, or torch.distributed hangs."""
+        fake_ray_cluster.use_node_ips("10.0.0.1", "10.0.0.2")
+        recorder = _LaunchRecorder()
+        spec = _make_spec(
+            "engine",
+            num_workers_per_cell=2,
+            port_infos=[
+                PortInfo(name="primary", static_port=8000, allow_dynamic=True),
+                PortInfo(name="dist_init", static_port=9000, mode="master", allow_dynamic=True),
+            ],
+            launch_command=recorder.command,
+        )
+        await _launch([spec])
+
+        peer = recorder.context_of(cell_index=0, worker_in_cell_index=1).self_addrs
+        assert peer["primary"].host == "10.0.0.2"
+        assert peer["dist_init"].host == "10.0.0.1"
+
+    async def test_a_master_addr_does_not_overwrite_a_workers_own_ports(self, fake_ray_cluster: FakeRayCluster):
+        """Merging the cell's master addr must leave each worker's own addresses intact."""
+        recorder = _LaunchRecorder()
+        spec = _make_spec(
+            "engine",
+            num_workers_per_cell=2,
+            port_infos=[
+                PortInfo(name="primary", static_port=8000, allow_dynamic=True),
+                PortInfo(name="dist_init", static_port=9000, mode="master", allow_dynamic=True),
+            ],
+            launch_command=recorder.command,
+        )
+        manager = await _launch([spec])
+
+        for worker_in_cell_index in range(2):
+            ctx = recorder.context_of(cell_index=0, worker_in_cell_index=worker_in_cell_index)
+            assert ctx.self_addrs["primary"] == manager.get_worker_addr(f"engine-0-{worker_in_cell_index}")
+
+
+class TestConcurrentPhases:
+    async def test_all_cells_of_a_phase_run_concurrently(self, fake_ray_cluster: FakeRayCluster):
+        """Cells are configured concurrently, so a large pool does not start up one cell at a time."""
+        from miles.utils.workers.ray_worker_manager import _CommandActorManager
+
+        original_alloc_ports = _CommandActorManager.alloc_ports
+        entered: list[int] = []
+        release = asyncio.Event()
+
+        async def gated_alloc(self) -> None:
+            entered.append(self.parent.cell_index)
+            if len(entered) == 3:
+                release.set()
+            await asyncio.wait_for(release.wait(), timeout=5)
+            await original_alloc_ports(self)
+
+        with pytest.MonkeyPatch.context() as patched:
+            patched.setattr(_CommandActorManager, "alloc_ports", gated_alloc)
+            manager = await asyncio.wait_for(_launch([_make_spec("engine", num_cells=3)]), timeout=10)
+
+        assert sorted(entered) == [0, 1, 2]
+        assert len({manager.get_worker_addr(f"engine-{index}-0").port for index in range(3)}) == 3