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/python/freetoken/scheduler/config.py b/python/freetoken/scheduler/config.py index b2bbe0aefe..0805513435 100644 --- a/python/freetoken/scheduler/config.py +++ b/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: 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/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index ce04b34102..f555baf31d 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 @@ -94,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 @@ -102,6 +107,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 +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) 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/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 ---- diff --git a/tests/scheduler/test_zmq_links.py b/tests/scheduler/test_zmq_links.py new file mode 100644 index 0000000000..4363c72b48 --- /dev/null +++ b/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://(?P127(?:\.\d{1,3}){3}):(?P\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] 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.