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

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

@@ -14,16 +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 import TableType, 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)
@@ -135,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):
@@ -147,14 +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()
-            # Pass schema to enable column projection
-            return table_utils.csv_to_dataframe(csv_bytes, schema=out_spec)
-        # 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

@@ -29,7 +29,7 @@ from typing import Any
 
 import numpy as np
 
-from mplang.core import (
+from mplang.v1.core import (
     ClusterSpec,
     InterpContext,
     InterpVar,
@@ -38,9 +38,9 @@ from mplang.core import (
     MPObject,
     MPType,
 )
-from mplang.core.expr.ast import Expr
-from mplang.kernels.value import TableValue, TensorValue
-from mplang.runtime.client import HttpExecutorClient
+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:

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/{runtime → v1/runtime}/server.py

@@ -30,14 +30,18 @@ from fastapi import (
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 
-from mplang.core import IrReader, TableType, TensorType
-from mplang.core.cluster import ClusterSpec
-from mplang.kernels.base import KernelContext
-from mplang.kernels.value import Value, decode_value, encode_value
-from mplang.protos.v1alpha1 import mpir_pb2
-from mplang.runtime.data_providers import DataProvider, ResolvedURI, register_provider
-from mplang.runtime.exceptions import InvalidRequestError, ResourceNotFound
-from mplang.runtime.session import (
+from mplang.v1.core import IrReader, TableType, TensorType
+from mplang.v1.core.cluster import ClusterSpec
+from mplang.v1.kernels.base import KernelContext
+from mplang.v1.kernels.value import Value, decode_value, encode_value
+from mplang.v1.protos.v1alpha1 import mpir_pb2
+from mplang.v1.runtime.data_providers import (
+    DataProvider,
+    ResolvedURI,
+    register_provider,
+)
+from mplang.v1.runtime.exceptions import InvalidRequestError, ResourceNotFound
+from mplang.v1.runtime.session import (
     Computation,
     Session,
     Symbol,
@@ -215,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
@@ -483,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/{runtime → v1/runtime}/session.py

@@ -25,51 +25,28 @@ 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
 
-from mplang.core.cluster import ClusterSpec
-from mplang.core.comm import ICommunicator
-from mplang.core.expr.ast import Expr
-from mplang.core.expr.evaluator import IEvaluator, create_evaluator
-from mplang.core.mask import Mask
-from mplang.kernels.context import RuntimeContext
-from mplang.kernels.spu import PFunction  # type: ignore
-from mplang.kernels.value import Value
-from mplang.runtime.communicator import HttpCommunicator
-from mplang.runtime.exceptions import ResourceNotFound
-from mplang.runtime.link_comm import LinkCommunicator
-from mplang.utils.spu_utils import parse_field, parse_protocol
+from mplang.v1.core.cluster import ClusterSpec
+from mplang.v1.core.comm import ICommunicator
+from mplang.v1.core.expr.ast import Expr
+from mplang.v1.core.expr.evaluator import IEvaluator, create_evaluator
+from mplang.v1.core.mask import Mask
+from mplang.v1.kernels.context import RuntimeContext
+from mplang.v1.kernels.spu import PFunction  # type: ignore
+from mplang.v1.kernels.value import Value
+from mplang.v1.runtime.communicator import HttpCommunicator
+from mplang.v1.runtime.exceptions import ResourceNotFound
+from mplang.v1.runtime.link_comm import LinkCommunicator
+from mplang.v1.utils.spu_utils import parse_field, parse_protocol
 
 if TYPE_CHECKING:  # pragma: no cover - import only for type checking
-    from mplang.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()
+    from mplang.v1.core.cluster import ClusterSpec, Node, RuntimeInfo
 
 
 @dataclass
@@ -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),

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

@@ -18,13 +18,14 @@ import concurrent.futures
 import faulthandler
 import logging
 import sys
+import threading
 import traceback
 from collections.abc import Sequence
 from typing import Any, cast
 
 import spu.libspu as libspu
 
-from mplang.core import (
+from mplang.v1.core import (
     ClusterSpec,
     CollectiveMixin,
     CommunicatorBase,
@@ -38,11 +39,11 @@ from mplang.core import (
     PFunction,  # for spu.seed_env kernel seeding
     TensorLike,
 )
-from mplang.core.expr.ast import Expr
-from mplang.core.expr.evaluator import IEvaluator, create_evaluator
-from mplang.kernels.context import RuntimeContext
-from mplang.runtime.link_comm import LinkCommunicator
-from mplang.utils.spu_utils import parse_field, parse_protocol
+from mplang.v1.core.expr.ast import Expr
+from mplang.v1.core.expr.evaluator import IEvaluator, create_evaluator
+from mplang.v1.kernels.context import RuntimeContext
+from mplang.v1.runtime.link_comm import LinkCommunicator
+from mplang.v1.utils.spu_utils import parse_field, parse_protocol
 
 
 class ThreadCommunicator(CommunicatorBase, CollectiveMixin):
@@ -129,16 +130,37 @@ class Simulator(InterpContext):
             comm.set_peers(self._comms)
 
         # Prepare link contexts for SPU parties (store for evaluator-time initialization)
-        spu_addrs = [f"P{spu_rank}" for spu_rank in spu_mask]
+        # Use Channels mode to reuse ThreadCommunicator instead of separate mem_link
         self._spu_link_ctxs: list[LinkCommunicator | None] = [None] * world_size
-        link_ctx_list = [
-            LinkCommunicator(idx, spu_addrs, mem_link=True)
-            for idx in range(spu_mask.num_parties())
+
+        # Create LinkCommunicators in parallel to avoid deadlock
+        # (create_with_channels does handshake via TestSend/TestRecv)
+        exceptions: dict[int, Exception] = {}
+
+        def create_link(g_rank: int) -> None:
+            try:
+                self._spu_link_ctxs[g_rank] = LinkCommunicator(
+                    rank=g_rank,
+                    comm=self._comms[g_rank],
+                    spu_mask=spu_mask,
+                )
+            except Exception as e:
+                exceptions[g_rank] = e
+
+        threads = [
+            threading.Thread(target=create_link, args=(g_rank,)) for g_rank in spu_mask
         ]
-        for g_rank in range(world_size):
-            if g_rank in spu_mask:
-                rel = Mask(spu_mask).global_to_relative_rank(g_rank)
-                self._spu_link_ctxs[g_rank] = link_ctx_list[rel]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        # Check for exceptions during link creation
+        if exceptions:
+            first_exc = next(iter(exceptions.values()))
+            raise RuntimeError(
+                f"Failed to create SPU link contexts for ranks {list(exceptions.keys())}"
+            ) from first_exc
 
         self._spu_runtime_cfg = libspu.RuntimeConfig(
             protocol=spu_protocol, field=spu_field