mplang-nightly 0.1.dev266__py3-none-any.whl → 0.1.dev267__py3-none-any.whl

mplang/v2/backends/channel.py

@@ -0,0 +1,217 @@
+# Copyright 2026 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""SPU IChannel implementation for MPLang v2.
+
+Bridges v2's simp_worker communicators (ThreadCommunicator/HttpCommunicator)
+to libspu's IChannel interface, enabling SPU to reuse existing communication
+infrastructure instead of creating separate BRPC connections.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Protocol
+
+import spu.libspu as libspu
+
+
+class CommunicatorProtocol(Protocol):
+    """Protocol for v2 communicators (duck typing).
+
+    Both ThreadCommunicator and HttpCommunicator implement this interface.
+    """
+
+    def send(self, to: int, key: str, data: bytes) -> None: ...
+    def recv(self, frm: int, key: str) -> bytes: ...
+
+
+class BaseChannel(libspu.link.IChannel):
+    """Bridge v2 communicator to SPU IChannel interface.
+
+    Supports both ThreadCommunicator and HttpCommunicator via duck typing.
+    Each BaseChannel represents a channel to ONE peer rank.
+
+    Communication Protocol:
+        - SPU calls send(tag, bytes) -> comm.send(peer, "spu:tag", bytes)
+        - SPU calls recv(tag) -> bytes <- comm.recv(peer, "spu:tag")
+
+    Tag Namespace:
+        All tags are prefixed with "spu:" to avoid collision with other
+        traffic on the same communicator.
+    """
+
+    def __init__(
+        self,
+        comm: CommunicatorProtocol,
+        local_rank: int,
+        peer_rank: int,
+        tag_prefix: str = "spu",
+    ):
+        """Initialize channel to a specific peer.
+
+        Args:
+            comm: v2 communicator (any object implementing send/recv)
+            local_rank: Global rank of this party
+            peer_rank: Global rank of the peer party
+            tag_prefix: Prefix for all tags (default: "spu")
+        """
+        super().__init__()
+        self._comm = comm
+        self._local_rank = local_rank
+        self._peer_rank = peer_rank
+        self._tag_prefix = tag_prefix
+
+        logging.debug(
+            f"BaseChannel initialized: local_rank={local_rank}, "
+            f"peer_rank={peer_rank}, tag_prefix={tag_prefix}"
+        )
+
+    def _make_key(self, tag: str) -> str:
+        """Create unique key for communicator.
+
+        Args:
+            tag: SPU-provided tag (e.g., "send_0")
+
+        Returns:
+            Prefixed key (e.g., "spu:send_0")
+        """
+        return f"{self._tag_prefix}:{tag}"
+
+    def Send(self, tag: str, data: bytes) -> None:
+        """Send bytes to peer.
+
+        Args:
+            tag: Message tag for matching send/recv pairs
+            data: Raw bytes to send
+        """
+        key = self._make_key(tag)
+        logging.debug(
+            f"BaseChannel.Send: {self._local_rank} -> {self._peer_rank}, "
+            f"tag={tag}, key={key}, size={len(data)}"
+        )
+
+        # Send raw bytes directly
+        # v2 communicators accept Any, bytes is valid
+        self._comm.send(self._peer_rank, key, data)
+
+    def Recv(self, tag: str) -> bytes:
+        """Receive bytes from peer (blocking).
+
+        Args:
+            tag: Message tag for matching send/recv pairs
+
+        Returns:
+            Raw bytes received
+
+        Raises:
+            TypeError: If received data is not bytes
+        """
+        key = self._make_key(tag)
+        logging.debug(
+            f"BaseChannel.Recv: {self._local_rank} <- {self._peer_rank}, "
+            f"tag={tag}, key={key}"
+        )
+
+        # Receive data (should be bytes)
+        data = self._comm.recv(self._peer_rank, key)
+
+        # Validate data type
+        if not isinstance(data, bytes):
+            raise TypeError(
+                f"Expected bytes from communicator, got {type(data).__name__}. "
+                f"Communicator must support raw bytes transmission for SPU channels."
+            )
+
+        logging.debug(
+            f"BaseChannel.Recv complete: {self._local_rank} <- {self._peer_rank}, "
+            f"tag={tag}, size={len(data)}"
+        )
+        return data
+
+    def SendAsync(self, tag: str, data: bytes) -> None:
+        """Async send.
+
+        For HttpCommunicator, underlying HTTP client is non-blocking.
+        For ThreadCommunicator, send is instant (memory transfer).
+
+        Args:
+            tag: Message tag
+            data: Raw bytes to send
+        """
+        self.Send(tag, data)
+
+    def SendAsyncThrottled(self, tag: str, data: bytes) -> None:
+        """Throttled async send.
+
+        Currently maps to regular SendAsync.
+
+        Args:
+            tag: Message tag
+            data: Raw bytes to send
+        """
+        self.SendAsync(tag, data)
+
+    def TestSend(self, timeout: int) -> None:
+        """Test if channel can send a dummy message to peer.
+
+        Uses fixed tag "__test__" for idempotency.
+
+        Args:
+            timeout: Timeout in milliseconds (informational)
+        """
+        test_data = b"\x00"  # Minimal 1-byte handshake
+        self.Send("__test__", test_data)
+
+    def TestRecv(self) -> None:
+        """Wait for dummy message from peer.
+
+        Timeout controlled by recv_timeout_ms in link descriptor.
+
+        Raises:
+            Warning if unexpected handshake data received
+        """
+        test_data = self.Recv("__test__")
+        if test_data != b"\x00":
+            logging.warning(
+                f"TestRecv: unexpected handshake from {self._peer_rank}, "
+                f"expected b'\\x00', got {test_data!r}"
+            )
+
+    def WaitLinkTaskFinish(self) -> None:
+        """Wait for all pending async tasks.
+
+        No-op for v2 communicators (handled automatically).
+        """
+
+    def Abort(self) -> None:
+        """Abort communication (cleanup).
+
+        Currently a no-op. Could be extended for resource cleanup.
+        """
+        logging.warning(f"BaseChannel.Abort: {self._local_rank} <-> {self._peer_rank}")
+
+    def SetThrottleWindowSize(self, size: int) -> None:
+        """Set throttle window size (no-op).
+
+        Args:
+            size: Window size (ignored)
+        """
+
+    def SetChunkParallelSendSize(self, size: int) -> None:
+        """Set chunk parallel send size (no-op).
+
+        Args:
+            size: Chunk size (ignored)
+        """

mplang/v2/backends/simp_worker/http.py

@@ -83,7 +83,7 @@ class HttpCommunicator:
         self.world_size = world_size
         self.endpoints = endpoints
         self.tracer = tracer
-        self._mailbox: dict[str, Any] = {}
+        self._mailbox: dict[tuple[int, str], Any] = {}
         self._cond = threading.Condition()
         self._send_executor = concurrent.futures.ThreadPoolExecutor(
             max_workers=world_size, thread_name_prefix=f"comm_send_{rank}"
@@ -100,8 +100,19 @@ class HttpCommunicator:
         """Perform the HTTP send."""
         url = f"{self.endpoints[to]}/comm/{key}"
         logger.debug(f"Rank {self.rank} sending to {to} key={key}")
-        # Use secure JSON serialization
-        payload = serde.dumps_b64(data)
+
+        # Detect SPU channel (tag prefix "spu:") and handle bytes
+        if key.startswith("spu:") and isinstance(data, bytes):
+            # Send raw bytes for SPU channels
+            import base64
+
+            payload = base64.b64encode(data).decode("ascii")
+            is_raw_bytes = True
+        else:
+            # Use secure JSON serialization
+            payload = serde.dumps_b64(data)
+            is_raw_bytes = False
+
         size_bytes = len(payload)
 
         # Log to profiler
@@ -116,7 +127,14 @@ class HttpCommunicator:
 
         try:
             t0 = time.time()
-            resp = self.client.put(url, json={"data": payload, "from_rank": self.rank})
+            resp = self.client.put(
+                url,
+                json={
+                    "data": payload,
+                    "from_rank": self.rank,
+                    "is_raw_bytes": is_raw_bytes,
+                },
+            )
             resp.raise_for_status()
             duration = time.time() - t0
             if self.tracer:
@@ -134,17 +152,21 @@ class HttpCommunicator:
     def recv(self, frm: int, key: str) -> Any:
         """Receive data from another rank (blocking)."""
         logger.debug(f"Rank {self.rank} waiting recv from {frm} key={key}")
+        mailbox_key = (frm, key)
         with self._cond:
-            while key not in self._mailbox:
+            while mailbox_key not in self._mailbox:
                 self._cond.wait(timeout=1.0)
-            return self._mailbox.pop(key)
+            return self._mailbox.pop(mailbox_key)
 
-    def on_receive(self, key: str, data: Any) -> None:
+    def on_receive(self, from_rank: int, key: str, data: Any) -> None:
         """Called when data is received from the HTTP endpoint."""
+        mailbox_key = (from_rank, key)
         with self._cond:
-            if key in self._mailbox:
-                logger.warning(f"Rank {self.rank} overwriting key={key}")
-            self._mailbox[key] = data
+            if mailbox_key in self._mailbox:
+                raise RuntimeError(
+                    f"Mailbox overflow: key {mailbox_key} already exists"
+                )
+            self._mailbox[mailbox_key] = data
             self._cond.notify_all()
 
     def wait_pending_sends(self) -> None:
@@ -176,6 +198,7 @@ class CommRequest(BaseModel):
 
     data: str
     from_rank: int
+    is_raw_bytes: bool = False  # NEW: indicates raw bytes (not serde)
 
 
 class FetchRequest(BaseModel):
@@ -279,9 +302,17 @@ def create_worker_app(
         """Receive communication data from another worker."""
         logger.debug(f"Worker {rank} received comm key={key} from {req.from_rank}")
         try:
-            # Use secure JSON deserialization
-            data = serde.loads_b64(req.data)
-            comm.on_receive(key, data)
+            # Handle raw bytes (SPU channels) vs serde data
+            if req.is_raw_bytes:
+                # Decode base64 to raw bytes
+                import base64
+
+                data = base64.b64decode(req.data)
+            else:
+                # Use secure JSON deserialization
+                data = serde.loads_b64(req.data)
+
+            comm.on_receive(req.from_rank, key, data)
             return {"status": "ok"}
         except Exception as e:
             logger.error(f"Worker {rank} comm failed: {e}")

mplang/v2/backends/simp_worker/mem.py

@@ -35,7 +35,8 @@ class ThreadCommunicator:
         self.world_size = world_size
         self.use_serde = use_serde
         self.peers: list[ThreadCommunicator] = []
-        self._mailbox: dict[str, Any] = {}
+        # Mailbox keyed by (from_rank, tag): each key has exactly one message
+        self._mailbox: dict[tuple[int, str], Any] = {}
         self._cond = threading.Condition()
         self._sent_events: dict[str, threading.Event] = {}
         self._shutdown = False
@@ -58,20 +59,22 @@ class ThreadCommunicator:
         self.peers[to]._on_receive(self.rank, key, data)
 
     def recv(self, frm: int, key: str) -> Any:
+        mailbox_key = (frm, key)
         with self._cond:
-            while key not in self._mailbox and not self._shutdown:
+            while mailbox_key not in self._mailbox and not self._shutdown:
                 self._cond.wait()
             if self._shutdown:
                 raise RuntimeError("Communicator shut down")
-            return self._mailbox.pop(key)
+            return self._mailbox.pop(mailbox_key)
 
     def _on_receive(self, frm: int, key: str, data: Any) -> None:
+        mailbox_key = (frm, key)
         with self._cond:
-            if key in self._mailbox:
+            if mailbox_key in self._mailbox:
                 raise RuntimeError(
-                    f"Mailbox overflow for key {key} at rank {self.rank}"
+                    f"Mailbox overflow: key {mailbox_key} already exists"
                 )
-            self._mailbox[key] = data
+            self._mailbox[mailbox_key] = data
             self._cond.notify_all()
 
 

mplang/v2/backends/spu_impl.py

@@ -26,7 +26,6 @@ import numpy as np
 import spu.api as spu_api
 import spu.libspu as libspu
 
-from mplang.v2.backends.simp_worker import SimpWorker
 from mplang.v2.backends.spu_state import SPUState
 from mplang.v2.backends.tensor_impl import TensorValue
 from mplang.v2.dialects import spu
@@ -161,6 +160,8 @@ def exec_impl(interpreter: Interpreter, op: Operation, *args: Any) -> Any:
     The SPU config must contain parties info to correctly map global rank
     to local SPU rank and determine SPU world size.
     """
+    from mplang.v2.backends.simp_worker.state import SimpWorker
+
     # Get SPU config from attrs (passed through from run_jax)
     config: spu.SPUConfig = op.attrs["config"]
 
@@ -193,9 +194,8 @@ def exec_impl(interpreter: Interpreter, op: Operation, *args: Any) -> Any:
         interpreter, "spu_endpoints", None
     )
     if spu_endpoints_map is None:
-        context = interpreter.get_dialect_state("simp")
-        if context is not None:
-            spu_endpoints_map = getattr(context, "spu_endpoints", None)
+        # Try getting from SimpWorker context (context is already SimpWorker)
+        spu_endpoints_map = getattr(context, "spu_endpoints", None)
 
     # Build ordered list of endpoints for SPU parties
     spu_endpoints: list[str] | None = None
@@ -209,6 +209,14 @@ def exec_impl(interpreter: Interpreter, op: Operation, *args: Any) -> Any:
                 )
             spu_endpoints.append(spu_endpoints_map[party_rank])
 
+    # Get communicator for Channels mode (reuse existing communication)
+    # If no BRPC endpoints configured, use Channels mode
+    communicator = None
+    if spu_endpoints is None:
+        # Use worker's communicator for channel reuse
+        # (SimpWorker already imported at function start)
+        communicator = context.communicator
+
     # Get or create SPUState for caching Runtime/Io
     spu_state = interpreter.get_dialect_state(SPUState.dialect_name)
     if not isinstance(spu_state, SPUState):
@@ -216,7 +224,12 @@ def exec_impl(interpreter: Interpreter, op: Operation, *args: Any) -> Any:
         interpreter.set_dialect_state(SPUState.dialect_name, spu_state)
 
     runtime, io = spu_state.get_or_create(
-        local_rank, spu_world_size, config, spu_endpoints
+        local_rank,
+        spu_world_size,
+        config,
+        spu_endpoints,
+        communicator=communicator,
+        parties=list(parties),
     )
 
     executable_code = op.attrs["executable"]

mplang/v2/backends/spu_state.py

@@ -20,7 +20,7 @@ multiple executions while binding to the Interpreter's lifecycle.
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import spu.api as spu_api
 import spu.libspu as libspu
@@ -56,6 +56,8 @@ class SPUState(DialectState):
         spu_world_size: int,
         config: spu.SPUConfig,
         spu_endpoints: list[str] | None = None,
+        communicator: object | None = None,
+        parties: list[int] | None = None,
     ) -> tuple[spu_api.Runtime, spu_api.Io]:
         """Get or create SPU Runtime and Io for the given configuration.
 
@@ -64,13 +66,24 @@ class SPUState(DialectState):
             spu_world_size: The number of parties in the SPU device.
             config: SPU configuration including protocol settings.
             spu_endpoints: Optional list of BRPC endpoints. If None, use mem link.
+            communicator: Optional v2 communicator (ThreadCommunicator/HttpCommunicator).
+                If provided, use Channels mode to reuse existing communication.
+            parties: Optional list of global ranks for SPU parties.
+                Required when communicator is provided.
 
         Returns:
             A tuple of (Runtime, Io) for this party.
         """
         from mplang.v2.backends.spu_impl import to_runtime_config
 
-        link_mode = "brpc" if spu_endpoints else "mem"
+        # Determine link mode
+        if communicator is not None:
+            link_mode = "channels"
+        elif spu_endpoints:
+            link_mode = "brpc"
+        else:
+            link_mode = "mem"
+
         cache_key = (
             local_rank,
             spu_world_size,
@@ -83,7 +96,13 @@ class SPUState(DialectState):
             return self._runtimes[cache_key]
 
         # Create Link
-        if spu_endpoints:
+        if communicator is not None:
+            if parties is None:
+                raise ValueError("parties required when using communicator")
+            link = self._create_channels_link(
+                local_rank, spu_world_size, communicator, parties
+            )
+        elif spu_endpoints:
             link = self._create_brpc_link(local_rank, spu_endpoints)
         else:
             link = self._create_mem_link(local_rank, spu_world_size)
@@ -106,6 +125,50 @@ class SPUState(DialectState):
             desc.add_party(f"P{i}", f"mem:{i}")
         return libspu.link.create_mem(desc, local_rank)
 
+    def _create_channels_link(
+        self,
+        local_rank: int,
+        spu_world_size: int,
+        communicator: Any,
+        parties: list[int],
+    ) -> libspu.link.Context:
+        """Create link using custom channels (reuse v2 communicator).
+
+        Args:
+            local_rank: SPU local rank (0-indexed, already converted from global)
+            spu_world_size: Number of SPU parties
+            communicator: v2 communicator (ThreadCommunicator/HttpCommunicator)
+            parties: List of global ranks for SPU parties (ordered by local rank)
+
+        Returns:
+            libspu link context using BaseChannel adapters
+        """
+        from mplang.v2.backends.channel import BaseChannel
+
+        # Get this worker's global rank
+        global_rank = parties[local_rank]
+
+        # Create channels list (world_size elements, self = None)
+        channels = []
+        for idx, peer_global_rank in enumerate(parties):
+            if idx == local_rank:
+                # Self channel must be None
+                channel = None
+            else:
+                # Create channel to peer
+                channel = BaseChannel(communicator, global_rank, peer_global_rank)
+            channels.append(channel)
+
+        # Create link descriptor
+        desc = libspu.link.Desc()  # type: ignore
+        desc.recv_timeout_ms = 100 * 1000  # 100 seconds
+
+        # Add party info (required for world_size inference)
+        for idx in range(spu_world_size):
+            desc.add_party(f"P{idx}", f"dummy_{parties[idx]}")
+
+        return libspu.link.create_with_channels(desc, local_rank, channels)
+
     def _create_brpc_link(
         self, local_rank: int, spu_endpoints: list[str]
     ) -> libspu.link.Context:

mplang/v2/libs/ml/sgb.py

@@ -1097,11 +1097,9 @@ def _update_tree_state(
 
         all_feats[party_idx] = simp.pcall_static(
             (party_rank,),
-            lambda pf=all_feats[party_idx],
-            bf=all_feats_level[party_idx],
-            ci=cur_indices_party,
-            op=owned_party_party,
-            il=is_leaf_party: tensor.run_jax(update_party_feats, pf, bf, ci, op, il),
+            lambda pf=all_feats[party_idx], bf=all_feats_level[party_idx], ci=cur_indices_party, op=owned_party_party, il=is_leaf_party: (
+                tensor.run_jax(update_party_feats, pf, bf, ci, op, il)
+            ),
         )
 
         def update_party_thresholds(
@@ -1123,21 +1121,17 @@ def _update_tree_state(
 
         all_thresholds[party_idx] = simp.pcall_static(
             (party_rank,),
-            lambda pt=all_thresholds[party_idx],
-            b=all_bins[party_idx],
-            bf=all_feats_level[party_idx],
-            bt_idx=all_threshs_level[party_idx],
-            ci=cur_indices_party,
-            op=owned_party_party,
-            il=is_leaf_party: tensor.run_jax(
-                update_party_thresholds,
-                pt,
-                b,
-                bf,
-                bt_idx,
-                ci,
-                op,
-                il,
+            lambda pt=all_thresholds[party_idx], b=all_bins[party_idx], bf=all_feats_level[party_idx], bt_idx=all_threshs_level[party_idx], ci=cur_indices_party, op=owned_party_party, il=is_leaf_party: (
+                tensor.run_jax(
+                    update_party_thresholds,
+                    pt,
+                    b,
+                    bf,
+                    bt_idx,
+                    ci,
+                    op,
+                    il,
+                )
             ),
         )
 
@@ -1152,13 +1146,8 @@ def _update_tree_state(
 
         tmp_bt = simp.pcall_static(
             (party_rank,),
-            lambda bi=all_bin_indices[party_idx],
-            bf=all_feats_level[party_idx],
-            bt_idx=all_threshs_level[party_idx],
-            bt_arr=bt_party,
-            bt_lv=bt_level_party,
-            il=is_leaf_party: tensor.run_jax(
-                update_bt, bt_arr, bt_lv, il, bi, bf, bt_idx
+            lambda bi=all_bin_indices[party_idx], bf=all_feats_level[party_idx], bt_idx=all_threshs_level[party_idx], bt_arr=bt_party, bt_lv=bt_level_party, il=is_leaf_party: (
+                tensor.run_jax(update_bt, bt_arr, bt_lv, il, bi, bf, bt_idx)
             ),
         )
 
@@ -1498,11 +1487,10 @@ def predict_tree(
     for i, rank in enumerate(all_ranks):
         mask = simp.pcall_static(
             (rank,),
-            lambda d=all_datas[i],
-            f=tree.feature[i],
-            t=tree.threshold[i],
-            idx=i: predict_tree_single_party(
-                d, f, t, tree.is_leaf, tree.owned_party_id, idx, n_nodes
+            lambda d=all_datas[i], f=tree.feature[i], t=tree.threshold[i], idx=i: (
+                predict_tree_single_party(
+                    d, f, t, tree.is_leaf, tree.owned_party_id, idx, n_nodes
+                )
             ),
         )
         # Transfer to AP

{mplang_nightly-0.1.dev266.dist-info → mplang_nightly-0.1.dev267.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev266
+Version: 0.1.dev267
 Summary: Multi-Party Programming Language
 Author-email: SecretFlow Team <secretflow-contact@service.alipay.com>
 License:                                  Apache License
@@ -219,7 +219,7 @@ Requires-Dist: pandas>=2.0.0
 Requires-Dist: protobuf<6.0,>=5.0
 Requires-Dist: pyarrow>=14.0.0
 Requires-Dist: pyyaml>=6.0
-Requires-Dist: spu>=0.10.0.dev20251208
+Requires-Dist: spu>=0.10.0.dev20251211
 Requires-Dist: sqlglot>=23.0.0
 Requires-Dist: tenseal==0.3.16
 Requires-Dist: typing-extensions