communication: vendored Rust transport (rust/) + RustZMQCommunicator (RFC #130 Step 1)

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

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

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_<rank>`` = base+100+rank).

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

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
+<https://www.maturin.rs>`_ (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
 ---------------
 

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,

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,

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__", "<pre-versioning>")
+            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)