fix(server): loopback-bind the TP rendezvous store, add --dist-port

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 |

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}"

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,29 +435,47 @@ 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 = (
                 config.max_forward_len * config.model_config.hidden_size * self.dtype.itemsize
             )
             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

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

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"

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