mplang-nightly 0.1.dev265__tar.gz → 0.1.dev267__tar.gz

{mplang_nightly-0.1.dev265 → mplang_nightly-0.1.dev267}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev265
+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

{mplang_nightly-0.1.dev265 → mplang_nightly-0.1.dev267}/mplang/v1/kernels/phe.py

@@ -637,7 +637,8 @@ def _phe_mul(pfunc: PFunction, ciphertext: CipherText, plaintext: TensorValue) -
             # Use numpy to create a properly broadcasted index mapping
             # Create a dummy array with same shape as ciphertext, fill with indices
             dummy_ct = (
-                np.arange(np.prod(ciphertext.semantic_shape))
+                np
+                .arange(np.prod(ciphertext.semantic_shape))
                 .reshape(ciphertext.semantic_shape)
                 .astype(np.int64)
             )
@@ -744,7 +745,8 @@ def _phe_add_ct2ct(ct1: CipherText, ct2: CipherText) -> CipherText:
     # Broadcast ct1 if needed
     if ct1.semantic_shape != result_shape:
         dummy_ct1 = (
-            np.arange(np.prod(ct1.semantic_shape))
+            np
+            .arange(np.prod(ct1.semantic_shape))
             .reshape(ct1.semantic_shape)
             .astype(np.int64)
         )
@@ -757,7 +759,8 @@ def _phe_add_ct2ct(ct1: CipherText, ct2: CipherText) -> CipherText:
     # Broadcast ct2 if needed
     if ct2.semantic_shape != result_shape:
         dummy_ct2 = (
-            np.arange(np.prod(ct2.semantic_shape))
+            np
+            .arange(np.prod(ct2.semantic_shape))
             .reshape(ct2.semantic_shape)
             .astype(np.int64)
         )
@@ -830,7 +833,8 @@ def _phe_add_ct2pt(ciphertext: CipherText, plaintext: TensorValue) -> CipherText
     # Broadcast ciphertext if needed
     if ciphertext.semantic_shape != result_shape:
         dummy_ct = (
-            np.arange(np.prod(ciphertext.semantic_shape))
+            np
+            .arange(np.prod(ciphertext.semantic_shape))
             .reshape(ciphertext.semantic_shape)
             .astype(np.int64)
         )

mplang_nightly-0.1.dev267/mplang/v1/runtime/channel.py

@@ -0,0 +1,230 @@
+# Copyright 2025 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 that bridges to MPLang CommunicatorBase.
+
+This module provides BaseChannel, which allows SPU to reuse MPLang's
+existing communication layer (ThreadCommunicator/HttpCommunicator) instead
+of creating separate BRPC connections.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import spu.libspu as libspu
+
+if TYPE_CHECKING:
+    from mplang.v1.core.comm import CommunicatorBase
+
+
+class BaseChannel(libspu.link.IChannel):
+    """Bridge MPLang CommunicatorBase to SPU IChannel interface.
+
+    This adapter allows SPU to use MPLang's existing communication layer
+    (ThreadCommunicator or HttpCommunicator) instead of creating separate
+    BRPC connections.
+
+    Each BaseChannel represents a channel to ONE peer rank.
+
+    Communication Protocol:
+        - SPU calls send(tag, bytes_data) -> MPLang comm.send(peer, key, bytes_data)
+        - SPU calls recv(tag) -> bytes_data <- MPLang comm.recv(peer, key)
+
+    Tag Namespace:
+        All tags are prefixed with "spu:" to avoid collision with other
+        MPLang traffic on the same communicator.
+    """
+
+    def __init__(
+        self,
+        comm: CommunicatorBase,
+        local_rank: int,
+        peer_rank: int,
+        tag_prefix: str = "spu",
+    ):
+        """Initialize channel to a specific peer.
+
+        Args:
+            comm: MPLang communicator instance (Thread/Http)
+            local_rank: Global rank of this party (for logging/debugging)
+            peer_rank: Global rank of the peer party
+            tag_prefix: Prefix for all tags to avoid collision (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 MPLang comm.
+
+        Prefixes the tag to avoid collision with non-SPU traffic.
+
+        Args:
+            tag: SPU-provided tag (e.g., "send_0", "recv_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 (synchronous in SPU semantics).
+
+        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
+        # Note: CommunicatorBase.send expects Any type, bytes is acceptable
+        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
+        """
+        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 (MPLang's send is already async at network layer).
+
+        For HttpCommunicator, the underlying httpx.put() is non-blocking
+        at the HTTP client level. For ThreadCommunicator, send is instant
+        (memory transfer).
+
+        Args:
+            tag: Message tag
+            data: Raw bytes to send
+        """
+        # Reuse synchronous send - it's already async underneath
+        self.Send(tag, data)
+
+    def SendAsyncThrottled(self, tag: str, data: bytes) -> None:
+        """Throttled async send.
+
+        Currently maps to regular SendAsync. Future optimization could
+        implement rate limiting if needed.
+
+        Args:
+            tag: Message tag
+            data: Raw bytes to send
+        """
+        self.SendAsync(tag, data)
+
+    def TestSend(self, timeout: int) -> None:
+        """Test if this channel can send a dummy msg to peer.
+
+        Uses fixed 0 seq_id as dummy msg's id to make this function reentrant.
+        ConnectToMesh will retry on this multiple times.
+
+        Args:
+            timeout: Timeout in milliseconds
+        """
+        # Send a handshake message to test connectivity
+        # Use fixed tag "__test__" to make this reentrant (idempotent)
+        test_data = b"\x00"  # Minimal 1-byte message with seq_id=0
+        self.Send("__test__", test_data)
+
+    def TestRecv(self) -> None:
+        """Wait for dummy msg from peer.
+
+        Timeout is controlled by recv_timeout_ms in link descriptor.
+        """
+        # Receive the handshake message from peer
+        # This blocks until message arrives (timeout from desc.recv_timeout_ms)
+        test_data = self.Recv("__test__")
+        # Validate it's the expected handshake message
+        if test_data != b"\x00":
+            logging.warning(
+                f"TestRecv: unexpected handshake data from {self._peer_rank}, "
+                f"expected b'\\x00', got {test_data!r}"
+            )
+
+    def WaitLinkTaskFinish(self) -> None:
+        """Wait for all pending async tasks.
+
+        For MPLang communicators:
+        - ThreadCommunicator: No-op (instant memory transfer)
+        - HttpCommunicator: No explicit wait needed (httpx handles it)
+
+        This is a no-op in current implementation.
+        """
+
+    def Abort(self) -> None:
+        """Abort communication (cleanup resources).
+
+        This could be extended to notify the communicator to drop pending
+        messages for this channel, but currently is a no-op.
+        """
+        logging.warning(
+            f"BaseChannel.Abort called: {self._local_rank} <-> {self._peer_rank}"
+        )
+        # Future: Could call comm.abort_session() if implemented
+
+    def SetThrottleWindowSize(self, size: int) -> None:
+        """Set throttle window size.
+
+        Not applicable to MPLang communicators. No-op.
+
+        Args:
+            size: Window size (ignored)
+        """
+
+    def SetChunkParallelSendSize(self, size: int) -> None:
+        """Set chunk parallel send size.
+
+        Not applicable to MPLang communicators. No-op.
+
+        Args:
+            size: Chunk size (ignored)
+        """

{mplang_nightly-0.1.dev265 → mplang_nightly-0.1.dev267}/mplang/v1/runtime/communicator.py

@@ -57,7 +57,12 @@ class HttpCommunicator(CommunicatorBase):
         return str(res)
 
     def send(self, to: int, key: str, data: Any) -> None:
-        """Sends data to a peer party by PUTing to its /comm/{key}/from/{from_rank} endpoint."""
+        """Sends data to a peer party by PUTing to its /comm/{key}/from/{from_rank} endpoint.
+
+        Supports two modes:
+        - SPU channel (key starts with "spu:"): sends raw bytes directly
+        - Normal channel: wraps data in Value envelope
+        """
         target_endpoint = self.endpoints[to]
         url = f"{target_endpoint}/sessions/{self.session_name}/comm/{key}/from/{self._rank}"
         logging.debug(
@@ -65,19 +70,20 @@ class HttpCommunicator(CommunicatorBase):
         )
 
         try:
-            # Serialize data using Value envelope.
-            if not isinstance(data, Value):
+            # SPU channel mode: send raw bytes directly
+            if key.startswith("spu:") and isinstance(data, bytes):
+                data_b64 = base64.b64encode(data).decode("utf-8")
+                request_data = {"data": data_b64, "is_raw_bytes": True}
+            # Normal mode: serialize using Value envelope
+            elif isinstance(data, Value):
+                data_bytes = encode_value(data)
+                data_b64 = base64.b64encode(data_bytes).decode("utf-8")
+                request_data = {"data": data_b64}
+            else:
                 raise TypeError(
                     f"Communicator requires Value instance, got {type(data).__name__}. "
                     "Wrap data in TensorValue or custom Value subclass."
                 )
-            data_bytes = encode_value(data)
-
-            data_b64 = base64.b64encode(data_bytes).decode("utf-8")
-
-            request_data = {
-                "data": data_b64,
-            }
 
             response = httpx.put(url, json=request_data, timeout=60)
             logging.debug(f"Send response: status={response.status_code}")
@@ -91,14 +97,32 @@ class HttpCommunicator(CommunicatorBase):
             raise OSError(f"Failed to send data to rank {to}") from e
 
     def recv(self, frm: int, key: str) -> Any:
-        """Wait until the key is set, returns the value. Override to add logging."""
+        """Wait until the key is set, returns the value.
+
+        Supports two modes:
+        - SPU channel (key starts with "spu:"): returns raw bytes
+        - Normal channel: returns deserialized Value
+        """
         logging.debug(
             f"Waiting to receive: from_rank={frm}, to_rank={self._rank}, key={key}"
         )
-        data_b64 = super().recv(frm, key)
+        received_data = super().recv(frm, key)
 
+        # Check if this is raw bytes (SPU channel)
+        if isinstance(received_data, dict) and received_data.get("is_raw_bytes"):
+            data_bytes = base64.b64decode(received_data["data"])
+            logging.debug(
+                f"Received raw bytes: from_rank={frm}, to_rank={self._rank}, key={key}, size={len(data_bytes)}"
+            )
+            return data_bytes
+
+        # Normal mode: deserialize Value envelope
+        data_b64 = (
+            received_data
+            if isinstance(received_data, str)
+            else received_data.get("data")
+        )
         data_bytes = base64.b64decode(data_b64)
-        # Deserialize using Value envelope
         result = decode_value(data_bytes)
 
         logging.debug(

mplang_nightly-0.1.dev267/mplang/v1/runtime/link_comm.py

@@ -0,0 +1,196 @@
+# Copyright 2025 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.
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import spu.libspu as libspu
+
+if TYPE_CHECKING:
+    from mplang.v1.core.comm import CommunicatorBase
+    from mplang.v1.core.mask import Mask
+
+
+class LinkCommunicator:
+    """Minimal wrapper for libspu link context.
+
+    Supports three modes:
+    1. BRPC: Production mode with separate BRPC ports (legacy)
+    2. Mem: In-memory links for testing (legacy)
+    3. Channels: Reuse MPLang communicator via IChannel bridge (NEW)
+
+    The mode is selected based on constructor arguments:
+    - If `comm` is provided: Channels mode (NEW)
+    - Elif `mem_link` is True: Mem mode
+    - Else: BRPC mode
+    """
+
+    def __init__(
+        self,
+        rank: int,
+        addrs: list[str] | None = None,
+        *,
+        mem_link: bool = False,
+        comm: CommunicatorBase | None = None,
+        spu_mask: Mask | None = None,
+    ):
+        """Initialize link communicator for SPU.
+
+        Args:
+            rank: Global rank of this party
+            addrs: List of addresses for all SPU parties (required for BRPC/Mem mode)
+            mem_link: If True, use in-memory link (Mem mode)
+            comm: MPLang communicator to reuse (Channels mode, NEW)
+            spu_mask: SPU parties mask (required for Channels mode)
+
+        Raises:
+            ValueError: If arguments are invalid for the selected mode
+        """
+        self._rank = rank
+
+        # Select initialization mode based on arguments
+        if comm is not None:
+            self._init_channels_mode(rank, comm, spu_mask)
+        elif mem_link:
+            self._init_mem_mode(rank, addrs)
+        else:
+            self._init_brpc_mode(rank, addrs)
+
+    def _init_channels_mode(
+        self, rank: int, comm: CommunicatorBase, spu_mask: Mask | None
+    ) -> None:
+        """Initialize Channels mode (reuse MPLang communicator).
+
+        Args:
+            rank: Global rank of this party
+            comm: MPLang communicator to reuse
+            spu_mask: SPU parties mask
+
+        Raises:
+            ValueError: If spu_mask is None or rank not in mask
+        """
+        if spu_mask is None:
+            raise ValueError("spu_mask required when using comm")
+        if rank not in spu_mask:
+            raise ValueError(f"rank {rank} not in spu_mask {spu_mask}")
+
+        # Lazy import to avoid circular dependency
+        from mplang.v1.runtime.channel import BaseChannel
+
+        # Create channels to ALL SPU parties (including self)
+        # libspu expects world_size channels, with self channel being None
+        channels = []
+        rel_rank = spu_mask.global_to_relative_rank(rank)
+
+        for _, peer_rank in enumerate(spu_mask):
+            if peer_rank == rank:
+                # For self, use None (won't be accessed by SPU)
+                channel = None
+            else:
+                channel = BaseChannel(comm, rank, peer_rank)
+            channels.append(channel)
+
+        # Create link context with custom channels
+        desc = libspu.link.Desc()  # type: ignore
+        desc.recv_timeout_ms = 100 * 1000  # 100 seconds
+
+        # Add party info to desc (required for world_size inference)
+        for idx, peer_rank in enumerate(spu_mask):
+            desc.add_party(f"P{idx}", f"dummy_{peer_rank}")
+
+        self.lctx = libspu.link.create_with_channels(desc, rel_rank, channels)
+        self._world_size = spu_mask.num_parties()
+
+        logging.info(
+            f"LinkCommunicator initialized with BaseChannel: "
+            f"rank={rank}, rel_rank={rel_rank}, spu_mask={spu_mask}, "
+            f"world_size={self._world_size}"
+        )
+
+    def _init_mem_mode(self, rank: int, addrs: list[str] | None) -> None:
+        """Initialize Mem mode (in-memory links for testing).
+
+        Args:
+            rank: Global rank of this party
+            addrs: List of addresses for all SPU parties
+
+        Raises:
+            ValueError: If addrs is None
+        """
+        if addrs is None:
+            raise ValueError("addrs required for Mem mode")
+
+        self._world_size = len(addrs)
+
+        desc = libspu.link.Desc()  # type: ignore
+        desc.recv_timeout_ms = 100 * 1000  # 100 seconds
+        desc.http_max_payload_size = 32 * 1024 * 1024  # 32M
+        for rank_idx, addr in enumerate(addrs):
+            desc.add_party(f"P{rank_idx}", addr)
+
+        self.lctx = libspu.link.create_mem(desc, self._rank)
+        logging.info(
+            f"LinkCommunicator initialized with Mem: "
+            f"rank={self._rank}, world_size={self._world_size}, addrs={addrs}"
+        )
+
+    def _init_brpc_mode(self, rank: int, addrs: list[str] | None) -> None:
+        """Initialize BRPC mode (production mode with separate BRPC ports).
+
+        Args:
+            rank: Global rank of this party
+            addrs: List of addresses for all SPU parties
+
+        Raises:
+            ValueError: If addrs is None
+        """
+        if addrs is None:
+            raise ValueError("addrs required for BRPC mode")
+
+        self._world_size = len(addrs)
+
+        desc = libspu.link.Desc()  # type: ignore
+        desc.recv_timeout_ms = 100 * 1000  # 100 seconds
+        desc.http_max_payload_size = 32 * 1024 * 1024  # 32M
+        for rank_idx, addr in enumerate(addrs):
+            desc.add_party(f"P{rank_idx}", addr)
+
+        self.lctx = libspu.link.create_brpc(desc, self._rank)
+        logging.info(
+            f"LinkCommunicator initialized with BRPC: "
+            f"rank={self._rank}, world_size={self._world_size}, addrs={addrs}"
+        )
+
+    @property
+    def rank(self) -> int:
+        """Get rank from underlying link context."""
+        return self.lctx.rank  # type: ignore[no-any-return]
+
+    @property
+    def world_size(self) -> int:
+        """Get world size from underlying link context."""
+        return self.lctx.world_size  # type: ignore[no-any-return]
+
+    def get_lctx(self) -> libspu.link.Context:
+        """Get the underlying libspu link context.
+
+        This is the primary interface - SPU runtime uses this context directly.
+        All communication and serialization is handled by libspu internally.
+
+        Returns:
+            The underlying libspu.link.Context instance.
+        """
+        return self.lctx

{mplang_nightly-0.1.dev265 → mplang_nightly-0.1.dev267}/mplang/v1/runtime/server.py

@@ -219,6 +219,7 @@ class SymbolResponse(BaseModel):
 
 class CommSendRequest(BaseModel):
     data: str  # Base64 encoded binary data
+    is_raw_bytes: bool = False  # True for SPU channel raw bytes
 
 
 # Response Models for enhanced status
@@ -487,6 +488,14 @@ def comm_send(
     # The receiver rank should be the rank of the server hosting this endpoint
     # We don't need to validate to_rank since the request is coming to this server
 
+    # For raw bytes (SPU channel), pass through as dict with flag
+    # For normal data, pass the base64 string directly
+    data_payload: str | dict[str, object]
+    if request.is_raw_bytes:
+        data_payload = {"data": request.data, "is_raw_bytes": True}
+    else:
+        data_payload = request.data
+
     # Use the proper onSent mechanism from CommunicatorBase
-    sess.communicator.onSent(from_rank, key, request.data)
+    sess.communicator.onSent(from_rank, key, data_payload)
     return {"status": "ok"}

{mplang_nightly-0.1.dev265 → mplang_nightly-0.1.dev267}/mplang/v1/runtime/session.py

@@ -25,12 +25,10 @@ Process-wide registries (sessions, global symbols) live in the server layer
 
 from __future__ import annotations
 
-import logging
 import time
 from dataclasses import dataclass, field
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, cast
-from urllib.parse import urlparse
 
 import spu.libspu as libspu
 
@@ -51,27 +49,6 @@ if TYPE_CHECKING:  # pragma: no cover - import only for type checking
     from mplang.v1.core.cluster import ClusterSpec, Node, RuntimeInfo
 
 
-class LinkCommFactory:
-    """Factory for creating and caching link communicators."""
-
-    def __init__(self) -> None:
-        self._cache: dict[tuple[int, tuple[str, ...]], LinkCommunicator] = {}
-
-    def create_link(self, rel_rank: int, addrs: list[str]) -> LinkCommunicator:
-        key = (rel_rank, tuple(addrs))
-        link = self._cache.get(key)
-        if link is not None:
-            return link
-        logging.info(f"LinkCommunicator created: rel_rank={rel_rank} addrs={addrs}")
-        link = LinkCommunicator(rel_rank, addrs)
-        self._cache[key] = link
-        return link
-
-
-# Shared link factory (module-local, not global registry of sessions)
-g_link_factory = LinkCommFactory()
-
-
 @dataclass
 class Symbol:
     name: str
@@ -184,23 +161,19 @@ class Session:
             return
 
         link_ctx = None
-        # TODO(jint): reuse same port for mplang and spu.
-        SPU_PORT_OFFSET = 100
 
         if self.is_spu_party:
-            # Build SPU address list across all endpoints for ranks in mask
-            spu_addrs: list[str] = []
-            for r, addr in enumerate(self.cluster_spec.endpoints):
-                if r in self.spu_mask:
-                    # TODO(oeqqwq): addr may contain other schema like grpc://
-                    if not addr.startswith(("http://", "https://")):
-                        addr = f"http://{addr}"
-                    parsed = urlparse(addr)
-                    assert isinstance(parsed.port, int)
-                    new_addr = f"{parsed.hostname}:{parsed.port + SPU_PORT_OFFSET}"
-                    spu_addrs.append(new_addr)
-            rel_index = sum(1 for r in range(self.rank) if r in self.spu_mask)
-            link_ctx = g_link_factory.create_link(rel_index, spu_addrs)
+            # Use Channels mode to reuse existing HttpCommunicator
+            # This eliminates the need for separate BRPC ports (SPU_PORT_OFFSET)
+            from mplang.v1.core.comm import CommunicatorBase
+
+            # Type assertion: ICommunicator is actually CommunicatorBase
+            comm = cast(CommunicatorBase, self.communicator)
+            link_ctx = LinkCommunicator(
+                rank=self.rank,
+                comm=comm,
+                spu_mask=self.spu_mask,
+            )
 
         spu_config = libspu.RuntimeConfig(
             protocol=parse_protocol(self.spu_protocol),