From b482b86221c7a3de88f7ed0347e11366d3d46882 Mon Sep 17 00:00:00 2001 From: Shikha Kathrani <99456505+shikhakath@users.noreply.github.com> Date: Sun, 13 Sep 2026 15:56:34 -0700 Subject: [PATCH] fix(server): loopback-bind the TP rendezvous store, add --dist-port --- docs/cli.md | 1 + python/freetoken/engine/config.py | 3 +- python/freetoken/engine/engine.py | 49 +++++++++++++------ python/freetoken/server/args.py | 28 +++++++++-- tests/engine/test_distributed_store.py | 66 ++++++++++++++++++++++++++ tests/server/test_dist_port_arg.py | 41 ++++++++++++++++ 6 files changed, 168 insertions(+), 20 deletions(-) create mode 100644 tests/engine/test_distributed_store.py create mode 100644 tests/server/test_dist_port_arg.py diff --git a/docs/cli.md b/docs/cli.md index e1eaa8aa3..144597bb0 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -40,6 +40,7 @@ parsers all resolve automatically from the checkpoint and the GPU. |---|---|---| | `--host` | 127.0.0.1 | Bind address | | `--port` | 1919 | Bind port | +| `--dist-port` | `--port` + 1 | Internal TP rendezvous port (loopback-only regardless of `--host`) | | `--gpu` | GPU 0 | GPU to run on: a UUID from `nvidia-smi -L` or an `nvidia-smi` index; see [below](#choosing-a-gpu) | | `--max-running-requests` | 4 | Max concurrently running requests | | `--max-output-tokens` | 32768 | Default output budget for requests that omit one | diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 7ab792f9e..0f16eeeca 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -80,6 +80,7 @@ class EngineConfig: # regardless of the full anchor; the ratio is the startup default and the fallback. swa_num_pages_override: int | None = None distributed_timeout: float = 60.0 + distributed_port: int = 2333 use_dummy_weight: bool = False use_pynccl: bool = True max_seq_len_override: int | None = None @@ -121,4 +122,4 @@ def max_forward_len(self) -> int: @property def distributed_addr(self) -> str: - return "tcp://127.0.0.1:2333" + return f"tcp://127.0.0.1:{self.distributed_port}" diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 006c19089..fd0d1c084 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -3,6 +3,7 @@ import gc import math import os +import socket from datetime import timedelta from typing import Any, Dict, Iterable, NamedTuple, Tuple @@ -434,15 +435,40 @@ def __init__(self, config: EngineConfig): # Prefill runs on the first comma part; warm its autotune cache. self._warmup_prefill() - def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup: - if config.tp_info.size == 1 or config.use_pynccl: - torch.distributed.init_process_group( - backend="gloo", - rank=config.tp_info.rank, - world_size=config.tp_info.size, - timeout=timedelta(seconds=config.distributed_timeout), - init_method=config.distributed_addr, + def _make_distributed_store(self, config: EngineConfig) -> torch.distributed.Store: + """The rendezvous store's C10d server ignores the host it's given and always listens + on every interface, so an unauthenticated TCPStore ends up reachable off-box even + when distributed_addr says 127.0.0.1. Rank 0 pre-binds the listening socket to + loopback itself and hands the fd to TCPStore (master_listen_fd) to force that.""" + timeout = timedelta(seconds=config.distributed_timeout) + if not config.tp_info.is_primary(): + return torch.distributed.TCPStore( + "127.0.0.1", config.distributed_port, config.tp_info.size, + is_master=False, timeout=timeout, multi_tenant=True, ) + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(("127.0.0.1", config.distributed_port)) + sock.listen(1024) + # TCPStore uses this fd directly rather than duplicating it, so it must outlive the + # store; kept alive on self, closed only when the process (and the store) exits. + self._distributed_listen_socket = sock + return torch.distributed.TCPStore( + "127.0.0.1", config.distributed_port, config.tp_info.size, + is_master=True, timeout=timeout, multi_tenant=True, master_listen_fd=sock.fileno(), + ) + + def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup: + store = self._make_distributed_store(config) + use_gloo = config.tp_info.size == 1 or config.use_pynccl + torch.distributed.init_process_group( + backend="gloo" if use_gloo else "nccl", + rank=config.tp_info.rank, + world_size=config.tp_info.size, + timeout=timedelta(seconds=config.distributed_timeout), + store=store, + ) + if use_gloo: tp_cpu_group = torch.distributed.group.WORLD assert tp_cpu_group is not None max_bytes = ( @@ -450,13 +476,6 @@ def _init_communication(self, config: EngineConfig) -> torch.distributed.Process ) enable_pynccl_distributed(config.tp_info, tp_cpu_group, max_bytes) else: - torch.distributed.init_process_group( - backend="nccl", - rank=config.tp_info.rank, - world_size=config.tp_info.size, - timeout=timedelta(seconds=config.distributed_timeout), - init_method=config.distributed_addr, - ) tp_cpu_group = torch.distributed.new_group(backend="gloo") assert tp_cpu_group is not None return tp_cpu_group diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 32e8d1266..c773367dc 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -93,10 +93,6 @@ def backend_create_detokenizer_link(self) -> bool: def frontend_create_tokenizer_link(self) -> bool: return not self.share_tokenizer - @property - def distributed_addr(self) -> str: - return f"tcp://127.0.0.1:{self.server_port + 1}" - def parse_args( args: List[str], @@ -143,6 +139,15 @@ def _positive_int(value: str) -> int: raise argparse.ArgumentTypeError("must be >= 1") return n + def _valid_port(value: str) -> int: + try: + n = int(value) + except ValueError as exc: + raise argparse.ArgumentTypeError("must be an integer") from exc + if not 0 <= n <= 65535: + raise argparse.ArgumentTypeError("must be between 0 and 65535") + return n + def _lazy_gpu_arg(value: str) -> tuple[str, ...]: from freetoken.gpu_select import gpu_arg @@ -338,6 +343,18 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="The port number for the server to listen on.", ) + parser.add_argument( + "--dist-port", + type=_valid_port, + dest="distributed_port", + default=None, + help=( + "Port for the internal TP rendezvous store, loopback-only regardless of --host. " + "Defaults to --port + 1; override when that collides with another instance or " + "service on the same host." + ), + ) + parser.add_argument( "--cuda-graph-max-bs", "--graph", @@ -721,6 +738,9 @@ def _infer_reasoning_parser(model_path: str) -> str | None: if entry: kwargs["quant_backend"] = entry + if kwargs["distributed_port"] is None: + kwargs["distributed_port"] = kwargs["server_port"] + 1 + if kwargs["model_path"].startswith("~"): kwargs["model_path"] = os.path.expanduser(kwargs["model_path"]) diff --git a/tests/engine/test_distributed_store.py b/tests/engine/test_distributed_store.py new file mode 100644 index 000000000..f70be1266 --- /dev/null +++ b/tests/engine/test_distributed_store.py @@ -0,0 +1,66 @@ +"""Regression test for #301: the TP rendezvous store must not fall back to a wildcard +bind. torch's tcp:// rendezvous handler ignores the host it's given when binding the +master's listening socket and always listens on every interface, so +init_method="tcp://127.0.0.1:PORT" looks loopback-only but isn't -- see +Engine._make_distributed_store. This drives the real store-construction path across two +spawned local ranks (CPU/gloo) and checks the actual bind address, not just the URL. +""" + +from __future__ import annotations + +import multiprocessing as mp +from datetime import timedelta +from types import SimpleNamespace + +import torch +import torch.distributed as dist + +from freetoken.distributed import DistributedInfo +from freetoken.engine.engine import Engine + +WORLD_SIZE = 2 +TIMEOUT = timedelta(seconds=15) + + +def _config(rank: int, port: int) -> SimpleNamespace: + return SimpleNamespace( + tp_info=DistributedInfo(rank, WORLD_SIZE), + distributed_timeout=TIMEOUT.total_seconds(), + distributed_port=port, + ) + + +def _run_rank(rank: int, port: int, result_q: mp.Queue) -> None: + engine = Engine.__new__(Engine) # only _make_distributed_store is under test + store = engine._make_distributed_store(_config(rank, port)) + dist.init_process_group( + backend="gloo", rank=rank, world_size=WORLD_SIZE, timeout=TIMEOUT, store=store + ) + total = torch.tensor([float(rank)]) + dist.all_reduce(total) + bind_host = engine._distributed_listen_socket.getsockname()[0] if rank == 0 else None + result_q.put((rank, total.item(), bind_host)) + dist.destroy_process_group() + + +def test_two_ranks_rendezvous_with_loopback_only_store(): + ctx = mp.get_context("spawn") + port = 29511 + result_q = ctx.Queue() + procs = [ctx.Process(target=_run_rank, args=(rank, port, result_q)) for rank in range(WORLD_SIZE)] + for p in procs: + p.start() + + results = {} + for _ in procs: + rank, total, bind_host = result_q.get(timeout=30) + results[rank] = (total, bind_host) + for p in procs: + p.join(timeout=30) + assert p.exitcode == 0 + + # both ranks completed the same all_reduce over the store -- rendezvous worked + assert results[0][0] == results[1][0] == 1.0 # sum(0, 1) + # the master's listening socket was pre-bound to loopback, not the wildcard + # address the C10d TCPStore server binds to on its own + assert results[0][1] == "127.0.0.1" diff --git a/tests/server/test_dist_port_arg.py b/tests/server/test_dist_port_arg.py new file mode 100644 index 000000000..a49def7b6 --- /dev/null +++ b/tests/server/test_dist_port_arg.py @@ -0,0 +1,41 @@ +"""Regression test for #301: the TP rendezvous port must be independently configurable +from --port, so two instances on adjacent ports don't have one's API port collide with +the other's rendezvous port. +""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest + +from freetoken.server.args import parse_args + +ANON_PATH = "/models/anon" + + +class _Config: + def to_dict(self) -> dict: + return {"architectures": ["LlamaForCausalLM"], "torch_dtype": "bfloat16"} + + +def _parse(argv: list[str]): + with patch("freetoken.utils.cached_load_hf_config", lambda _path: _Config()): + return parse_args(["--model", ANON_PATH, *argv]) + + +def test_dist_port_defaults_to_server_port_plus_one(): + args, _ = _parse(["--port", "8081"]) + assert args.distributed_port == 8082 + assert args.distributed_addr == "tcp://127.0.0.1:8082" + + +def test_dist_port_override_is_independent_of_server_port(): + args, _ = _parse(["--port", "8082", "--dist-port", "9000"]) + assert args.distributed_port == 9000 + assert args.distributed_addr == "tcp://127.0.0.1:9000" + + +def test_dist_port_rejects_out_of_range_value(): + with pytest.raises(SystemExit): + _parse(["--dist-port", "70000"])