mplang-nightly 0.1.dev158__py3-none-any.whl → 0.1.dev268__py3-none-any.whl

mplang/{runtime → v1/runtime}/communicator.py

@@ -21,20 +21,33 @@ import base64
 import logging
 from typing import Any
 
-import cloudpickle as pickle
 import httpx
 
-from mplang.core.comm import CommunicatorBase
+from mplang.v1.core.comm import CommunicatorBase
+from mplang.v1.kernels.value import Value, decode_value, encode_value
 
 
 class HttpCommunicator(CommunicatorBase):
     def __init__(self, session_name: str, rank: int, endpoints: list[str]):
+        # Validate endpoints
+        if not endpoints:
+            raise ValueError("endpoints cannot be empty")
+
+        if not all(endpoint for endpoint in endpoints):
+            raise ValueError("endpoints cannot contain empty elements")
+
         super().__init__(rank, len(endpoints))
         self.session_name = session_name
-        self.endpoints = endpoints
+        # Ensure all endpoints have protocol prefix
+        self.endpoints = [
+            endpoint
+            if endpoint.startswith(("http://", "https://"))
+            else f"http://{endpoint}"
+            for endpoint in endpoints
+        ]
         self._counter = 0
         logging.info(
-            f"HttpCommunicator initialized: session={session_name}, rank={rank}, endpoints={endpoints}"
+            f"HttpCommunicator initialized: session={session_name}, rank={rank}, endpoints={self.endpoints}"
         )
 
     # override
@@ -44,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(
@@ -52,13 +70,20 @@ class HttpCommunicator(CommunicatorBase):
         )
 
         try:
-            # Use cloudpickle for robust serialization of complex Python objects
-            data_bytes = pickle.dumps(data)
-            data_b64 = base64.b64encode(data_bytes).decode("utf-8")
-
-            request_data = {
-                "data": data_b64,
-            }
+            # 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."
+                )
 
             response = httpx.put(url, json=request_data, timeout=60)
             logging.debug(f"Send response: status={response.status_code}")
@@ -72,14 +97,33 @@ 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)
-        result = pickle.loads(data_bytes)
+        result = decode_value(data_bytes)
 
         logging.debug(
             f"Received data: from_rank={frm}, to_rank={self._rank}, key={key}"

mplang/{runtime → v1/runtime}/data_providers.py

@@ -14,17 +14,24 @@
 
 from __future__ import annotations
 
+import pathlib
+import struct
 from dataclasses import dataclass
 from typing import Any
 from urllib.parse import ParseResult, urlparse
 
 import numpy as np
-import pandas as pd
 
-from mplang.core.table import TableType
-from mplang.core.tensor import TensorType
-from mplang.kernels.base import KernelContext
-from mplang.utils import table_utils
+from mplang.v1.core import TableLike, TableType, TensorType
+from mplang.v1.kernels.base import KernelContext
+from mplang.v1.kernels.value import (
+    TableValue,
+    TensorValue,
+    Value,
+    decode_value,
+    encode_value,
+)
+from mplang.v1.utils import table_utils
 
 
 @dataclass(frozen=True)
@@ -136,6 +143,11 @@ def get_provider(scheme: str) -> DataProvider | None:
 
 
 # ---------------- Default Providers ----------------
+MAGIC_MPLANG = b"MPLG"
+MAGIC_PARQUET = b"PAR1"
+MAGIC_ORC = b"ORC"
+MAGIC_NUMPY = b"\x93NUMPY"
+VERSION = 0x01
 
 
 class FileProvider(DataProvider):
@@ -148,13 +160,52 @@ class FileProvider(DataProvider):
     def read(
         self, uri: ResolvedURI, out_spec: TensorType | TableType, *, ctx: KernelContext
     ) -> Any:
-        path = uri.local_path or uri.raw
-        if isinstance(out_spec, TableType):
-            with open(path, "rb") as f:
-                csv_bytes = f.read()
-            return table_utils.csv_to_dataframe(csv_bytes)
-        # tensor path
-        return np.load(path)
+        path = pathlib.Path(uri.local_path or uri.raw)
+        # try load by magic
+        with path.open("rb") as f:
+            # this is the maximum length needed to detect all supported formats
+            # (numpy requires 6 bytes: '\x93NUMPY').
+            MAGIC_BYTES_LEN_MAX = 6
+            magic = f.read(MAGIC_BYTES_LEN_MAX)
+            f.seek(0)
+            if magic.startswith(MAGIC_MPLANG):
+                MPLANG_HEADER_LEN = len(MAGIC_MPLANG) + 1
+                header = f.read(MPLANG_HEADER_LEN)
+                _, version = struct.unpack(">4sB", header)
+                if version != VERSION:
+                    raise ValueError(f"unsupported mplang version {version}")
+                payload = f.read()
+                return decode_value(payload)
+            elif magic.startswith(MAGIC_PARQUET):
+                if not isinstance(out_spec, TableType):
+                    raise ValueError(
+                        f"PARQUET files require TableType output spec, got {type(out_spec).__name__}"
+                    )
+                return table_utils.read_table(
+                    f, format="parquet", columns=list(out_spec.column_names())
+                )
+            elif magic.startswith(MAGIC_ORC):
+                if not isinstance(out_spec, TableType):
+                    raise ValueError(
+                        f"ORC files require TableType output spec, got {type(out_spec).__name__}"
+                    )
+                return table_utils.read_table(
+                    f, format="orc", columns=list(out_spec.column_names())
+                )
+            elif magic.startswith(MAGIC_NUMPY):
+                if not isinstance(out_spec, TensorType):
+                    raise ValueError(
+                        f"NumPy files require TensorType output spec, got {type(out_spec).__name__}"
+                    )
+                return np.load(f)
+
+            # Fallback: open the file for CSV or NumPy loading.
+            if isinstance(out_spec, TableType):
+                return table_utils.read_table(
+                    f, format="csv", columns=list(out_spec.column_names())
+                )
+            else:
+                return np.load(f)
 
     def write(self, uri: ResolvedURI, value: Any, *, ctx: KernelContext) -> None:
         import os
@@ -163,14 +214,24 @@ class FileProvider(DataProvider):
         dir_name = os.path.dirname(path)
         if dir_name:
             os.makedirs(dir_name, exist_ok=True)
-        # Table-like to CSV bytes
-        if hasattr(value, "__dataframe__") or isinstance(value, pd.DataFrame):
-            csv_bytes = table_utils.dataframe_to_csv(value)  # type: ignore
+
+        if not isinstance(value, Value):
+            value = (
+                TableValue(value)
+                if isinstance(value, TableLike)
+                else TensorValue(value)
+            )
+
+        if isinstance(value, TableValue):
+            table_utils.write_table(value.to_arrow(), path, format="parquet")
+        elif isinstance(value, TensorValue):
+            with open(path, "wb") as f:
+                np.save(f, value.to_numpy())
+        else:
+            payload = encode_value(value)
             with open(path, "wb") as f:
-                f.write(csv_bytes)
-            return
-        # Tensor-like via numpy
-        np.save(path, np.asarray(value))
+                f.write(struct.pack(">4sB", MAGIC_MPLANG, VERSION))
+                f.write(payload)
 
 
 class MemProvider(DataProvider):

mplang/{runtime → v1/runtime}/driver.py

@@ -15,8 +15,8 @@
 """
 HTTP-based driver implementation for distributed execution.
 
-This module provides an HTTP-based alternative to the gRPC Driver,
-using REST APIs for distributed multi-party computation coordination.
+This module provides an HTTP-based driver, using REST APIs
+for distributed multi-party computation coordination.
 """
 
 from __future__ import annotations
@@ -27,14 +27,20 @@ import uuid
 from collections.abc import Sequence
 from typing import Any
 
-from mplang.core.cluster import ClusterSpec
-from mplang.core.expr.ast import Expr
-from mplang.core.interp import InterpContext, InterpVar
-from mplang.core.mask import Mask
-from mplang.core.mpir import Writer
-from mplang.core.mpobject import MPObject
-from mplang.core.mptype import MPType
-from mplang.runtime.client import HttpExecutorClient
+import numpy as np
+
+from mplang.v1.core import (
+    ClusterSpec,
+    InterpContext,
+    InterpVar,
+    IrWriter,
+    Mask,
+    MPObject,
+    MPType,
+)
+from mplang.v1.core.expr.ast import Expr
+from mplang.v1.kernels.value import TableValue, TensorValue
+from mplang.v1.runtime.client import HttpExecutorClient
 
 
 def new_uuid() -> str:
@@ -195,7 +201,7 @@ class Driver(InterpContext):
 
         var_name_mapping = dict(zip(var_names, party_symbol_names, strict=True))
 
-        writer = Writer(var_name_mapping)
+        writer = IrWriter(var_name_mapping)
         program_proto = writer.dumps(expr)
 
         output_symbols = [self.new_name() for _ in range(expr.num_outputs)]
@@ -257,7 +263,19 @@ class Driver(InterpContext):
             try:
                 # The results will be in the same order as the clients (ranks)
                 results = await asyncio.gather(*tasks)
-                return list(results)
+                converted: list[Any] = []
+                for value in results:
+                    if isinstance(value, TensorValue):
+                        arr = value.to_numpy()
+                        if isinstance(arr, np.ndarray) and arr.size == 1:
+                            converted.append(arr.item())
+                        else:
+                            converted.append(arr)
+                    elif isinstance(value, TableValue):
+                        converted.append(value.to_pandas())
+                    else:
+                        converted.append(value)
+                return converted
             except RuntimeError as e:
                 raise RuntimeError(
                     f"Failed to fetch symbol from one or more parties: {e}"

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