From b99cf087cf42b8e39f6f273c45750dd81a43604c Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Tue, 28 Jul 2026 10:10:18 +0800 Subject: [PATCH] Make start_engines async and await engine init inside it start_engines now awaits its engines' init and returns only the new engine indices; recover drops its separate gather, and start_rollout_servers submits one start future per group and collects them with wait_futures. Actor creation, port allocation and state marking all sit before the first await point, so concurrently running starts cannot double-allocate a slot, and port allocation stays serialized because alloc's blocking ray.get never yields the loop. Squashed commits: - Make start_engines async and await engine init inside it - Apply black auto-fix --- miles/ray/rollout/rollout_server.py | 10 +-- miles/ray/rollout/server_group.py | 22 +++---- .../rollout/real_ray/test_fault_tolerance.py | 19 +++--- .../rollout/real_ray/test_rollout_server.py | 19 +++--- .../ray/rollout/real_ray/test_server_group.py | 63 +++++++++---------- .../real_ray/test_server_group_teardown.py | 5 +- tests/fast/ray/rollout/test_addr_allocator.py | 5 +- tests/fast/ray/rollout/test_config_matrix.py | 4 +- 8 files changed, 68 insertions(+), 79 deletions(-) diff --git a/miles/ray/rollout/rollout_server.py b/miles/ray/rollout/rollout_server.py index 8236643059e..9aa29edbad7 100644 --- a/miles/ray/rollout/rollout_server.py +++ b/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) diff --git a/miles/ray/rollout/server_group.py b/miles/ray/rollout/server_group.py index c7488b101aa..456b0c98bd9 100644 --- a/miles/ray/rollout/server_group.py +++ b/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 = [] diff --git a/tests/fast/ray/rollout/real_ray/test_fault_tolerance.py b/tests/fast/ray/rollout/real_ray/test_fault_tolerance.py index ba8cd38dd4b..41fb4762a65 100644 --- a/tests/fast/ray/rollout/real_ray/test_fault_tolerance.py +++ b/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: diff --git a/tests/fast/ray/rollout/real_ray/test_rollout_server.py b/tests/fast/ray/rollout/real_ray/test_rollout_server.py index d1895b7c73b..70cd123ef46 100644 --- a/tests/fast/ray/rollout/real_ray/test_rollout_server.py +++ b/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() diff --git a/tests/fast/ray/rollout/real_ray/test_server_group.py b/tests/fast/ray/rollout/real_ray/test_server_group.py index 196be7f43d0..64ed1b6db96 100644 --- a/tests/fast/ray/rollout/real_ray/test_server_group.py +++ b/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,7 +172,7 @@ 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, @@ -184,9 +180,8 @@ def test_real_allocator_assigns_distinct_ports_via_remote_calls( 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,14 +221,14 @@ 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) @@ -241,12 +236,10 @@ def test_real_allocator_advances_cursor_across_sequential_groups( 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()) diff --git a/tests/fast/ray/rollout/real_ray/test_server_group_teardown.py b/tests/fast/ray/rollout/real_ray/test_server_group_teardown.py index b15f907a012..3ba60fc30ad 100644 --- a/tests/fast/ray/rollout/real_ray/test_server_group_teardown.py +++ b/tests/fast/ray/rollout/real_ray/test_server_group_teardown.py @@ -46,11 +46,10 @@ def _is_dead(actor_handle, *, timeout: float = 60.0) -> bool: class TestTeardownIsTerminal: - def test_a_failing_shutdown_still_kills_the_actor(self, patched_sglang_engine, placement_group_factory): + async def test_a_failing_shutdown_still_kills_the_actor(self, patched_sglang_engine, placement_group_factory): """A graceful shutdown that raises must not leave the actor and its server process behind.""" group = _build_group(pg_tuple=placement_group_factory(1)) - handles, _ = group.start_engines(PortAllocator.empty()) - ray.get(handles) + await group.start_engines(PortAllocator.empty()) actor_handle = flatten_cells(group.cells)[0].actor_handle ray.get(actor_handle.set_fault.remote("shutdown", RuntimeError("shutdown blew up"))) diff --git a/tests/fast/ray/rollout/test_addr_allocator.py b/tests/fast/ray/rollout/test_addr_allocator.py index 04469262c42..bd40cbd7790 100644 --- a/tests/fast/ray/rollout/test_addr_allocator.py +++ b/tests/fast/ray/rollout/test_addr_allocator.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio from unittest.mock import patch from tests.fast.ray.rollout.conftest import chunk_engines_into_cells, fake_actor_handle, fake_engine, make_args @@ -39,13 +40,15 @@ def _start_engines_and_collect_addressing( for index, slot in enumerate(slots): if rank_offset + index not in requested: slot.mark_allocated_uninitialized(fake_actor_handle()) + for engine in requested.values(): + engine.init.remote.side_effect = lambda **kwargs: asyncio.sleep(0) started_cell_indices = sorted({(rank - rank_offset) // nodes_per_engine for rank in requested}) def _launch(*, global_rank, **kwargs): return requested[global_rank] with patch.object(server_group_module, "launch_sglang_ray_actor", side_effect=_launch): - group.start_engines(port_allocator, start_cell_indices=started_cell_indices) + asyncio.run(group.start_engines(port_allocator, start_cell_indices=started_cell_indices)) return {rank: dict(engine.init.remote.call_args.kwargs) for rank, engine in requested.items()} diff --git a/tests/fast/ray/rollout/test_config_matrix.py b/tests/fast/ray/rollout/test_config_matrix.py index c3e53b9609a..665a9c155eb 100644 --- a/tests/fast/ray/rollout/test_config_matrix.py +++ b/tests/fast/ray/rollout/test_config_matrix.py @@ -165,7 +165,7 @@ def test_sglang_config_aggregates_across_models(self): class TestRolloutExternalPath: - def test_starting_engines_in_external_mode_is_not_implemented(self): + async def test_starting_engines_in_external_mode_is_not_implemented(self): """The external allocator was removed; starting engines must fail loudly until the replacement lands.""" from miles.ray.rollout.addr_allocator import PortAllocator from miles.ray.rollout.server_cell import ServerCell @@ -181,4 +181,4 @@ def test_starting_engines_in_external_mode_is_not_implemented(self): has_new_engines=False, ) with pytest.raises(NotImplementedError): - group.start_engines(PortAllocator.empty()) + await group.start_engines(PortAllocator.empty())