diff --git a/areno/api/backend/cuda/backend.py b/areno/api/backend/cuda/backend.py index 1b181708..119860af 100644 --- a/areno/api/backend/cuda/backend.py +++ b/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, diff --git a/areno/engine/parallel/context.py b/areno/engine/parallel/context.py index 47ab42c4..5c1f7ebe 100644 --- a/areno/engine/parallel/context.py +++ b/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, ) diff --git a/areno/engine/protocol.py b/areno/engine/protocol.py index dc6cbeae..a99d70e2 100644 --- a/areno/engine/protocol.py +++ b/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() diff --git a/tests/test_parallel_partition_cpu.py b/tests/test_parallel_partition_cpu.py index 26fa4c19..05b38591 100644 --- a/tests/test_parallel_partition_cpu.py +++ b/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 diff --git a/tests/test_policy_tensor_sync_cpu.py b/tests/test_policy_tensor_sync_cpu.py index 12b26abb..82d2a7c9 100644 --- a/tests/test_policy_tensor_sync_cpu.py +++ b/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) ]