From fc32af0e5c9ad2732c01a21403f8e9adea6752a4 Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Mon, 17 Aug 2026 01:18:20 -0500 Subject: [PATCH 1/2] fix(atom): relay KV events from data-parallel ranks Signed-off-by: Cam Quilici --- .../engine/atom/hooks/kv_event_bootstrap.py | 7 +- infera/engine/atom/hooks/kv_events.py | 56 +++++--- infera/engine/atom/kv_event_proxy.py | 99 ++++++++++++++ infera/engine/atom/worker.py | 88 ++++++------ tests/unit/engine/test_atom_kv_event_proxy.py | 128 ++++++++++++++++++ 5 files changed, 312 insertions(+), 66 deletions(-) create mode 100644 infera/engine/atom/kv_event_proxy.py create mode 100644 tests/unit/engine/test_atom_kv_event_proxy.py diff --git a/infera/engine/atom/hooks/kv_event_bootstrap.py b/infera/engine/atom/hooks/kv_event_bootstrap.py index 5c4d22b3..c75e65b2 100644 --- a/infera/engine/atom/hooks/kv_event_bootstrap.py +++ b/infera/engine/atom/hooks/kv_event_bootstrap.py @@ -14,8 +14,8 @@ The work is gated on ``INFERA_ATOM_KV_EVENTS_ENDPOINT`` so that unrelated Python invocations (pip, debug shells, the launcher itself before it spawns -ATOM) pay nothing. The infera ATOM launcher sets that env var (to the ZMQ -bind address) only for the ATOM subprocess it spawns. +ATOM) pay nothing. The Infera ATOM launcher sets that env var to the local +relay ingress and enables connect mode only for the ATOM subprocess it spawns. """ from __future__ import annotations @@ -30,7 +30,8 @@ def _activate() -> None: try: from infera.engine.atom.hooks.kv_events import arm_kv_event_hooks - arm_kv_event_hooks(endpoint) + connect = os.environ.get("INFERA_ATOM_KV_EVENTS_CONNECT") == "1" + arm_kv_event_hooks(endpoint, connect=connect) except Exception: # A sitecustomize/.pth hook must never take down an unrelated # interpreter; surface the traceback and continue. diff --git a/infera/engine/atom/hooks/kv_events.py b/infera/engine/atom/hooks/kv_events.py index a0b406fa..ad827ce8 100644 --- a/infera/engine/atom/hooks/kv_events.py +++ b/infera/engine/atom/hooks/kv_events.py @@ -23,11 +23,11 @@ change is needed — ATOM workers light up KV-aware routing exactly like vLLM/SGLang ones. -The hooks are installed inside the ATOM ``EngineCore`` subprocess (which is -where ``BlockManager`` actually lives) via -:mod:`infera.engine.atom.hooks.kv_event_bootstrap`. The PUB socket binds -lazily on the first ``BlockManager`` instantiation, so only the process that -owns the block manager binds the port. +The hooks are installed inside every ATOM ``EngineCore`` subprocess (where +each data-parallel rank owns a ``BlockManager``) via +:mod:`infera.engine.atom.hooks.kv_event_bootstrap`. Each PUB socket connects +to the worker-local relay, which multiplexes all ranks onto the single endpoint +advertised for the logical ATOM worker. """ from __future__ import annotations @@ -101,15 +101,16 @@ def _check_block_manager_compat(BlockManager) -> str | None: class _Publisher: """Thin ZMQ PUB wrapper that encodes batches with the router's msgspec - structs (guaranteeing byte-for-byte wire compatibility) and binds lazily. + structs (guaranteeing byte-for-byte wire compatibility) and opens lazily. All publishing happens from the single EngineCore scheduler thread that drives ``BlockManager``, so a plain (non-async, non-thread-safe) PUB socket is safe. """ - def __init__(self, bind_endpoint: str) -> None: - self._bind_endpoint = bind_endpoint + def __init__(self, endpoint: str, *, connect: bool = False) -> None: + self._endpoint = endpoint + self._connect = connect self._sock = None self._encoder = None self._events_mod = None @@ -130,9 +131,16 @@ def ensure_bound(self) -> None: # the router rebuilds its view from a fresh subscription if it ever # falls behind; we never want a slow subscriber to block the engine. sock.setsockopt(zmq.SNDHWM, 100_000) - sock.bind(self._bind_endpoint) + if self._connect: + sock.connect(self._endpoint) + else: + sock.bind(self._endpoint) self._sock = sock - logger.info("ATOM kv-events: PUB bound at %s", self._bind_endpoint) + logger.info( + "ATOM kv-events: PUB %s %s", + "connected to" if self._connect else "bound at", + self._endpoint, + ) def _send_batch(self, events: list) -> None: if self._sock is None or not events: @@ -174,9 +182,9 @@ def block_removed(self, block_hash: int) -> None: def _patch_block_manager(BlockManager, publisher: _Publisher) -> None: """Wrap ``BlockManager`` methods to publish KV cache events. - Idempotent. The PUB socket only binds when a ``BlockManager`` is actually - instantiated (i.e. inside the EngineCore subprocess), so this is safe to - apply in any process that imports the class. + Idempotent. The PUB socket is only opened when a ``BlockManager`` is + actually instantiated (i.e. inside an EngineCore subprocess), so this is + safe to apply in any process that imports the class. """ if getattr(BlockManager, "_infera_kv_patched", False): return @@ -285,7 +293,7 @@ def patched_allocate_block(self, block_id): BlockManager.hash_blocks = patched_hash_blocks BlockManager._allocate_block = patched_allocate_block BlockManager._infera_kv_patched = True - logger.info("ATOM kv-events: BlockManager hooks installed (bind=%s)", publisher._bind_endpoint) + logger.info("ATOM kv-events: BlockManager hooks installed (endpoint=%s)", publisher._endpoint) class _DeferredPatchFinder: @@ -300,8 +308,9 @@ class _DeferredPatchFinder: point ``site`` init is long done. """ - def __init__(self, endpoint: str) -> None: + def __init__(self, endpoint: str, *, connect: bool = False) -> None: self._endpoint = endpoint + self._connect = connect self._done = False def find_spec(self, fullname, path, target=None): @@ -314,13 +323,14 @@ def find_spec(self, fullname, path, target=None): return None self._done = True endpoint = self._endpoint + connect = self._connect loader = spec.loader orig_exec = loader.exec_module - def exec_module(module, _orig=orig_exec, _ep=endpoint): + def exec_module(module, _orig=orig_exec, _ep=endpoint, _connect=connect): _orig(module) try: - _patch_block_manager(module.BlockManager, _Publisher(_ep)) + _patch_block_manager(module.BlockManager, _Publisher(_ep, connect=_connect)) except Exception: traceback.print_exc() @@ -328,7 +338,7 @@ def exec_module(module, _orig=orig_exec, _ep=endpoint): return spec -def arm_kv_event_hooks(bind_endpoint: str) -> None: +def arm_kv_event_hooks(endpoint: str, *, connect: bool = False) -> None: """Arrange for ATOM's ``BlockManager`` to publish KV cache events. If the block-manager module is already imported, patch immediately; @@ -338,9 +348,13 @@ def arm_kv_event_hooks(bind_endpoint: str) -> None: """ mod = sys.modules.get(_BLOCK_MANAGER_MODULE) if mod is not None: - _patch_block_manager(mod.BlockManager, _Publisher(bind_endpoint)) + _patch_block_manager(mod.BlockManager, _Publisher(endpoint, connect=connect)) return if any(isinstance(f, _DeferredPatchFinder) for f in sys.meta_path): return - sys.meta_path.insert(0, _DeferredPatchFinder(bind_endpoint)) - logger.info("ATOM kv-events: armed deferred BlockManager patch (bind=%s)", bind_endpoint) + sys.meta_path.insert(0, _DeferredPatchFinder(endpoint, connect=connect)) + logger.info( + "ATOM kv-events: armed deferred BlockManager patch (endpoint=%s connect=%s)", + endpoint, + connect, + ) diff --git a/infera/engine/atom/kv_event_proxy.py b/infera/engine/atom/kv_event_proxy.py new file mode 100644 index 00000000..8b7f219b --- /dev/null +++ b/infera/engine/atom/kv_event_proxy.py @@ -0,0 +1,99 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Process-safe relay for ATOM KV-cache events. + +ATOM creates one ``EngineCore`` process per data-parallel rank. Each process +owns an independent ``BlockManager`` and therefore needs to publish its own KV +events, while Infera advertises one logical KV-event endpoint per ATOM worker. +The relay gives all EngineCore publishers a local XSUB ingress and forwards +their streams through one externally advertised XPUB endpoint. +""" + +from __future__ import annotations + +import logging +import os +import threading +import uuid + +import zmq + +logger = logging.getLogger(__name__) + + +class AtomKvEventProxy: + """Relay multiple child PUB sockets through one advertised endpoint.""" + + def __init__(self, external_bind_endpoint: str) -> None: + suffix = f"{os.getpid()}-{uuid.uuid4().hex[:8]}" + self.external_bind_endpoint = external_bind_endpoint + self.ingress_endpoint = f"ipc:///tmp/infera-atom-kv-{suffix}.sock" + self._control_endpoint = f"inproc://infera-atom-kv-control-{suffix}" + self._ready = threading.Event() + self._error: BaseException | None = None + self._thread: threading.Thread | None = None + self._controller: zmq.Socket | None = None + + def start(self, timeout: float = 10.0) -> None: + if self._thread is not None: + return + self._thread = threading.Thread( + target=self._run, + name="atom-kv-event-proxy", + daemon=True, + ) + self._thread.start() + if not self._ready.wait(timeout): + raise TimeoutError("ATOM KV-event proxy did not start") + if self._error is not None: + raise RuntimeError("ATOM KV-event proxy failed to start") from self._error + + controller = zmq.Context.instance().socket(zmq.PAIR) + controller.setsockopt(zmq.LINGER, 0) + controller.connect(self._control_endpoint) + self._controller = controller + + def stop(self, timeout: float = 5.0) -> None: + if self._thread is None: + return + if self._controller is not None: + try: + self._controller.send(b"TERMINATE") + except zmq.ZMQError: + logger.exception("failed to stop ATOM KV-event proxy cleanly") + self._controller.close(linger=0) + self._controller = None + self._thread.join(timeout) + if self._thread.is_alive(): + logger.error("ATOM KV-event proxy did not stop within %.1fs", timeout) + self._thread = None + + def _run(self) -> None: + context = zmq.Context.instance() + ingress = context.socket(zmq.XSUB) + egress = context.socket(zmq.XPUB) + control = context.socket(zmq.PAIR) + for socket in (ingress, egress, control): + socket.setsockopt(zmq.LINGER, 0) + try: + ingress.bind(self.ingress_endpoint) + egress.bind(self.external_bind_endpoint) + control.bind(self._control_endpoint) + logger.info( + "ATOM kv-events: proxy ingress=%s advertise-bind=%s", + self.ingress_endpoint, + self.external_bind_endpoint, + ) + self._ready.set() + zmq.proxy_steerable(ingress, egress, None, control) + except BaseException as exc: + self._error = exc + self._ready.set() + logger.exception("ATOM KV-event proxy failed") + finally: + ingress.close(linger=0) + egress.close(linger=0) + control.close(linger=0) diff --git a/infera/engine/atom/worker.py b/infera/engine/atom/worker.py index bc45b033..bd609d73 100644 --- a/infera/engine/atom/worker.py +++ b/infera/engine/atom/worker.py @@ -11,14 +11,13 @@ :class:`EngineConfig` back to the launcher for etcd registration. KV-aware routing is **off by default**. ATOM has no native KV-event stream, -so when the operator opts in with ``--enable-kv-events`` infera installs a +so when the operator opts in with ``--enable-kv-events`` Infera installs a BlockManager hook (see :mod:`infera.engine.atom.hooks.kv_events`) inside the -ATOM subprocess that republishes ATOM's prefix-cache index on a ZMQ PUB -socket in the router's wire format: the launcher passes the bind endpoint to -the subprocess via ``INFERA_ATOM_KV_EVENTS_ENDPOINT`` and advertises the -reachable endpoint + block size in :class:`EngineConfig`. When not enabled -(the default) the config carries no ``kv_events_endpoint`` and the router -routes the worker round-robin. +ATOM subprocess that republishes ATOM's prefix-cache index on ZMQ PUB sockets +in the router's wire format. A worker-local proxy combines every data-parallel +rank into the one reachable endpoint advertised in :class:`EngineConfig`. +When not enabled (the default), the config carries no ``kv_events_endpoint`` +and the router routes the worker round-robin. """ from __future__ import annotations @@ -34,6 +33,7 @@ import httpx from infera.common.worker_pool import DisaggMode, EngineType +from infera.engine.atom.kv_event_proxy import AtomKvEventProxy from infera.engine.base import BaseEngine, EngineConfig logger = logging.getLogger(__name__) @@ -60,13 +60,15 @@ def __init__( self.port = port self.advertise_host = advertise_host or host # ``kv_events_endpoint`` is what the router connects to (advertised); - # ``kv_events_bind`` is what the ATOM subprocess binds (``tcp://*:port``). + # ``kv_events_bind`` is bound by a local relay. Every data-parallel + # EngineCore connects its own publisher to the relay ingress. self.kv_events_endpoint = kv_events_endpoint self.kv_events_bind = kv_events_bind self.kv_block_size = kv_block_size self.disagg_mode = disagg_mode self.disagg_meta: dict[str, Any] = dict(disagg_meta or {}) self._proc: subprocess.Popen | None = None + self._kv_event_proxy: AtomKvEventProxy | None = None async def start(self) -> EngineConfig: cmd = [ @@ -77,15 +79,19 @@ async def start(self) -> EngineConfig: ] logger.info("spawning atom subprocess: %s", " ".join(cmd)) env = os.environ.copy() - # Hand the ZMQ bind endpoint to the ATOM subprocess. The site-startup - # hook (infera.engine.atom.hooks.kv_event_bootstrap, run via a .pth in - # every interpreter — including ATOM's spawned EngineCore) reads this - # and installs the BlockManager KV-event publisher. + # A TP+DP ATOM worker owns one BlockManager per EngineCore process. + # Relay all of their PUB streams through the one endpoint advertised + # for this logical worker; otherwise every DP rank races to bind the + # same TCP port and all but one EngineCore fail during startup. if self.kv_events_bind: - env["INFERA_ATOM_KV_EVENTS_ENDPOINT"] = self.kv_events_bind + self._kv_event_proxy = AtomKvEventProxy(self.kv_events_bind) + self._kv_event_proxy.start() + env["INFERA_ATOM_KV_EVENTS_ENDPOINT"] = self._kv_event_proxy.ingress_endpoint + env["INFERA_ATOM_KV_EVENTS_CONNECT"] = "1" logger.info( - "ATOM kv-events enabled: bind=%s advertise=%s block_size=%s", + "ATOM kv-events enabled: relay=%s ingress=%s advertise=%s block_size=%s", self.kv_events_bind, + self._kv_event_proxy.ingress_endpoint, self.kv_events_endpoint, self.kv_block_size, ) @@ -114,36 +120,34 @@ async def start(self) -> EngineConfig: async def stop(self) -> None: logger.info("ATOM engine stopping") - if self._proc is None: - return - # The subprocess is a session leader (start_new_session=True), so its - # pid doubles as the process-group id for the whole ATOM tree. - pgid = self._proc.pid - if self._proc.poll() is None: + if self._proc is not None: + # The subprocess is a session leader (start_new_session=True), so + # its pid doubles as the process-group id for the whole ATOM tree. + pgid = self._proc.pid + if self._proc.poll() is None: + try: + os.killpg(pgid, signal.SIGTERM) + except ProcessLookupError: + pass + else: + # Graceful window kept short so we always reach SIGKILL + # within the launcher's own teardown budget. + for _ in range(15): + if self._proc.poll() is not None: + break + await asyncio.sleep(1) + # ATOM/AITER helpers can outlive the leader and retain GPU VRAM. try: - os.killpg(pgid, signal.SIGTERM) + os.killpg(pgid, signal.SIGKILL) except ProcessLookupError: - return - # Graceful window kept short so we always reach the SIGKILL sweep - # within the launcher's own teardown budget (the parent harness - # SIGKILLs this process group ~25s after SIGTERM). - for _ in range(15): - if self._proc.poll() is not None: - break - await asyncio.sleep(1) - # The leader (ATOM's openai_server) exiting does NOT mean the whole - # group is gone: ATOM/AITER helpers — e.g. the shared-memory broadcast - # worker — can hang and ignore SIGTERM, holding GPU VRAM until the - # container dies and starving the next worker (HIP-OOM). Always sweep - # the group with SIGKILL so VRAM is released on teardown. - try: - os.killpg(pgid, signal.SIGKILL) - except ProcessLookupError: - pass - try: - self._proc.wait(timeout=10) - except Exception: - pass + pass + try: + self._proc.wait(timeout=10) + except Exception: + pass + if self._kv_event_proxy is not None: + self._kv_event_proxy.stop() + self._kv_event_proxy = None async def _wait_ready(self, timeout: float | None = None) -> None: if timeout is None: diff --git a/tests/unit/engine/test_atom_kv_event_proxy.py b/tests/unit/engine/test_atom_kv_event_proxy.py new file mode 100644 index 00000000..11dadd10 --- /dev/null +++ b/tests/unit/engine/test_atom_kv_event_proxy.py @@ -0,0 +1,128 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### + +from __future__ import annotations + +import socket +import time +from unittest.mock import AsyncMock + +import msgspec +import pytest +import zmq + +from infera.engine.atom.hooks.kv_events import _Publisher +from infera.engine.atom.kv_event_proxy import AtomKvEventProxy +from infera.engine.atom.worker import AtomEngine +from infera.router.kv_event.events import BlockRemoved, KVEventBatch + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def test_proxy_relays_events_from_multiple_engine_core_publishers() -> None: + """Every DP rank may publish without racing to bind the worker endpoint.""" + context = zmq.Context.instance() + external = f"tcp://127.0.0.1:{_free_port()}" + proxy = AtomKvEventProxy(external) + subscriber = context.socket(zmq.SUB) + subscriber.setsockopt(zmq.SUBSCRIBE, b"kv-events") + subscriber.setsockopt(zmq.RCVTIMEO, 100) + publishers = [ + _Publisher(proxy.ingress_endpoint, connect=True), + _Publisher(proxy.ingress_endpoint, connect=True), + ] + try: + proxy.start() + subscriber.connect(external) + for publisher in publishers: + publisher.ensure_bound() + + observed: set[int] = set() + deadline = time.monotonic() + 3 + while time.monotonic() < deadline and observed != {101, 202}: + publishers[0].block_removed(101) + publishers[1].block_removed(202) + try: + topic, payload = subscriber.recv_multipart() + except zmq.Again: + continue + assert topic == b"kv-events" + batch = msgspec.msgpack.decode(payload, type=KVEventBatch) + assert isinstance(batch.events[0], BlockRemoved) + observed.add(batch.events[0].block_hashes[0]) + + assert observed == {101, 202} + finally: + for publisher in publishers: + if publisher._sock is not None: + publisher._sock.close(linger=0) + subscriber.close(linger=0) + proxy.stop() + + +@pytest.mark.asyncio +async def test_atom_engine_connects_children_to_worker_proxy(monkeypatch) -> None: + captured: dict[str, object] = {} + + class FakeProxy: + ingress_endpoint = "ipc:///tmp/infera-atom-kv-test.sock" + + def __init__(self, endpoint: str) -> None: + captured["proxy_endpoint"] = endpoint + + def start(self) -> None: + captured["proxy_started"] = True + + def stop(self) -> None: + captured["proxy_stopped"] = True + + class FakeProcess: + pid = 12345 + returncode = 0 + + def poll(self) -> int: + return 0 + + def wait(self, timeout: float) -> None: + captured["wait_timeout"] = timeout + + def fake_popen(cmd, *, env, start_new_session, stdout, stderr): + captured["cmd"] = cmd + captured["env"] = env + captured["start_new_session"] = start_new_session + return FakeProcess() + + monkeypatch.setattr("infera.engine.atom.worker.AtomKvEventProxy", FakeProxy) + monkeypatch.setattr("infera.engine.atom.worker.subprocess.Popen", fake_popen) + monkeypatch.setattr("infera.engine.atom.worker.os.killpg", lambda *_: None) + + engine = AtomEngine( + atom_argv=["--tensor-parallel-size", "8"], + model_name="test/model", + host="0.0.0.0", + port=6100, + advertise_host="10.0.0.1", + kv_events_endpoint="tcp://10.0.0.1:7000", + kv_events_bind="tcp://*:7000", + kv_block_size=16, + ) + engine._wait_ready = AsyncMock() + + config = await engine.start() + env = captured["env"] + assert isinstance(env, dict) + assert captured["proxy_endpoint"] == "tcp://*:7000" + assert captured["proxy_started"] is True + assert env["INFERA_ATOM_KV_EVENTS_ENDPOINT"] == FakeProxy.ingress_endpoint + assert env["INFERA_ATOM_KV_EVENTS_CONNECT"] == "1" + assert config.kv_events_endpoint == "tcp://10.0.0.1:7000" + + await engine.stop() + assert captured["proxy_stopped"] is True From 503dcbd1b60229a5a508fd3da45994c56a2b40ec Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Mon, 17 Aug 2026 01:22:54 -0500 Subject: [PATCH 2/2] refactor(atom): preserve engine teardown semantics Signed-off-by: Cam Quilici --- infera/engine/atom/worker.py | 33 ++++++++++++++++++++------------- 1 file changed, 20 insertions(+), 13 deletions(-) diff --git a/infera/engine/atom/worker.py b/infera/engine/atom/worker.py index bd609d73..0ca379a2 100644 --- a/infera/engine/atom/worker.py +++ b/infera/engine/atom/worker.py @@ -120,7 +120,9 @@ async def start(self) -> EngineConfig: async def stop(self) -> None: logger.info("ATOM engine stopping") - if self._proc is not None: + try: + if self._proc is None: + return # The subprocess is a session leader (start_new_session=True), so # its pid doubles as the process-group id for the whole ATOM tree. pgid = self._proc.pid @@ -128,15 +130,19 @@ async def stop(self) -> None: try: os.killpg(pgid, signal.SIGTERM) except ProcessLookupError: - pass - else: - # Graceful window kept short so we always reach SIGKILL - # within the launcher's own teardown budget. - for _ in range(15): - if self._proc.poll() is not None: - break - await asyncio.sleep(1) - # ATOM/AITER helpers can outlive the leader and retain GPU VRAM. + return + # Graceful window kept short so we always reach the SIGKILL + # sweep within the launcher's own teardown budget (the parent + # harness SIGKILLs this process group ~25s after SIGTERM). + for _ in range(15): + if self._proc.poll() is not None: + break + await asyncio.sleep(1) + # The leader (ATOM's openai_server) exiting does NOT mean the whole + # group is gone: ATOM/AITER helpers — e.g. the shared-memory + # broadcast worker — can hang and ignore SIGTERM, holding GPU VRAM + # until the container dies and starving the next worker (HIP-OOM). + # Always sweep the group with SIGKILL so VRAM is released. try: os.killpg(pgid, signal.SIGKILL) except ProcessLookupError: @@ -145,9 +151,10 @@ async def stop(self) -> None: self._proc.wait(timeout=10) except Exception: pass - if self._kv_event_proxy is not None: - self._kv_event_proxy.stop() - self._kv_event_proxy = None + finally: + if self._kv_event_proxy is not None: + self._kv_event_proxy.stop() + self._kv_event_proxy = None async def _wait_ready(self, timeout: float | None = None) -> None: if timeout is None: