diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 227a7b8e0..10e7f888d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -17,4 +17,35 @@ 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 (rust/): cargo tests, then the pyzmq interop + # pytest against a freshly built extension. + 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 + # 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 diff --git a/README.md b/README.md index d6ddd2c21..f70a51207 100644 --- a/README.md +++ b/README.md @@ -140,6 +140,17 @@ 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` (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). + ## 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..5957fab13 --- /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`` + - ``AUTO`` + - Transport selection for the ZeroMQ control mesh (see + :func:`mstar.communication.communicator.make_communicator`). + ``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 + 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 --------------- 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..e953ae303 100644 --- a/mstar/communication/communicator.py +++ b/mstar/communication/communicator.py @@ -9,6 +9,19 @@ 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" + TCP = "TCP" + RDMA = "RDMA" + SHM = "SHM" + class BaseCommunicator(ABC): @abstractmethod @@ -22,18 +35,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,28 +116,17 @@ 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 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 @@ -123,8 +144,18 @@ 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. `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: try: # zmq.NOBLOCK means zmq doesn't wait for a new message to be @@ -141,3 +172,49 @@ 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`` selects it (see ``docs/environment_variables.rst``): + + * ``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", "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": + 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/mstar/communication/rust_communicator.py b/mstar/communication/rust_communicator.py new file mode 100644 index 000000000..f69f80cf5 --- /dev/null +++ b/mstar/communication/rust_communicator.py @@ -0,0 +1,179 @@ +"""``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 + +from mstar.communication.communicator import BaseCommunicator, CommProtocol +from mstar.communication.event import EventWakeup + +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:: + + 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 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.""" + + def __init__( + self, + my_id: str, + push_ids: list[str], + protocol: CommProtocol = CommProtocol.IPC, + ipc_socket_path_prefix: str = "/tmp/mstar/", + 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.codec = 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: `_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: + 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) -> 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: + return # work is already waiting + self._poll_once(timeout_ms) + + def poll_for_messages(self, timeout_ms: int = 20) -> bool: + """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) + 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.codec.encode(msg)) + + 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. `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(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] + self._buffered.clear() + # 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/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, 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..f41870dab --- /dev/null +++ b/rust/Cargo.lock @@ -0,0 +1,625 @@ +# 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 = "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 = [ + "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" +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 = "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" +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..721746a2d --- /dev/null +++ b/rust/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "mstar-rust" +version = "0.1.0" +edition = "2021" +license = "Apache-2.0" + +[lib] +name = "mstar_rust" +# 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"] } +rmp-serde = "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..c0b4a7c1d --- /dev/null +++ b/rust/src/communicator.rs @@ -0,0 +1,710 @@ +//! 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 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 ([`MsgpackCodec`] by default — the language-neutral wire 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(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)")] + 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. +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. `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(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, + /// 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, +} + +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::WakeFdError, RecvEvent::WakeFdError) => 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. +/// +/// 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). + // 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 + // 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 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(()) + } + + /// Non-blocking: next inbound frame, or None. + pub fn try_recv(&self) -> Option { + let pull = self.pull.lock().expect("pull lock"); + pull.recv_msg(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_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 { + 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_msg(zmq::DONTWAIT) { + return RecvEvent::Message(b); + } + } + // 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().contains(zmq::POLLERR)) + { + return RecvEvent::WakeFdError; + } + 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_msg(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 for language-neutral ones). The +/// transport never looks inside the bytes. +pub trait Codec { + fn encode(msg: &M) -> Result, CommError>; + /// 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 +/// 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 MsgpackCodec { + fn encode(msg: &M) -> Result, CommError> { + rmp_serde::to_vec_named(msg).map_err(CommError::Serialize) + } + fn decode(bytes: &[u8]) -> Result { + rmp_serde::from_slice(bytes).map_err(|e| CommError::Decode { + len: bytes.len(), + detail: e.to_string(), + }) + } +} + +/// A typed mailbox: [`RawZmqCommunicator`] + a [`Codec`]. `M` is the entity's +/// message type; the default codec is MessagePack. +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. `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) -> 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) -> Result, CommError> { + self.raw + .recv_timeout(timeout) + .map(|b| C::decode(&b)) + .transpose() + } + + /// Drain all currently-queued inbound messages. + pub fn drain(&self) -> Result, CommError> { + self.raw + .drain() + .into_iter() + .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().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().unwrap()); + 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().unwrap() { + 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().unwrap()); + 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().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().unwrap()); + 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]) -> 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] + 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..9847321af --- /dev/null +++ b/rust/src/lib.rs @@ -0,0 +1,112 @@ +//! 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/. + +pub 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); + } + + /// 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())) + } + + fn try_recv<'py>(&self, py: Python<'py>) -> Option> { + 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)) + } + + /// ("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, + ) -> PyResult<(&'static str, Option>)> { + let ev = py.allow_threads(|| self.inner.recv_or_wake(Duration::from_millis(timeout_ms))); + 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 new file mode 100644 index 000000000..2206d3f0d --- /dev/null +++ b/test/rust/test_rust_communicator.py @@ -0,0 +1,123 @@ +"""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 +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"] + + +@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") + # 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) + + +@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 + + +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}]