From 31d4af4c580a901972eb5e133d786ad472901e9f Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 27 Jul 2026 16:13:14 +0800 Subject: [PATCH] Rename PortCursors to PortAllocator and key it by the node's ip The cursor dict was keyed by ``node_index = local_rank // num_engines_per_node`` -- an arithmetic proxy for "which machine", derived from an assumed rank layout and restarting at zero for every group. The node ip is the machine's actual identity and needs no such assumption. Nothing misbehaves today: the dict is rebuilt on every call, and groups are kept apart by ``next_base_port()``, one high-water mark taken across all nodes. That guard is safe but coarse -- it drags every node's cursor up to the busiest node's. The rekey is what makes a per-node cursor mean anything once the allocator becomes shared and long-lived: keyed by a per-group index, one shared dict would alias two groups' distinct machines onto a single cursor, and conversely split one machine into two cursors whenever the two groups computed different indices for it. Reading the ip costs one extra ``_get_current_node_ip_and_free_port`` round trip per node. --- miles/ray/rollout/addr_allocator.py | 23 ++++++----- miles/ray/rollout/inference_controller.py | 6 +-- miles/ray/rollout/rollout_server.py | 10 ++--- miles/ray/rollout/server_group.py | 16 ++++---- tests/fast/ray/rollout/real_ray/conftest.py | 4 +- .../rollout/real_ray/test_fault_tolerance.py | 18 ++++----- .../real_ray/test_inference_controller.py | 4 +- .../rollout/real_ray/test_rollout_server.py | 4 +- .../ray/rollout/real_ray/test_server_group.py | 26 ++++++------ .../real_ray/test_server_group_teardown.py | 4 +- tests/fast/ray/rollout/test_addr_allocator.py | 40 +++++++++---------- 11 files changed, 78 insertions(+), 77 deletions(-) diff --git a/miles/ray/rollout/addr_allocator.py b/miles/ray/rollout/addr_allocator.py index 7900a6dfea1..50b5e39a645 100644 --- a/miles/ray/rollout/addr_allocator.py +++ b/miles/ray/rollout/addr_allocator.py @@ -7,14 +7,14 @@ @dataclass -class PortCursors: - _values: dict[int, int] +class PortAllocator: + _values: dict[str, int] @staticmethod - def empty() -> "PortCursors": - return PortCursors(_values={}) + def empty() -> "PortAllocator": + return PortAllocator(_values={}) - def assign(self, other: "PortCursors"): + def assign(self, other: "PortAllocator"): self._values = other._values.copy() def next_base_port(self) -> int: @@ -43,7 +43,7 @@ def allocate_rollout_engine_addr_and_ports_normal( # Track per-node port cursors so that different server groups (called # sequentially) never race for the same ports on a given node. - node_port_cursor: dict[int, int] = {} + node_port_cursor: dict[str, int] = {} visited_nodes = set() for rank, engine in rollout_engines: @@ -56,10 +56,10 @@ def allocate_rollout_engine_addr_and_ports_normal( # e.g. for 8 gpus, if we are restarting engine on gpu 3, we will set port for engine 3,4,5,6,7 on this node. num_engines_on_this_node = num_engines_per_node - (local_rank % num_engines_per_node) - def get_addr_and_ports(engine, node_idx): + def get_addr_and_ports(engine, node_ip): # use small ports to prevent ephemeral port between 32768 and 65536. # also, ray uses port 10002-19999, thus we avoid near-10002 to avoid racing condition - start_port = node_port_cursor.get(node_idx, base_port) + start_port = node_port_cursor.get(node_ip, base_port) def port(consecutive=1): nonlocal start_port @@ -70,7 +70,7 @@ def port(consecutive=1): ) ) start_port = port + consecutive - node_port_cursor[node_idx] = start_port + node_port_cursor[node_ip] = start_port return port def addr(): @@ -79,7 +79,8 @@ def addr(): return addr, port - get_addr, get_port = get_addr_and_ports(engine, node_index) + node_ip, _ = ray.get(engine._get_current_node_ip_and_free_port.remote()) + get_addr, get_port = get_addr_and_ports(engine, node_ip) for i in range(num_engines_on_this_node): current_rank = rank + i @@ -109,7 +110,7 @@ def addr(): assert key in addr_and_ports[i], f"Engine {i} {key} is not set." logger.info(f"Ports for engine {i}: {addr_and_ports[i]}") - return addr_and_ports, PortCursors(_values=node_port_cursor) + return addr_and_ports, PortAllocator(_values=node_port_cursor) def allocate_rollout_engine_addr_and_ports_external(args, rollout_engines): diff --git a/miles/ray/rollout/inference_controller.py b/miles/ray/rollout/inference_controller.py index e99ae23fbcb..f0daa4bb88f 100644 --- a/miles/ray/rollout/inference_controller.py +++ b/miles/ray/rollout/inference_controller.py @@ -6,7 +6,7 @@ from miles.backends.sglang_utils.sglang_api_client import SGLangApiClient from miles.dashboard import hooks as dashboard_hooks -from miles.ray.rollout.addr_allocator import PortCursors +from miles.ray.rollout.addr_allocator import PortAllocator from miles.ray.rollout.eval_fleet import EvalFleet from miles.ray.rollout.rollout_server import RolloutServer, start_rollout_servers from miles.ray.rollout.router_manager import start_session_server @@ -136,10 +136,10 @@ def _get_updatable_server(self) -> RolloutServer | None: # -------------------------- external start/stop ----------------------------- async def start_cell(self, cell_id: int): - port_cursors = PortCursors.empty() + port_allocator = PortAllocator.empty() idx = get_cell_indexer_of_id_map(self.servers)[cell_id] group = self.servers[idx.srv_key].server_groups[idx.group_index] - await group.recover(port_cursors=port_cursors, filter_cell_indices=[idx.cell_index]) + await group.recover(port_allocator=port_allocator, filter_cell_indices=[idx.cell_index]) async def stop_cell(self, cell_id: int): idx = get_cell_indexer_of_id_map(self.servers)[cell_id] diff --git a/miles/ray/rollout/rollout_server.py b/miles/ray/rollout/rollout_server.py index bea23778310..8236643059e 100644 --- a/miles/ray/rollout/rollout_server.py +++ b/miles/ray/rollout/rollout_server.py @@ -6,7 +6,7 @@ from miles.backends.sglang_utils.arguments import collect_eval_sglang_overrides from miles.backends.sglang_utils.sglang_config import ModelConfig, ServerGroupConfig, SglangConfig -from miles.ray.rollout.addr_allocator import PortCursors +from miles.ray.rollout.addr_allocator import PortAllocator from miles.ray.rollout.router_manager import start_router from miles.ray.rollout.server_cell import ServerCell from miles.ray.rollout.server_engine import ServerEngine @@ -43,7 +43,7 @@ def start_rollout_servers(args, pg) -> dict[str, "RolloutServer"]: server_groups: list[ServerGroup] = [] all_init_handles: list = [] new_engine_indices_per_group: list[list[int]] = [] - port_cursors = PortCursors.empty() + port_allocator = PortAllocator.empty() for group_cfg in model_cfg.server_groups: gpus_per_engine = group_cfg.num_gpus_per_engine @@ -84,7 +84,7 @@ 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_cursors) + handles, new_engine_indices = group.start_engines(port_allocator) all_init_handles.extend(handles) server_groups.append(group) new_engine_indices_per_group.append(new_engine_indices) @@ -271,8 +271,8 @@ async def probe_and_mark_dead(self): async def recover(self): """Recover dead engines across all active groups, overlapping init.""" - port_cursors = PortCursors.empty() - await asyncio.gather(*[g.recover(port_cursors=port_cursors) for g in self.server_groups]) + port_allocator = PortAllocator.empty() + await asyncio.gather(*[g.recover(port_allocator=port_allocator) for g in self.server_groups]) async def offload(self, tags: list[str] | None = None): per_group = await asyncio.gather(*[g.offload(tags=tags) for g in self.server_groups]) diff --git a/miles/ray/rollout/server_group.py b/miles/ray/rollout/server_group.py index 90cd87c0547..03e6cb3372a 100644 --- a/miles/ray/rollout/server_group.py +++ b/miles/ray/rollout/server_group.py @@ -8,7 +8,7 @@ from miles.backends.sglang_utils.sglang_engine import build_server_url from miles.backends.sglang_utils.sglang_router_api_client import SGLangRouterApiClient, use_legacy_router_api from miles.ray.rollout.addr_allocator import ( - PortCursors, + PortAllocator, allocate_rollout_engine_addr_and_ports_external, allocate_rollout_engine_addr_and_ports_normal, ) @@ -59,11 +59,11 @@ def engines(self) -> list[ServerEngine]: return [cell.engines[0] for cell in self.cells] def start_engines( - self, port_cursors: PortCursors, start_cell_indices: list[int] | None = None + 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. - Mutates ``port_cursors`` in place to advance past any newly assigned ports. + 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. @@ -121,8 +121,8 @@ def start_engines( args=self.args, rollout_engines=new_engines ) else: - base_port = port_cursors.next_base_port() - addr_and_ports, next_port_cursors = allocate_rollout_engine_addr_and_ports_normal( + base_port = port_allocator.next_base_port() + addr_and_ports, next_port_allocator = allocate_rollout_engine_addr_and_ports_normal( args=self.args, rollout_engines=new_engines, worker_type=self.worker_type, @@ -130,7 +130,7 @@ def start_engines( rank_offset=self.rank_offset, base_port=base_port, ) - port_cursors.assign(next_port_cursors) + port_allocator.assign(next_port_allocator) for index, _ in new_engines: engine_addr_and_ports = addr_and_ports[index] @@ -204,7 +204,7 @@ def stop_engines(self, cell_indices: list[int]): for cell_index in sorted(set(cell_indices)): self.cells[cell_index].stop() - async def recover(self, port_cursors: PortCursors, filter_cell_indices: list[int] | None = None): + async def recover(self, port_allocator: PortAllocator, filter_cell_indices: list[int] | None = None): if filter_cell_indices is None: filter_cell_indices = [ cell_index @@ -212,7 +212,7 @@ async def recover(self, port_cursors: PortCursors, filter_cell_indices: list[int if any(not engine.is_allocated for engine in cell.engines) ] - handles, new_engine_indices = self.start_engines(port_cursors, start_cell_indices=filter_cell_indices) + handles, new_engine_indices = self.start_engines(port_allocator, start_cell_indices=filter_cell_indices) await asyncio.gather(*handles) all_engines = flatten_cells(self.cells) diff --git a/tests/fast/ray/rollout/real_ray/conftest.py b/tests/fast/ray/rollout/real_ray/conftest.py index a05cb419e6d..1022947ba3f 100644 --- a/tests/fast/ray/rollout/real_ray/conftest.py +++ b/tests/fast/ray/rollout/real_ray/conftest.py @@ -69,7 +69,7 @@ def patched_sglang_engine(monkeypatch, mock_engine_class, mock_engine_http_serve monkeypatch.setattr(cell_mod, "SGLangEngine", mock_engine_class) - from miles.ray.rollout.addr_allocator import PortCursors + from miles.ray.rollout.addr_allocator import PortAllocator def _fake_alloc(*args, **kwargs): engines = kwargs["rollout_engines"] @@ -83,7 +83,7 @@ def _fake_alloc(*args, **kwargs): engine_info_bootstrap_port=32000 + rank, dist_init_addr=f"127.0.0.1:{33000 + rank}", ) - return addr_and_ports, PortCursors(_values={0: 34000}) + return addr_and_ports, PortAllocator(_values={"127.0.0.1": 34000}) monkeypatch.setattr(mod, "allocate_rollout_engine_addr_and_ports_normal", _fake_alloc) 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 5f5932c4fb6..13647dd554d 100644 --- a/tests/fast/ray/rollout/real_ray/test_fault_tolerance.py +++ b/tests/fast/ray/rollout/real_ray/test_fault_tolerance.py @@ -9,7 +9,7 @@ import ray from tests.fast.ray.rollout.conftest import chunk_engines_into_cells, make_args -from miles.ray.rollout.addr_allocator import PortCursors +from miles.ray.rollout.addr_allocator import PortAllocator from miles.ray.rollout.server_cell import flatten_cells from miles.ray.rollout.server_engine import ServerEngine from miles.ray.rollout.server_group import ServerGroup @@ -39,7 +39,7 @@ def _build_group( def _start(group: ServerGroup) -> None: - handles, indices = group.start_engines(PortCursors.empty()) + handles, indices = group.start_engines(PortAllocator.empty()) ray.get(handles) group.mark_alive(indices) @@ -76,7 +76,7 @@ async def test_recover_creates_new_actor_after_real_kill( flatten_cells(group.cells)[0].mark_stopped() try: - await group.recover(port_cursors=PortCursors.empty(), filter_cell_indices=[0]) + await group.recover(port_allocator=PortAllocator.empty(), filter_cell_indices=[0]) # New actor for slot 0 assert flatten_cells(group.cells)[0].is_allocated assert flatten_cells(group.cells)[0].actor_handle is not original_handles[0] @@ -106,7 +106,7 @@ async def test_recover_default_filter_picks_all_dead_slots( flatten_cells(group.cells)[i].mark_stopped() try: - await group.recover(port_cursors=PortCursors.empty()) + await group.recover(port_allocator=PortAllocator.empty()) for i in (0, 2): assert flatten_cells(group.cells)[i].is_allocated assert flatten_cells(group.cells)[i].actor_handle is not old[i] @@ -143,7 +143,7 @@ async def remove_worker(self, **kwargs): try: with patch.object(ServerGroup, "_router_api_client", property(lambda self: _Recorder())): - await group.recover(port_cursors=PortCursors.empty(), filter_cell_indices=[0]) + await group.recover(port_allocator=PortAllocator.empty(), filter_cell_indices=[0]) assert [event["worker_url"] for event in events] == [flatten_cells(group.cells)[0].addr_info.server_url] assert flatten_cells(group.cells)[0].is_alive @@ -168,7 +168,7 @@ async def test_recover_with_offload_calls_release_then_resume( flatten_cells(group.cells)[0].mark_stopped() try: - await group.recover(port_cursors=PortCursors.empty(), filter_cell_indices=[0]) + await group.recover(port_allocator=PortAllocator.empty(), filter_cell_indices=[0]) calls = ray.get(flatten_cells(group.cells)[0].actor_handle.get_calls.remote()) assert "init" in [c[0] for c in calls] @@ -231,8 +231,8 @@ async def test_two_groups_recover_in_parallel_completes_without_deadlock( try: # Real concurrent recover via asyncio.gather await asyncio.gather( - a.recover(port_cursors=PortCursors.empty(), filter_cell_indices=[0]), - b.recover(port_cursors=PortCursors.empty(), filter_cell_indices=[0]), + a.recover(port_allocator=PortAllocator.empty(), filter_cell_indices=[0]), + b.recover(port_allocator=PortAllocator.empty(), filter_cell_indices=[0]), ) assert flatten_cells(a.cells)[0].is_allocated assert flatten_cells(b.cells)[0].is_allocated @@ -283,7 +283,7 @@ async def test_recover_releases_and_resumes_only_on_node0( assert group.nodes_per_engine == 2 try: - await group.recover(port_cursors=PortCursors.empty()) + await group.recover(port_allocator=PortAllocator.empty()) node0_paths = mock_engine_http_servers.for_rank(0).paths node1_paths = mock_engine_http_servers.for_rank(1).paths diff --git a/tests/fast/ray/rollout/real_ray/test_inference_controller.py b/tests/fast/ray/rollout/real_ray/test_inference_controller.py index acfb7a4fdad..2b9fe0740a5 100644 --- a/tests/fast/ray/rollout/real_ray/test_inference_controller.py +++ b/tests/fast/ray/rollout/real_ray/test_inference_controller.py @@ -39,7 +39,7 @@ def patch_low_level(monkeypatch, mock_engine_http_servers): import miles.ray.rollout.rollout_server as rsrv import miles.ray.rollout.server_cell as scell import miles.ray.rollout.server_group as sg - from miles.ray.rollout.addr_allocator import PortCursors + from miles.ray.rollout.addr_allocator import PortAllocator from miles.utils.test_utils.mock_sglang_engine import MockSGLangEngine monkeypatch.setattr(scell, "SGLangEngine", MockSGLangEngine.__ray_actor_class__) @@ -64,7 +64,7 @@ def _fake_alloc(*args, **kwargs): ) for rank, _ in engines }, - PortCursors(_values={0: 34000}), + PortAllocator(_values={"127.0.0.1": 34000}), ) monkeypatch.setattr(sg, "allocate_rollout_engine_addr_and_ports_normal", _fake_alloc) 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 51d1ccffa8b..83e16279376 100644 --- a/tests/fast/ray/rollout/real_ray/test_rollout_server.py +++ b/tests/fast/ray/rollout/real_ray/test_rollout_server.py @@ -4,7 +4,7 @@ import ray from tests.fast.ray.rollout.conftest import chunk_engines_into_cells, make_args -from miles.ray.rollout.addr_allocator import PortCursors +from miles.ray.rollout.addr_allocator import PortAllocator from miles.ray.rollout.rollout_server import RolloutServer from miles.ray.rollout.server_cell import flatten_cells from miles.ray.rollout.server_engine import ServerEngine @@ -36,7 +36,7 @@ def _build_group( def _start_group(group: ServerGroup) -> None: - handles, _ = group.start_engines(PortCursors.empty()) + handles, _ = group.start_engines(PortAllocator.empty()) ray.get(handles) 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 56c7e23dd6f..83dd68f8cda 100644 --- a/tests/fast/ray/rollout/real_ray/test_server_group.py +++ b/tests/fast/ray/rollout/real_ray/test_server_group.py @@ -4,7 +4,7 @@ import ray from tests.fast.ray.rollout.conftest import chunk_engines_into_cells, make_args -from miles.ray.rollout.addr_allocator import PortCursors +from miles.ray.rollout.addr_allocator import PortAllocator from miles.ray.rollout.server_cell import flatten_cells from miles.ray.rollout.server_engine import ServerEngine from miles.ray.rollout.server_group import ServerGroup @@ -42,7 +42,7 @@ 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(PortCursors.empty()) + handles, indices = group.start_engines(PortAllocator.empty()) assert handles == [] and indices == [] assert group.has_new_engines is False for e in flatten_cells(group.cells): @@ -52,7 +52,7 @@ 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(PortCursors.empty()) + handles, indices = group.start_engines(PortAllocator.empty()) assert handles == [] and indices == [] assert group.has_new_engines is False @@ -68,7 +68,7 @@ def test_creates_real_actors_and_init_runs( pg = placement_group_factory(2) group = _build_group(pg_tuple=pg, num_engines=2) - handles, indices = group.start_engines(PortCursors.empty()) + handles, indices = 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. @@ -92,7 +92,7 @@ def test_start_cell_indices_filters_to_subset(self, patched_sglang_engine, place pg = placement_group_factory(4) group = _build_group(pg_tuple=pg, num_engines=4) - handles, indices = group.start_engines(PortCursors.empty(), start_cell_indices=[1, 3]) + handles, indices = group.start_engines(PortAllocator.empty(), start_cell_indices=[1, 3]) assert sorted(indices) == [1, 3] ray.get(handles) @@ -111,12 +111,12 @@ def test_already_allocated_slot_is_skipped(self, patched_sglang_engine, placemen group = _build_group(pg_tuple=pg, num_engines=2) # First call: allocates both slots. - handles, _ = group.start_engines(PortCursors.empty()) + handles, _ = group.start_engines(PortAllocator.empty()) ray.get(handles) 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(PortCursors.empty()) + handles2, indices2 = group.start_engines(PortAllocator.empty()) assert handles2 == [] and indices2 == [] for first, e in zip(first_handles, flatten_cells(group.cells), strict=True): assert e.actor_handle is first # still the same actor @@ -136,7 +136,7 @@ class TestStopEnginesRealKill: 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(PortCursors.empty()) + handles, _ = group.start_engines(PortAllocator.empty()) ray.get(handles) actors = [e.actor_handle for e in flatten_cells(group.cells)] @@ -158,7 +158,7 @@ def test_stop_handles_shutdown_failure_gracefully(self, patched_sglang_engine, p 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(PortCursors.empty()) + handles, _ = group.start_engines(PortAllocator.empty()) ray.get(handles) # Plant a one-shot shutdown failure on engine 1. @@ -193,7 +193,7 @@ 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(PortCursors.empty()) + handles, indices = group.start_engines(PortAllocator.empty()) assert sorted(indices) == [0, 1] ray.get(handles) @@ -241,7 +241,7 @@ def test_real_allocator_advances_cursor_across_sequential_groups( placement_group_factory, ): """Two sequentially-started groups on independent PGs both invoke the - real allocator. ``start_engines`` mutates the passed-in PortCursors + real allocator. ``start_engines`` mutates the passed-in PortAllocator in place (via ``assign``); reusing it for B must shift B's ports past A's — that's the cursor's job.""" pg_a = placement_group_factory(2) @@ -249,7 +249,7 @@ def test_real_allocator_advances_cursor_across_sequential_groups( a = _build_group(pg_tuple=pg_a, num_engines=2) b = _build_group(pg_tuple=pg_b, num_engines=2) - cursors = PortCursors.empty() + cursors = PortAllocator.empty() handles_a, _ = a.start_engines(cursors) ray.get(handles_a) # `cursors` now carries the next-free-port state from group A. @@ -307,4 +307,4 @@ def test_host_or_port_override_is_rejected(self, patched_sglang_engine, placemen group.sglang_overrides = overrides with pytest.raises(AssertionError, match="must not override host/port"): - group.start_engines(PortCursors.empty()) + 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 13a92464703..b15f907a012 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 @@ -7,7 +7,7 @@ from tests.fast.ray.rollout.conftest import chunk_engines_into_cells, make_args import miles.ray.rollout.server_cell as server_cell_module -from miles.ray.rollout.addr_allocator import PortCursors +from miles.ray.rollout.addr_allocator import PortAllocator from miles.ray.rollout.server_cell import flatten_cells from miles.ray.rollout.server_engine import ServerEngine from miles.ray.rollout.server_group import ServerGroup @@ -49,7 +49,7 @@ class TestTeardownIsTerminal: 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(PortCursors.empty()) + handles, _ = group.start_engines(PortAllocator.empty()) ray.get(handles) 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 088b677982d..e7314eef2ca 100644 --- a/tests/fast/ray/rollout/test_addr_allocator.py +++ b/tests/fast/ray/rollout/test_addr_allocator.py @@ -5,37 +5,37 @@ from tests.fast.ray.rollout.conftest import fake_engine, make_args from miles.ray.rollout.addr_allocator import ( - PortCursors, + PortAllocator, allocate_rollout_engine_addr_and_ports_external, allocate_rollout_engine_addr_and_ports_normal, ) -class TestPortCursors: +class TestPortAllocator: def test_empty_has_no_values(self): - c = PortCursors.empty() + c = PortAllocator.empty() assert c._values == {} def test_next_base_port_default_when_empty(self): - assert PortCursors.empty().next_base_port() == 15000 + assert PortAllocator.empty().next_base_port() == 15000 def test_next_base_port_returns_max_value(self): - c = PortCursors(_values={0: 17000, 1: 16500, 2: 18000}) + c = PortAllocator(_values={"10.0.0.1": 17000, "10.0.0.2": 16500, "10.0.0.3": 18000}) assert c.next_base_port() == 18000 def test_assign_copies_values(self): - a = PortCursors.empty() - b = PortCursors(_values={0: 19000, 1: 19500}) + a = PortAllocator.empty() + b = PortAllocator(_values={"10.0.0.1": 19000, "10.0.0.2": 19500}) a.assign(b) - assert a._values == {0: 19000, 1: 19500} + assert a._values == {"10.0.0.1": 19000, "10.0.0.2": 19500} def test_assign_is_decoupled(self): """After assign, mutating source must not bleed into target.""" - a = PortCursors.empty() - b = PortCursors(_values={0: 19000}) + a = PortAllocator.empty() + b = PortAllocator(_values={"10.0.0.1": 19000}) a.assign(b) - b._values[0] = 99999 - assert a._values == {0: 19000}, "assign must deep-copy the inner dict" + b._values["10.0.0.1"] = 99999 + assert a._values == {"10.0.0.1": 19000}, "assign must deep-copy the inner dict" def _all_ports(addr_and_ports: dict) -> list[int]: @@ -80,11 +80,11 @@ def test_single_node_8_cards_tp1(self, patch_ray_get): } assert len(same_rank_ports) == 4, f"rank {rank} reused a port: {addr_and_ports[rank]}" - # Cursor must reflect the *node*'s next free port (single-node here → key 0). - assert isinstance(cursors, PortCursors) - assert set(cursors._values.keys()) == {0} + # Cursor must reflect the *node*'s next free port (single node → its ip). + assert isinstance(cursors, PortAllocator) + assert set(cursors._values.keys()) == {"10.0.0.1"} # And it must sit past every port we handed out. - assert cursors._values[0] >= max(_all_ports(addr_and_ports)) + 1 + assert cursors._values["10.0.0.1"] >= max(_all_ports(addr_and_ports)) + 1 # Cross-rank: every numeric port across all 8 engines must be unique. all_ports = _all_ports(addr_and_ports) @@ -180,7 +180,7 @@ def test_base_port_propagates_into_cursor(self, patch_ray_get): # also reserves consecutive blocks for dist_init_addr that aren't all # visible in the output, so we can't pin to max_issued + 1). max_issued = max(_all_ports(addr_and_ports)) - assert cursors._values[0] > max_issued + assert cursors._values["10.0.0.1"] > max_issued # And the lowest port must be >= base_port (allocator never went below it). assert min(_all_ports(addr_and_ports)) >= 22000 @@ -207,13 +207,13 @@ def test_ipv4_addr_split_is_consistent(self): assert isinstance(result[0]["port"], int) -class TestSharedPortCursorsAcrossGroups: - """Two ``ServerGroup``s sharing one ``PortCursors`` must produce disjoint +class TestSharedPortAllocatorAcrossGroups: + """Two ``ServerGroup``s sharing one ``PortAllocator`` must produce disjoint port allocations across nodes — required for parallel recover.""" def test_sequential_groups_share_cursor_and_avoid_overlap(self, patch_ray_get): args = make_args(num_gpus_per_node=8, sglang_dp_size=1) - cursors = PortCursors.empty() + cursors = PortAllocator.empty() engines_a = [(rank, fake_engine(port_seed=0)) for rank in range(4)] addrs_a, next_a = allocate_rollout_engine_addr_and_ports_normal(