fix(engine): allocate distributed init ports outside the ephemeral range

areno/api/backend/cuda/backend.py

@@ -178,7 +178,6 @@ def initialize(self, ctx: Context):
         from areno.engine.protocol import (
             ClusterPartition,
             DistributedWorldSpec,
-            find_free_port,
             start_partitioned_clusters,
         )
 
@@ -204,7 +203,9 @@ def initialize(self, ctx: Context):
         )
         world_spec = DistributedWorldSpec(
             master_addr="127.0.0.1",
-            master_port=find_free_port(),
+            # Placeholder: start_partitioned_clusters creates a coordinator-held
+            # TCPStore with port=0 and fills in the resolved port before spawn.
+            master_port=0,
             global_world_size=world_size + len(rollout_devices),
             train=train_partition,
             rollout=rollout_partition,

areno/engine/parallel/context.py

@@ -113,9 +113,13 @@ def init_process_group(
 
     resolved_global_rank = rank if global_rank is None else global_rank
     resolved_global_world_size = world_size if global_world_size is None else global_world_size
+    # The coordinator holds the server-side rendezvous store on `master_port`
+    # (created with port=0 in protocol.py), so every worker joins as a TCPStore
+    # client instead of racing to bind the port itself (#517).
+    store = dist.TCPStore(master_addr, master_port, world_size=resolved_global_world_size, is_master=False)
     dist.init_process_group(
         backend=backend,
-        init_method=f"tcp://{master_addr}:{master_port}",
+        store=store,
         rank=resolved_global_rank,
         world_size=resolved_global_world_size,
     )

areno/engine/protocol.py

@@ -13,10 +13,9 @@
 import asyncio
 import multiprocessing as mp
 import queue
-import socket
 import threading
 import traceback
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from enum import Enum, auto
 from itertools import count
 from typing import Any, Literal
@@ -235,12 +234,20 @@ def result(self, timeout: float | None = None) -> list[Any]:
         return self._pending.results
 
 
-def find_free_port() -> int:
-    """Reserve an available localhost TCP port for torch distributed init."""
+def _create_rendezvous_store(master_addr: str, world_size: int):
+    """Create and retain a coordinator-side TCPStore with an OS-assigned port.
 
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-        sock.bind(("127.0.0.1", 0))
-        return int(sock.getsockname()[1])
+    The store binds ``port=0`` and keeps the listening socket open for the
+    coordinator's lifetime, so no other process (including outbound traffic
+    reusing ephemeral ports) can steal the resolved port before the workers
+    connect as client stores. This replaces the bind-close-bind probe of the
+    old ``find_free_port``, which raced with outbound connections on busy
+    hosts and failed worker startup with EADDRINUSE (#517).
+    """
+
+    import torch.distributed as dist
+
+    return dist.TCPStore(master_addr, 0, world_size=world_size, is_master=True, wait_for_workers=False)
 
 
 def _rollout_payload_count(payload: RolloutPayload) -> int:
@@ -292,6 +299,9 @@ def __init__(
             raise ValueError("world_spec and partition must be provided together")
         self.world_spec = world_spec
         self.partition = partition
+        # Coordinator-side rendezvous store; retained for the cluster lifetime
+        # so the resolved master port stays reserved for the workers (#517).
+        self._rendezvous_store = None
         # `spawn` start method is required by CUDA-aware workers; do not
         # inherit fds/CUDA state from the parent.
         self.ctx = mp.get_context("spawn")
@@ -314,9 +324,12 @@ def start(self) -> None:
         if self.world_spec is not None:
             raise RuntimeError("partitioned clusters must be started with start_partitioned_clusters()")
         world_size = self.config.tp_size * int(self.config.dp_size)
+        # The coordinator holds the rendezvous store open so the resolved port
+        # is genuinely reserved until the workers connect (see #517).
+        self._rendezvous_store = _create_rendezvous_store("127.0.0.1", world_size)
         world_spec = DistributedWorldSpec(
             master_addr="127.0.0.1",
-            master_port=find_free_port(),
+            master_port=int(self._rendezvous_store.port),
             global_world_size=world_size,
             train=ClusterPartition(
                 role="train",
@@ -710,6 +723,13 @@ def start_partitioned_clusters(
     clusters = (train_cluster, rollout_cluster)
     if train_cluster.world_spec != world_spec or rollout_cluster.world_spec != world_spec:
         raise ValueError("both clusters must use the supplied world_spec")
+    # One coordinator-held store serves the combined world; the resolved port
+    # is rebuilt into the frozen world_spec before any worker spawns (#517).
+    store = _create_rendezvous_store(world_spec.master_addr, world_spec.global_world_size)
+    resolved_spec = replace(world_spec, master_port=int(store.port))
+    for cluster in clusters:
+        cluster.world_spec = resolved_spec
+        cluster._rendezvous_store = store
     try:
         for cluster in clusters:
             cluster._spawn_workers()

tests/test_parallel_partition_cpu.py

@@ -2,11 +2,12 @@
 
 import multiprocessing as mp
 
+import pytest
 import torch
 
 from areno.engine.parallel import context
 from areno.engine.parallel.collectives import broadcast_object
-from areno.engine.protocol import find_free_port
+from areno.engine.protocol import _create_rendezvous_store
 
 
 def _run_offset_tp_broadcast(global_rank: int, port: int, output_queue) -> None:
@@ -37,6 +38,9 @@ def _run_offset_tp_broadcast(global_rank: int, port: int, output_queue) -> None:
 def _mock_distributed(monkeypatch):
     calls = []
     monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    # These tests exercise group construction only, not the rendezvous, so
+    # the client TCPStore and the process-group init are both faked.
+    monkeypatch.setattr(context.dist, "TCPStore", lambda *args, **kwargs: object())
     monkeypatch.setattr(context.dist, "init_process_group", lambda **kwargs: calls.append(("init", kwargs)))
 
     def new_group(*, ranks):
@@ -127,7 +131,10 @@ def test_single_engine_context_creates_no_policy_publisher_group(monkeypatch) ->
 def test_real_gloo_tp_broadcast_uses_partition_global_root() -> None:
     spawn = mp.get_context("spawn")
     output_queue = spawn.Queue()
-    port = find_free_port()
+    # The coordinator holds the server store (port=0) so the resolved port is
+    # genuinely reserved before the workers join as client stores.
+    store = _create_rendezvous_store("127.0.0.1", 4)
+    port = int(store.port)
     processes = [
         spawn.Process(target=_run_offset_tp_broadcast, args=(global_rank, port, output_queue))
         for global_rank in range(4)
@@ -145,3 +152,102 @@ def test_real_gloo_tp_broadcast_uses_partition_global_root() -> None:
         2: "rollout-root",
         3: "rollout-root",
     }
+
+
+def test_rendezvous_store_resolves_and_holds_its_port() -> None:
+    """The coordinator store resolves a real port and keeps it bound.
+
+    The resolved port must be immediately usable by client stores, and binding
+    the same port again must fail while the server store is alive (the socket
+    is genuinely reserved, unlike the old bind-close-bind probe, #517).
+    """
+
+    import socket
+
+    import torch.distributed as dist
+
+    store = _create_rendezvous_store("127.0.0.1", 2)
+    port = int(store.port)
+    assert port > 0
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        with pytest.raises(OSError):
+            sock.bind(("127.0.0.1", port))
+    # A client store connects to the retained server.
+    client = dist.TCPStore("127.0.0.1", port, world_size=2, is_master=False)
+    store.set("k", "v")
+    assert client.get("k") == b"v"
+
+
+def test_start_partitioned_clusters_resolves_frozen_world_spec() -> None:
+    """`start_partitioned_clusters` must fill in the resolved rendezvous port
+    without mutating the frozen `world_spec`.
+
+    `DistributedWorldSpec` is ``dataclass(frozen=True)``, so assigning
+    ``world_spec.master_port`` raises ``FrozenInstanceError``. The
+    coordinator-held store's resolved port is instead carried into a rebuilt
+    spec before any worker spawns (see #517).
+    """
+
+    from areno.engine.protocol import (
+        ClusterPartition,
+        DistributedWorldSpec,
+        start_partitioned_clusters,
+    )
+
+    class FakeCluster:
+        """Stand-in for TPCluster exposing exactly what the helper touches."""
+
+        def __init__(self, partition):
+            self.world_spec = None
+            self.partition = partition
+            self._rendezvous_store = None
+            self.started = False
+
+        def _spawn_workers(self) -> None:
+            pass
+
+        def _wait_for_worker_ready(self, ranks) -> None:
+            pass
+
+        def _abort_start(self) -> None:
+            pass
+
+        def _start_result_pump(self) -> None:
+            pass
+
+    # Mirror backend.py: a placeholder port that the helper must resolve.
+    world_spec = DistributedWorldSpec(
+        master_addr="127.0.0.1",
+        master_port=0,
+        global_world_size=4,
+        train=ClusterPartition(
+            role="train",
+            global_rank_offset=0,
+            local_world_size=2,
+            tp_size=2,
+            devices=(0, 1),
+        ),
+        rollout=ClusterPartition(
+            role="rollout",
+            global_rank_offset=2,
+            local_world_size=2,
+            tp_size=2,
+            devices=(2, 3),
+        ),
+    )
+    train = FakeCluster(world_spec.train)
+    rollout = FakeCluster(world_spec.rollout)
+    train.world_spec = world_spec
+    rollout.world_spec = world_spec
+
+    start_partitioned_clusters(train, rollout, world_spec)
+
+    # The shared world_spec's port is resolved and both partitions agree.
+    assert train.started and rollout.started
+    assert train.world_spec is rollout.world_spec
+    assert train.world_spec.master_port > 0
+    assert train.world_spec.master_port == rollout.world_spec.master_port
+    # The coordinator store is retained on every cluster.
+    assert train._rendezvous_store is rollout._rendezvous_store
+    # The frozen input spec must not have been mutated in place.
+    assert world_spec.master_port == 0

tests/test_policy_tensor_sync_cpu.py

@@ -22,7 +22,7 @@
     build_adapter_policy_plan,
     transfer_policy_weights,
 )
-from areno.engine.protocol import PolicySyncPayload, find_free_port
+from areno.engine.protocol import PolicySyncPayload, _create_rendezvous_store
 
 
 def _set_rank(rank: int, world_size: int) -> None:
@@ -353,7 +353,9 @@ def _gloo_policy_sync_worker(global_rank: int, port: int, output_queue) -> None:
 def test_real_gloo_collectives_reshard_train_tp2_to_rollout_tp1() -> None:
     ctx = mp.get_context("spawn")
     output_queue = ctx.Queue()
-    port = find_free_port()
+    # Coordinator-held store mirrors production: workers join as client stores.
+    store = _create_rendezvous_store("127.0.0.1", 3)
+    port = int(store.port)
     processes = [ctx.Process(target=_gloo_policy_sync_worker, args=(rank, port, output_queue)) for rank in range(3)]
     for process in processes:
         process.start()
@@ -402,7 +404,9 @@ def _gloo_policy_sync_reverse_worker(global_rank: int, port: int, output_queue)
 def test_real_gloo_collectives_reshard_train_tp1_to_rollout_tp2() -> None:
     ctx = mp.get_context("spawn")
     output_queue = ctx.Queue()
-    port = find_free_port()
+    # Coordinator-held store mirrors production: workers join as client stores.
+    store = _create_rendezvous_store("127.0.0.1", 3)
+    port = int(store.port)
     processes = [
         ctx.Process(target=_gloo_policy_sync_reverse_worker, args=(rank, port, output_queue)) for rank in range(3)
     ]