Make start_engines async and await engine init inside it

miles/ray/rollout/rollout_server.py

@@ -41,8 +41,7 @@ def start_rollout_servers(args, pg) -> dict[str, "RolloutServer"]:
             args.sglang_router_port = router_port
 
         server_groups: list[ServerGroup] = []
-        all_init_handles: list = []
-        new_engine_indices_per_group: list[list[int]] = []
+        start_futures: list = []
         port_allocator = PortAllocator.empty()
 
         for group_cfg in model_cfg.server_groups:
@@ -84,16 +83,13 @@ def start_rollout_servers(args, pg) -> dict[str, "RolloutServer"]:
                 router_port=router_port,
                 update_weights=model_cfg.update_weights,
             )
-            handles, new_engine_indices = group.start_engines(port_allocator)
-            all_init_handles.extend(handles)
+            start_futures.append(async_utils.submit(group.start_engines(port_allocator)))
             server_groups.append(group)
-            new_engine_indices_per_group.append(new_engine_indices)
 
             engine_offset += num_engines
             gpu_offset += group_cfg.num_gpus
 
-        if all_init_handles:
-            ray.get(all_init_handles)
+        new_engine_indices_per_group = async_utils.wait_futures(start_futures)
 
         for group, new_engine_indices in zip(server_groups, new_engine_indices_per_group, strict=True):
             group.mark_alive(engine_indices=new_engine_indices)

miles/ray/rollout/server_group.py

@@ -56,15 +56,15 @@ def engines(self) -> list[ServerEngine]:
         """Node-0 engines only (for multi-node serving)."""
         return [cell.engines[0] for cell in self.cells]
 
-    def start_engines(
+    async def start_engines(
         self, port_allocator: PortAllocator, start_cell_indices: list[int] | None = None
-    ) -> tuple[list, list[int]]:
-        """Create Ray actors, allocate ports, and fire ``engine.init()`` without waiting.
+    ) -> list[int]:
+        """Create Ray actors, allocate ports, and run ``engine.init()`` on every new engine.
 
         Mutates ``port_allocator`` in place to advance past any newly assigned ports.
-        Returns ``(init_handles, new_engine_indices)`` where *init_handles* is a list
-        of Ray ObjectRefs (one per newly created engine) and *new_engine_indices* is
-        the list of indices into the group's flat engine list that were just allocated.
+        Returns the list of indices into the group's flat engine list that were just
+        allocated. Actor creation, port allocation and state marking all happen before
+        the first await point, so concurrent callers cannot double-start a slot.
         """
         assert not ({"host", "port"} & set(self.sglang_overrides)), (
             f"sglang_overrides must not override host/port ({self.sglang_overrides=}): the rollout process derives "
@@ -73,7 +73,7 @@ def start_engines(
 
         if self.args.debug_train_only or self.worker_type == "placeholder":
             self.has_new_engines = False
-            return [], []
+            return []
 
         if self.args.rollout_external:
             raise NotImplementedError(
@@ -117,7 +117,7 @@ def start_engines(
         self.has_new_engines |= curr_num_new_engines > 0
 
         if curr_num_new_engines == 0:
-            return [], []
+            return []
 
         addr_and_ports: dict[int, dict[str, Any]] = {}
         for cell_index in sorted({index // self.nodes_per_engine for index in new_engine_indices}):
@@ -153,7 +153,8 @@ def start_engines(
             )
 
         init_handles = [engine.init.remote(**addr_and_ports[index]) for index, engine in new_engines]
-        return init_handles, new_engine_indices
+        await asyncio.gather(*init_handles)
+        return new_engine_indices
 
     async def register_workers(self, engine_indices: list[int]) -> None:
         if self.args.rollout_external or not (self.router_ip and self.router_port):
@@ -221,8 +222,7 @@ async def recover(self, port_allocator: PortAllocator, filter_cell_indices: list
                 if any(not engine.is_allocated for engine in cell.engines)
             ]
 
-        handles, new_engine_indices = self.start_engines(port_allocator, start_cell_indices=filter_cell_indices)
-        await asyncio.gather(*handles)
+        new_engine_indices = await self.start_engines(port_allocator, start_cell_indices=filter_cell_indices)
 
         all_engines = flatten_cells(self.cells)
         release_handles = []

tests/fast/ray/rollout/real_ray/test_fault_tolerance.py

@@ -38,9 +38,8 @@ def _build_group(
     )
 
 
-def _start(group: ServerGroup) -> None:
-    handles, indices = group.start_engines(PortAllocator.empty())
-    ray.get(handles)
+async def _start(group: ServerGroup) -> None:
+    indices = await group.start_engines(PortAllocator.empty())
     group.mark_alive(indices)
 
 
@@ -67,7 +66,7 @@ async def test_recover_creates_new_actor_after_real_kill(
         and the surviving engine is untouched."""
         pg = placement_group_factory(2)
         group = _build_group(pg_tuple=pg, num_engines=2)
-        _start(group)
+        await _start(group)
 
         original_handles = [e.actor_handle for e in flatten_cells(group.cells)]
         # Real fault: kill engine 0 + mark its slot stopped (production code's
@@ -98,7 +97,7 @@ async def test_recover_default_filter_picks_all_dead_slots(
         only 0 and 2 to be re-created."""
         pg = placement_group_factory(3)
         group = _build_group(pg_tuple=pg, num_engines=3)
-        _start(group)
+        await _start(group)
 
         old = [e.actor_handle for e in flatten_cells(group.cells)]
         for i in (0, 2):
@@ -136,7 +135,7 @@ async def remove_worker(self, **kwargs):
         pg = placement_group_factory(1)
         group = _build_group(pg_tuple=pg, num_engines=1)
         group.router_ip, group.router_port = "10.0.0.9", 9000
-        _start(group)
+        await _start(group)
         ray.kill(flatten_cells(group.cells)[0].actor_handle)
         flatten_cells(group.cells)[0].mark_stopped()
 
@@ -159,7 +158,7 @@ async def test_recover_with_offload_calls_release_then_resume(
         Verify by reading the recovered engine's mock HTTP server log."""
         pg = placement_group_factory(2)
         group = _build_group(pg_tuple=pg, num_engines=2, needs_offload=True, update_weights=True)
-        _start(group)
+        await _start(group)
         old = [e.actor_handle for e in flatten_cells(group.cells)]
 
         ray.kill(old[0])
@@ -221,8 +220,8 @@ async def test_two_groups_recover_in_parallel_completes_without_deadlock(
         pg_b = placement_group_factory(2)
         a = _build_group(pg_tuple=pg_a, num_engines=2)
         b = _build_group(pg_tuple=pg_b, num_engines=2)
-        _start(a)
-        _start(b)
+        await _start(a)
+        await _start(b)
 
         # Kill one engine in each group
         for g in (a, b):
@@ -261,7 +260,7 @@ async def test_simulate_crash_then_health_check_still_returns(
     ):
         pg = placement_group_factory(1)
         group = _build_group(pg_tuple=pg, num_engines=1)
-        _start(group)
+        await _start(group)
         actor = flatten_cells(group.cells)[0].actor_handle
 
         try:

tests/fast/ray/rollout/real_ray/test_rollout_server.py

@@ -35,9 +35,8 @@ def _build_group(
     )
 
 
-def _start_group(group: ServerGroup) -> None:
-    handles, _ = group.start_engines(PortAllocator.empty())
-    ray.get(handles)
+async def _start_group(group: ServerGroup) -> None:
+    await group.start_engines(PortAllocator.empty())
 
 
 def _kill_group(group: ServerGroup) -> None:
@@ -63,8 +62,8 @@ async def test_aggregates_across_groups_via_real_asyncio_gather(
         pg_b = placement_group_factory(3)
         a = _build_group(pg_tuple=pg_a, num_engines=2)
         b = _build_group(pg_tuple=pg_b, num_engines=3, rank_offset=2)
-        _start_group(a)
-        _start_group(b)
+        await _start_group(a)
+        await _start_group(b)
         a.mark_alive([0, 1])
         b.mark_alive([0, 1, 2])
 
@@ -100,8 +99,8 @@ async def test_offload_and_onload_reach_every_engine_of_every_group(
         pg_b = placement_group_factory(3)
         a = _build_group(pg_tuple=pg_a, num_engines=2, needs_offload=True)
         b = _build_group(pg_tuple=pg_b, num_engines=3, rank_offset=2, needs_offload=True)
-        _start_group(a)
-        _start_group(b)
+        await _start_group(a)
+        await _start_group(b)
         a.mark_alive([0, 1])
         b.mark_alive([0, 1, 2])
 
@@ -143,8 +142,8 @@ async def test_a_group_that_does_not_need_offload_is_skipped(
         pg_b = placement_group_factory(2)
         offloading = _build_group(pg_tuple=pg_a, num_engines=2, needs_offload=True)
         resident = _build_group(pg_tuple=pg_b, num_engines=2, rank_offset=2, needs_offload=False)
-        _start_group(offloading)
-        _start_group(resident)
+        await _start_group(offloading)
+        await _start_group(resident)
         offloading.mark_alive([0, 1])
         resident.mark_alive([0, 1])
 
@@ -166,7 +165,7 @@ async def test_a_dead_engine_is_not_addressed(
         """Offload must not block forever on an engine the group already gave up on."""
         pg = placement_group_factory(2)
         group = _build_group(pg_tuple=pg, num_engines=2, needs_offload=True)
-        _start_group(group)
+        await _start_group(group)
         group.mark_alive([0, 1])
         flatten_cells(group.cells)[1].mark_stopped()
 

tests/fast/ray/rollout/real_ray/test_server_group.py

@@ -39,22 +39,22 @@ def _build_group(
 class TestStartEnginesShortCircuits:
     """Branches that bail before hitting the PG / actor creation path."""
 
-    def test_debug_train_only_returns_immediately(self, placement_group_factory):
+    async def test_debug_train_only_returns_immediately(self, placement_group_factory):
         # PG made but unused — start_engines should bail before scheduling.
         pg = placement_group_factory(2)
         group = _build_group(pg_tuple=pg, num_engines=2, debug_train_only=True)
-        handles, indices = group.start_engines(PortAllocator.empty())
-        assert handles == [] and indices == []
+        indices = await group.start_engines(PortAllocator.empty())
+        assert indices == []
         assert group.has_new_engines is False
         for e in flatten_cells(group.cells):
             assert not e.is_allocated
 
-    def test_placeholder_worker_short_circuits(self, placement_group_factory):
+    async def test_placeholder_worker_short_circuits(self, placement_group_factory):
         # PG is unused in this short-circuit path; min size 1 keeps Ray happy.
         pg = placement_group_factory(1)
         group = _build_group(pg_tuple=pg, num_engines=0, worker_type="placeholder")
-        handles, indices = group.start_engines(PortAllocator.empty())
-        assert handles == [] and indices == []
+        indices = await group.start_engines(PortAllocator.empty())
+        assert indices == []
         assert group.has_new_engines is False
 
 
@@ -63,15 +63,13 @@ class TestStartEnginesRealActors:
     real Ray actors (via ``get_calls()`` round-trip) and that ``init`` was
     invoked with the addr/port kwargs from the allocator."""
 
-    def test_creates_real_actors_and_init_runs(self, patched_sglang_engine, placement_group_factory):
+    async def test_creates_real_actors_and_init_runs(self, patched_sglang_engine, placement_group_factory):
         pg = placement_group_factory(2)
         group = _build_group(pg_tuple=pg, num_engines=2)
 
-        handles, indices = group.start_engines(PortAllocator.empty())
+        indices = await group.start_engines(PortAllocator.empty())
         assert sorted(indices) == [0, 1]
         assert group.has_new_engines is True
-        # Wait for init.remote() to actually complete on each actor.
-        ray.get(handles)
 
         for e in flatten_cells(group.cells):
             assert e.is_allocated
@@ -86,13 +84,12 @@ def test_creates_real_actors_and_init_runs(self, patched_sglang_engine, placemen
         for e in flatten_cells(group.cells):
             ray.kill(e.actor_handle)
 
-    def test_start_cell_indices_filters_to_subset(self, patched_sglang_engine, placement_group_factory):
+    async def test_start_cell_indices_filters_to_subset(self, patched_sglang_engine, placement_group_factory):
         pg = placement_group_factory(4)
         group = _build_group(pg_tuple=pg, num_engines=4)
 
-        handles, indices = group.start_engines(PortAllocator.empty(), start_cell_indices=[1, 3])
+        indices = await group.start_engines(PortAllocator.empty(), start_cell_indices=[1, 3])
         assert sorted(indices) == [1, 3]
-        ray.get(handles)
 
         assert not flatten_cells(group.cells)[0].is_allocated
         assert flatten_cells(group.cells)[1].is_allocated
@@ -102,20 +99,19 @@ def test_start_cell_indices_filters_to_subset(self, patched_sglang_engine, place
         for i in (1, 3):
             ray.kill(flatten_cells(group.cells)[i].actor_handle)
 
-    def test_already_allocated_slot_is_skipped(self, patched_sglang_engine, placement_group_factory):
+    async def test_already_allocated_slot_is_skipped(self, patched_sglang_engine, placement_group_factory):
         """A second start_engines() call must NOT replace an already-allocated
         actor — the existing handle is preserved verbatim."""
         pg = placement_group_factory(2)
         group = _build_group(pg_tuple=pg, num_engines=2)
 
         # First call: allocates both slots.
-        handles, _ = group.start_engines(PortAllocator.empty())
-        ray.get(handles)
+        await group.start_engines(PortAllocator.empty())
         first_handles = [e.actor_handle for e in flatten_cells(group.cells)]
 
         # Second call with no start_cell_indices: should skip both.
-        handles2, indices2 = group.start_engines(PortAllocator.empty())
-        assert handles2 == [] and indices2 == []
+        indices2 = await group.start_engines(PortAllocator.empty())
+        assert indices2 == []
         for first, e in zip(first_handles, flatten_cells(group.cells), strict=True):
             assert e.actor_handle is first  # still the same actor
 
@@ -131,11 +127,12 @@ class TestStopEnginesRealKill:
     """``ray.kill`` is the real thing here — we verify the actor is actually
     dead by issuing a follow-up ``.remote()`` and expecting RayActorError."""
 
-    def test_stop_marks_engines_stopped_and_actor_truly_dies(self, patched_sglang_engine, placement_group_factory):
+    async def test_stop_marks_engines_stopped_and_actor_truly_dies(
+        self, patched_sglang_engine, placement_group_factory
+    ):
         pg = placement_group_factory(2)
         group = _build_group(pg_tuple=pg, num_engines=2)
-        handles, _ = group.start_engines(PortAllocator.empty())
-        ray.get(handles)
+        await group.start_engines(PortAllocator.empty())
 
         actors = [e.actor_handle for e in flatten_cells(group.cells)]
         group.stop_engines(cell_indices=[0, 1])
@@ -149,15 +146,14 @@ def test_stop_marks_engines_stopped_and_actor_truly_dies(self, patched_sglang_en
             with pytest.raises((ray.exceptions.RayActorError, ray.exceptions.RayTaskError)):
                 ray.get(actor.get_calls.remote(), timeout=10.0)
 
-    def test_stop_handles_shutdown_failure_gracefully(self, patched_sglang_engine, placement_group_factory):
+    async def test_stop_handles_shutdown_failure_gracefully(self, patched_sglang_engine, placement_group_factory):
         """If ``shutdown`` raises on the actor, ``stop_engines`` must still
         mark the engine stopped (and ray.kill is still called).
 
         We use ``set_fault`` to make shutdown raise on its next invocation."""
         pg = placement_group_factory(2)
         group = _build_group(pg_tuple=pg, num_engines=2)
-        handles, _ = group.start_engines(PortAllocator.empty())
-        ray.get(handles)
+        await group.start_engines(PortAllocator.empty())
 
         # Plant a one-shot shutdown failure on engine 1.
         ray.get(
@@ -176,17 +172,16 @@ class TestStartEnginesRealAllocator:
     """Drive ``start_engines`` with real actors so that the actor → driver port
     round-trip via ``_get_current_node_ip_and_free_port.remote`` actually runs."""
 
-    def test_real_allocator_assigns_distinct_ports_via_remote_calls(
+    async def test_real_allocator_assigns_distinct_ports_via_remote_calls(
         self,
         patched_sglang_engine,
         placement_group_factory,
     ):
         pg = placement_group_factory(2)
         group = _build_group(pg_tuple=pg, num_engines=2)
 
-        handles, indices = group.start_engines(PortAllocator.empty())
+        indices = await group.start_engines(PortAllocator.empty())
         assert sorted(indices) == [0, 1]
-        ray.get(handles)
 
         # init kwargs == the addr_and_ports map produced by the real allocator
         kwargs0, kwargs1 = ray.get(
@@ -226,27 +221,25 @@ def test_real_allocator_assigns_distinct_ports_via_remote_calls(
         for e in flatten_cells(group.cells):
             ray.kill(e.actor_handle)
 
-    def test_real_allocator_advances_cursor_across_sequential_groups(
+    async def test_real_allocator_advances_cursor_across_sequential_groups(
         self,
         patched_sglang_engine,
         placement_group_factory,
     ):
         """Two sequentially-started groups on independent PGs both invoke the
         real allocator. ``start_engines`` mutates the passed-in PortAllocator
-        in place (via ``assign``); reusing it for B must shift B's ports past
+        in place; reusing it for B must shift B's ports past
         A's — that's the cursor's job."""
         pg_a = placement_group_factory(2)
         pg_b = placement_group_factory(2)
         a = _build_group(pg_tuple=pg_a, num_engines=2)
         b = _build_group(pg_tuple=pg_b, num_engines=2)
 
         cursors = PortAllocator.empty()
-        handles_a, _ = a.start_engines(cursors)
-        ray.get(handles_a)
+        await a.start_engines(cursors)
         # `cursors` now carries the next-free-port state from group A.
 
-        handles_b, _ = b.start_engines(cursors)
-        ray.get(handles_b)
+        await b.start_engines(cursors)
 
         kwargs_a = ray.get([e.actor_handle.get_init_kwargs.remote() for e in flatten_cells(a.cells)])
         kwargs_b = ray.get([e.actor_handle.get_init_kwargs.remote() for e in flatten_cells(b.cells)])
@@ -292,10 +285,10 @@ def test_an_engineless_group_is_exempt(self, placement_group_factory):
 
 class TestRejectedConfigurations:
     @pytest.mark.parametrize("overrides", [{"port": 40000}, {"host": "10.9.9.9"}, {"host": "10.9.9.9", "port": 40000}])
-    def test_host_or_port_override_is_rejected(self, patched_sglang_engine, placement_group_factory, overrides):
+    async def test_host_or_port_override_is_rejected(self, patched_sglang_engine, placement_group_factory, overrides):
         """An override of host or port would make the rollout process address the wrong endpoint."""
         group = _build_group(pg_tuple=placement_group_factory(1), num_engines=1)
         group.sglang_overrides = overrides
 
         with pytest.raises(AssertionError, match="must not override host/port"):
-            group.start_engines(PortAllocator.empty())
+            await group.start_engines(PortAllocator.empty())