mplang-nightly 0.1.dev158__py3-none-any.whl → 0.1.dev268__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)
+        """