Rename PortCursors to PortAllocator and key it by the node's ip

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):

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]

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])

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,16 +121,16 @@ 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,
                 num_gpus_per_engine=self.num_gpus_per_engine,
                 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,15 +204,15 @@ 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
                 for cell_index, cell in enumerate(self.cells)
                 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)

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)
 

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

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)

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)