Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions areno/api/backend/cuda/backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -178,7 +178,6 @@ def initialize(self, ctx: Context):
from areno.engine.protocol import (
ClusterPartition,
DistributedWorldSpec,
find_free_port,
start_partitioned_clusters,
)

Expand All @@ -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,
Expand Down
6 changes: 5 additions & 1 deletion areno/engine/parallel/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down
34 changes: 27 additions & 7 deletions areno/engine/protocol.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@
import asyncio
import multiprocessing as mp
import queue
import socket
import threading
import traceback
from dataclasses import dataclass
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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")
Expand All @@ -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",
Expand Down Expand Up @@ -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 for the combined train + rollout world; the
# resolved port replaces the placeholder in the shared world_spec before
# any worker spawns (see #517).
store = _create_rendezvous_store(world_spec.master_addr, world_spec.global_world_size)
world_spec.master_port = int(store.port)
for cluster in clusters:
cluster._rendezvous_store = store
try:
for cluster in clusters:
cluster._spawn_workers()
Expand Down
35 changes: 33 additions & 2 deletions tests/test_parallel_partition_cpu.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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)
Expand All @@ -145,3 +152,27 @@ 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"
10 changes: 7 additions & 3 deletions tests/test_policy_tensor_sync_cpu.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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)
]
Expand Down
Loading