Probe every node of a cell concurrently to avoid serialized initialization

miles/ray/rollout/server_cell.py

@@ -108,10 +108,17 @@ async def start_engines(self, port_allocator: PortAllocator) -> None:
 
         global_ranks = [self.rank_offset + local_index for local_index in range(self.num_nodes)]
 
+        node_ips = [
+            node_ip
+            for node_ip, _ in await asyncio.gather(
+                *[actor._get_current_node_ip_and_free_port.remote() for actor in actor_handles]
+            )
+        ]
+
         addr_and_ports: dict[int, dict[str, Any]] = {}
         dist_init_addr = None
         for local_index, (rank, actor) in enumerate(zip(global_ranks, actor_handles, strict=True)):
-            node_ip, _ = ray.get(actor._get_current_node_ip_and_free_port.remote())
+            node_ip = node_ips[local_index]
             alloc = functools.partial(port_allocator.alloc, engine=actor, node_ip=node_ip)
 
             if local_index == 0:

tests/fast/ray/rollout/conftest.py

@@ -290,9 +290,10 @@ def fake_engine(host: str = "10.0.0.1", port_seed: int = 30000) -> MagicMock:
 
     Mocks ``_get_current_node_ip_and_free_port.remote(start_port, consecutive)``
     with a deterministic ``max(seq, start_port)`` counter so allocator tests
-    can predict and assert on port assignment. It also passes
-    ``isinstance(x, ray.actor.ActorHandle)`` so it can be handed to
-    ``ServerEngine.mark_allocated_uninitialized`` (see ``fake_actor_handle``)."""
+    can predict and assert on port assignment. The argument-less form is the
+    node-ip probe, which the cell awaits, so it returns an awaitable just like a
+    real ``ObjectRef``. It also passes ``isinstance(x, ray.actor.ActorHandle)``
+    so it can be handed to ``mark_allocated_uninitialized``."""
     e = MagicMock()
     e._spec_class = ray.actor.ActorHandle
     e._port_cursor = port_seed
@@ -302,7 +303,10 @@ def _alloc(start_port: int = 15000, consecutive: int = 1):
         e._port_cursor = port + consecutive
         return (host, port)
 
-    e._get_current_node_ip_and_free_port.remote.side_effect = lambda **kw: _alloc(**kw)
+    async def _probe():
+        return _alloc()
+
+    e._get_current_node_ip_and_free_port.remote.side_effect = lambda **kw: _alloc(**kw) if kw else _probe()
     return e
 
 

tests/fast/ray/rollout/test_addr_allocator.py

@@ -240,3 +240,41 @@ def test_sequential_batches_share_cursor_and_avoid_overlap(self, patch_ray_get):
         ports_a = {addrs_a[r]["port"] for r in addrs_a} | {addrs_a[r]["nccl_port"] for r in addrs_a}
         ports_b = {addrs_b[r]["port"] for r in addrs_b} | {addrs_b[r]["nccl_port"] for r in addrs_b}
         assert ports_a.isdisjoint(ports_b), f"port overlap A={ports_a} B={ports_b}"
+
+
+class TestConcurrentNodeProbes:
+    async def test_a_cell_probes_all_of_its_nodes_concurrently(self, patch_ray_get):
+        """Serializing the node probes would make cell startup scale with the node count."""
+        events: list[tuple[str, int]] = []
+        num_nodes = 3
+
+        def _instrumented(index: int):
+            engine = fake_engine(host=f"10.0.0.{index + 1}", port_seed=0)
+            engine.__class__ = ray.actor.ActorHandle
+            engine.init.remote.side_effect = lambda **kwargs: asyncio.sleep(0)
+            alloc = engine._get_current_node_ip_and_free_port.remote.side_effect
+
+            async def _probe():
+                events.append(("enter", index))
+                await asyncio.sleep(0.05)
+                events.append(("exit", index))
+                return alloc(start_port=15000, consecutive=1)
+
+            engine._get_current_node_ip_and_free_port.remote.side_effect = lambda **kw: alloc(**kw) if kw else _probe()
+            return engine
+
+        actors = {rank: _instrumented(rank) for rank in range(num_nodes)}
+        cell = ServerCell(
+            args=make_args(num_gpus_per_node=8, sglang_dp_size=1),
+            num_nodes=num_nodes,
+            worker_type="regular",
+            cell_id="cell-0",
+        )
+        with patch.object(
+            server_cell_module,
+            "launch_sglang_ray_actor",
+            side_effect=lambda *, global_rank, **kw: actors[global_rank],
+        ):
+            await cell.start_engines(PortAllocator())
+
+        assert [kind for kind, _ in events[:num_nodes]] == ["enter"] * num_nodes, events