From 82e8365036e980e778be9b50238119151b0d6ef1 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Tue, 29 Sep 2026 03:44:09 +0300 Subject: [PATCH 1/7] fix(windows): serve over loopback ZMQ and a selector event loop --- python/freetoken/scheduler/config.py | 34 ++++++++++++--- python/freetoken/server/api_server.py | 24 +++++++++-- python/freetoken/server/args.py | 4 +- tests/scheduler/test_zmq_links.py | 59 +++++++++++++++++++++++++++ 4 files changed, 110 insertions(+), 11 deletions(-) create mode 100644 tests/scheduler/test_zmq_links.py diff --git a/python/freetoken/scheduler/config.py b/python/freetoken/scheduler/config.py index b2bbe0aefe..c46e28cce2 100644 --- a/python/freetoken/scheduler/config.py +++ b/python/freetoken/scheduler/config.py @@ -5,10 +5,32 @@ from freetoken.engine import EngineConfig -def _get_pid_suffix() -> str: +ZMQ_LINK_COUNT = 5 + + +def _free_loopback_ports(count: int) -> list[int]: + import socket + + held = [socket.socket(socket.AF_INET, socket.SOCK_STREAM) for _ in range(count)] + try: + for sock in held: + sock.bind(("127.0.0.1", 0)) + return [sock.getsockname()[1] for sock in held] + finally: + for sock in held: + sock.close() + + +def _choose_zmq_links() -> tuple[str, ...]: + """The ZMQ endpoints, chosen once in the parent so every spawned worker reads the same ones; + loopback TCP ports on Windows, where libzmq has no ipc transport.""" import os - return f".pid={os.getpid()}" + import zmq + + if zmq.has("ipc"): + return tuple(f"ipc:///tmp/freetoken_{link}.pid={os.getpid()}" for link in range(ZMQ_LINK_COUNT)) + return tuple(f"tcp://127.0.0.1:{port}" for port in _free_loopback_ports(ZMQ_LINK_COUNT)) @dataclass(frozen=True) @@ -20,19 +42,19 @@ class SchedulerConfig(EngineConfig): special_token_ckpt: bool = False # networking config - _unique_suffix: str = field(default_factory=_get_pid_suffix) + _zmq_links: tuple[str, ...] = field(default_factory=_choose_zmq_links) @property def zmq_backend_addr(self) -> str: - return "ipc:///tmp/freetoken_0" + self._unique_suffix + return self._zmq_links[0] @property def zmq_detokenizer_addr(self) -> str: - return "ipc:///tmp/freetoken_1" + self._unique_suffix + return self._zmq_links[1] @property def zmq_scheduler_broadcast_addr(self) -> str: - return "ipc:///tmp/freetoken_2" + self._unique_suffix + return self._zmq_links[2] @property def max_forward_len(self) -> int: diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index ce04b34102..87d9cb21cf 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -7,6 +7,7 @@ import json import os import signal +import sys import threading import time import uuid @@ -61,6 +62,8 @@ # signal handler). The backend supervisor reads this so a worker exiting AS PART of that # shutdown is treated as expected — no ERROR log, no "failed" latch. See run_backend_supervisor. _SHUTTING_DOWN = threading.Event() +# zmq.asyncio needs add_reader, which the Proactor loop uvicorn picks on Windows does not implement. +UVICORN_LOOP = "asyncio:SelectorEventLoop" if sys.platform == "win32" else "auto" BACKEND_DEATH_EXIT_GRACE_S = 10.0 @@ -102,6 +105,18 @@ def _stop() -> None: return timer +def _listener_ended(task: asyncio.Task) -> None: + """No reply reaches any request once the tokenizer listener is gone, so its end takes the serve down.""" + if task.cancelled() or _SHUTTING_DOWN.is_set(): + return + error = task.exception() + logger.error("Frontend listener stopped: %r; no request can be answered", error) + if _GLOBAL_STATE is not None: + _GLOBAL_STATE.fatal_error = f"frontend listener stopped: {error!r}" + _GLOBAL_STATE.maintenance_state = "failed" + _exit_after_backend_death(0.0) + + def _reap_backend_workers(processes: List[Any], timeout: float = 5.0) -> None: """Wait out a preceding ``_terminate_backend_workers`` and SIGKILL whatever is still standing. Only the shell path needs this: it owns the process lifetime end to end (no @@ -193,6 +208,8 @@ class FrontendManager: # future resolution back onto the loop (asyncio Futures are not thread-safe). None until the # first send_one starts the listener. _loop: Any = None + # The listener task, held so its failure is observed rather than dropped with the task. + _listener: Any = None # Frontend-side tokenizer for /v1/messages/count_tokens, built lazily on the first count # and cached for the process lifetime. It is the SAME TokenizeManager the engine's # tokenizer worker runs (chat template + generation prompt, DSV4/GGUF handling), so a @@ -335,7 +352,8 @@ def _resolve_all() -> None: def _create_listener_once(self): if not self.initialized: self._loop = asyncio.get_running_loop() - asyncio.create_task(self.listen()) + self._listener = asyncio.create_task(self.listen()) + self._listener.add_done_callback(_listener_ended) self.initialized = True async def send_one(self, msg: BaseTokenizerMsg): @@ -991,7 +1009,7 @@ def _serve_and_run_shell(host: str, port: int, api_key: str | None) -> None: netloc = f"[{host}]:{port}" if ":" in host else f"{host}:{port}" origin = resolve_server_url(f"http://{netloc}").origin - server = uvicorn.Server(uvicorn.Config(app, host=host, port=port, access_log=False)) + server = uvicorn.Server(uvicorn.Config(app, host=host, port=port, access_log=False, loop=UVICORN_LOOP)) thread = threading.Thread(target=server.run, name="freetoken-uvicorn", daemon=True) thread.start() _install_shell_stop_handlers() @@ -1132,4 +1150,4 @@ def _on_meta(meta: dict) -> None: _serve_and_run_shell(host, port, config.api_key) return # uvicorn stays on the main thread (signal handling unchanged); ^C reaches the worker group. - uvicorn.run(app, host=host, port=port) + uvicorn.run(app, host=host, port=port, loop=UVICORN_LOOP) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 3d75143e2e..924d6eef08 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -84,13 +84,13 @@ def share_tokenizer(self) -> bool: @property def zmq_frontend_addr(self) -> str: - return "ipc:///tmp/freetoken_3" + self._unique_suffix + return self._zmq_links[3] @property def zmq_tokenizer_addr(self) -> str: if self.share_tokenizer: return self.zmq_detokenizer_addr - result = "ipc:///tmp/freetoken_4" + self._unique_suffix + result = self._zmq_links[4] assert result != self.zmq_detokenizer_addr return result diff --git a/tests/scheduler/test_zmq_links.py b/tests/scheduler/test_zmq_links.py new file mode 100644 index 0000000000..b26cbc58ff --- /dev/null +++ b/tests/scheduler/test_zmq_links.py @@ -0,0 +1,59 @@ +"""The ZMQ endpoints every worker process connects through: one set per server, on either transport.""" + +from __future__ import annotations + +import os +import pickle +import re + +import pytest +import torch +import zmq + +from freetoken.distributed import DistributedInfo +from freetoken.scheduler.config import ZMQ_LINK_COUNT, SchedulerConfig, _choose_zmq_links + +LOOPBACK = re.compile(r"^tcp://127\.0\.0\.1:(\d+)$") + + +@pytest.fixture(params=[True, False], ids=["ipc", "no-ipc"]) +def transport(request, monkeypatch) -> bool: + monkeypatch.setattr(zmq, "has", lambda capability: request.param if capability == "ipc" else True) + return request.param + + +def test_every_link_is_distinct(transport): + links = _choose_zmq_links() + assert len(links) == ZMQ_LINK_COUNT + assert len(set(links)) == ZMQ_LINK_COUNT + + +def test_without_ipc_every_link_is_a_loopback_port(monkeypatch): + monkeypatch.setattr(zmq, "has", lambda capability: capability != "ipc") + ports = [LOOPBACK.match(link) for link in _choose_zmq_links()] + assert all(ports), "a link that is not loopback TCP cannot bind where libzmq lacks ipc" + assert len({int(port.group(1)) for port in ports}) == ZMQ_LINK_COUNT + + +def test_with_ipc_the_links_are_per_process_socket_paths(monkeypatch): + monkeypatch.setattr(zmq, "has", lambda capability: True) + assert _choose_zmq_links() == tuple( + f"ipc:///tmp/freetoken_{link}.pid={os.getpid()}" for link in range(ZMQ_LINK_COUNT) + ) + + +def test_a_pickled_config_keeps_its_links(transport): + config = SchedulerConfig( + model_path="unused", tp_info=DistributedInfo(rank=0, size=1), dtype=torch.float16 + ) + spawned = pickle.loads(pickle.dumps(config)) + assert (spawned.zmq_backend_addr, spawned.zmq_detokenizer_addr, spawned.zmq_scheduler_broadcast_addr) == ( + config.zmq_backend_addr, + config.zmq_detokenizer_addr, + config.zmq_scheduler_broadcast_addr, + ) + + +def test_this_platform_picks_the_transport_libzmq_has(): + link = _choose_zmq_links()[0] + assert link.startswith("ipc://") == zmq.has("ipc") From 5a4ff5ffd7556ef39a9e59601c3f1c176f452c28 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Thu, 1 Oct 2026 06:25:17 +0300 Subject: [PATCH 2/7] test(scheduler): drop the ZMQ link tests the transport tests cover The two transport tests already assert five distinct links each, and the platform test restated the branch it ran through. --- tests/scheduler/test_zmq_links.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/tests/scheduler/test_zmq_links.py b/tests/scheduler/test_zmq_links.py index b26cbc58ff..ae4c2198df 100644 --- a/tests/scheduler/test_zmq_links.py +++ b/tests/scheduler/test_zmq_links.py @@ -22,12 +22,6 @@ def transport(request, monkeypatch) -> bool: return request.param -def test_every_link_is_distinct(transport): - links = _choose_zmq_links() - assert len(links) == ZMQ_LINK_COUNT - assert len(set(links)) == ZMQ_LINK_COUNT - - def test_without_ipc_every_link_is_a_loopback_port(monkeypatch): monkeypatch.setattr(zmq, "has", lambda capability: capability != "ipc") ports = [LOOPBACK.match(link) for link in _choose_zmq_links()] @@ -52,8 +46,3 @@ def test_a_pickled_config_keeps_its_links(transport): config.zmq_detokenizer_addr, config.zmq_scheduler_broadcast_addr, ) - - -def test_this_platform_picks_the_transport_libzmq_has(): - link = _choose_zmq_links()[0] - assert link.startswith("ipc://") == zmq.has("ipc") From b2d61d531f121bb979199b9ab7defd076a205bc9 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Thu, 1 Oct 2026 12:12:20 +0300 Subject: [PATCH 3/7] docs(windows): describe the ZMQ links by the capability that picks them --- python/freetoken/scheduler/config.py | 2 +- python/freetoken/server/api_server.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/freetoken/scheduler/config.py b/python/freetoken/scheduler/config.py index c46e28cce2..52e4540a6b 100644 --- a/python/freetoken/scheduler/config.py +++ b/python/freetoken/scheduler/config.py @@ -23,7 +23,7 @@ def _free_loopback_ports(count: int) -> list[int]: def _choose_zmq_links() -> tuple[str, ...]: """The ZMQ endpoints, chosen once in the parent so every spawned worker reads the same ones; - loopback TCP ports on Windows, where libzmq has no ipc transport.""" + loopback TCP ports where libzmq has no ipc transport, as on Windows.""" import os import zmq diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 87d9cb21cf..92e84cd61f 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -1034,7 +1034,7 @@ def run_api_server(config: ServerArgs, start_backend: Callable[[], "Any"], run_s Run the frontend API server (FastAPI + uvicorn) and wire it to the tokenizer process via ZMQ. Args: - config: Server configuration (host/port, ZMQ IPC addresses, etc). + config: Server configuration (host/port, ZMQ addresses, etc). start_backend: Callback that launches the backend worker processes (TP schedulers + tokenizer/detokenizer). run_shell: If True, also attach the interactive terminal shell to the served API. From 78d41befa0b6a7c1148df37f09d06b48eb27df00 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Fri, 2 Oct 2026 20:53:35 +0300 Subject: [PATCH 4/7] fix(windows): stop on backend death through the server's own SIGTERM handler --- python/freetoken/server/api_server.py | 4 +++- tests/server/test_supervisor.py | 24 ++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 92e84cd61f..f555baf31d 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -97,7 +97,9 @@ def _stop() -> None: if _SHUTTING_DOWN.is_set(): return # an external stop got here first logger.error("Backend worker is gone and cannot be restarted; stopping the API server") - os.kill(os.getpid(), signal.SIGTERM) + # Raised in-process so uvicorn's handler runs the lifespan shutdown that stops the workers; + # on Windows, os.kill on our own pid is TerminateProcess and runs no handler. + signal.raise_signal(signal.SIGTERM) timer = threading.Timer(grace_s, _stop) timer.daemon = True diff --git a/tests/server/test_supervisor.py b/tests/server/test_supervisor.py index 3568d90505..2a96e2722a 100644 --- a/tests/server/test_supervisor.py +++ b/tests/server/test_supervisor.py @@ -1,6 +1,8 @@ from __future__ import annotations import queue +import subprocess +import sys from queue import Empty as _Empty @@ -239,6 +241,28 @@ def is_alive(self) -> bool: assert "failure" not in seen # silenced: expected exit during shutdown +def test_backend_death_stops_the_server_through_its_own_signal_handler(): + """The stop has to run the server's SIGTERM handler: its lifespan shutdown is what terminates + the remaining workers. On Windows os.kill on our own pid is TerminateProcess, which runs no + handler and leaves the workers running.""" + # Python runs the handler on the main thread; the wakeup socket wakes it out of select when the + # signal arrives. + probe = ( + "import os, select, signal, socket\n" + "from freetoken.server.api_server import _exit_after_backend_death\n" + "signal.signal(signal.SIGTERM, lambda signum, frame: os._exit(0))\n" + "reader, writer = socket.socketpair()\n" + "writer.setblocking(False)\n" + "signal.set_wakeup_fd(writer.fileno())\n" + "_exit_after_backend_death(0.0)\n" + "if select.select([reader], [], [], 60)[0]:\n" + " reader.recv(1)\n" + "os._exit(1)\n" + ) + result = subprocess.run([sys.executable, "-c", probe], capture_output=True, text=True, timeout=300) + assert result.returncode == 0, result.stdout + result.stderr + + # --------------------------------------------------------------------------- # progress sink: the byte_bar -> set_progress_sink pipe drain_ready consumes, and the # phase_slug normalization that labels the three serve bars. From 6035cba9eb60ec466c7ad443b101590a93a31781 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Tue, 13 Oct 2026 12:43:48 +0300 Subject: [PATCH 5/7] fix(engine): set the allocator config through torch's non-deprecated binding --- python/freetoken/engine/engine.py | 3 ++- tests/engine/test_cache_budget.py | 12 ++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 5365cd659f..eb109550bb 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -1532,7 +1532,8 @@ def _ensure_expandable_segments() -> None: if os.environ.get("PYTORCH_ALLOC_CONF") or os.environ.get("PYTORCH_CUDA_ALLOC_CONF"): return try: - torch.cuda.memory._set_allocator_settings("expandable_segments:True") + # torch.cuda.memory._set_allocator_settings is a deprecated wrapper of this binding + torch._C._accelerator_setAllocatorSettings("expandable_segments:True") except Exception as exc: # pragma: no cover - depends on torch build logger.info_rank0(f"Could not enable expandable_segments ({exc}); continuing") return diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index 5c1f93b348..401d60234a 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -506,6 +506,18 @@ def test_adjust_config_rope_gate_exempts_dsv4(): _adjust_config(cfg) # must not raise +def test_expandable_segments_are_requested_without_a_deprecation_warning(monkeypatch, recwarn): + from freetoken.engine.engine import _ensure_expandable_segments + + settings: list[str] = [] + monkeypatch.setattr(torch._C, "_accelerator_setAllocatorSettings", settings.append) + monkeypatch.delenv("PYTORCH_ALLOC_CONF", raising=False) + monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF", raising=False) + _ensure_expandable_segments() + assert settings == ["expandable_segments:True"] + assert [str(warning.message) for warning in recwarn] == [] + + # ---- _pin_budget_bytes: host bytes already pinned outside the expert banks ---- From 760af395e811a7cd6929126fd7c9ad08bf745036 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Wed, 4 Nov 2026 13:10:03 +0200 Subject: [PATCH 6/7] fix(scheduler): bind the scheduler's links before the engine loads --- python/freetoken/scheduler/io.py | 8 ++++-- python/freetoken/scheduler/scheduler.py | 5 ++-- tests/scheduler/test_zmq_links.py | 37 +++++++++++++++++++++++++ 3 files changed, 45 insertions(+), 5 deletions(-) diff --git a/python/freetoken/scheduler/io.py b/python/freetoken/scheduler/io.py index 37557a97fc..e365fed447 100644 --- a/python/freetoken/scheduler/io.py +++ b/python/freetoken/scheduler/io.py @@ -22,11 +22,15 @@ class SchedulerIOMixin: receive_msg: Function to receive messages from the tokenizer. send_result: Function to send results back to the tokenizer. sync_all_ranks: Function to synchronize all ranks on CPU side. + + The links open before the engine loads and ``tp_cpu_group`` is set once the engine forms it: a port bound after + the load is free for all of it, and a peer that connects meanwhile can meet another socket there and give up. """ - def __init__(self, config: SchedulerConfig, tp_cpu_group: torch.distributed.ProcessGroup): + tp_cpu_group: torch.distributed.ProcessGroup + + def __init__(self, config: SchedulerConfig): tp_info = config.tp_info - self.tp_cpu_group: Final = tp_cpu_group if config.offline_mode: self.receive_msg = self.offline_receive_msg self.send_result = self.offline_send_result diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 4bdf2631d8..80d3bd7cac 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -68,7 +68,9 @@ class Scheduler(SchedulerIOMixin): def __init__(self, config: SchedulerConfig): from freetoken.engine import Engine + super().__init__(config) self.engine = Engine(config) + self.tp_cpu_group = self.engine.tp_cpu_group # use another stream to overlap metadata processing with computation self.device = self.engine.device @@ -159,9 +161,6 @@ def __init__(self, config: SchedulerConfig): moe_stats=self._moe_stats_line if config.moe_collect_stats and moe_cache is not None else None, ) - # Initialize the I/O mixin - super().__init__(config, self.engine.tp_cpu_group) - def _moe_stats_line(self) -> str | None: # on the engine stream, so the read and the zeroing fall between decode replays with torch.cuda.stream(self.engine.stream): diff --git a/tests/scheduler/test_zmq_links.py b/tests/scheduler/test_zmq_links.py index ae4c2198df..325d445c53 100644 --- a/tests/scheduler/test_zmq_links.py +++ b/tests/scheduler/test_zmq_links.py @@ -5,6 +5,7 @@ import os import pickle import re +import socket import pytest import torch @@ -46,3 +47,39 @@ def test_a_pickled_config_keeps_its_links(transport): config.zmq_detokenizer_addr, config.zmq_scheduler_broadcast_addr, ) + + +class _EngineLoading(Exception): + """Ends the scheduler's start inside the engine's load, once the load has looked at the links.""" + + +def test_the_scheduler_binds_its_links_before_the_engine_loads(monkeypatch): + # A port nothing holds for a whole load is a port a connecting worker can meet anything on, itself included, + # and a ZMQ connect that meets a socket of the wrong type disconnects and never retries. + import freetoken.engine + from freetoken.scheduler.scheduler import Scheduler + + monkeypatch.setattr(zmq, "has", lambda capability: capability != "ipc") + config = SchedulerConfig(model_path="unused", tp_info=DistributedInfo(rank=0, size=1), dtype=torch.float16) + held_during_the_load: list[bool] = [] + + class LoadingEngine: + def __init__(self, _config): + port = int(LOOPBACK.match(config.zmq_backend_addr).group(1)) + with socket.socket() as stranger: + try: + stranger.bind(("127.0.0.1", port)) + held_during_the_load.append(False) + except OSError: + held_during_the_load.append(True) + raise _EngineLoading + + monkeypatch.setattr(freetoken.engine, "Engine", LoadingEngine) + scheduler = Scheduler.__new__(Scheduler) + with pytest.raises(_EngineLoading): + scheduler.__init__(config) + for link in ("_recv_from_tokenizer", "_send_into_tokenizer"): + queue = getattr(scheduler, link, None) + if queue is not None: + queue.stop() + assert held_during_the_load == [True] From f75278c422da1b2e27172b99882a8a57df8d37e4 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Wed, 4 Nov 2026 16:06:20 +0200 Subject: [PATCH 7/7] fix(windows): bind the zmq links on 127.0.0.2, so no connect to one can be its own peer --- python/freetoken/scheduler/config.py | 7 +++-- tests/scheduler/test_zmq_links.py | 39 ++++++++++++++++++---------- 2 files changed, 30 insertions(+), 16 deletions(-) diff --git a/python/freetoken/scheduler/config.py b/python/freetoken/scheduler/config.py index 52e4540a6b..0805513435 100644 --- a/python/freetoken/scheduler/config.py +++ b/python/freetoken/scheduler/config.py @@ -6,6 +6,9 @@ ZMQ_LINK_COUNT = 5 +# A connect to 127.0.0.2 leaves from 127.0.0.1, so it can never be handed the address it connects to; on 127.0.0.1 +# a connect made before its link binds can get the link's own port, connect to itself, and ZMQ never retries that. +ZMQ_LINK_HOST = "127.0.0.2" def _free_loopback_ports(count: int) -> list[int]: @@ -14,7 +17,7 @@ def _free_loopback_ports(count: int) -> list[int]: held = [socket.socket(socket.AF_INET, socket.SOCK_STREAM) for _ in range(count)] try: for sock in held: - sock.bind(("127.0.0.1", 0)) + sock.bind((ZMQ_LINK_HOST, 0)) return [sock.getsockname()[1] for sock in held] finally: for sock in held: @@ -30,7 +33,7 @@ def _choose_zmq_links() -> tuple[str, ...]: if zmq.has("ipc"): return tuple(f"ipc:///tmp/freetoken_{link}.pid={os.getpid()}" for link in range(ZMQ_LINK_COUNT)) - return tuple(f"tcp://127.0.0.1:{port}" for port in _free_loopback_ports(ZMQ_LINK_COUNT)) + return tuple(f"tcp://{ZMQ_LINK_HOST}:{port}" for port in _free_loopback_ports(ZMQ_LINK_COUNT)) @dataclass(frozen=True) diff --git a/tests/scheduler/test_zmq_links.py b/tests/scheduler/test_zmq_links.py index 325d445c53..4363c72b48 100644 --- a/tests/scheduler/test_zmq_links.py +++ b/tests/scheduler/test_zmq_links.py @@ -14,7 +14,7 @@ from freetoken.distributed import DistributedInfo from freetoken.scheduler.config import ZMQ_LINK_COUNT, SchedulerConfig, _choose_zmq_links -LOOPBACK = re.compile(r"^tcp://127\.0\.0\.1:(\d+)$") +LOOPBACK = re.compile(r"^tcp://(?P127(?:\.\d{1,3}){3}):(?P\d+)$") @pytest.fixture(params=[True, False], ids=["ipc", "no-ipc"]) @@ -25,9 +25,21 @@ def transport(request, monkeypatch) -> bool: def test_without_ipc_every_link_is_a_loopback_port(monkeypatch): monkeypatch.setattr(zmq, "has", lambda capability: capability != "ipc") - ports = [LOOPBACK.match(link) for link in _choose_zmq_links()] - assert all(ports), "a link that is not loopback TCP cannot bind where libzmq lacks ipc" - assert len({int(port.group(1)) for port in ports}) == ZMQ_LINK_COUNT + links = [LOOPBACK.match(link) for link in _choose_zmq_links()] + assert all(links), "a link that is not loopback TCP cannot bind where libzmq lacks ipc" + assert len({int(link["port"]) for link in links}) == ZMQ_LINK_COUNT + + +def test_a_connect_to_a_link_leaves_from_another_address(monkeypatch): + # A connect made before its link binds can be handed the link's port as its own and connect to itself, a peer of + # its own socket type, which ZMQ never retries; a source address other than the link's rules that out. + monkeypatch.setattr(zmq, "has", lambda capability: capability != "ipc") + link = LOOPBACK.match(_choose_zmq_links()[0]) + with socket.socket() as listener, socket.socket() as client: + listener.bind((link["host"], int(link["port"]))) + listener.listen(1) + client.connect((link["host"], int(link["port"]))) + assert client.getsockname()[0] != link["host"] def test_with_ipc_the_links_are_per_process_socket_paths(monkeypatch): @@ -49,34 +61,33 @@ def test_a_pickled_config_keeps_its_links(transport): ) -class _EngineLoading(Exception): - """Ends the scheduler's start inside the engine's load, once the load has looked at the links.""" - - def test_the_scheduler_binds_its_links_before_the_engine_loads(monkeypatch): - # A port nothing holds for a whole load is a port a connecting worker can meet anything on, itself included, - # and a ZMQ connect that meets a socket of the wrong type disconnects and never retries. + # A port nothing holds for a whole load is one another program can take, and a ZMQ connect that meets a socket of + # the wrong type there never retries. import freetoken.engine from freetoken.scheduler.scheduler import Scheduler + class EngineLoading(Exception): + pass + monkeypatch.setattr(zmq, "has", lambda capability: capability != "ipc") config = SchedulerConfig(model_path="unused", tp_info=DistributedInfo(rank=0, size=1), dtype=torch.float16) + backend = LOOPBACK.match(config.zmq_backend_addr) held_during_the_load: list[bool] = [] class LoadingEngine: def __init__(self, _config): - port = int(LOOPBACK.match(config.zmq_backend_addr).group(1)) with socket.socket() as stranger: try: - stranger.bind(("127.0.0.1", port)) + stranger.bind((backend["host"], int(backend["port"]))) held_during_the_load.append(False) except OSError: held_during_the_load.append(True) - raise _EngineLoading + raise EngineLoading monkeypatch.setattr(freetoken.engine, "Engine", LoadingEngine) scheduler = Scheduler.__new__(Scheduler) - with pytest.raises(_EngineLoading): + with pytest.raises(EngineLoading): scheduler.__init__(config) for link in ("_recv_from_tokenizer", "_send_into_tokenizer"): queue = getattr(scheduler, link, None)