From ddb8a86ec501278e394d3b86a0b62a19fcb3e92e Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Wed, 15 Jul 2026 00:27:20 +0000 Subject: [PATCH 01/21] communication: Rust transport (vendored rust/) + RustZMQCommunicator (RFC #130 Step 1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Vendors the mstar-rs ZMQ transport into rust/ (communicator.rs: opaque byte frames over PUSH/PULL, ipc + tcp endpoints, wakeup-fd polling, and a typed Codec layer as the pickle->msgpack migration seam) with a PyO3 module (mstar_rust; maturin develop in rust/) — no external dependency. RustZMQCommunicator is a drop-in for ZMQCommunicator over that module: same constructor/methods/endpoints (ipc prefix + the MSTAR_ZMQ_* TCP scheme), wire-compatible with unwrapped pyzmq entities in both directions (pickle-default codec, no added framing), eventfd wakeup forwarded to the Rust poller with unchanged drain semantics, and a buffering bridge so a readiness poll never drops or reorders a frame. test/modular/test_rust_communicator.py (skips unless mstar_rust is built): pickle interop both directions on one mesh; an EventWakeup fire cuts wait_for_work to ~50 ms against a 2 s timeout; FIFO-lossless readiness polls. All 3 pass against the vendored extension. --- mstar/communication/rust_communicator.py | 150 ++++++ rust/.gitignore | 1 + rust/Cargo.lock | 606 ++++++++++++++++++++++ rust/Cargo.toml | 23 + rust/pyproject.toml | 13 + rust/src/communicator.rs | 630 +++++++++++++++++++++++ rust/src/lib.rs | 95 ++++ test/modular/test_rust_communicator.py | 61 +++ 8 files changed, 1579 insertions(+) create mode 100644 mstar/communication/rust_communicator.py create mode 100644 rust/.gitignore create mode 100644 rust/Cargo.lock create mode 100644 rust/Cargo.toml create mode 100644 rust/pyproject.toml create mode 100644 rust/src/communicator.rs create mode 100644 rust/src/lib.rs create mode 100644 test/modular/test_rust_communicator.py diff --git a/mstar/communication/rust_communicator.py b/mstar/communication/rust_communicator.py new file mode 100644 index 000000000..cc6d22f90 --- /dev/null +++ b/mstar/communication/rust_communicator.py @@ -0,0 +1,150 @@ +"""DRAFT (RFC #130 Step 1): ``ZMQCommunicator`` as a thin wrapper over the +mstar-rs Rust communicator (``mstar_rs._core.ZmqCommunicator``). + +Drop-in for :class:`mstar.communication.communicator.ZMQCommunicator` — same +constructor, same methods, same semantics — with the transport moved to Rust: + +* **Wire-compatible with the existing pyzmq communicator.** The Rust transport + moves opaque byte frames (no added framing) over the same endpoints + (``ipc://{prefix}/{id}.ipc`` / the same TCP host+port scheme), and the + default codec is pickle — so a wrapped entity talks to unwrapped pyzmq + entities in both directions. Migration can proceed one process at a time. +* **encode/decode seam.** ``codec`` is a ``(dumps, loads)`` pair defaulting to + pickle (today's wire). Swapping to msgpack later is a codec change on both + ends of an edge, never a transport change. +* **eventfd wakeup.** ``register_event_for_poll`` forwards the ``EventWakeup`` + fd to the Rust poller (``register_wakeup_fd``): a completed compute future + wakes ``wait_for_work`` / ``poll_for_messages`` immediately, not on the poll + timeout. Drain semantics are identical (the event is drained here, in + Python, exactly as before). +* **Readiness without consumption.** The pyzmq poller reports "a message is + readable" without receiving it; the Rust ``recv_or_wake`` consumes. The + wrapper bridges the two with an internal deque: a message consumed during a + poll is buffered and handed out by the next ``get_all_new_messages`` — no + message is ever dropped or reordered. + +Requires the vendored extension: ``maturin develop`` (or ``pip install .``) in ``rust/``. +""" + +from __future__ import annotations + +import logging +import os +import pickle +from collections import deque + +from mstar_rust import ZmqCommunicator as _RustZmq + +from mstar.communication.communicator import BaseCommunicator, CommProtocol +from mstar.communication.event import EventWakeup + +logger = logging.getLogger(__name__) + +#: The encode/decode seam. Pickle matches today's ``send_pyobj`` wire, so a +#: wrapped entity interoperates with unwrapped pyzmq entities. Migrate an edge +#: to msgpack by giving both endpoints a msgpack codec. +PickleCodec = (pickle.dumps, pickle.loads) + + +class RustZMQCommunicator(BaseCommunicator): + """The pyzmq ``ZMQCommunicator`` surface over the Rust transport.""" + + def __init__( + self, + my_id: str, + push_ids: list[str], + protocol: CommProtocol = CommProtocol.IPC, + ipc_socket_path_prefix: str = "/tmp/mstar/", + codec=PickleCodec, + ): + transport = os.getenv("MSTAR_ZMQ_TRANSPORT", protocol.value).upper() + self.protocol = CommProtocol(transport) + self.my_id = my_id + self.ipc_socket_path_prefix = ipc_socket_path_prefix + self._dumps, self._loads = codec + self.event: EventWakeup | None = None + # Messages consumed by a readiness poll, awaiting get_all_new_messages. + self._buffered: deque = deque() + + if self.protocol == CommProtocol.IPC: + os.makedirs(ipc_socket_path_prefix, exist_ok=True) + self._inner = _RustZmq(my_id, ipc_socket_path_prefix) + elif self.protocol == CommProtocol.TCP: + self._inner = _RustZmq.bind_endpoint(my_id, self._endpoint(my_id)) + # TCP has no directory scheme: register every peer explicitly + # (lazily extended in send() for peers not known up front). + self._registered: set[str] = set() + for peer in push_ids: + if peer != my_id: + self._register(peer) + else: + raise NotImplementedError(f"Protocol {protocol} not yet supported yet") + + # -- endpoint scheme (identical to the pyzmq communicator's) ------------ + + def _endpoint(self, entity_id: str) -> str: + if self.protocol == CommProtocol.IPC: + return f"ipc://{self.ipc_socket_path_prefix}/{entity_id}.ipc" + host = os.getenv("MSTAR_ZMQ_TCP_HOST", "127.0.0.1") + return f"tcp://{host}:{self._tcp_port(entity_id)}" + + @staticmethod + def _tcp_port(entity_id: str) -> int: + base_port = int(os.getenv("MSTAR_ZMQ_TCP_BASE_PORT", "19000")) + if entity_id == "api_server": + return base_port + if entity_id == "conductor": + return base_port + 1 + if entity_id == "api_server_preprocess_worker": + return base_port + 2 + if entity_id.startswith("worker_"): + rank = entity_id.removeprefix("worker_") + if rank.isdigit(): + return base_port + 100 + int(rank) + return base_port + 1000 + (sum(entity_id.encode("utf-8")) % 1000) + + def _register(self, entity_id: str) -> None: + if entity_id not in self._registered: + self._inner.register_peer(entity_id, self._endpoint(entity_id)) + self._registered.add(entity_id) + + # -- the wrapped surface ------------------------------------------------- + + def register_event_for_poll(self, event: EventWakeup) -> None: + self._inner.register_wakeup_fd(event.fd) + self.event = event + + def _poll_once(self, timeout_ms: int) -> None: + """One wake-aware poll. A consumed message goes to the buffer; a + wakeup drains the event (same place the pyzmq path drains it).""" + kind, payload = self._inner.recv_or_wake(timeout_ms) + if kind == "msg": + self._buffered.append(payload) + elif kind == "wake" and self.event is not None: + self.event.drain() + + def wait_for_work(self, timeout_ms: int = 50) -> None: + if self._buffered: + return # work is already waiting + self._poll_once(timeout_ms) + + def poll_for_messages(self, timeout_ms: int = 20) -> bool: + """Block up to ``timeout_ms`` for a readable message; True when one is + available (buffered here, delivered by ``get_all_new_messages``).""" + if self._buffered: + return True + self._poll_once(timeout_ms) + return bool(self._buffered) + + def send(self, entity_id: str, msg) -> None: + logger.debug("%s to send a message %s to entity %s", self.my_id, str(msg), entity_id) + if self.protocol == CommProtocol.TCP: + self._register(entity_id) + self._inner.send(entity_id, self._dumps(msg)) + + def get_all_new_messages(self, blocking: bool = False) -> list: + messages = [self._loads(b) for b in self._buffered] + self._buffered.clear() + while (b := self._inner.try_recv()) is not None: + messages.append(self._loads(b)) + return messages diff --git a/rust/.gitignore b/rust/.gitignore new file mode 100644 index 000000000..2f7896d1d --- /dev/null +++ b/rust/.gitignore @@ -0,0 +1 @@ +target/ diff --git a/rust/Cargo.lock b/rust/Cargo.lock new file mode 100644 index 000000000..1a9ffae7d --- /dev/null +++ b/rust/Cargo.lock @@ -0,0 +1,606 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "bincode" +version = "1.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad" +dependencies = [ + "serde", +] + +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + +[[package]] +name = "cc" +version = "1.2.67" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17dd265a7d0f31ef544e1b20e03add05d3b45b491b633b10d67145d2acc1a38" +dependencies = [ + "find-msvc-tools", + "jobserver", + "libc", + "shlex", +] + +[[package]] +name = "cfg-expr" +version = "0.15.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d067ad48b8650848b989a59a86c6c36a995d02d2bf778d45c3c5d57bc2718f02" +dependencies = [ + "smallvec", + "target-lexicon", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "crossbeam" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1137cd7e7fc0fb5d3c5a8678be38ec56e819125d8d7907411fe24ccb943faca8" +dependencies = [ + "crossbeam-channel", + "crossbeam-deque", + "crossbeam-epoch", + "crossbeam-queue", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-channel" +version = "0.5.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d85363c37faeca707aef026efa9f3b34d077bce547e48f770770625c6013679e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-queue" +version = "0.3.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "803d13fb3b09d88be9f4dbc29062c66b19bf7170867ceb746d2a8689bf6c7a26" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "dircpy" +version = "0.3.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebcbec2b9a580ddee352ac38523d2ecd4dcaad53532957034394556909e27f4b" +dependencies = [ + "jwalk", + "log", + "walkdir", +] + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown", +] + +[[package]] +name = "indoc" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +dependencies = [ + "rustversion", +] + +[[package]] +name = "jobserver" +version = "0.1.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" +dependencies = [ + "getrandom", + "libc", +] + +[[package]] +name = "jwalk" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2735847566356cd2179a2a38264839308f7079fa96e6bd5a42d740460e003c56" +dependencies = [ + "crossbeam", + "rayon", +] + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + +[[package]] +name = "mstar-rust" +version = "0.1.0" +dependencies = [ + "bincode", + "libc", + "pyo3", + "serde", + "thiserror", + "zmq", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pyo3" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872" +dependencies = [ + "cfg-if", + "indoc", + "libc", + "memoffset", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "unindent", +] + +[[package]] +name = "pyo3-build-config" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb" +dependencies = [ + "once_cell", + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_spanned" +version = "0.6.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf41e0cfaf7226dca15e8197172c295a782857fcb97fad1808a166870dee75a3" +dependencies = [ + "serde", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "system-deps" +version = "6.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3e535eb8dded36d55ec13eddacd30dec501792ff23a0b1682c38601b8cf2349" +dependencies = [ + "cfg-expr", + "heck", + "pkg-config", + "toml", + "version-compare", +] + +[[package]] +name = "target-lexicon" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "toml" +version = "0.8.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1beb996b9d83529a9e75c17a1686767d148d70663143c7854d8b4a09ced362" +dependencies = [ + "serde", + "serde_spanned", + "toml_datetime", + "toml_edit", +] + +[[package]] +name = "toml_datetime" +version = "0.6.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22cddaf88f4fbc13c51aebbf5f8eceb5c7c5a9da2ac40a13519eb5b0a0e8f11c" +dependencies = [ + "serde", +] + +[[package]] +name = "toml_edit" +version = "0.22.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41fe8c660ae4257887cf66394862d21dbca4a6ddd26f04a3560410406a2f819a" +dependencies = [ + "indexmap", + "serde", + "serde_spanned", + "toml_datetime", + "winnow", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unindent" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + +[[package]] +name = "version-compare" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03c2856837ef78f57382f06b2b8563a2f512f7185d732608fd9176cb3b8edf0e" + +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "winnow" +version = "0.7.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945" +dependencies = [ + "memchr", +] + +[[package]] +name = "zeromq-src" +version = "0.2.6+4.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc120b771270365d5ed0dfb4baf1005f2243ae1ae83703265cb3504070f4160b" +dependencies = [ + "cc", + "dircpy", +] + +[[package]] +name = "zmq" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd3091dd571fb84a9b3e5e5c6a807d186c411c812c8618786c3c30e5349234e7" +dependencies = [ + "bitflags", + "libc", + "zmq-sys", +] + +[[package]] +name = "zmq-sys" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e8351dc72494b4d7f5652a681c33634063bbad58046c1689e75270908fdc864" +dependencies = [ + "libc", + "system-deps", + "zeromq-src", +] diff --git a/rust/Cargo.toml b/rust/Cargo.toml new file mode 100644 index 000000000..bfb581949 --- /dev/null +++ b/rust/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "mstar-rust" +version = "0.1.0" +edition = "2021" +license = "Apache-2.0" + +[lib] +name = "mstar_rust" +crate-type = ["cdylib"] + +[dependencies] +serde = { version = "1", features = ["derive"] } +bincode = "1.3" +thiserror = "2" +zmq = "0.10" +pyo3 = { version = "0.23", features = ["extension-module"] } + +[dev-dependencies] +libc = "0.2" + +[profile.release] +opt-level = 3 +lto = true diff --git a/rust/pyproject.toml b/rust/pyproject.toml new file mode 100644 index 000000000..7c3301073 --- /dev/null +++ b/rust/pyproject.toml @@ -0,0 +1,13 @@ +[build-system] +requires = ["maturin>=1.0,<2.0"] +build-backend = "maturin" + +[project] +name = "mstar-rust" +version = "0.1.0" +description = "Rust transport for mstar (RFC #130): ZMQ PUSH/PULL mesh with wakeup-fd polling" +requires-python = ">=3.10" + +[tool.maturin] +module-name = "mstar_rust" +manifest-path = "Cargo.toml" diff --git a/rust/src/communicator.rs b/rust/src/communicator.rs new file mode 100644 index 000000000..8b636b12d --- /dev/null +++ b/rust/src/communicator.rs @@ -0,0 +1,630 @@ +//! Per-entity mailbox over ZeroMQ PUSH/PULL — the direct analogue of mstar's +//! `ZMQCommunicator`. +//! +//! Layered as **transport + codec**, so the wire format is a seam rather than +//! a baked-in choice (mstar migrates pickle → language-neutral encodings by +//! swapping the codec, not the transport): +//! +//! * [`RawZmqCommunicator`] — the transport. Sends/receives **opaque byte +//! frames** (one zmq frame = exactly the payload, no added framing), so a +//! Python pickle blob, a msgpack blob, or bincode all pass through +//! untouched. Also owns the delivery machinery: ipc *and* tcp endpoints, +//! lazily-cached PUSH sockets, and **wakeup fds** (poll an eventfd alongside +//! the PULL socket so an external event — e.g. a completed compute future — +//! wakes the receive loop immediately instead of on the poll timeout). +//! * [`ZmqCommunicator`] — a typed wrapper: `Codec` encodes/decodes `M` +//! to bytes ([`BincodeCodec`] by default, matching the previous behavior for +//! Rust-internal messaging). +//! +//! Each entity binds one **PULL** inbox (default `ipc:///.ipc`; +//! or any zmq endpoint via [`RawZmqCommunicator::bind_endpoint`], e.g. +//! `tcp://0.0.0.0:5701` for the multi-node path) and connects a lazily-cached +//! **PUSH** socket per peer. PUSH/PULL gives fire-and-forget, ordered, +//! load-balanced delivery; libzmq queues to a not-yet-bound peer and +//! transparently reconnects when a peer restarts, so there is no +//! "unreachable" error to handle. + +use std::collections::HashMap; +use std::marker::PhantomData; +use std::os::fd::RawFd; +use std::path::{Path, PathBuf}; +use std::sync::Mutex; +use std::time::Duration; + +use serde::{de::DeserializeOwned, Serialize}; +use thiserror::Error; + +#[derive(Debug, Error)] +pub enum CommError { + #[error("zmq: {0}")] + Zmq(#[from] zmq::Error), + #[error("serialize: {0}")] + Serialize(bincode::Error), + #[error("io: {0}")] + Io(#[from] std::io::Error), + #[error("no endpoint for peer '{0}' (register_peer it, or bind with an ipc dir)")] + UnknownPeer(String), +} + +/// The default ipc endpoint an entity binds/connects to. +fn ipc_endpoint(dir: &Path, id: &str) -> String { + format!("ipc://{}/{}.ipc", dir.display(), id) +} + +/// The backing socket file (for clearing a stale one before bind). +fn sock_path(dir: &Path, id: &str) -> PathBuf { + dir.join(format!("{id}.ipc")) +} + +/// What a wake-aware receive returned. +#[derive(Debug, PartialEq, Eq)] +pub enum RecvEvent { + /// An inbound message frame. + Message(Vec), + /// A registered wakeup fd is readable (the registrant reads/clears it — + /// e.g. `read(2)` on the eventfd — the transport only polls it). + Wake, + /// The timeout elapsed with neither. + Timeout, +} + +/// Byte-frame transport: PUSH/PULL mailbox with pluggable endpoints and +/// wakeup-fd polling. This is the language-neutral seam — payloads are opaque. +/// +/// Field order matters for `Drop`: the sockets must close before the +/// `Context` is dropped (`zmq_ctx_term` blocks until its sockets are gone). +pub struct RawZmqCommunicator { + my_id: String, + /// ipc directory for the default peer-endpoint scheme (None when bound via + /// an explicit endpoint and peers must be `register_peer`ed). + dir: Option, + // PULL inbox. Behind a Mutex because a zmq `Socket` is `!Sync` (and must be + // used from one thread at a time); the drive loop is the single consumer. + pull: Mutex, + // peer id -> connected PUSH socket (created on first send to that peer). + peers: Mutex>, + // peer id -> explicit endpoint (tcp or ipc), overriding the dir scheme. + peer_endpoints: Mutex>, + // External fds polled alongside the PULL socket (mstar's worker registers + // an eventfd so a completed compute future wakes the message loop + // immediately — "done wrong, this silently stalls the async worker"). + wakeup_fds: Mutex>, + ctx: zmq::Context, +} + +impl RawZmqCommunicator { + /// Bind this entity's PULL inbox at `ipc:///.ipc` (the default + /// single-node scheme: peers resolve by id within the same dir). + pub fn bind(my_id: impl Into, dir: impl Into) -> Result { + let my_id = my_id.into(); + let dir = dir.into(); + std::fs::create_dir_all(&dir)?; + let _ = std::fs::remove_file(sock_path(&dir, &my_id)); // clear a stale socket + let endpoint = ipc_endpoint(&dir, &my_id); + Self::bind_inner(my_id, Some(dir), &endpoint) + } + + /// Bind this entity's PULL inbox at an explicit zmq endpoint — e.g. + /// `tcp://0.0.0.0:5701` for the multi-node path, or `tcp://127.0.0.1:*` + /// for an OS-assigned port (query it with [`Self::last_endpoint`]). + /// Peers have no implicit scheme here: `register_peer` each one. + pub fn bind_endpoint( + my_id: impl Into, + endpoint: &str, + ) -> Result { + Self::bind_inner(my_id.into(), None, endpoint) + } + + fn bind_inner( + my_id: String, + dir: Option, + endpoint: &str, + ) -> Result { + let ctx = zmq::Context::new(); + let pull = ctx.socket(zmq::PULL)?; + pull.set_linger(0)?; // don't block on close + pull.bind(endpoint)?; + Ok(Self { + my_id, + dir, + pull: Mutex::new(pull), + peers: Mutex::new(HashMap::new()), + peer_endpoints: Mutex::new(HashMap::new()), + wakeup_fds: Mutex::new(Vec::new()), + ctx, + }) + } + + pub fn id(&self) -> &str { + &self.my_id + } + + /// The bound endpoint as zmq reports it — with `tcp://…:*` this carries + /// the OS-assigned port, which is what peers must `register_peer`. + pub fn last_endpoint(&self) -> Result { + let pull = self.pull.lock().expect("pull lock"); + Ok(pull + .get_last_endpoint()? + .unwrap_or_else(|_| String::new())) + } + + /// Map `peer_id` to an explicit endpoint (tcp or ipc). Overrides the + /// ipc-dir scheme; required for peers reached over tcp. Replacing a + /// mapping drops the cached socket so the next send reconnects. + pub fn register_peer(&self, peer_id: &str, endpoint: &str) { + self.peer_endpoints + .lock() + .expect("peer endpoints lock") + .insert(peer_id.to_string(), endpoint.to_string()); + self.peers.lock().expect("peers lock").remove(peer_id); + } + + /// Poll an external fd alongside the PULL socket: when it becomes readable + /// the receive loop wakes immediately ([`RecvEvent::Wake`]) instead of on + /// the poll timeout. The registrant owns clearing it (e.g. reading the + /// eventfd); until cleared, wake-aware receives keep returning `Wake` + /// (level-triggered). + pub fn register_wakeup_fd(&self, fd: RawFd) { + self.wakeup_fds.lock().expect("wakeup fds lock").push(fd); + } + + fn resolve(&self, peer_id: &str) -> Result { + if let Some(ep) = self + .peer_endpoints + .lock() + .expect("peer endpoints lock") + .get(peer_id) + { + return Ok(ep.clone()); + } + match &self.dir { + Some(dir) => Ok(ipc_endpoint(dir, peer_id)), + None => Err(CommError::UnknownPeer(peer_id.to_string())), + } + } + + /// Send one opaque byte frame to `peer_id` (fire-and-forget). Queues if + /// the peer isn't bound yet and reconnects transparently if it restarted. + pub fn send(&self, peer_id: &str, payload: &[u8]) -> Result<(), CommError> { + let mut peers = self.peers.lock().expect("peers lock"); + if !peers.contains_key(peer_id) { + let endpoint = self.resolve(peer_id)?; + let push = self.ctx.socket(zmq::PUSH)?; + push.set_linger(0)?; + push.connect(&endpoint)?; + peers.insert(peer_id.to_string(), push); + } + peers + .get(peer_id) + .expect("just inserted") + .send(payload, 0)?; + Ok(()) + } + + /// Non-blocking: next inbound frame, or None. + pub fn try_recv(&self) -> Option> { + let pull = self.pull.lock().expect("pull lock"); + pull.recv_bytes(zmq::DONTWAIT).ok() + } + + /// Block until the next inbound frame (ignores wakeup fds). + pub fn recv(&self) -> Option> { + let pull = self.pull.lock().expect("pull lock"); + pull.recv_bytes(0).ok() + } + + /// Block up to `timeout` for the next inbound frame. Wakeup fds cut the + /// wait short (returns None early); use [`Self::recv_or_wake`] to + /// distinguish a wake from a timeout. + pub fn recv_timeout(&self, timeout: Duration) -> Option> { + match self.recv_or_wake(timeout) { + RecvEvent::Message(b) => Some(b), + _ => None, + } + } + + /// Wake-aware receive: a message frame, a wakeup-fd trip, or a timeout — + /// whichever comes first. Messages win a simultaneous wake (the wake is + /// level-triggered, so it is not lost: the next call reports it). + pub fn recv_or_wake(&self, timeout: Duration) -> RecvEvent { + let pull = self.pull.lock().expect("pull lock"); + let ms = timeout.as_millis().min(i64::MAX as u128) as i64; + let fds: Vec = self.wakeup_fds.lock().expect("wakeup fds lock").clone(); + + let mut items = Vec::with_capacity(1 + fds.len()); + items.push(pull.as_poll_item(zmq::POLLIN)); + for &fd in &fds { + items.push(zmq::PollItem::from_fd(fd, zmq::POLLIN)); + } + let n = zmq::poll(&mut items, ms).unwrap_or(0); + if n <= 0 { + return RecvEvent::Timeout; + } + if items[0].is_readable() { + if let Ok(b) = pull.recv_bytes(zmq::DONTWAIT) { + return RecvEvent::Message(b); + } + } + if items[1..].iter().any(|it| it.is_readable()) { + return RecvEvent::Wake; + } + RecvEvent::Timeout + } + + /// Drain all currently-queued inbound frames. + pub fn drain(&self) -> Vec> { + let pull = self.pull.lock().expect("pull lock"); + let mut out = Vec::new(); + while let Ok(b) = pull.recv_bytes(zmq::DONTWAIT) { + out.push(b); + } + out + } +} + +impl Drop for RawZmqCommunicator { + fn drop(&mut self) { + // Sockets (pull + peers) close first via field-drop order; zmq unlinks + // the bound ipc file on close, but remove it defensively too. + if let Some(dir) = &self.dir { + let _ = std::fs::remove_file(sock_path(dir, &self.my_id)); + } + } +} + +// --------------------------------------------------------------------------- +// Codec seam + typed wrapper +// --------------------------------------------------------------------------- + +/// Message encoding — the seam mstar migrates across (pickle for a +/// Python-to-Python mesh, msgpack/bincode for language-neutral ones). The +/// transport never looks inside the bytes. +pub trait Codec { + fn encode(msg: &M) -> Result, CommError>; + fn decode(bytes: &[u8]) -> Option; +} + +/// The default codec for Rust-internal messaging. +pub struct BincodeCodec; + +impl Codec for BincodeCodec { + fn encode(msg: &M) -> Result, CommError> { + bincode::serialize(msg).map_err(CommError::Serialize) + } + fn decode(bytes: &[u8]) -> Option { + bincode::deserialize(bytes).ok() + } +} + +/// A typed mailbox: [`RawZmqCommunicator`] + a [`Codec`]. `M` is the entity's +/// message type; the default codec is bincode (the previous behavior). +pub struct ZmqCommunicator { + raw: RawZmqCommunicator, + _marker: PhantomData (M, C)>, +} + +impl ZmqCommunicator +where + M: Send + 'static, + C: Codec, +{ + /// Bind this entity's PULL inbox at `ipc:///.ipc`. + pub fn bind(my_id: impl Into, dir: impl Into) -> Result { + Ok(Self { + raw: RawZmqCommunicator::bind(my_id, dir)?, + _marker: PhantomData, + }) + } + + /// Bind at an explicit zmq endpoint (e.g. `tcp://0.0.0.0:5701`). + pub fn bind_endpoint(my_id: impl Into, endpoint: &str) -> Result { + Ok(Self { + raw: RawZmqCommunicator::bind_endpoint(my_id, endpoint)?, + _marker: PhantomData, + }) + } + + /// The underlying byte transport (peer registration, wakeup fds, …). + pub fn raw(&self) -> &RawZmqCommunicator { + &self.raw + } + + pub fn id(&self) -> &str { + self.raw.id() + } + + /// Send `msg` to peer `peer_id` (fire-and-forget). + pub fn send(&self, peer_id: &str, msg: &M) -> Result<(), CommError> { + self.raw.send(peer_id, &C::encode(msg)?) + } + + /// Non-blocking: next inbound message, or None. + pub fn try_recv(&self) -> Option { + self.raw.try_recv().and_then(|b| C::decode(&b)) + } + + /// Block until the next inbound message. + pub fn recv(&self) -> Option { + self.raw.recv().and_then(|b| C::decode(&b)) + } + + /// Block up to `timeout` for the next inbound message (a wakeup fd cuts + /// the wait short — see [`RawZmqCommunicator::recv_or_wake`]). + pub fn recv_timeout(&self, timeout: Duration) -> Option { + self.raw.recv_timeout(timeout).and_then(|b| C::decode(&b)) + } + + /// Drain all currently-queued inbound messages. + pub fn drain(&self) -> Vec { + self.raw + .drain() + .into_iter() + .filter_map(|b| C::decode(&b)) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde::Deserialize; + + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] + enum Msg { + Hello(String), + Batch { id: u64, node: String }, + } + + fn tmpdir(tag: &str) -> PathBuf { + // Unique per test via the tag + thread id (no Date/rand available). + let t = format!("{:?}", std::thread::current().id()); + let dir = std::env::temp_dir().join(format!("mstar_comm_{tag}_{t}")); + let _ = std::fs::remove_dir_all(&dir); + dir + } + + fn wait_for) -> Option>( + mb: &ZmqCommunicator, + f: F, + ) -> Msg + where + M: Send + 'static, + C: Codec, + { + for _ in 0..500 { + if let Some(m) = f(mb) { + return m; + } + std::thread::sleep(Duration::from_millis(4)); + } + panic!("timed out waiting for message"); + } + + #[test] + fn two_entities_exchange_messages() { + let dir = tmpdir("exch"); + let conductor: ZmqCommunicator = ZmqCommunicator::bind("conductor", &dir).unwrap(); + let worker: ZmqCommunicator = ZmqCommunicator::bind("worker_0", &dir).unwrap(); + + conductor + .send("worker_0", &Msg::Batch { + id: 1, + node: "LLM".into(), + }) + .unwrap(); + let got = wait_for(&worker, |w| w.try_recv()); + assert_eq!(got, Msg::Batch { id: 1, node: "LLM".into() }); + + worker.send("conductor", &Msg::Hello("done".into())).unwrap(); + let got = wait_for(&conductor, |c| c.try_recv()); + assert_eq!(got, Msg::Hello("done".into())); + } + + #[test] + fn many_messages_arrive_in_order_from_one_sender() { + let dir = tmpdir("order"); + let a: ZmqCommunicator = ZmqCommunicator::bind("a", &dir).unwrap(); + let b: ZmqCommunicator = ZmqCommunicator::bind("b", &dir).unwrap(); + for i in 0..100 { + a.send("b", &Msg::Batch { id: i, node: "n".into() }).unwrap(); + } + // PUSH->PULL over one connection preserves FIFO order. + let mut seen = 0u64; + for _ in 0..1000 { + for m in b.drain() { + if let Msg::Batch { id, .. } = m { + assert_eq!(id, seen); + seen += 1; + } + } + if seen == 100 { + break; + } + std::thread::sleep(Duration::from_millis(2)); + } + assert_eq!(seen, 100); + } + + #[test] + fn send_before_peer_binds_is_queued() { + // Unlike the old UDS transport (which errored on a missing peer), zmq + // PUSH queues to a not-yet-bound endpoint and delivers once it binds. + let dir = tmpdir("queue"); + let a: ZmqCommunicator = ZmqCommunicator::bind("a", &dir).unwrap(); + a.send("late", &Msg::Hello("queued".into())).unwrap(); // no peer yet — no error + let late: ZmqCommunicator = ZmqCommunicator::bind("late", &dir).unwrap(); + let got = wait_for(&late, |m| m.try_recv()); + assert_eq!(got, Msg::Hello("queued".into())); + } + + #[test] + fn reconnects_after_peer_restart() { + let dir = tmpdir("restart"); + let a: ZmqCommunicator = ZmqCommunicator::bind("a", &dir).unwrap(); + { + let b: ZmqCommunicator = ZmqCommunicator::bind("b", &dir).unwrap(); + a.send("b", &Msg::Hello("1".into())).unwrap(); + wait_for(&b, |b| b.try_recv()); + } // b drops; zmq unlinks its inbox + std::thread::sleep(Duration::from_millis(20)); + // New b at the same id: a's cached PUSH auto-reconnects to it. + let b2: ZmqCommunicator = ZmqCommunicator::bind("b", &dir).unwrap(); + a.send("b", &Msg::Hello("2".into())).unwrap(); + let got = wait_for(&b2, |b| b.try_recv()); + assert_eq!(got, Msg::Hello("2".into())); + } + + // ---- raw transport: byte passthrough, tcp, wakeup fds ---------------- + + #[test] + fn raw_frames_pass_through_unmodified() { + // The wire is the payload — no framing added. A foreign blob (e.g. a + // Python pickle) must arrive byte-identical. + let dir = tmpdir("raw"); + let a = RawZmqCommunicator::bind("a", &dir).unwrap(); + let b = RawZmqCommunicator::bind("b", &dir).unwrap(); + let blob: Vec = (0..=255).collect(); + a.send("b", &blob).unwrap(); + for _ in 0..500 { + if let Some(got) = b.try_recv() { + assert_eq!(got, blob); + return; + } + std::thread::sleep(Duration::from_millis(4)); + } + panic!("frame never arrived"); + } + + #[test] + fn tcp_endpoint_exchange() { + // Multi-node shape: bind on an OS-assigned tcp port, peer registers + // the reported endpoint explicitly (no ipc dir involved). + let a = RawZmqCommunicator::bind_endpoint("a", "tcp://127.0.0.1:*").unwrap(); + let b = RawZmqCommunicator::bind_endpoint("b", "tcp://127.0.0.1:*").unwrap(); + let a_ep = a.last_endpoint().unwrap(); + let b_ep = b.last_endpoint().unwrap(); + assert!(a_ep.starts_with("tcp://"), "{a_ep}"); + a.register_peer("b", &b_ep); + b.register_peer("a", &a_ep); + + a.send("b", b"over tcp").unwrap(); + for _ in 0..500 { + if let Some(got) = b.try_recv() { + assert_eq!(got, b"over tcp"); + // and the reverse direction + b.send("a", b"ack").unwrap(); + for _ in 0..500 { + if let Some(back) = a.try_recv() { + assert_eq!(back, b"ack"); + return; + } + std::thread::sleep(Duration::from_millis(4)); + } + panic!("ack never arrived"); + } + std::thread::sleep(Duration::from_millis(4)); + } + panic!("tcp frame never arrived"); + } + + #[test] + fn unknown_peer_without_dir_errors() { + let a = RawZmqCommunicator::bind_endpoint("a", "tcp://127.0.0.1:*").unwrap(); + assert!(matches!( + a.send("nowhere", b"x"), + Err(CommError::UnknownPeer(_)) + )); + } + + #[cfg(target_os = "linux")] + #[test] + fn wakeup_fd_cuts_recv_wait_short() { + // mstar's worker registers an eventfd alongside the PULL socket so a + // completed compute future wakes the loop immediately — not on the + // poll timeout. Fire the eventfd from a thread mid-wait and require + // the wake to arrive far sooner than the 2 s timeout. + let dir = tmpdir("wake"); + let a = RawZmqCommunicator::bind("a", &dir).unwrap(); + let efd = unsafe { libc::eventfd(0, 0) }; + assert!(efd >= 0); + a.register_wakeup_fd(efd); + + let t = std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(50)); + let one: u64 = 1; + let n = unsafe { + libc::write(efd, &one as *const u64 as *const libc::c_void, 8) + }; + assert_eq!(n, 8); + }); + + let start = std::time::Instant::now(); + let ev = a.recv_or_wake(Duration::from_secs(2)); + let waited = start.elapsed(); + t.join().unwrap(); + assert_eq!(ev, RecvEvent::Wake); + assert!( + waited < Duration::from_millis(500), + "wake took {waited:?} — the fd poll isn't cutting the wait short" + ); + + // Level-triggered until the registrant clears it... + assert_eq!(a.recv_or_wake(Duration::from_millis(10)), RecvEvent::Wake); + let mut buf = 0u64; + unsafe { + libc::read(efd, &mut buf as *mut u64 as *mut libc::c_void, 8); + } + // ...and quiet after the read (times out). + assert_eq!( + a.recv_or_wake(Duration::from_millis(10)), + RecvEvent::Timeout + ); + + // A message still wins while the fd is quiet. + let b = RawZmqCommunicator::bind("b", &dir).unwrap(); + b.send("a", b"msg").unwrap(); + for _ in 0..500 { + match a.recv_or_wake(Duration::from_millis(20)) { + RecvEvent::Message(m) => { + assert_eq!(m, b"msg"); + unsafe { libc::close(efd) }; + return; + } + _ => continue, + } + } + panic!("message never arrived"); + } + + // ---- codec seam -------------------------------------------------------- + + /// A toy foreign codec (length-prefixed utf8) proving the transport is + /// format-agnostic — the pickle/msgpack seam mstar wants. + struct TextCodec; + impl Codec for TextCodec { + fn encode(msg: &String) -> Result, CommError> { + Ok(msg.as_bytes().to_vec()) + } + fn decode(bytes: &[u8]) -> Option { + String::from_utf8(bytes.to_vec()).ok() + } + } + + #[test] + fn custom_codec_over_same_transport() { + let dir = tmpdir("codec"); + let a: ZmqCommunicator = ZmqCommunicator::bind("a", &dir).unwrap(); + // The peer reads RAW bytes: what TextCodec sent is exactly the utf8 — + // no transport framing in between. + let b = RawZmqCommunicator::bind("b", &dir).unwrap(); + a.send("b", &"hello seam".to_string()).unwrap(); + for _ in 0..500 { + if let Some(bytes) = b.try_recv() { + assert_eq!(bytes, b"hello seam"); + return; + } + std::thread::sleep(Duration::from_millis(4)); + } + panic!("frame never arrived"); + } +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs new file mode 100644 index 000000000..c774900ae --- /dev/null +++ b/rust/src/lib.rs @@ -0,0 +1,95 @@ +//! mstar's Rust transport (RFC #130 Step 1): the ZMQ PUSH/PULL mesh, vendored +//! from mstar-rs. `communicator.rs` is the transport + codec split; this file +//! is the PyO3 surface (`mstar_rust.ZmqCommunicator`) the Python +//! `RustZMQCommunicator` wrapper drives. Build: `maturin develop` in rust/. + +mod communicator; + +use std::time::Duration; + +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; +use pyo3::types::PyBytes; + +use communicator::{RawZmqCommunicator, RecvEvent}; + +/// Opaque byte frames over ZMQ PUSH/PULL: ipc or tcp endpoints, lazily-cached +/// peers, and wakeup-fd polling (an eventfd wakes `recv_or_wake` instantly). +/// Encoding is the caller's (pickle today; msgpack by swapping the codec). +#[pyclass(name = "ZmqCommunicator")] +struct PyZmqCommunicator { + inner: RawZmqCommunicator, +} + +#[pymethods] +impl PyZmqCommunicator { + /// Bind the PULL inbox at `ipc://{dir}/{my_id}.ipc`. + #[new] + fn new(my_id: &str, dir: &str) -> PyResult { + Ok(Self { + inner: RawZmqCommunicator::bind(my_id, dir) + .map_err(|e| PyRuntimeError::new_err(e.to_string()))?, + }) + } + + /// Bind at an explicit zmq endpoint (e.g. `tcp://0.0.0.0:5701`). + #[staticmethod] + fn bind_endpoint(my_id: &str, endpoint: &str) -> PyResult { + Ok(Self { + inner: RawZmqCommunicator::bind_endpoint(my_id, endpoint) + .map_err(|e| PyRuntimeError::new_err(e.to_string()))?, + }) + } + + fn last_endpoint(&self) -> PyResult { + self.inner + .last_endpoint() + .map_err(|e| PyRuntimeError::new_err(e.to_string())) + } + + fn register_peer(&self, peer_id: &str, endpoint: &str) { + self.inner.register_peer(peer_id, endpoint); + } + + /// Poll `fd` (an eventfd) alongside the inbox; when readable, + /// `recv_or_wake` returns ("wake", None) immediately. The registrant + /// reads/clears the fd (level-triggered). + fn register_wakeup_fd(&self, fd: i32) { + self.inner.register_wakeup_fd(fd); + } + + fn send(&self, peer_id: &str, data: &[u8]) -> PyResult<()> { + self.inner + .send(peer_id, data) + .map_err(|e| PyRuntimeError::new_err(e.to_string())) + } + + fn try_recv<'py>(&self, py: Python<'py>) -> Option> { + self.inner.try_recv().map(|b| PyBytes::new(py, &b)) + } + + fn recv_timeout<'py>(&self, py: Python<'py>, timeout_ms: u64) -> Option> { + py.allow_threads(|| self.inner.recv_timeout(Duration::from_millis(timeout_ms))) + .map(|b| PyBytes::new(py, &b)) + } + + /// ("msg", bytes) | ("wake", None) | ("timeout", None). + fn recv_or_wake<'py>( + &self, + py: Python<'py>, + timeout_ms: u64, + ) -> (&'static str, Option>) { + let ev = py.allow_threads(|| self.inner.recv_or_wake(Duration::from_millis(timeout_ms))); + match ev { + RecvEvent::Message(b) => ("msg", Some(PyBytes::new(py, &b))), + RecvEvent::Wake => ("wake", None), + RecvEvent::Timeout => ("timeout", None), + } + } +} + +#[pymodule] +fn mstar_rust(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + Ok(()) +} diff --git a/test/modular/test_rust_communicator.py b/test/modular/test_rust_communicator.py new file mode 100644 index 000000000..5b0d46bb8 --- /dev/null +++ b/test/modular/test_rust_communicator.py @@ -0,0 +1,61 @@ +"""RustZMQCommunicator interop with the pyzmq ZMQCommunicator (RFC #130 +Step 1): same mesh, pickle wire, eventfd wakeup, lossless readiness polls. +Skipped unless the mstar-rs wheel is installed.""" +import os +import tempfile +import threading +import time + +import pytest + +pytest.importorskip("mstar_rust") + +from mstar.communication.communicator import ZMQCommunicator +from mstar.communication.event import EventWakeup +from mstar.communication.rust_communicator import RustZMQCommunicator + + +def _wait_msgs(comm, n=1, timeout=5.0): + out = [] + deadline = time.time() + timeout + while len(out) < n and time.time() < deadline: + out.extend(comm.get_all_new_messages()) + time.sleep(0.005) + return out + + +@pytest.fixture() +def pair(): + prefix = tempfile.mkdtemp(prefix="mstar_wrap_test_") + orig = ZMQCommunicator("orig", push_ids=["rust"], ipc_socket_path_prefix=prefix) + rust = RustZMQCommunicator("rust", push_ids=["orig"], ipc_socket_path_prefix=prefix) + return orig, rust + + +def test_pickle_interop_both_directions(pair): + orig, rust = pair + payload = {"op": "execute", "rids": [1, 2, 3], "nested": {"f": 1.5}} + orig.send("rust", payload) + assert _wait_msgs(rust) == [payload] + rust.send("orig", ("done", 42)) + assert _wait_msgs(orig) == [("done", 42)] + + +def test_eventfd_wakeup_cuts_wait_short(pair): + _, rust = pair + ev = EventWakeup() + rust.register_event_for_poll(ev) + threading.Thread(target=lambda: (time.sleep(0.05), + os.eventfd_write(ev.fd, 1))).start() + t0 = time.time() + rust.wait_for_work(timeout_ms=2000) + assert time.time() - t0 < 0.5, "wake must beat the poll timeout" + + +def test_readiness_poll_never_drops_or_reorders(pair): + orig, rust = pair + orig.send("rust", "first") + assert any(rust.poll_for_messages(timeout_ms=10) for _ in range(500)) + orig.send("rust", "second") + time.sleep(0.1) + assert rust.get_all_new_messages() == ["first", "second"] From e6966d34341ce4071a6a9f7bc8b8eca6644270e7 Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Wed, 15 Jul 2026 02:40:36 +0000 Subject: [PATCH 02/21] communication: MSTAR_RUST_ZMQ=1 opt-in for the Rust transport MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit make_communicator() selects the transport at the four construction sites (worker, conductor, api_server, data worker): default is the pyzmq ZMQCommunicator, byte-identical behavior; MSTAR_RUST_ZMQ=1 opts a process into RustZMQCommunicator. The two are wire-compatible (same endpoints, same pickle frames), so the flag is per-process — one entity can be switched and A/B'd while the rest of the mesh stays on pyzmq. The default flip stays a later step, after perf gates and packaging (building the extension in CI/wheels). --- mstar/api_server/data_worker.py | 6 +++--- mstar/api_server/entrypoint.py | 4 ++-- mstar/communication/communicator.py | 18 ++++++++++++++++++ mstar/conductor/conductor.py | 4 ++-- mstar/worker/worker.py | 4 ++-- 5 files changed, 27 insertions(+), 9 deletions(-) diff --git a/mstar/api_server/data_worker.py b/mstar/api_server/data_worker.py index d983d0810..462cc8c61 100644 --- a/mstar/api_server/data_worker.py +++ b/mstar/api_server/data_worker.py @@ -22,7 +22,7 @@ ResultChunk, ResultTensors, ) -from mstar.communication.communicator import CommProtocol, ZMQCommunicator +from mstar.communication.communicator import BaseCommunicator, CommProtocol, make_communicator from mstar.communication.tensors import NameToTensorList, create_tensor_communication_manager from mstar.model.base import Model from mstar.profile.format import InputInfo, RxInfo, TxInfo @@ -74,7 +74,7 @@ def __init__( # The socket is only *used* from the worker thread, but owning the # tensor manager here lets the main thread read its tx/rx profiling # directly once a request is done (no cross-thread queue / race). - self.communicator = ZMQCommunicator( + self.communicator = make_communicator( my_id="api_server_preprocess_worker", push_ids=["conductor"], ipc_socket_path_prefix=socket_path_prefix, @@ -233,7 +233,7 @@ def __init__( abort_request_queue: queue.Queue, discard_tensor_queue: queue.Queue, stop_event: threading.Event, - communicator: ZMQCommunicator, + communicator: BaseCommunicator, tensor_manager, device: str = "cpu", model: Model | None = None, diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index 6adf0dd82..cdc19939e 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -23,7 +23,7 @@ from mstar.api_server.data_worker import PreprocessWorker from mstar.api_server.request_types import APIServerMessage, PreprocessInput, ResultChunk -from mstar.communication.communicator import CommProtocol, ZMQCommunicator +from mstar.communication.communicator import CommProtocol, make_communicator from mstar.model.registry import HF_MODELS from mstar.profile.display import pretty_print_profile from mstar.profile.format import OutputInfo, RequestProfile, RequestTiming @@ -207,7 +207,7 @@ def __init__( self.running = True # ZMQ channel shared with conductor / workers - self.communicator = ZMQCommunicator( + self.communicator = make_communicator( my_id="api_server", push_ids=["conductor"], ipc_socket_path_prefix=socket_path_prefix, diff --git a/mstar/communication/communicator.py b/mstar/communication/communicator.py index 8975c91c7..b60f3fefe 100644 --- a/mstar/communication/communicator.py +++ b/mstar/communication/communicator.py @@ -141,3 +141,21 @@ def get_all_new_messages(self, blocking=False) -> list: # zmq.Again actually means no messages left to read break return messages + + +def make_communicator(*args, **kwargs) -> BaseCommunicator: + """Construct the process's communicator, selecting the transport. + + ``MSTAR_RUST_ZMQ=1`` opts a process into the Rust-backed + ``RustZMQCommunicator`` (vendored ``rust/`` extension; see + ``communication/rust_communicator.py``). Default is the pyzmq + ``ZMQCommunicator``, byte-identical behavior. The two are + wire-compatible (same endpoints, same pickle frames), so the flag can + be set per-process — one entity at a time — while the rest of the + mesh stays on pyzmq. + """ + if os.getenv("MSTAR_RUST_ZMQ", "0") == "1": + from mstar.communication.rust_communicator import RustZMQCommunicator + + return RustZMQCommunicator(*args, **kwargs) + return ZMQCommunicator(*args, **kwargs) diff --git a/mstar/conductor/conductor.py b/mstar/conductor/conductor.py index e90144424..80ad6788d 100644 --- a/mstar/conductor/conductor.py +++ b/mstar/conductor/conductor.py @@ -13,7 +13,7 @@ import yaml from mstar.api_server.request_types import APIServerMessage, RequestComplete -from mstar.communication.communicator import CommProtocol, ZMQCommunicator +from mstar.communication.communicator import CommProtocol, make_communicator from mstar.conductor.request_info import ( CurrentForwardConductorMetadata, CurrentForwardPassInfo, @@ -306,7 +306,7 @@ def __init__( self._derive_worker_info() self._launch_workers() - self.communicator = ZMQCommunicator( + self.communicator = make_communicator( my_id="conductor", push_ids=self.worker_ids + ["api_server", "api_server_preprocess_worker"], ipc_socket_path_prefix=socket_path_prefix, diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 0ee5c3da2..82c731541 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -13,7 +13,7 @@ import torch from mstar.api_server.request_types import APIServerMessage, ResultTensors -from mstar.communication.communicator import CommProtocol, ZMQCommunicator +from mstar.communication.communicator import CommProtocol, make_communicator from mstar.communication.event import EventWakeup from mstar.communication.tensors import NameToTensorList, create_tensor_communication_manager from mstar.conductor.request_info import CurrentForwardPassInfo @@ -168,7 +168,7 @@ def __init__( for node_name in section.get_nodes(): node_to_partition[node_name] = pdef.name - self.communicator = ZMQCommunicator( + self.communicator = make_communicator( my_id=worker_id, push_ids=worker_ids + ["conductor", "api_server", "api_server_preprocess_worker"], ipc_socket_path_prefix=socket_path_prefix, From 58dceb62e6e1f9314aebb981066b53ebb5a9932e Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Wed, 15 Jul 2026 03:32:13 +0000 Subject: [PATCH 03/21] ci: build the vendored Rust transport and run its tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A path-filtered workflow (rust/, communication/, its test file): cargo tests for the transport semantics (raw frames, tcp endpoints, wakeup-fd polling), then maturin-build the mstar_rust extension and run the pyzmq <-> Rust interop pytest against it — so the importorskip guard never silently skips in CI when the transport is what changed. --- .github/workflows/rust.yml | 41 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) create mode 100644 .github/workflows/rust.yml diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml new file mode 100644 index 000000000..afacd14cd --- /dev/null +++ b/.github/workflows/rust.yml @@ -0,0 +1,41 @@ +name: Rust transport + +# Builds the vendored rust/ extension and runs its tests — the cargo unit +# tests (transport semantics: raw frames, tcp endpoints, wakeup-fd polling) +# and the pyzmq <-> Rust interop pytest (pickle wire both directions, +# eventfd wake latency, lossless readiness polls). Path-filtered so it only +# runs when the transport or its wrapper changes. + +on: + pull_request: + branches: [main] + paths: + - "rust/**" + - "mstar/communication/**" + - "test/modular/test_rust_communicator.py" + - ".github/workflows/rust.yml" + +jobs: + rust-transport: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install libzmq + run: sudo apt-get update && sudo apt-get install -y libzmq3-dev + - name: Cargo tests (transport semantics) + working-directory: rust + run: cargo test --release + - name: Build and install the extension + working-directory: rust + run: | + python -m pip install --upgrade pip maturin + maturin build --release + pip install target/wheels/*.whl + - name: Interop tests (pyzmq <-> Rust) + run: | + pip install pytest pyzmq + pytest test/modular/test_rust_communicator.py -v From 3befdeb925bcb0ac1027b543a0faf0b50318810b Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Wed, 15 Jul 2026 03:40:09 +0000 Subject: [PATCH 04/21] ci: move the transport interop test out of test/modular MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit test/modular/conftest.py imports torch, which the transport test neither needs nor should pull into its CI job — test/rust/ has no conftest, so the job stays pyzmq-only. --- .github/workflows/rust.yml | 4 ++-- test/{modular => rust}/test_rust_communicator.py | 0 2 files changed, 2 insertions(+), 2 deletions(-) rename test/{modular => rust}/test_rust_communicator.py (100%) diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index afacd14cd..68ebac581 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -12,7 +12,7 @@ on: paths: - "rust/**" - "mstar/communication/**" - - "test/modular/test_rust_communicator.py" + - "test/rust/test_rust_communicator.py" - ".github/workflows/rust.yml" jobs: @@ -38,4 +38,4 @@ jobs: - name: Interop tests (pyzmq <-> Rust) run: | pip install pytest pyzmq - pytest test/modular/test_rust_communicator.py -v + pytest test/rust/test_rust_communicator.py -v diff --git a/test/modular/test_rust_communicator.py b/test/rust/test_rust_communicator.py similarity index 100% rename from test/modular/test_rust_communicator.py rename to test/rust/test_rust_communicator.py From 7d238599b14b2d3d88691b32ca94d8a7c89e6330 Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Wed, 15 Jul 2026 03:41:33 +0000 Subject: [PATCH 05/21] ci: fold the Rust-transport job into ci.yml MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One workflow, two jobs (ruff + rust-transport). Path filtering is workflow-level only, so the transport job now runs on every PR — with Swatinem cargo caching it stays around a minute warm, and the transport is core enough to earn that. --- .github/workflows/ci.yml | 32 ++++++++++++++++++++++++++++- .github/workflows/rust.yml | 41 -------------------------------------- 2 files changed, 31 insertions(+), 42 deletions(-) delete mode 100644 .github/workflows/rust.yml diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 227a7b8e0..a315439a9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -17,4 +17,34 @@ jobs: pip install ruff - name: Run Ruff # Use GitHub’s annotation format so lint errors show inline - run: ruff check --output-format=github . \ No newline at end of file + run: ruff check --output-format=github . + # Vendored Rust transport: cargo tests (raw frames, tcp endpoints, + # wakeup-fd polling), then build the mstar_rust extension and run the + # pyzmq <-> Rust interop pytest against it. Cargo caching keeps this to + # ~a minute on warm runs. + rust-transport: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + - uses: Swatinem/rust-cache@v2 + with: + workspaces: rust + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install libzmq + run: sudo apt-get update && sudo apt-get install -y libzmq3-dev + - name: Cargo tests (transport semantics) + working-directory: rust + run: cargo test --release + - name: Build and install the extension + working-directory: rust + run: | + python -m pip install --upgrade pip maturin + maturin build --release + pip install target/wheels/*.whl + - name: Interop tests (pyzmq <-> Rust) + run: | + pip install pytest pyzmq + pytest test/rust/test_rust_communicator.py -v diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml deleted file mode 100644 index 68ebac581..000000000 --- a/.github/workflows/rust.yml +++ /dev/null @@ -1,41 +0,0 @@ -name: Rust transport - -# Builds the vendored rust/ extension and runs its tests — the cargo unit -# tests (transport semantics: raw frames, tcp endpoints, wakeup-fd polling) -# and the pyzmq <-> Rust interop pytest (pickle wire both directions, -# eventfd wake latency, lossless readiness polls). Path-filtered so it only -# runs when the transport or its wrapper changes. - -on: - pull_request: - branches: [main] - paths: - - "rust/**" - - "mstar/communication/**" - - "test/rust/test_rust_communicator.py" - - ".github/workflows/rust.yml" - -jobs: - rust-transport: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - uses: dtolnay/rust-toolchain@stable - - uses: actions/setup-python@v5 - with: - python-version: "3.12" - - name: Install libzmq - run: sudo apt-get update && sudo apt-get install -y libzmq3-dev - - name: Cargo tests (transport semantics) - working-directory: rust - run: cargo test --release - - name: Build and install the extension - working-directory: rust - run: | - python -m pip install --upgrade pip maturin - maturin build --release - pip install target/wheels/*.whl - - name: Interop tests (pyzmq <-> Rust) - run: | - pip install pytest pyzmq - pytest test/rust/test_rust_communicator.py -v From b0e29f235865ee46ff64c285f4bf9facada2be16 Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Wed, 15 Jul 2026 03:58:46 +0000 Subject: [PATCH 06/21] ci: install mstar (no deps) before the interop pytest MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The runner had no importable mstar package; --no-deps keeps the job lean — the transport test chain needs only pyzmq and the stdlib. --- .github/workflows/ci.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a315439a9..c2ab1a9e8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -47,4 +47,7 @@ jobs: - name: Interop tests (pyzmq <-> Rust) run: | pip install pytest pyzmq + # the package itself, without its (heavy) deps — the transport + # test chain needs only pyzmq + stdlib + pip install --no-deps -e . pytest test/rust/test_rust_communicator.py -v From c46d426dab5391ecaca274753cd63046c34e8f6c Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Wed, 15 Jul 2026 08:16:47 +0000 Subject: [PATCH 07/21] rust: build the crate as a library too (fixes dead-code warnings) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The typed ZmqCommunicator / Codec layer is the API the migration's later steps (Rust conductor, API server) consume, and the cargo tests exercise it — but with only a cdylib target and a private module, rustc saw it as unreachable and warned. Exposing it (crate-type rlib + pub mod communicator) makes the crate a real Rust library and the build warning-free. --- rust/Cargo.toml | 5 ++++- rust/src/lib.rs | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/rust/Cargo.toml b/rust/Cargo.toml index bfb581949..8e674daef 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -6,7 +6,10 @@ license = "Apache-2.0" [lib] name = "mstar_rust" -crate-type = ["cdylib"] +# cdylib: the Python extension module. rlib: the same crate as a Rust +# library, so the migration's later steps (Rust conductor / API server) +# consume the typed `ZmqCommunicator` layer directly. +crate-type = ["cdylib", "rlib"] [dependencies] serde = { version = "1", features = ["derive"] } diff --git a/rust/src/lib.rs b/rust/src/lib.rs index c774900ae..67595a374 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -3,7 +3,7 @@ //! is the PyO3 surface (`mstar_rust.ZmqCommunicator`) the Python //! `RustZMQCommunicator` wrapper drives. Build: `maturin develop` in rust/. -mod communicator; +pub mod communicator; use std::time::Duration; From 0e158afead9b35b1dd5891ce3671af00a35b7926 Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Wed, 15 Jul 2026 08:16:47 +0000 Subject: [PATCH 08/21] comm: honor blocking receives; share the endpoint scheme; Codec class Review follow-ups: - get_all_new_messages(blocking=True) now waits for at least one message before draining, on BOTH communicators (the pyzmq one had the same latent bug: the flag was accepted and ignored). A registered wakeup event also ends the wait, mirroring wait_for_work. - _endpoint/_tcp_port were copied between ZMQCommunicator and RustZMQCommunicator; they live on BaseCommunicator now. - The codec seam is a small Codec class (encode/decode), mirroring the Rust trait, with PickleCodec as the wire default. - poll_for_messages docstring states the actual contract: message, wake, or timeout, whichever comes first. - MSTAR_RUST_ZMQ is tri-state: 0 (default, pyzmq), 1 (Rust, raise if the extension is missing), AUTO (Rust when importable). Tests cover the blocking receive on both transports and the factory selection. --- mstar/communication/communicator.py | 106 ++++++++++++++--------- mstar/communication/rust_communicator.py | 99 ++++++++++++--------- test/rust/test_rust_communicator.py | 26 ++++++ 3 files changed, 150 insertions(+), 81 deletions(-) diff --git a/mstar/communication/communicator.py b/mstar/communication/communicator.py index b60f3fefe..8f0808af6 100644 --- a/mstar/communication/communicator.py +++ b/mstar/communication/communicator.py @@ -10,6 +10,13 @@ logger = logging.getLogger(__name__) +class CommProtocol(Enum): + IPC = "IPC" + TCP = "TCP" + RDMA = "RDMA" + SHM = "SHM" + + class BaseCommunicator(ABC): @abstractmethod def send(self, entity_id: str, msg): @@ -22,18 +29,37 @@ def send(self, entity_id: str, msg): def get_all_new_messages(self) -> list: pass + # -- endpoint scheme (shared by every ZMQ-based communicator) ------------ + # Subclasses set ``self.protocol`` and ``self.ipc_socket_path_prefix``. + + def _endpoint(self, entity_id: str) -> str: + if self.protocol == CommProtocol.IPC: + return f"ipc://{self.ipc_socket_path_prefix}/{entity_id}.ipc" + if self.protocol == CommProtocol.TCP: + host = os.getenv("MSTAR_ZMQ_TCP_HOST", "127.0.0.1") + return f"tcp://{host}:{self._tcp_port(entity_id)}" + raise NotImplementedError(f"Protocol {self.protocol} not yet supported yet") + + @staticmethod + def _tcp_port(entity_id: str) -> int: + base_port = int(os.getenv("MSTAR_ZMQ_TCP_BASE_PORT", "19000")) + if entity_id == "api_server": + return base_port + if entity_id == "conductor": + return base_port + 1 + if entity_id == "api_server_preprocess_worker": + return base_port + 2 + if entity_id.startswith("worker_"): + rank = entity_id.removeprefix("worker_") + if rank.isdigit(): + return base_port + 100 + int(rank) + return base_port + 1000 + (sum(entity_id.encode("utf-8")) % 1000) + # @abstractmethod # def get_session_id(self) -> str: # pass -class CommProtocol(Enum): - IPC = "IPC" - TCP = "TCP" - RDMA = "RDMA" - SHM = "SHM" - - class ZMQCommunicator(BaseCommunicator): def __init__( self, @@ -84,29 +110,6 @@ def wait_for_work(self, timeout_ms=50): if self.event.fd in events: self.event.drain() - def _endpoint(self, entity_id: str) -> str: - if self.protocol == CommProtocol.IPC: - return f"ipc://{self.ipc_socket_path_prefix}/{entity_id}.ipc" - if self.protocol == CommProtocol.TCP: - host = os.getenv("MSTAR_ZMQ_TCP_HOST", "127.0.0.1") - return f"tcp://{host}:{self._tcp_port(entity_id)}" - raise NotImplementedError(f"Protocol {self.protocol} not yet supported yet") - - @staticmethod - def _tcp_port(entity_id: str) -> int: - base_port = int(os.getenv("MSTAR_ZMQ_TCP_BASE_PORT", "19000")) - if entity_id == "api_server": - return base_port - if entity_id == "conductor": - return base_port + 1 - if entity_id == "api_server_preprocess_worker": - return base_port + 2 - if entity_id.startswith("worker_"): - rank = entity_id.removeprefix("worker_") - if rank.isdigit(): - return base_port + 100 + int(rank) - return base_port + 1000 + (sum(entity_id.encode("utf-8")) % 1000) - # def get_session_id(self) -> str: # return self.session_id @@ -125,6 +128,14 @@ def send(self, entity_id: str, msg): def get_all_new_messages(self, blocking=False) -> list: messages = [] + if blocking: + # Wait until the pull socket is readable before draining. A + # registered wakeup event also ends the wait (and is drained + # here, exactly as in wait_for_work), so a completed compute + # future can interrupt a blocking receive. + events = dict(self.poller.poll()) + if self.event is not None and self.event.fd in events: + self.event.drain() while True: try: # zmq.NOBLOCK means zmq doesn't wait for a new message to be @@ -146,16 +157,29 @@ def get_all_new_messages(self, blocking=False) -> list: def make_communicator(*args, **kwargs) -> BaseCommunicator: """Construct the process's communicator, selecting the transport. - ``MSTAR_RUST_ZMQ=1`` opts a process into the Rust-backed - ``RustZMQCommunicator`` (vendored ``rust/`` extension; see - ``communication/rust_communicator.py``). Default is the pyzmq - ``ZMQCommunicator``, byte-identical behavior. The two are - wire-compatible (same endpoints, same pickle frames), so the flag can - be set per-process — one entity at a time — while the rest of the - mesh stays on pyzmq. - """ - if os.getenv("MSTAR_RUST_ZMQ", "0") == "1": - from mstar.communication.rust_communicator import RustZMQCommunicator + ``MSTAR_RUST_ZMQ`` selects it (see ``docs/environment_variables.rst``): + + * ``0`` (default) — the pyzmq ``ZMQCommunicator``. + * ``1`` — the Rust-backed ``RustZMQCommunicator`` (vendored ``rust/`` + extension; see ``communication/rust_communicator.py``); raises if the + extension is not installed. + * ``AUTO`` — the Rust communicator when the extension imports + successfully, pyzmq otherwise. - return RustZMQCommunicator(*args, **kwargs) + The two are wire-compatible (same endpoints, same pickle frames), so the + flag can be set per-process — one entity at a time — while the rest of + the mesh stays on pyzmq. + """ + choice = os.getenv("MSTAR_RUST_ZMQ", "0").upper() + if choice not in ("0", "1", "AUTO"): + raise ValueError(f"MSTAR_RUST_ZMQ must be 0, 1, or AUTO; got {choice!r}") + if choice != "0": + try: + from mstar.communication.rust_communicator import RustZMQCommunicator + except ImportError: + if choice == "1": + raise + logger.debug("MSTAR_RUST_ZMQ=AUTO: mstar_rust not installed, using pyzmq") + else: + return RustZMQCommunicator(*args, **kwargs) return ZMQCommunicator(*args, **kwargs) diff --git a/mstar/communication/rust_communicator.py b/mstar/communication/rust_communicator.py index cc6d22f90..67865cb36 100644 --- a/mstar/communication/rust_communicator.py +++ b/mstar/communication/rust_communicator.py @@ -9,9 +9,10 @@ (``ipc://{prefix}/{id}.ipc`` / the same TCP host+port scheme), and the default codec is pickle — so a wrapped entity talks to unwrapped pyzmq entities in both directions. Migration can proceed one process at a time. -* **encode/decode seam.** ``codec`` is a ``(dumps, loads)`` pair defaulting to - pickle (today's wire). Swapping to msgpack later is a codec change on both - ends of an edge, never a transport change. +* **encode/decode seam.** ``codec`` is a :class:`Codec` (mirroring the Rust + ``Codec`` trait) defaulting to :class:`PickleCodec` (today's wire). + Swapping to msgpack later is a codec change on both ends of an edge, + never a transport change. * **eventfd wakeup.** ``register_event_for_poll`` forwards the ``EventWakeup`` fd to the Rust poller (``register_wakeup_fd``): a completed compute future wakes ``wait_for_work`` / ``poll_for_messages`` immediately, not on the poll @@ -23,7 +24,9 @@ poll is buffered and handed out by the next ``get_all_new_messages`` — no message is ever dropped or reordered. -Requires the vendored extension: ``maturin develop`` (or ``pip install .``) in ``rust/``. +Requires the vendored extension — from ``rust/``: ``maturin develop --release`` +(or ``pip install .``). Use ``--release``: debug builds cost real latency on +the hot receive path. See ``docs/installation.rst``. """ from __future__ import annotations @@ -40,10 +43,34 @@ logger = logging.getLogger(__name__) -#: The encode/decode seam. Pickle matches today's ``send_pyobj`` wire, so a -#: wrapped entity interoperates with unwrapped pyzmq entities. Migrate an edge -#: to msgpack by giving both endpoints a msgpack codec. -PickleCodec = (pickle.dumps, pickle.loads) +class Codec: + """The encode/decode seam, mirroring the Rust ``Codec`` trait:: + + pub trait Codec { + fn encode(msg: &M) -> Result, CommError>; + fn decode(bytes: &[u8]) -> Option; + } + + The transport never looks inside the bytes; migrating an edge to another + encoding (e.g. msgpack) means giving both endpoints that codec — never a + transport change. + """ + + @staticmethod + def encode(msg) -> bytes: + raise NotImplementedError + + @staticmethod + def decode(data: bytes): + raise NotImplementedError + + +class PickleCodec(Codec): + """Pickle matches today's ``send_pyobj`` wire, so a wrapped entity + interoperates with unwrapped pyzmq entities in both directions.""" + + encode = staticmethod(pickle.dumps) + decode = staticmethod(pickle.loads) class RustZMQCommunicator(BaseCommunicator): @@ -55,13 +82,13 @@ def __init__( push_ids: list[str], protocol: CommProtocol = CommProtocol.IPC, ipc_socket_path_prefix: str = "/tmp/mstar/", - codec=PickleCodec, + codec: type[Codec] = PickleCodec, ): transport = os.getenv("MSTAR_ZMQ_TRANSPORT", protocol.value).upper() self.protocol = CommProtocol(transport) self.my_id = my_id self.ipc_socket_path_prefix = ipc_socket_path_prefix - self._dumps, self._loads = codec + self.codec = codec self.event: EventWakeup | None = None # Messages consumed by a readiness poll, awaiting get_all_new_messages. self._buffered: deque = deque() @@ -80,28 +107,8 @@ def __init__( else: raise NotImplementedError(f"Protocol {protocol} not yet supported yet") - # -- endpoint scheme (identical to the pyzmq communicator's) ------------ - - def _endpoint(self, entity_id: str) -> str: - if self.protocol == CommProtocol.IPC: - return f"ipc://{self.ipc_socket_path_prefix}/{entity_id}.ipc" - host = os.getenv("MSTAR_ZMQ_TCP_HOST", "127.0.0.1") - return f"tcp://{host}:{self._tcp_port(entity_id)}" - - @staticmethod - def _tcp_port(entity_id: str) -> int: - base_port = int(os.getenv("MSTAR_ZMQ_TCP_BASE_PORT", "19000")) - if entity_id == "api_server": - return base_port - if entity_id == "conductor": - return base_port + 1 - if entity_id == "api_server_preprocess_worker": - return base_port + 2 - if entity_id.startswith("worker_"): - rank = entity_id.removeprefix("worker_") - if rank.isdigit(): - return base_port + 100 + int(rank) - return base_port + 1000 + (sum(entity_id.encode("utf-8")) % 1000) + # endpoint scheme: `_endpoint` / `_tcp_port` come from BaseCommunicator + # (shared with the pyzmq ZMQCommunicator). def _register(self, entity_id: str) -> None: if entity_id not in self._registered: @@ -114,14 +121,16 @@ def register_event_for_poll(self, event: EventWakeup) -> None: self._inner.register_wakeup_fd(event.fd) self.event = event - def _poll_once(self, timeout_ms: int) -> None: - """One wake-aware poll. A consumed message goes to the buffer; a - wakeup drains the event (same place the pyzmq path drains it).""" + def _poll_once(self, timeout_ms: int) -> str: + """One wake-aware poll; returns ``"msg"`` / ``"wake"`` / ``"timeout"``. + A consumed message goes to the buffer; a wakeup drains the event + (same place the pyzmq path drains it).""" kind, payload = self._inner.recv_or_wake(timeout_ms) if kind == "msg": self._buffered.append(payload) elif kind == "wake" and self.event is not None: self.event.drain() + return kind def wait_for_work(self, timeout_ms: int = 50) -> None: if self._buffered: @@ -129,8 +138,11 @@ def wait_for_work(self, timeout_ms: int = 50) -> None: self._poll_once(timeout_ms) def poll_for_messages(self, timeout_ms: int = 20) -> bool: - """Block up to ``timeout_ms`` for a readable message; True when one is - available (buffered here, delivered by ``get_all_new_messages``).""" + """Block until a message is readable, a registered wakeup event + fires, or ``timeout_ms`` elapses — whichever comes first. True when + a message is available (buffered here, delivered by + ``get_all_new_messages``); a wakeup ends the poll early with False + (the event is drained, exactly as in ``wait_for_work``).""" if self._buffered: return True self._poll_once(timeout_ms) @@ -140,11 +152,18 @@ def send(self, entity_id: str, msg) -> None: logger.debug("%s to send a message %s to entity %s", self.my_id, str(msg), entity_id) if self.protocol == CommProtocol.TCP: self._register(entity_id) - self._inner.send(entity_id, self._dumps(msg)) + self._inner.send(entity_id, self.codec.encode(msg)) def get_all_new_messages(self, blocking: bool = False) -> list: - messages = [self._loads(b) for b in self._buffered] + if blocking and not self._buffered: + # Wait until a message is readable before draining — same + # semantics as the pyzmq communicator: a registered wakeup + # event also ends the wait, so a completed compute future can + # interrupt a blocking receive. + while self._poll_once(50) == "timeout": + pass + messages = [self.codec.decode(b) for b in self._buffered] self._buffered.clear() while (b := self._inner.try_recv()) is not None: - messages.append(self._loads(b)) + messages.append(self.codec.decode(b)) return messages diff --git a/test/rust/test_rust_communicator.py b/test/rust/test_rust_communicator.py index 5b0d46bb8..6125ae9ba 100644 --- a/test/rust/test_rust_communicator.py +++ b/test/rust/test_rust_communicator.py @@ -59,3 +59,29 @@ def test_readiness_poll_never_drops_or_reorders(pair): orig.send("rust", "second") time.sleep(0.1) assert rust.get_all_new_messages() == ["first", "second"] + + +@pytest.mark.parametrize("receiver", ["rust", "orig"]) +def test_blocking_receive_waits_for_a_message(pair, receiver): + """get_all_new_messages(blocking=True) waits instead of returning [] — + on both communicators (the pyzmq one had the same latent bug).""" + orig, rust = pair + dst, src = (rust, orig) if receiver == "rust" else (orig, rust) + threading.Thread(target=lambda: (time.sleep(0.1), + src.send(receiver, "late"))).start() + assert dst.get_all_new_messages(blocking=True) == ["late"] + + +def test_make_communicator_flag(monkeypatch, tmp_path): + from mstar.communication.communicator import make_communicator + + def make(value): + monkeypatch.setenv("MSTAR_RUST_ZMQ", value) + return make_communicator( + f"m_{value}", push_ids=[], ipc_socket_path_prefix=str(tmp_path)) + + assert isinstance(make("0"), ZMQCommunicator) + assert isinstance(make("1"), RustZMQCommunicator) + assert isinstance(make("AUTO"), RustZMQCommunicator) # extension installed + with pytest.raises(ValueError): + make("yes") From 3fc075324ebf053271bfa793dfd4a58e9176a6cd Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Wed, 15 Jul 2026 08:17:02 +0000 Subject: [PATCH 09/21] docs: environment-variables reference; Rust transport setup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - docs/environment_variables.rst starts the documented-env-vars precedent, seeded with the communication variables (MSTAR_RUST_ZMQ, MSTAR_ZMQ_TRANSPORT, MSTAR_ZMQ_TCP_HOST, MSTAR_ZMQ_TCP_BASE_PORT). - installation.rst + README: how to build the optional rust/ extension. maturin defaults to a debug build, so the instructions say 'maturin develop --release' explicitly — debug costs real latency on the hot receive path. --- README.md | 8 ++++++++ docs/environment_variables.rst | 37 ++++++++++++++++++++++++++++++++++ docs/index.rst | 1 + docs/installation.rst | 24 ++++++++++++++++++++++ 4 files changed, 70 insertions(+) create mode 100644 docs/environment_variables.rst diff --git a/README.md b/README.md index d6ddd2c21..6ac9df16e 100644 --- a/README.md +++ b/README.md @@ -140,6 +140,14 @@ single-GPU or fully disaggregated by changing only the YAML `node_groups`. Four primitives — `Sequential`, `Parallel`, `Loop`, and a cross-partition `StreamingGraphEdge` — express every model family above. See the [paper](https://arxiv.org/abs/2606.12688) for the full design. +### Optional: Rust ZMQ transport + +The ZeroMQ control mesh can run over a Rust transport (vendored in [`rust/`](rust/)) instead of +pyzmq — wire-compatible, selectable per process with `MSTAR_RUST_ZMQ=1` (or `AUTO`). Build it with +[maturin](https://www.maturin.rs): `uv pip install maturin && maturin develop --release -m rust/Cargo.toml`. +See the [installation docs](https://m-star.org/mstar/installation.html) and +[environment variables](https://m-star.org/mstar/environment_variables.html). + ## Performance Across every model we benchmark, M\* matches or beats the system specialized for that family — unified diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst new file mode 100644 index 000000000..81851442f --- /dev/null +++ b/docs/environment_variables.rst @@ -0,0 +1,37 @@ +Environment variables +===================== + +Runtime knobs M* reads from the environment. New variables should be +documented here as they are introduced. + +Communication +------------- + +.. list-table:: + :header-rows: 1 + :widths: 28 14 58 + + * - Variable + - Default + - Meaning + * - ``MSTAR_RUST_ZMQ`` + - ``0`` + - Transport selection for the ZeroMQ control mesh (see + :func:`mstar.communication.communicator.make_communicator`). + ``0``: the pyzmq ``ZMQCommunicator``. ``1``: the Rust-backed + ``RustZMQCommunicator`` (requires the vendored ``rust/`` extension; + raises if it is not installed). ``AUTO``: Rust when the extension + imports successfully, pyzmq otherwise. The two transports are + wire-compatible, so this can be set per-process while the rest of + the mesh stays on pyzmq. + * - ``MSTAR_ZMQ_TRANSPORT`` + - constructor's protocol + - Overrides the communicator protocol (``IPC`` or ``TCP``) for a + process, e.g. to run entities on separate hosts. + * - ``MSTAR_ZMQ_TCP_HOST`` + - ``127.0.0.1`` + - Host used to build peer endpoints when the protocol is ``TCP``. + * - ``MSTAR_ZMQ_TCP_BASE_PORT`` + - ``19000`` + - Base of the deterministic entity-id → TCP port map (``api_server`` + = base, ``conductor`` = base+1, ``worker_`` = base+100+rank). diff --git a/docs/index.rst b/docs/index.rst index 68c95d3d3..6358f3ee7 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -39,6 +39,7 @@ vision-language-action policies, and world models — through a **Python SDK**, architecture models api + environment_variables .. toctree:: :maxdepth: 2 diff --git a/docs/installation.rst b/docs/installation.rst index e02e0b388..4009c0022 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -248,6 +248,30 @@ Verify the install mstar --help mstar-serve --help +Optional: the Rust ZMQ transport +-------------------------------- + +The ZeroMQ control mesh can run over a Rust transport (vendored in ``rust/``) +instead of pyzmq — same endpoints, same wire format, selectable per process +with ``MSTAR_RUST_ZMQ`` (see :doc:`environment_variables`). It is optional: +without it, everything runs on pyzmq as before. + +Build the extension into your environment with `maturin +`_ (needs a Rust toolchain; ``rustup`` works): + +.. code-block:: bash + + uv pip install maturin + maturin develop --release -m rust/Cargo.toml + +Build with ``--release`` — an unoptimized debug build (maturin's default) +costs real latency on the hot receive path. Verify with: + +.. code-block:: bash + + python -c "import mstar_rust; print('mstar_rust OK')" + pytest test/rust/test_rust_communicator.py + Troubleshooting --------------- From 825fcd331ff2d55e6c553f269fb07b3a76c78fed Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Wed, 15 Jul 2026 08:36:26 +0000 Subject: [PATCH 10/21] docs: drop the stale DRAFT wording; fix the extension module name MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The wrapper predates the vendoring: it said DRAFT and pointed at the prototype's module path (mstar_rs._core). It is neither — the extension is mstar_rust, built from rust/ in this repo. --- mstar/communication/rust_communicator.py | 4 ++-- rust/src/lib.rs | 8 ++++---- test/rust/test_rust_communicator.py | 3 ++- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/mstar/communication/rust_communicator.py b/mstar/communication/rust_communicator.py index 67865cb36..a28dd5434 100644 --- a/mstar/communication/rust_communicator.py +++ b/mstar/communication/rust_communicator.py @@ -1,5 +1,5 @@ -"""DRAFT (RFC #130 Step 1): ``ZMQCommunicator`` as a thin wrapper over the -mstar-rs Rust communicator (``mstar_rs._core.ZmqCommunicator``). +"""RFC #130 Step 1: ``ZMQCommunicator`` as a thin wrapper over the Rust +communicator vendored in ``rust/`` (``mstar_rust.ZmqCommunicator``). Drop-in for :class:`mstar.communication.communicator.ZMQCommunicator` — same constructor, same methods, same semantics — with the transport moved to Rust: diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 67595a374..cb50376aa 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -1,7 +1,7 @@ -//! mstar's Rust transport (RFC #130 Step 1): the ZMQ PUSH/PULL mesh, vendored -//! from mstar-rs. `communicator.rs` is the transport + codec split; this file -//! is the PyO3 surface (`mstar_rust.ZmqCommunicator`) the Python -//! `RustZMQCommunicator` wrapper drives. Build: `maturin develop` in rust/. +//! mstar's Rust transport (RFC #130 Step 1): the ZMQ PUSH/PULL mesh. +//! `communicator.rs` is the transport + codec split; this file is the PyO3 +//! surface (`mstar_rust.ZmqCommunicator`) the Python `RustZMQCommunicator` +//! wrapper drives. Build: `maturin develop --release` in rust/. pub mod communicator; diff --git a/test/rust/test_rust_communicator.py b/test/rust/test_rust_communicator.py index 6125ae9ba..94ac15900 100644 --- a/test/rust/test_rust_communicator.py +++ b/test/rust/test_rust_communicator.py @@ -1,6 +1,7 @@ """RustZMQCommunicator interop with the pyzmq ZMQCommunicator (RFC #130 Step 1): same mesh, pickle wire, eventfd wakeup, lossless readiness polls. -Skipped unless the mstar-rs wheel is installed.""" +Skipped unless the ``mstar_rust`` extension (built from ``rust/``) is +installed.""" import os import tempfile import threading From 2f1870a5b8cfe287c5000156915f55a797f3a559 Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Wed, 15 Jul 2026 08:58:08 +0000 Subject: [PATCH 11/21] rust: default typed codec bincode -> MessagePack (rmp-serde) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit bincode is a governance risk (1.x frozen; upstream has left GitHub and disavows further updates there) — and MessagePack is what this migration standardizes on anyway: the codec seam's stated target wire, readable by Python peers with the msgpack package (to_vec_named = field-name maps, i.e. Python dicts). rmp-serde is actively maintained. The transport is unaffected (opaque frames); only the Rust-internal typed default changes. --- rust/Cargo.lock | 39 +++++++++++++++++++++++++++++---------- rust/Cargo.toml | 2 +- rust/src/communicator.rs | 24 +++++++++++++----------- 3 files changed, 43 insertions(+), 22 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 1a9ffae7d..f41870dab 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -8,15 +8,6 @@ version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" -[[package]] -name = "bincode" -version = "1.3.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad" -dependencies = [ - "serde", -] - [[package]] name = "bitflags" version = "1.3.2" @@ -229,14 +220,23 @@ dependencies = [ name = "mstar-rust" version = "0.1.0" dependencies = [ - "bincode", "libc", "pyo3", + "rmp-serde", "serde", "thiserror", "zmq", ] +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -362,6 +362,25 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "rmp" +version = "0.8.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ba8be72d372b2c9b35542551678538b562e7cf86c3315773cae48dfbfe7790c" +dependencies = [ + "num-traits", +] + +[[package]] +name = "rmp-serde" +version = "1.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f81bee8c8ef9b577d1681a70ebbc962c232461e397b22c208c43c04b67a155" +dependencies = [ + "rmp", + "serde", +] + [[package]] name = "rustversion" version = "1.0.23" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 8e674daef..721746a2d 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -13,7 +13,7 @@ crate-type = ["cdylib", "rlib"] [dependencies] serde = { version = "1", features = ["derive"] } -bincode = "1.3" +rmp-serde = "1.3" thiserror = "2" zmq = "0.10" pyo3 = { version = "0.23", features = ["extension-module"] } diff --git a/rust/src/communicator.rs b/rust/src/communicator.rs index 8b636b12d..55c4db541 100644 --- a/rust/src/communicator.rs +++ b/rust/src/communicator.rs @@ -7,13 +7,13 @@ //! //! * [`RawZmqCommunicator`] — the transport. Sends/receives **opaque byte //! frames** (one zmq frame = exactly the payload, no added framing), so a -//! Python pickle blob, a msgpack blob, or bincode all pass through +//! Python pickle blob or a msgpack blob passes through //! untouched. Also owns the delivery machinery: ipc *and* tcp endpoints, //! lazily-cached PUSH sockets, and **wakeup fds** (poll an eventfd alongside //! the PULL socket so an external event — e.g. a completed compute future — //! wakes the receive loop immediately instead of on the poll timeout). //! * [`ZmqCommunicator`] — a typed wrapper: `Codec` encodes/decodes `M` -//! to bytes ([`BincodeCodec`] by default, matching the previous behavior for +//! to bytes ([`MsgpackCodec`] by default — the language-neutral wire for //! Rust-internal messaging). //! //! Each entity binds one **PULL** inbox (default `ipc:///.ipc`; @@ -39,7 +39,7 @@ pub enum CommError { #[error("zmq: {0}")] Zmq(#[from] zmq::Error), #[error("serialize: {0}")] - Serialize(bincode::Error), + Serialize(rmp_serde::encode::Error), #[error("io: {0}")] Io(#[from] std::io::Error), #[error("no endpoint for peer '{0}' (register_peer it, or bind with an ipc dir)")] @@ -277,28 +277,30 @@ impl Drop for RawZmqCommunicator { // --------------------------------------------------------------------------- /// Message encoding — the seam mstar migrates across (pickle for a -/// Python-to-Python mesh, msgpack/bincode for language-neutral ones). The +/// Python-to-Python mesh, msgpack for language-neutral ones). The /// transport never looks inside the bytes. pub trait Codec { fn encode(msg: &M) -> Result, CommError>; fn decode(bytes: &[u8]) -> Option; } -/// The default codec for Rust-internal messaging. -pub struct BincodeCodec; +/// The default codec: MessagePack, the language-neutral encoding the +/// migration standardizes on — Python peers read the same frames with +/// `msgpack` (`to_vec_named`: maps with field names, like Python dicts). +pub struct MsgpackCodec; -impl Codec for BincodeCodec { +impl Codec for MsgpackCodec { fn encode(msg: &M) -> Result, CommError> { - bincode::serialize(msg).map_err(CommError::Serialize) + rmp_serde::to_vec_named(msg).map_err(CommError::Serialize) } fn decode(bytes: &[u8]) -> Option { - bincode::deserialize(bytes).ok() + rmp_serde::from_slice(bytes).ok() } } /// A typed mailbox: [`RawZmqCommunicator`] + a [`Codec`]. `M` is the entity's -/// message type; the default codec is bincode (the previous behavior). -pub struct ZmqCommunicator { +/// message type; the default codec is MessagePack. +pub struct ZmqCommunicator { raw: RawZmqCommunicator, _marker: PhantomData (M, C)>, } From 930da03bf3821274224af762a743862e37d25178 Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Wed, 15 Jul 2026 10:20:27 +0000 Subject: [PATCH 12/21] rust: release the GIL on send; batch receives with drain - send: a PUSH send blocks when the peer sits at its high-water mark; doing that with the GIL held freezes every Python thread in the process. Release it around the send, as pyzmq does. - get_all_new_messages picks up the whole queued batch with one drain() call instead of one try_recv FFI round-trip per message. --- mstar/communication/rust_communicator.py | 4 ++-- rust/src/lib.rs | 15 ++++++++++++--- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/mstar/communication/rust_communicator.py b/mstar/communication/rust_communicator.py index a28dd5434..172fc5ea1 100644 --- a/mstar/communication/rust_communicator.py +++ b/mstar/communication/rust_communicator.py @@ -164,6 +164,6 @@ def get_all_new_messages(self, blocking: bool = False) -> list: pass messages = [self.codec.decode(b) for b in self._buffered] self._buffered.clear() - while (b := self._inner.try_recv()) is not None: - messages.append(self.codec.decode(b)) + # One FFI call for the whole queued batch (vs try_recv per message). + messages.extend(self.codec.decode(b) for b in self._inner.drain()) return messages diff --git a/rust/src/lib.rs b/rust/src/lib.rs index cb50376aa..f1650e57f 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -58,9 +58,11 @@ impl PyZmqCommunicator { self.inner.register_wakeup_fd(fd); } - fn send(&self, peer_id: &str, data: &[u8]) -> PyResult<()> { - self.inner - .send(peer_id, data) + /// Released-GIL send: a PUSH send blocks when the peer is at its + /// high-water mark, and blocking with the GIL held would freeze every + /// Python thread in the process (pyzmq releases it here too). + fn send(&self, py: Python<'_>, peer_id: &str, data: &[u8]) -> PyResult<()> { + py.allow_threads(|| self.inner.send(peer_id, data)) .map_err(|e| PyRuntimeError::new_err(e.to_string())) } @@ -68,6 +70,13 @@ impl PyZmqCommunicator { self.inner.try_recv().map(|b| PyBytes::new(py, &b)) } + /// Drain all currently-queued frames in one call — one GIL round-trip + /// per batch instead of one `try_recv` per message. + fn drain<'py>(&self, py: Python<'py>) -> Vec> { + let frames = py.allow_threads(|| self.inner.drain()); + frames.iter().map(|b| PyBytes::new(py, b)).collect() + } + fn recv_timeout<'py>(&self, py: Python<'py>, timeout_ms: u64) -> Option> { py.allow_threads(|| self.inner.recv_timeout(Duration::from_millis(timeout_ms))) .map(|b| PyBytes::new(py, &b)) From 83e1f731deae508d4c291e3eb864643daf571751 Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Wed, 15 Jul 2026 10:31:17 +0000 Subject: [PATCH 13/21] rust: one payload copy on receive (zmq::Message end to end) recv_bytes copied every frame twice: zmq's buffer -> Vec, then Vec -> PyBytes (or the codec's decode). Receives now hand the zmq::Message itself through RecvEvent/try_recv/drain (Deref), so the single copy happens at the consumer boundary and nowhere else. --- rust/src/communicator.rs | 47 ++++++++++++++++++++++++++-------------- 1 file changed, 31 insertions(+), 16 deletions(-) diff --git a/rust/src/communicator.rs b/rust/src/communicator.rs index 55c4db541..b74df8044 100644 --- a/rust/src/communicator.rs +++ b/rust/src/communicator.rs @@ -56,11 +56,13 @@ fn sock_path(dir: &Path, id: &str) -> PathBuf { dir.join(format!("{id}.ipc")) } -/// What a wake-aware receive returned. -#[derive(Debug, PartialEq, Eq)] +/// What a wake-aware receive returned. `Message` carries the zmq frame +/// itself (`Deref`), so bytes are copied exactly once — +/// at the consumer's decode/PyBytes boundary, never in the transport. +#[derive(Debug)] pub enum RecvEvent { /// An inbound message frame. - Message(Vec), + Message(zmq::Message), /// A registered wakeup fd is readable (the registrant reads/clears it — /// e.g. `read(2)` on the eventfd — the transport only polls it). Wake, @@ -68,6 +70,19 @@ pub enum RecvEvent { Timeout, } +impl PartialEq for RecvEvent { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (RecvEvent::Message(a), RecvEvent::Message(b)) => a[..] == b[..], + (RecvEvent::Wake, RecvEvent::Wake) => true, + (RecvEvent::Timeout, RecvEvent::Timeout) => true, + _ => false, + } + } +} + +impl Eq for RecvEvent {} + /// Byte-frame transport: PUSH/PULL mailbox with pluggable endpoints and /// wakeup-fd polling. This is the language-neutral seam — payloads are opaque. /// @@ -202,21 +217,21 @@ impl RawZmqCommunicator { } /// Non-blocking: next inbound frame, or None. - pub fn try_recv(&self) -> Option> { + pub fn try_recv(&self) -> Option { let pull = self.pull.lock().expect("pull lock"); - pull.recv_bytes(zmq::DONTWAIT).ok() + pull.recv_msg(zmq::DONTWAIT).ok() } /// Block until the next inbound frame (ignores wakeup fds). - pub fn recv(&self) -> Option> { + pub fn recv(&self) -> Option { let pull = self.pull.lock().expect("pull lock"); - pull.recv_bytes(0).ok() + pull.recv_msg(0).ok() } /// Block up to `timeout` for the next inbound frame. Wakeup fds cut the /// wait short (returns None early); use [`Self::recv_or_wake`] to /// distinguish a wake from a timeout. - pub fn recv_timeout(&self, timeout: Duration) -> Option> { + pub fn recv_timeout(&self, timeout: Duration) -> Option { match self.recv_or_wake(timeout) { RecvEvent::Message(b) => Some(b), _ => None, @@ -241,7 +256,7 @@ impl RawZmqCommunicator { return RecvEvent::Timeout; } if items[0].is_readable() { - if let Ok(b) = pull.recv_bytes(zmq::DONTWAIT) { + if let Ok(b) = pull.recv_msg(zmq::DONTWAIT) { return RecvEvent::Message(b); } } @@ -252,10 +267,10 @@ impl RawZmqCommunicator { } /// Drain all currently-queued inbound frames. - pub fn drain(&self) -> Vec> { + pub fn drain(&self) -> Vec { let pull = self.pull.lock().expect("pull lock"); let mut out = Vec::new(); - while let Ok(b) = pull.recv_bytes(zmq::DONTWAIT) { + while let Ok(b) = pull.recv_msg(zmq::DONTWAIT) { out.push(b); } out @@ -489,7 +504,7 @@ mod tests { a.send("b", &blob).unwrap(); for _ in 0..500 { if let Some(got) = b.try_recv() { - assert_eq!(got, blob); + assert_eq!(&got[..], &blob[..]); return; } std::thread::sleep(Duration::from_millis(4)); @@ -512,12 +527,12 @@ mod tests { a.send("b", b"over tcp").unwrap(); for _ in 0..500 { if let Some(got) = b.try_recv() { - assert_eq!(got, b"over tcp"); + assert_eq!(&got[..], b"over tcp"); // and the reverse direction b.send("a", b"ack").unwrap(); for _ in 0..500 { if let Some(back) = a.try_recv() { - assert_eq!(back, b"ack"); + assert_eq!(&back[..], b"ack"); return; } std::thread::sleep(Duration::from_millis(4)); @@ -588,7 +603,7 @@ mod tests { for _ in 0..500 { match a.recv_or_wake(Duration::from_millis(20)) { RecvEvent::Message(m) => { - assert_eq!(m, b"msg"); + assert_eq!(&m[..], b"msg"); unsafe { libc::close(efd) }; return; } @@ -622,7 +637,7 @@ mod tests { a.send("b", &"hello seam".to_string()).unwrap(); for _ in 0..500 { if let Some(bytes) = b.try_recv() { - assert_eq!(bytes, b"hello seam"); + assert_eq!(&bytes[..], b"hello seam"); return; } std::thread::sleep(Duration::from_millis(4)); From 255fa7f5905617b37a6989ab0d2803409c613147 Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Wed, 15 Jul 2026 11:36:20 +0000 Subject: [PATCH 14/21] ci: trim the rust-transport job comment --- .github/workflows/ci.yml | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c2ab1a9e8..10e7f888d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -18,10 +18,8 @@ jobs: - name: Run Ruff # Use GitHub’s annotation format so lint errors show inline run: ruff check --output-format=github . - # Vendored Rust transport: cargo tests (raw frames, tcp endpoints, - # wakeup-fd polling), then build the mstar_rust extension and run the - # pyzmq <-> Rust interop pytest against it. Cargo caching keeps this to - # ~a minute on warm runs. + # Vendored Rust transport (rust/): cargo tests, then the pyzmq interop + # pytest against a freshly built extension. rust-transport: runs-on: ubuntu-latest steps: From ce38a5bf8853c9b924c0d8217cc731ec7b487d8c Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Thu, 16 Jul 2026 01:25:35 +0000 Subject: [PATCH 15/21] comm: default MSTAR_RUST_ZMQ to AUTO (review request) Unset now means: the Rust transport when the vendored extension imports, pyzmq otherwise. Explicit 0/1 behave as before. Test covers the default. --- docs/environment_variables.rst | 14 +++++++------- mstar/communication/communicator.py | 13 ++++++------- test/rust/test_rust_communicator.py | 6 ++++++ 3 files changed, 19 insertions(+), 14 deletions(-) diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index 81851442f..5957fab13 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -15,15 +15,15 @@ Communication - Default - Meaning * - ``MSTAR_RUST_ZMQ`` - - ``0`` + - ``AUTO`` - Transport selection for the ZeroMQ control mesh (see :func:`mstar.communication.communicator.make_communicator`). - ``0``: the pyzmq ``ZMQCommunicator``. ``1``: the Rust-backed - ``RustZMQCommunicator`` (requires the vendored ``rust/`` extension; - raises if it is not installed). ``AUTO``: Rust when the extension - imports successfully, pyzmq otherwise. The two transports are - wire-compatible, so this can be set per-process while the rest of - the mesh stays on pyzmq. + ``AUTO``: the Rust-backed ``RustZMQCommunicator`` when the vendored + ``rust/`` extension imports successfully, pyzmq otherwise. + ``1``: the Rust communicator, raising if the extension is missing. + ``0``: always pyzmq. The two transports are wire-compatible, so + this can be set per-process while the rest of the mesh stays on + pyzmq. * - ``MSTAR_ZMQ_TRANSPORT`` - constructor's protocol - Overrides the communicator protocol (``IPC`` or ``TCP``) for a diff --git a/mstar/communication/communicator.py b/mstar/communication/communicator.py index 8f0808af6..cfd605068 100644 --- a/mstar/communication/communicator.py +++ b/mstar/communication/communicator.py @@ -159,18 +159,17 @@ def make_communicator(*args, **kwargs) -> BaseCommunicator: ``MSTAR_RUST_ZMQ`` selects it (see ``docs/environment_variables.rst``): - * ``0`` (default) — the pyzmq ``ZMQCommunicator``. - * ``1`` — the Rust-backed ``RustZMQCommunicator`` (vendored ``rust/`` - extension; see ``communication/rust_communicator.py``); raises if the - extension is not installed. - * ``AUTO`` — the Rust communicator when the extension imports - successfully, pyzmq otherwise. + * ``AUTO`` (default) — the Rust-backed ``RustZMQCommunicator`` (vendored + ``rust/`` extension; see ``communication/rust_communicator.py``) when + the extension imports successfully, pyzmq otherwise. + * ``1`` — the Rust communicator; raises if the extension is missing. + * ``0`` — always the pyzmq ``ZMQCommunicator``. The two are wire-compatible (same endpoints, same pickle frames), so the flag can be set per-process — one entity at a time — while the rest of the mesh stays on pyzmq. """ - choice = os.getenv("MSTAR_RUST_ZMQ", "0").upper() + choice = os.getenv("MSTAR_RUST_ZMQ", "AUTO").upper() if choice not in ("0", "1", "AUTO"): raise ValueError(f"MSTAR_RUST_ZMQ must be 0, 1, or AUTO; got {choice!r}") if choice != "0": diff --git a/test/rust/test_rust_communicator.py b/test/rust/test_rust_communicator.py index 94ac15900..42faa3287 100644 --- a/test/rust/test_rust_communicator.py +++ b/test/rust/test_rust_communicator.py @@ -86,3 +86,9 @@ def make(value): assert isinstance(make("AUTO"), RustZMQCommunicator) # extension installed with pytest.raises(ValueError): make("yes") + # default is AUTO: with the extension installed, unset -> Rust + monkeypatch.delenv("MSTAR_RUST_ZMQ", raising=False) + assert isinstance( + make_communicator("m_default", push_ids=[], + ipc_socket_path_prefix=str(tmp_path)), + RustZMQCommunicator) From efead2864b520d92a39dfa5c6ae6e8164be7bed8 Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Thu, 16 Jul 2026 06:53:04 +0000 Subject: [PATCH 16/21] comm: bounded blocking receives; MsgpackCodec; tighter module doc MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups: get_all_new_messages(blocking=True) takes an optional timeout_s on BOTH communicators (None = wait indefinitely; on expiry, drain whatever arrived) so a future blocking caller cannot wait forever. MsgpackCodec joins PickleCodec — msgpack is language-neutral, so edges using it can terminate in (future) Rust processes. Module docstring condensed. --- mstar/communication/communicator.py | 8 ++- mstar/communication/rust_communicator.py | 70 +++++++++++++----------- test/rust/test_rust_communicator.py | 11 ++++ 3 files changed, 53 insertions(+), 36 deletions(-) diff --git a/mstar/communication/communicator.py b/mstar/communication/communicator.py index cfd605068..f1fb85ce6 100644 --- a/mstar/communication/communicator.py +++ b/mstar/communication/communicator.py @@ -126,14 +126,16 @@ def send(self, entity_id: str, msg): self.push_sockets[entity_id] = sock self.push_sockets[entity_id].send_pyobj(msg) - def get_all_new_messages(self, blocking=False) -> list: + def get_all_new_messages(self, blocking=False, timeout_s=None) -> list: messages = [] if blocking: # Wait until the pull socket is readable before draining. A # registered wakeup event also ends the wait (and is drained # here, exactly as in wait_for_work), so a completed compute - # future can interrupt a blocking receive. - events = dict(self.poller.poll()) + # future can interrupt a blocking receive. `timeout_s` bounds + # the wait (None = indefinitely); on expiry, drain what's there. + timeout_ms = None if timeout_s is None else int(timeout_s * 1000) + events = dict(self.poller.poll(timeout=timeout_ms)) if self.event is not None and self.event.fd in events: self.event.drain() while True: diff --git a/mstar/communication/rust_communicator.py b/mstar/communication/rust_communicator.py index 172fc5ea1..505ae903f 100644 --- a/mstar/communication/rust_communicator.py +++ b/mstar/communication/rust_communicator.py @@ -1,39 +1,19 @@ -"""RFC #130 Step 1: ``ZMQCommunicator`` as a thin wrapper over the Rust -communicator vendored in ``rust/`` (``mstar_rust.ZmqCommunicator``). - -Drop-in for :class:`mstar.communication.communicator.ZMQCommunicator` — same -constructor, same methods, same semantics — with the transport moved to Rust: - -* **Wire-compatible with the existing pyzmq communicator.** The Rust transport - moves opaque byte frames (no added framing) over the same endpoints - (``ipc://{prefix}/{id}.ipc`` / the same TCP host+port scheme), and the - default codec is pickle — so a wrapped entity talks to unwrapped pyzmq - entities in both directions. Migration can proceed one process at a time. -* **encode/decode seam.** ``codec`` is a :class:`Codec` (mirroring the Rust - ``Codec`` trait) defaulting to :class:`PickleCodec` (today's wire). - Swapping to msgpack later is a codec change on both ends of an edge, - never a transport change. -* **eventfd wakeup.** ``register_event_for_poll`` forwards the ``EventWakeup`` - fd to the Rust poller (``register_wakeup_fd``): a completed compute future - wakes ``wait_for_work`` / ``poll_for_messages`` immediately, not on the poll - timeout. Drain semantics are identical (the event is drained here, in - Python, exactly as before). -* **Readiness without consumption.** The pyzmq poller reports "a message is - readable" without receiving it; the Rust ``recv_or_wake`` consumes. The - wrapper bridges the two with an internal deque: a message consumed during a - poll is buffered and handed out by the next ``get_all_new_messages`` — no - message is ever dropped or reordered. - -Requires the vendored extension — from ``rust/``: ``maturin develop --release`` -(or ``pip install .``). Use ``--release``: debug builds cost real latency on -the hot receive path. See ``docs/installation.rst``. -""" +"""RFC #130 Step 1: ``ZMQCommunicator`` over the Rust transport vendored in +``rust/`` (``mstar_rust.ZmqCommunicator``) — drop-in for the pyzmq class: +same constructor, methods, endpoints, and pickle wire, so wrapped and +unwrapped entities interoperate and migration can proceed one process at a +time. The codec is a :class:`Codec` (pickle default; msgpack for +cross-language edges); ``register_event_for_poll`` forwards the eventfd to +the Rust poller; a readiness poll buffers any consumed message so nothing is +dropped or reordered. Build: ``maturin develop --release`` in ``rust/`` +(see ``docs/installation.rst``).""" from __future__ import annotations import logging import os import pickle +import time from collections import deque from mstar_rust import ZmqCommunicator as _RustZmq @@ -73,6 +53,24 @@ class PickleCodec(Codec): decode = staticmethod(pickle.loads) +class MsgpackCodec(Codec): + """msgpack is language-neutral: edges using it can terminate in (future) + Rust processes — the migration's target wire. Both endpoints of an edge + must use the same codec. Requires the ``msgpack`` package.""" + + @staticmethod + def encode(msg) -> bytes: + import msgpack + + return msgpack.packb(msg, default=str) + + @staticmethod + def decode(data: bytes): + import msgpack + + return msgpack.unpackb(data, raw=False) + + class RustZMQCommunicator(BaseCommunicator): """The pyzmq ``ZMQCommunicator`` surface over the Rust transport.""" @@ -154,14 +152,20 @@ def send(self, entity_id: str, msg) -> None: self._register(entity_id) self._inner.send(entity_id, self.codec.encode(msg)) - def get_all_new_messages(self, blocking: bool = False) -> list: + def get_all_new_messages( + self, blocking: bool = False, timeout_s: float | None = None, + ) -> list: if blocking and not self._buffered: # Wait until a message is readable before draining — same # semantics as the pyzmq communicator: a registered wakeup # event also ends the wait, so a completed compute future can - # interrupt a blocking receive. + # interrupt a blocking receive. `timeout_s` bounds the wait + # (None = wait indefinitely); on expiry, drain whatever is there. + deadline = None if timeout_s is None else ( + time.monotonic() + timeout_s) while self._poll_once(50) == "timeout": - pass + if deadline is not None and time.monotonic() >= deadline: + break messages = [self.codec.decode(b) for b in self._buffered] self._buffered.clear() # One FFI call for the whole queued batch (vs try_recv per message). diff --git a/test/rust/test_rust_communicator.py b/test/rust/test_rust_communicator.py index 42faa3287..d6a1b2cfe 100644 --- a/test/rust/test_rust_communicator.py +++ b/test/rust/test_rust_communicator.py @@ -92,3 +92,14 @@ def make(value): make_communicator("m_default", push_ids=[], ipc_socket_path_prefix=str(tmp_path)), RustZMQCommunicator) + + +@pytest.mark.parametrize("receiver", ["rust", "orig"]) +def test_blocking_receive_timeout_bounds_the_wait(pair, receiver): + import time as _t + + orig, rust = pair + dst = rust if receiver == "rust" else orig + t0 = _t.monotonic() + assert dst.get_all_new_messages(blocking=True, timeout_s=0.15) == [] + assert 0.1 < _t.monotonic() - t0 < 2.0 From b9b0c8024f2ccbb7c93a4a883607c52c7481cb94 Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Thu, 16 Jul 2026 07:15:14 +0000 Subject: [PATCH 17/21] docs: describe the transport by what it is, not by migration step --- mstar/communication/rust_communicator.py | 2 +- rust/src/lib.rs | 2 +- test/rust/test_rust_communicator.py | 3 +-- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/mstar/communication/rust_communicator.py b/mstar/communication/rust_communicator.py index 505ae903f..fe607c7c2 100644 --- a/mstar/communication/rust_communicator.py +++ b/mstar/communication/rust_communicator.py @@ -1,4 +1,4 @@ -"""RFC #130 Step 1: ``ZMQCommunicator`` over the Rust transport vendored in +"""``ZMQCommunicator`` over the Rust transport vendored in ``rust/`` (``mstar_rust.ZmqCommunicator``) — drop-in for the pyzmq class: same constructor, methods, endpoints, and pickle wire, so wrapped and unwrapped entities interoperate and migration can proceed one process at a diff --git a/rust/src/lib.rs b/rust/src/lib.rs index f1650e57f..9a6232f79 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -1,4 +1,4 @@ -//! mstar's Rust transport (RFC #130 Step 1): the ZMQ PUSH/PULL mesh. +//! mstar's Rust transport: the ZMQ PUSH/PULL control mesh. //! `communicator.rs` is the transport + codec split; this file is the PyO3 //! surface (`mstar_rust.ZmqCommunicator`) the Python `RustZMQCommunicator` //! wrapper drives. Build: `maturin develop --release` in rust/. diff --git a/test/rust/test_rust_communicator.py b/test/rust/test_rust_communicator.py index d6a1b2cfe..c270faaa1 100644 --- a/test/rust/test_rust_communicator.py +++ b/test/rust/test_rust_communicator.py @@ -1,5 +1,4 @@ -"""RustZMQCommunicator interop with the pyzmq ZMQCommunicator (RFC #130 -Step 1): same mesh, pickle wire, eventfd wakeup, lossless readiness polls. +"""RustZMQCommunicator interop with the pyzmq ZMQCommunicator: same mesh, pickle wire, eventfd wakeup, lossless readiness polls. Skipped unless the ``mstar_rust`` extension (built from ``rust/``) is installed.""" import os From 15fc614a2448032619ec054a3812526d343338fa Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Thu, 16 Jul 2026 07:20:39 +0000 Subject: [PATCH 18/21] docs: wrap the interop test docstring (line length) --- test/rust/test_rust_communicator.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/rust/test_rust_communicator.py b/test/rust/test_rust_communicator.py index c270faaa1..3a1be0a9f 100644 --- a/test/rust/test_rust_communicator.py +++ b/test/rust/test_rust_communicator.py @@ -1,4 +1,5 @@ -"""RustZMQCommunicator interop with the pyzmq ZMQCommunicator: same mesh, pickle wire, eventfd wakeup, lossless readiness polls. +"""RustZMQCommunicator interop with the pyzmq ZMQCommunicator: same +mesh, pickle wire, eventfd wakeup, lossless readiness polls. Skipped unless the ``mstar_rust`` extension (built from ``rust/``) is installed.""" import os From bc17818cd6b6f8eb3e638f7a9d0a67fa9c4472f2 Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Thu, 16 Jul 2026 07:44:51 +0000 Subject: [PATCH 19/21] comm: name the blocking-wait slice constant MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The 50 ms is a signal-responsiveness bound, not a latency knob (events end a slice immediately) — worth saying at a named constant instead of a literal in the loop. --- mstar/communication/rust_communicator.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/mstar/communication/rust_communicator.py b/mstar/communication/rust_communicator.py index fe607c7c2..f69f80cf5 100644 --- a/mstar/communication/rust_communicator.py +++ b/mstar/communication/rust_communicator.py @@ -23,6 +23,12 @@ logger = logging.getLogger(__name__) +#: Slice bound for indefinite blocking waits. Events (message / wakeup) end a +#: slice immediately, so this is not a latency knob — it bounds how long one +#: FFI poll can hold the thread, keeping Python signal handling and the +#: ``timeout_s`` deadline check responsive. Matches ``wait_for_work``'s tick. +BLOCKING_POLL_SLICE_MS = 50 + class Codec: """The encode/decode seam, mirroring the Rust ``Codec`` trait:: @@ -163,7 +169,7 @@ def get_all_new_messages( # (None = wait indefinitely); on expiry, drain whatever is there. deadline = None if timeout_s is None else ( time.monotonic() + timeout_s) - while self._poll_once(50) == "timeout": + while self._poll_once(BLOCKING_POLL_SLICE_MS) == "timeout": if deadline is not None and time.monotonic() >= deadline: break messages = [self.codec.decode(b) for b in self._buffered] From 838fb63a24a19fa8d0090e851dd887e380fdfb3a Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Fri, 17 Jul 2026 17:27:05 +0000 Subject: [PATCH 20/21] =?UTF-8?q?comm:=20review=20fixes=20=E2=80=94=20lock?= =?UTF-8?q?=20scope,=20loud=20decode,=20parity,=20version,=20fd=20guard?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Five review items on the transport: - send() no longer holds the peers-map mutex across the socket send: a PUSH send blocks at the peer's high-water mark, so one stalled consumer would freeze sends to every other peer (and registration) process-wide once anything drives this layer multi-threaded. The map now stores Arc>: the map lock covers only lookup/insert, the per-socket lock serializes sends to one peer (per-peer FIFO kept), and an HWM-blocked send stalls only that peer. - The typed layer's Codec::decode returns Result: a frame that fails to decode is a codec/version mismatch with the peer, and silently conflating it with "no message" would turn skew into requests that hang with no log line in the later Rust consumers of this API. try_recv/recv/recv_timeout return Result>, drain Result>; new test proves garbage surfaces as CommError::Decode, never as None. - poll_for_messages now exists on BOTH communicators (the pyzmq twin is the poller without consuming), restoring the factory's drop-in guarantee; parity test covers the contract on each. - Transport observability: the factory logs the chosen transport (and extension version) at INFO on construction; mstar_rust exports __version__, and a mismatch against the tree's expected version warns loudly — a stale wheel can no longer take over the mesh silently under AUTO. README notes what AUTO-by-default means operationally. - recv_or_wake distinguishes a closed/invalid wakeup fd (POLLERR/POLLNVAL) from a timeout and fails loudly instead of degrading every blocking wait into an instant-timeout CPU spin. --- README.md | 5 +- mstar/communication/communicator.py | 34 ++++++++ rust/src/communicator.rs | 129 +++++++++++++++++++++------- rust/src/lib.rs | 16 +++- test/rust/test_rust_communicator.py | 18 ++++ 5 files changed, 164 insertions(+), 38 deletions(-) diff --git a/README.md b/README.md index 6ac9df16e..f70a51207 100644 --- a/README.md +++ b/README.md @@ -143,7 +143,10 @@ primitives — `Sequential`, `Parallel`, `Loop`, and a cross-partition ### Optional: Rust ZMQ transport The ZeroMQ control mesh can run over a Rust transport (vendored in [`rust/`](rust/)) instead of -pyzmq — wire-compatible, selectable per process with `MSTAR_RUST_ZMQ=1` (or `AUTO`). Build it with +pyzmq — wire-compatible, selectable per process with `MSTAR_RUST_ZMQ` (default `AUTO`). +Note what AUTO-by-default means: on any machine where the extension is built, the whole +mesh switches to the Rust transport with no configuration change — each process logs its +choice at startup (`control mesh transport: ...`), and `MSTAR_RUST_ZMQ=0` pins pyzmq. Build it with [maturin](https://www.maturin.rs): `uv pip install maturin && maturin develop --release -m rust/Cargo.toml`. See the [installation docs](https://m-star.org/mstar/installation.html) and [environment variables](https://m-star.org/mstar/environment_variables.html). diff --git a/mstar/communication/communicator.py b/mstar/communication/communicator.py index f1fb85ce6..e953ae303 100644 --- a/mstar/communication/communicator.py +++ b/mstar/communication/communicator.py @@ -9,6 +9,12 @@ logger = logging.getLogger(__name__) +#: The ``mstar_rust`` extension version this tree expects (the vendored +#: ``rust/`` crate's version). Under ``MSTAR_RUST_ZMQ=AUTO`` a mismatching +#: install - e.g. a stale wheel after an upgrade - takes over the mesh +#: silently, so the factory warns when the imported version differs. +EXPECTED_MSTAR_RUST_VERSION = "0.1.0" + class CommProtocol(Enum): IPC = "IPC" @@ -110,6 +116,18 @@ def wait_for_work(self, timeout_ms=50): if self.event.fd in events: self.event.drain() + def poll_for_messages(self, timeout_ms=20): + """Block until a message is readable, a registered wakeup event + fires, or ``timeout_ms`` elapses — whichever comes first. True when + a message is available (left queued for ``get_all_new_messages``); + a wakeup ends the poll early with False (the event is drained, + exactly as in ``wait_for_work``). Mirrors the Rust communicator's + method so call sites work against either transport.""" + events = dict(self.poller.poll(timeout=timeout_ms)) + if self.event is not None and self.event.fd in events: + self.event.drain() + return self.pull_socket in events + # def get_session_id(self) -> str: # return self.session_id @@ -176,11 +194,27 @@ def make_communicator(*args, **kwargs) -> BaseCommunicator: raise ValueError(f"MSTAR_RUST_ZMQ must be 0, 1, or AUTO; got {choice!r}") if choice != "0": try: + import mstar_rust + from mstar.communication.rust_communicator import RustZMQCommunicator except ImportError: if choice == "1": raise logger.debug("MSTAR_RUST_ZMQ=AUTO: mstar_rust not installed, using pyzmq") else: + # A support bundle must be able to tell what a mesh was running, + # and an old wheel left in an env must not silently take over + # the whole mesh under AUTO after an upgrade. + version = getattr(mstar_rust, "__version__", "") + logger.info( + "control mesh transport: rust %s (MSTAR_RUST_ZMQ=%s)", + version, choice) + if version != EXPECTED_MSTAR_RUST_VERSION: + logger.warning( + "mstar_rust version %s does not match this tree's " + "expected %s — a stale wheel may be shadowing the " + "vendored rust/ build (rebuild with `maturin develop " + "--release`)", version, EXPECTED_MSTAR_RUST_VERSION) return RustZMQCommunicator(*args, **kwargs) + logger.info("control mesh transport: pyzmq (MSTAR_RUST_ZMQ=%s)", choice) return ZMQCommunicator(*args, **kwargs) diff --git a/rust/src/communicator.rs b/rust/src/communicator.rs index b74df8044..368696db2 100644 --- a/rust/src/communicator.rs +++ b/rust/src/communicator.rs @@ -44,6 +44,9 @@ pub enum CommError { Io(#[from] std::io::Error), #[error("no endpoint for peer '{0}' (register_peer it, or bind with an ipc dir)")] UnknownPeer(String), + #[error("decode failed on a {len}-byte frame: {detail} — codec/version \ + mismatch with the peer?")] + Decode { len: usize, detail: String }, } /// The default ipc endpoint an entity binds/connects to. @@ -66,6 +69,10 @@ pub enum RecvEvent { /// A registered wakeup fd is readable (the registrant reads/clears it — /// e.g. `read(2)` on the eventfd — the transport only polls it). Wake, + /// A registered wakeup fd is closed/invalid (POLLERR/POLLNVAL): the + /// registrant's bug — surfaced loudly instead of degrading every + /// blocking wait into an instant-timeout CPU spin. + WakeFdError, /// The timeout elapsed with neither. Timeout, } @@ -75,6 +82,7 @@ impl PartialEq for RecvEvent { match (self, other) { (RecvEvent::Message(a), RecvEvent::Message(b)) => a[..] == b[..], (RecvEvent::Wake, RecvEvent::Wake) => true, + (RecvEvent::WakeFdError, RecvEvent::WakeFdError) => true, (RecvEvent::Timeout, RecvEvent::Timeout) => true, _ => false, } @@ -97,7 +105,12 @@ pub struct RawZmqCommunicator { // used from one thread at a time); the drive loop is the single consumer. pull: Mutex, // peer id -> connected PUSH socket (created on first send to that peer). - peers: Mutex>, + // Arc> so the MAP lock is held only for lookup/insert: + // a PUSH send blocks at the peer's high-water mark, and holding the map + // across it would let one stalled peer freeze sends to every other peer + // (and registration) process-wide. The per-socket lock serializes sends + // to one peer (a zmq socket is single-user), preserving per-peer FIFO. + peers: Mutex>>>, // peer id -> explicit endpoint (tcp or ipc), overriding the dir scheme. peer_endpoints: Mutex>, // External fds polled alongside the PULL socket (mstar's worker registers @@ -201,18 +214,22 @@ impl RawZmqCommunicator { /// Send one opaque byte frame to `peer_id` (fire-and-forget). Queues if /// the peer isn't bound yet and reconnects transparently if it restarted. pub fn send(&self, peer_id: &str, payload: &[u8]) -> Result<(), CommError> { - let mut peers = self.peers.lock().expect("peers lock"); - if !peers.contains_key(peer_id) { - let endpoint = self.resolve(peer_id)?; - let push = self.ctx.socket(zmq::PUSH)?; - push.set_linger(0)?; - push.connect(&endpoint)?; - peers.insert(peer_id.to_string(), push); - } - peers - .get(peer_id) - .expect("just inserted") - .send(payload, 0)?; + let socket = { + let mut peers = self.peers.lock().expect("peers lock"); + match peers.get(peer_id) { + Some(s) => s.clone(), + None => { + let endpoint = self.resolve(peer_id)?; + let push = self.ctx.socket(zmq::PUSH)?; + push.set_linger(0)?; + push.connect(&endpoint)?; + let s = std::sync::Arc::new(Mutex::new(push)); + peers.insert(peer_id.to_string(), s.clone()); + s + } + } + }; // map lock dropped: an HWM-blocked send stalls only THIS peer + socket.lock().expect("peer socket lock").send(payload, 0)?; Ok(()) } @@ -260,6 +277,16 @@ impl RawZmqCommunicator { return RecvEvent::Message(b); } } + // A closed/invalid wakeup fd makes poll return POLLNVAL instantly: + // without this check every blocking wait would degrade into a + // silent 100%-CPU spin. Fail loudly instead. + if items[1..] + .iter() + .any(|it| it.get_revents().intersects( + zmq::POLLERR | zmq::PollEvents::from_bits_truncate(0x20))) + { + return RecvEvent::WakeFdError; + } if items[1..].iter().any(|it| it.is_readable()) { return RecvEvent::Wake; } @@ -296,7 +323,11 @@ impl Drop for RawZmqCommunicator { /// transport never looks inside the bytes. pub trait Codec { fn encode(msg: &M) -> Result, CommError>; - fn decode(bytes: &[u8]) -> Option; + /// Errors are LOUD by contract: a frame that fails to decode is a + /// codec/version mismatch with the peer, and silently dropping it turns + /// into requests that hang with no log line. The typed receive methods + /// propagate the error instead of conflating it with "no message". + fn decode(bytes: &[u8]) -> Result; } /// The default codec: MessagePack, the language-neutral encoding the @@ -308,8 +339,11 @@ impl Codec for MsgpackCodec { fn encode(msg: &M) -> Result, CommError> { rmp_serde::to_vec_named(msg).map_err(CommError::Serialize) } - fn decode(bytes: &[u8]) -> Option { - rmp_serde::from_slice(bytes).ok() + fn decode(bytes: &[u8]) -> Result { + rmp_serde::from_slice(bytes).map_err(|e| CommError::Decode { + len: bytes.len(), + detail: e.to_string(), + }) } } @@ -355,28 +389,32 @@ where self.raw.send(peer_id, &C::encode(msg)?) } - /// Non-blocking: next inbound message, or None. - pub fn try_recv(&self) -> Option { - self.raw.try_recv().and_then(|b| C::decode(&b)) + /// Non-blocking: next inbound message. `Ok(None)` = nothing queued; + /// `Err` = a frame arrived but failed to decode (codec/version skew). + pub fn try_recv(&self) -> Result, CommError> { + self.raw.try_recv().map(|b| C::decode(&b)).transpose() } /// Block until the next inbound message. - pub fn recv(&self) -> Option { - self.raw.recv().and_then(|b| C::decode(&b)) + pub fn recv(&self) -> Result, CommError> { + self.raw.recv().map(|b| C::decode(&b)).transpose() } /// Block up to `timeout` for the next inbound message (a wakeup fd cuts /// the wait short — see [`RawZmqCommunicator::recv_or_wake`]). - pub fn recv_timeout(&self, timeout: Duration) -> Option { - self.raw.recv_timeout(timeout).and_then(|b| C::decode(&b)) + pub fn recv_timeout(&self, timeout: Duration) -> Result, CommError> { + self.raw + .recv_timeout(timeout) + .map(|b| C::decode(&b)) + .transpose() } /// Drain all currently-queued inbound messages. - pub fn drain(&self) -> Vec { + pub fn drain(&self) -> Result, CommError> { self.raw .drain() .into_iter() - .filter_map(|b| C::decode(&b)) + .map(|b| C::decode(&b)) .collect() } } @@ -429,11 +467,11 @@ mod tests { node: "LLM".into(), }) .unwrap(); - let got = wait_for(&worker, |w| w.try_recv()); + let got = wait_for(&worker, |w| w.try_recv().unwrap()); assert_eq!(got, Msg::Batch { id: 1, node: "LLM".into() }); worker.send("conductor", &Msg::Hello("done".into())).unwrap(); - let got = wait_for(&conductor, |c| c.try_recv()); + let got = wait_for(&conductor, |c| c.try_recv().unwrap()); assert_eq!(got, Msg::Hello("done".into())); } @@ -448,7 +486,7 @@ mod tests { // PUSH->PULL over one connection preserves FIFO order. let mut seen = 0u64; for _ in 0..1000 { - for m in b.drain() { + for m in b.drain().unwrap() { if let Msg::Batch { id, .. } = m { assert_eq!(id, seen); seen += 1; @@ -470,7 +508,7 @@ mod tests { let a: ZmqCommunicator = ZmqCommunicator::bind("a", &dir).unwrap(); a.send("late", &Msg::Hello("queued".into())).unwrap(); // no peer yet — no error let late: ZmqCommunicator = ZmqCommunicator::bind("late", &dir).unwrap(); - let got = wait_for(&late, |m| m.try_recv()); + let got = wait_for(&late, |m| m.try_recv().unwrap()); assert_eq!(got, Msg::Hello("queued".into())); } @@ -481,13 +519,13 @@ mod tests { { let b: ZmqCommunicator = ZmqCommunicator::bind("b", &dir).unwrap(); a.send("b", &Msg::Hello("1".into())).unwrap(); - wait_for(&b, |b| b.try_recv()); + wait_for(&b, |b| b.try_recv().unwrap()); } // b drops; zmq unlinks its inbox std::thread::sleep(Duration::from_millis(20)); // New b at the same id: a's cached PUSH auto-reconnects to it. let b2: ZmqCommunicator = ZmqCommunicator::bind("b", &dir).unwrap(); a.send("b", &Msg::Hello("2".into())).unwrap(); - let got = wait_for(&b2, |b| b.try_recv()); + let got = wait_for(&b2, |b| b.try_recv().unwrap()); assert_eq!(got, Msg::Hello("2".into())); } @@ -622,9 +660,34 @@ mod tests { fn encode(msg: &String) -> Result, CommError> { Ok(msg.as_bytes().to_vec()) } - fn decode(bytes: &[u8]) -> Option { - String::from_utf8(bytes.to_vec()).ok() + fn decode(bytes: &[u8]) -> Result { + String::from_utf8(bytes.to_vec()).map_err(|e| CommError::Decode { + len: bytes.len(), + detail: e.to_string(), + }) + } + } + + #[test] + fn decode_failure_is_loud_not_silent() { + // A mis-codec'd frame must surface as Err (codec/version skew with + // the peer), never be conflated with "no message". + let dir = tmpdir("badframe"); + let a = RawZmqCommunicator::bind("a", &dir).unwrap(); + let b: ZmqCommunicator = ZmqCommunicator::bind("b", &dir).unwrap(); + a.send("b", b"\xc1 definitely not msgpack").unwrap(); + for _ in 0..500 { + match b.try_recv() { + Ok(None) => std::thread::sleep(Duration::from_millis(4)), + Ok(Some(_)) => panic!("garbage decoded as a message"), + Err(CommError::Decode { len, .. }) => { + assert!(len > 0); + return; + } + Err(e) => panic!("wrong error: {e}"), + } } + panic!("frame never arrived"); } #[test] diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 9a6232f79..9847321af 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -82,23 +82,31 @@ impl PyZmqCommunicator { .map(|b| PyBytes::new(py, &b)) } - /// ("msg", bytes) | ("wake", None) | ("timeout", None). + /// ("msg", bytes) | ("wake", None) | ("timeout", None). Raises if a + /// registered wakeup fd is closed/invalid (the registrant's bug — loud, + /// instead of degrading blocking waits into an instant-timeout spin). fn recv_or_wake<'py>( &self, py: Python<'py>, timeout_ms: u64, - ) -> (&'static str, Option>) { + ) -> PyResult<(&'static str, Option>)> { let ev = py.allow_threads(|| self.inner.recv_or_wake(Duration::from_millis(timeout_ms))); - match ev { + Ok(match ev { RecvEvent::Message(b) => ("msg", Some(PyBytes::new(py, &b))), RecvEvent::Wake => ("wake", None), + RecvEvent::WakeFdError => { + return Err(PyRuntimeError::new_err( + "a registered wakeup fd is closed/invalid (POLLERR/\ + POLLNVAL): fix the EventWakeup lifetime")); + } RecvEvent::Timeout => ("timeout", None), - } + }) } } #[pymodule] fn mstar_rust(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add("__version__", env!("CARGO_PKG_VERSION"))?; m.add_class::()?; Ok(()) } diff --git a/test/rust/test_rust_communicator.py b/test/rust/test_rust_communicator.py index 3a1be0a9f..2206d3f0d 100644 --- a/test/rust/test_rust_communicator.py +++ b/test/rust/test_rust_communicator.py @@ -103,3 +103,21 @@ def test_blocking_receive_timeout_bounds_the_wait(pair, receiver): t0 = _t.monotonic() assert dst.get_all_new_messages(blocking=True, timeout_s=0.15) == [] assert 0.1 < _t.monotonic() - t0 < 2.0 + + +def test_poll_for_messages_parity(tmp_path): + """Both transports expose poll_for_messages with the same contract, so a + call site written against the factory works on either (the drop-in + guarantee): True leaves the message queued for get_all_new_messages.""" + from mstar.communication.communicator import ZMQCommunicator + from mstar.communication.rust_communicator import RustZMQCommunicator + + for cls, tag in ((ZMQCommunicator, "py"), (RustZMQCommunicator, "rs")): + a = cls(f"pfm_{tag}_a", [f"pfm_{tag}_b"], + ipc_socket_path_prefix=str(tmp_path) + "/") + b = cls(f"pfm_{tag}_b", [f"pfm_{tag}_a"], + ipc_socket_path_prefix=str(tmp_path) + "/") + assert b.poll_for_messages(timeout_ms=10) is False + a.send(f"pfm_{tag}_b", {"hello": tag}) + assert b.poll_for_messages(timeout_ms=2000) is True + assert b.get_all_new_messages() == [{"hello": tag}] From 5b18c1db52c7881d40a71d5d5024976b71ac9278 Mon Sep 17 00:00:00 2001 From: npuichigo <418121364@qq.com> Date: Fri, 17 Jul 2026 19:32:11 +0000 Subject: [PATCH 21/21] comm: drop the dead half of the wakeup-fd error guard PollEvents::from_bits_truncate(0x20) truncates to empty (0x20 is not a defined flag in the zmq crate); the guard worked because libzmq folds a raw fd's POLLNVAL/POLLERR into ZMQ_POLLERR, which the other half already catches. Keep only the live half, with the explanation. --- rust/src/communicator.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/rust/src/communicator.rs b/rust/src/communicator.rs index 368696db2..c0b4a7c1d 100644 --- a/rust/src/communicator.rs +++ b/rust/src/communicator.rs @@ -277,13 +277,13 @@ impl RawZmqCommunicator { return RecvEvent::Message(b); } } - // A closed/invalid wakeup fd makes poll return POLLNVAL instantly: - // without this check every blocking wait would degrade into a - // silent 100%-CPU spin. Fail loudly instead. + // A closed/invalid wakeup fd makes poll return instantly with an + // error event: without this check every blocking wait would degrade + // into a silent 100%-CPU spin. libzmq folds a raw fd's POLLNVAL and + // POLLERR into ZMQ_POLLERR, so this one flag covers both. if items[1..] .iter() - .any(|it| it.get_revents().intersects( - zmq::POLLERR | zmq::PollEvents::from_bits_truncate(0x20))) + .any(|it| it.get_revents().contains(zmq::POLLERR)) { return RecvEvent::WakeFdError; }