fix(windows): serve over loopback ZMQ and a selector event loop

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

python/freetoken/scheduler/config.py

@@ -5,10 +5,35 @@
 from freetoken.engine import EngineConfig
 
 
-def _get_pid_suffix() -> str:
+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]:
+    import socket
+
+    held = [socket.socket(socket.AF_INET, socket.SOCK_STREAM) for _ in range(count)]
+    try:
+        for sock in held:
+            sock.bind((ZMQ_LINK_HOST, 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 where libzmq has no ipc transport, as on Windows."""
     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://{ZMQ_LINK_HOST}:{port}" for port in _free_loopback_ports(ZMQ_LINK_COUNT))
 
 
 @dataclass(frozen=True)
@@ -20,19 +45,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:

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

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):

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
 
 
@@ -94,14 +97,28 @@ 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
     timer.start()
     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 +210,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 +354,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 +1011,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()
@@ -1016,7 +1036,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.
@@ -1132,4 +1152,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)

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
 

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

tests/scheduler/test_zmq_links.py

@@ -0,0 +1,96 @@
+"""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 socket
+
+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://(?P<host>127(?:\.\d{1,3}){3}):(?P<port>\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_without_ipc_every_link_is_a_loopback_port(monkeypatch):
+    monkeypatch.setattr(zmq, "has", lambda capability: capability != "ipc")
+    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):
+    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_the_scheduler_binds_its_links_before_the_engine_loads(monkeypatch):
+    # 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):
+            with socket.socket() as stranger:
+                try:
+                    stranger.bind((backend["host"], int(backend["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]

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.