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

mplang/ops/crypto.py

@@ -1,109 +0,0 @@
-# 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.
-
-"""
-Crypto frontend operations: operation signatures, types, and high-level semantics.
-
-Scope and contracts:
-- This module defines portable API shapes; it does not implement cryptography.
-- Backends execute the operations and must meet the security semantics required
-    by the deployment (confidentiality, authenticity, correctness, etc.).
-- The enc/dec API in this frontend uses a conventional 12-byte nonce prefix
-    (ciphertext = nonce || payload), and dec expects that format. Other security
-    properties (e.g., AEAD) are backend responsibilities.
-"""
-
-from __future__ import annotations
-
-from mplang.core.dtype import UINT8
-from mplang.core.tensor import TensorType
-from mplang.ops.base import stateless_mod
-
-_CRYPTO_MOD = stateless_mod("crypto")
-
-
-@_CRYPTO_MOD.simple_op()
-def keygen(*, length: int = 32) -> TensorType:
-    """Generate random bytes for symmetric keys or generic randomness.
-
-    API: keygen(length: int = 32) -> key: u8[length]
-
-    Notes:
-    - Frontend defines the type/shape; backend provides randomness.
-    - Raises ValueError when length <= 0.
-    """
-    if length <= 0:
-        raise ValueError("length must be > 0")
-    return TensorType(UINT8, (length,))
-
-
-@_CRYPTO_MOD.simple_op()
-def enc(plaintext: TensorType, key: TensorType) -> TensorType:
-    """Symmetric encryption.
-
-    API: enc(plaintext: u8[N], key: u8[M]) -> ciphertext: u8[N + 12]
-    """
-    pt_ty = plaintext
-    if pt_ty.dtype != UINT8:
-        raise TypeError("enc expects UINT8 plaintext")
-    if len(pt_ty.shape) != 1:
-        raise TypeError("enc expects 1-D plaintext")
-    length = pt_ty.shape[0]
-    if length >= 0:
-        return TensorType(UINT8, (length + 12,))
-    return TensorType(UINT8, (-1,))
-
-
-@_CRYPTO_MOD.simple_op()
-def dec(ciphertext: TensorType, key: TensorType) -> TensorType:
-    """Symmetric decryption.
-
-    API: dec(ciphertext: u8[N + 12], key: u8[M]) -> plaintext: u8[N]
-    """
-    ct_ty = ciphertext
-    if ct_ty.dtype != UINT8:
-        raise TypeError("dec expects UINT8 ciphertext")
-    if len(ct_ty.shape) != 1:
-        raise TypeError("dec expects 1-D ciphertext with nonce")
-    length = ct_ty.shape[0]
-    if length >= 0 and length < 12:
-        raise TypeError("dec expects 1-D ciphertext with nonce")
-    if length >= 0:
-        return TensorType(UINT8, (length - 12,))
-    return TensorType(UINT8, (-1,))
-
-
-@_CRYPTO_MOD.simple_op()
-def kem_keygen(*, suite: str = "x25519") -> tuple[TensorType, TensorType]:
-    """KEM-style keypair generation: returns (sk, pk) bytes."""
-    sk_ty = TensorType(UINT8, (32,))
-    pk_ty = TensorType(UINT8, (32,))
-    return sk_ty, pk_ty
-
-
-@_CRYPTO_MOD.simple_op()
-def kem_derive(
-    sk: TensorType, peer_pk: TensorType, *, suite: str = "x25519"
-) -> TensorType:
-    """KEM-style shared secret derivation: returns secret bytes."""
-    _ = sk
-    _ = peer_pk
-    return TensorType(UINT8, (32,))
-
-
-@_CRYPTO_MOD.simple_op()
-def hkdf(secret: TensorType, *, info: str) -> TensorType:
-    """HKDF-style key derivation: returns a 32-byte key."""
-    _ = secret
-    return TensorType(UINT8, (32,))

mplang/ops/ibis_cc.py

@@ -1,139 +0,0 @@
-# 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.
-
-
-import inspect
-from collections.abc import Callable
-from typing import Any
-
-import ibis
-from jax.tree_util import PyTreeDef, tree_flatten
-
-from mplang.core import dtype
-from mplang.core.mpobject import MPObject
-from mplang.core.pfunc import PFunction
-from mplang.core.table import TableType
-from mplang.ops.base import FeOperation, stateless_mod
-from mplang.utils.func_utils import normalize_fn
-
-
-def ibis2sql(
-    expr: ibis.Table,
-    in_schemas: list[ibis.Schema],
-    in_names: list[str],
-    fn_name: str = "",
-) -> PFunction:
-    """
-    Compile a ibis expr to sql and return the PFunction.
-
-    Args:
-        expr: ibis expr.
-        in_schemas: the input table schemas
-        in_names: the input table names, If there is only one table, it is usually defaulted to "table"
-    Return:
-        PFunction: The compiled PFunction
-    """
-    assert len(in_schemas) == len(in_names), (
-        f"length of input table names and schemas mismatch. {len(in_schemas)}!={len(in_names)}"
-    )
-
-    def _convert(s: ibis.Schema) -> TableType:
-        return TableType.from_pairs([
-            (name, dtype.from_numpy(dt.to_numpy())) for name, dt in s.fields.items()
-        ])
-
-    ins_info = [_convert(s) for s in in_schemas]
-    outs_info = [_convert(expr.schema())]
-
-    sql = ibis.to_sql(expr, dialect="duckdb")
-    # Emit generic sql.run op; runtime maps to backend-specific kernel.
-    pfn = PFunction(
-        fn_type="sql.run",
-        fn_name=fn_name,
-        fn_text=sql,
-        ins_info=tuple(ins_info),
-        outs_info=tuple(outs_info),
-        in_names=tuple(in_names),
-        dialect="duckdb",
-    )
-    return pfn
-
-
-def is_ibis_function(func: Callable) -> bool:
-    """
-    Verify whether a function is an ibis function.
-    The func signature should like def foo(t0:ibis.Table, t1:ibis.Table)->ibis.Table
-    """
-    try:
-        sig = inspect.signature(func)
-    except (ValueError, TypeError):
-        return False
-
-    ret_anno = sig.return_annotation
-    if ret_anno is ibis.Table:
-        return True
-
-    for param in sig.parameters.values():
-        par_anno = param.annotation
-        if par_anno is ibis.Table:
-            return True
-
-    return False
-
-
-_IBIS_MOD = stateless_mod("ibis")
-
-
-class IbisCompiler(FeOperation):
-    """Ibis compiler frontend operation."""
-
-    def trace(
-        self, func: Callable, *args: Any, **kwargs: Any
-    ) -> tuple[PFunction, list[MPObject], PyTreeDef]:
-        """Compile an Ibis function to SQL format.
-
-        Args:
-            func: The Ibis function to compile
-            *args: Positional arguments to the function
-            **kwargs: Keyword arguments to the function
-
-        Returns:
-            tuple[PFunction, list[MPObject], Any]: The compiled PFunction, input variables, and output tree
-        """
-
-        def is_variable(arg: Any) -> bool:
-            return isinstance(arg, MPObject)
-
-        normalized_fn, in_vars = normalize_fn(func, args, kwargs, is_variable)
-
-        in_args, in_schemas, in_names = [], [], []
-        idx = 0
-        for arg in in_vars:
-            columns = [(p[0], p[1].to_numpy()) for p in arg.schema.columns]
-            schema = ibis.schema(columns)
-            name = f"table{idx}"
-            table = ibis.table(schema=schema, name=name)
-            in_args.append(table)
-            in_schemas.append(schema)
-            in_names.append(name)
-            idx += 1
-
-        result = normalized_fn(in_args)
-        assert isinstance(result, ibis.Table)
-        pfunc = ibis2sql(result, in_schemas, in_names, func.__name__)
-        _, treedef = tree_flatten(result)
-        return pfunc, in_vars, treedef
-
-
-ibis_compile = IbisCompiler(_IBIS_MOD, "compile")

mplang/ops/sql.py

@@ -1,61 +0,0 @@
-# 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 jax.tree_util import PyTreeDef, tree_flatten
-
-from mplang.core.mpobject import MPObject
-from mplang.core.pfunc import PFunction
-from mplang.core.table import TableType
-from mplang.ops.base import FeOperation, stateless_mod
-
-_SQL_MOD = stateless_mod("sql")
-
-
-class SqlFE(FeOperation):
-    def __init__(self, dialect: str = "duckdb"):
-        # Bind to sql module with a stable op name for registry/dispatch
-        super().__init__(_SQL_MOD, "run")
-        self._dialect = dialect
-
-    def trace(
-        self,
-        sql: str,
-        out_type: TableType,
-        in_tables: dict[str, MPObject] | None = None,
-    ) -> tuple[PFunction, list[MPObject], PyTreeDef]:
-        in_names: list[str] = []
-        ins_info: list[TableType] = []
-        in_vars: list[MPObject] = []
-        if in_tables:
-            for name, tbl in in_tables.items():
-                assert isinstance(tbl, MPObject)
-                assert tbl.schema is not None
-                in_names.append(name)
-                ins_info.append(tbl.schema)
-                in_vars.append(tbl)
-
-        pfn = PFunction(
-            fn_type="sql.run",
-            fn_name="",
-            fn_text=sql,
-            ins_info=tuple(ins_info),
-            outs_info=(out_type,),
-            in_names=tuple(in_names),
-            dialect=self._dialect,
-        )
-        _, treedef = tree_flatten(out_type)
-        return pfn, in_vars, treedef
-
-
-sql_run = SqlFE("duckdb")

mplang/protos/v1alpha1/mpir_pb2_grpc.py

@@ -1,3 +0,0 @@
-# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
-"""Client and server classes corresponding to protobuf-defined services."""
-import grpc

mplang/runtime/link_comm.py

@@ -1,131 +0,0 @@
-# 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 Any
-
-import cloudpickle as pickle
-import spu.libspu as libspu
-
-from mplang.core.comm import ICollective, ICommunicator
-
-
-class LinkCommunicator(ICommunicator, ICollective):
-    """Wraps libspu link communicator for distributed communication"""
-
-    def __init__(self, rank: int, addrs: list[str], *, mem_link: bool = False):
-        self._rank = rank
-        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  # Default set link payload to 32M
-        for rank, addr in enumerate(addrs):
-            desc.add_party(f"P{rank}", addr)
-
-        if mem_link:
-            self.lctx = libspu.link.create_mem(desc, self._rank)
-        else:
-            self.lctx = libspu.link.create_brpc(desc, self._rank)
-
-        logging.info(
-            f"LinkCommunicator initialized with rank={self._rank}, world_size={self._world_size}, addrs={addrs}",
-        )
-
-        self._counter = 0
-
-    @property
-    def rank(self) -> int:
-        return self.lctx.rank  # type: ignore[no-any-return]
-
-    @property
-    def world_size(self) -> int:
-        return self.lctx.world_size  # type: ignore[no-any-return]
-
-    def get_lctx(self) -> libspu.link.Context:
-        """Get the link context"""
-        return self.lctx
-
-    # override
-    def new_id(self) -> str:
-        res = self._counter
-        self._counter += 1
-        return str(res)
-
-    def wrap(self, obj: Any) -> str:
-        data = pickle.dumps(obj)
-        return data.hex()  # type: ignore[no-any-return]
-
-    def unwrap(self, obj: str) -> Any:
-        data = bytes.fromhex(obj)
-        return pickle.loads(data)  # type: ignore[no-any-return]
-
-    def send(self, to: int, key: str, data: Any) -> None:
-        serialized = pickle.dumps((key, data))
-        self.lctx.send(to, serialized.hex())
-
-    def recv(self, frm: int, key: str) -> Any:
-        serialized = self.lctx.recv(frm)
-        rkey, data = pickle.loads(bytes.fromhex(serialized.decode()))
-        assert key == rkey, f"recv key {key} != {rkey}"
-        return data  # type: ignore[no-any-return]
-
-    def p2p(self, frm: int, to: int, data: Any) -> Any:
-        assert 0 <= frm < self.world_size
-        assert 0 <= to < self.world_size
-
-        # TODO: link handles cid internally?
-        cid = self.new_id()
-
-        if self.rank == frm:
-            self.send(to, cid, data)
-            return None
-        elif self.rank == to:
-            return self.recv(frm, cid)
-        else:
-            return None
-
-    def gather(self, root: int, data: Any) -> list[Any]:
-        assert 0 <= root < self.world_size
-        rets = self.lctx.gather(self.wrap(data), root)
-        return [self.unwrap(ret) for ret in rets]
-
-    def scatter(self, root: int, args: list[Any]) -> Any:
-        assert 0 <= root < self.world_size
-        assert len(args) == self.world_size, f"{len(args)} != {self.world_size}"
-        ret = self.lctx.scatter([self.wrap(arg) for arg in args], root)
-        return self.unwrap(ret)
-
-    def allgather(self, arg: Any) -> list[Any]:
-        rets = self.lctx.all_gather(self.wrap(arg))
-        return [self.unwrap(ret) for ret in rets]
-
-    def bcast(self, root: int, arg: Any) -> Any:
-        assert 0 <= root < self.world_size
-        ret = self.lctx.broadcast(self.wrap(arg), root)
-        return self.unwrap(ret)
-
-    def gather_m(self, pmask: int, root: int, data: Any) -> list[Any]:
-        raise ValueError("Not supported by LinkCommunicator")
-
-    def scatter_m(self, pmask: int, root: int, args: list[Any]) -> Any:
-        raise ValueError("Not supported by LinkCommunicator")
-
-    def allgather_m(self, pmask: int, arg: Any) -> list[Any]:
-        raise ValueError("Not supported by LinkCommunicator")
-
-    def bcast_m(self, pmask: int, root: int, arg: Any) -> Any:
-        raise ValueError("Not supported by LinkCommunicator")

mplang/simp/smpc.py

@@ -1,201 +0,0 @@
-# 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 abc import ABC, abstractmethod
-from collections.abc import Callable
-from enum import Enum
-from functools import wraps
-from typing import Any
-
-from jax.tree_util import tree_unflatten
-
-from mplang.core import Mask, MPObject, Rank, peval, psize
-from mplang.core.context_mgr import cur_ctx
-from mplang.ops import spu
-from mplang.simp import mpi
-
-
-class SecureAPI(ABC):
-    """Base class for secure APIs."""
-
-    @abstractmethod
-    def seal(self, obj: MPObject, frm_mask: Mask | None) -> list[MPObject]: ...
-
-    @abstractmethod
-    def sealFrom(self, obj: MPObject, root: Rank) -> MPObject: ...
-
-    @abstractmethod
-    def seval(self, fe_type: str, pyfn: Callable, *args: Any, **kwargs: Any) -> Any:
-        """Run a function in the secure environment."""
-
-    @abstractmethod
-    def reveal(self, obj: MPObject, to_mask: Mask) -> MPObject: ...
-
-    @abstractmethod
-    def revealTo(self, obj: MPObject, to_rank: Rank) -> MPObject: ...
-
-
-class Delegation(SecureAPI):
-    """Delegate to a trusted third-party to perform secure operations."""
-
-    def seal(self, obj: MPObject, frm_mask: Mask | None = None) -> list[MPObject]:
-        raise NotImplementedError("TODO")
-
-    def sealFrom(self, obj: MPObject, root: Rank) -> MPObject:
-        raise NotImplementedError("TODO")
-
-    def seval(self, fe_type: str, pyfn: Callable, *args: Any, **kwargs: Any) -> Any:
-        raise NotImplementedError("TODO")
-
-    def reveal(self, obj: MPObject, to_mask: Mask) -> MPObject:
-        raise NotImplementedError("TODO")
-
-    def revealTo(self, obj: MPObject, to_rank: Rank) -> MPObject:
-        raise NotImplementedError("TODO")
-
-
-class SPU(SecureAPI):
-    """Use SPU to  perform secure operations."""
-
-    def get_spu_mask(self) -> Mask:
-        spu_devices = cur_ctx().cluster_spec.get_devices_by_kind("SPU")
-        if not spu_devices:
-            raise ValueError("No SPU device found in the cluster specification")
-        if len(spu_devices) > 1:
-            raise ValueError("Multiple SPU devices found in the cluster specification")
-        spu_device = spu_devices[0]
-        spu_mask = Mask.from_ranks([member.rank for member in spu_device.members])
-        return spu_mask
-
-    def seal(self, obj: MPObject, frm_mask: Mask | None = None) -> list[MPObject]:
-        spu_mask: Mask = self.get_spu_mask()
-        if obj.pmask is None:
-            if frm_mask is None:
-                # NOTE: The length of the return list is statically determined by obj_mask,
-                # so only static masks are supported here.
-                raise ValueError("Seal does not support dynamic masks.")
-            else:
-                # Force seal from the given mask, the runtime will raise error if the mask
-                # does not match obj.pmask.
-                # TODO(jint): add set_pmask primitive.
-                pass
-        else:
-            if frm_mask is None:
-                frm_mask = obj.pmask
-            else:
-                if not Mask(frm_mask).is_subset(obj.pmask):
-                    raise ValueError(f"Cannot seal from {frm_mask} to {obj.pmask}, ")
-
-        # Get the world_size from spu_mask (number of parties in SPU computation)
-        world_size = Mask(spu_mask).num_parties()
-        pfunc, ins, _ = spu.makeshares(
-            obj, world_size=world_size, visibility=spu.Visibility.SECRET
-        )
-        assert len(ins) == 1
-        shares = peval(pfunc, ins, frm_mask)
-
-        # scatter the shares to each party.
-        return [mpi.scatter_m(spu_mask, rank, shares) for rank in Mask(frm_mask)]
-
-    def sealFrom(self, obj: MPObject, root: Rank) -> MPObject:
-        results = seal(obj, frm_mask=Mask.from_ranks(root))
-        assert len(results) == 1, f"Expected one result, got {len(results)}"
-        return results[0]
-
-    def seval(self, fe_type: str, pyfn: Callable, *args: Any, **kwargs: Any) -> Any:
-        if fe_type != "jax":
-            raise ValueError(f"Unsupported fe_type: {fe_type}")
-
-        spu_mask = self.get_spu_mask()
-        pfunc, in_vars, out_tree = spu.jax_compile(pyfn, *args, **kwargs)
-        assert all(var.pmask == spu_mask for var in in_vars), in_vars
-        out_flat = peval(pfunc, in_vars, spu_mask)
-        return tree_unflatten(out_tree, out_flat)
-
-    def reveal(self, obj: MPObject, to_mask: Mask) -> MPObject:
-        spu_mask = self.get_spu_mask()
-
-        assert obj.pmask == spu_mask, (obj.pmask, spu_mask)
-
-        # (n_parties, n_shares)
-        shares = [mpi.bcast_m(to_mask, rank, obj) for rank in Mask(spu_mask)]
-        assert len(shares) == Mask(spu_mask).num_parties(), (shares, spu_mask)
-        assert all(share.pmask == to_mask for share in shares)
-
-        # Reconstruct the original object from shares
-        pfunc, ins, _ = spu.reconstruct(*shares)
-        return peval(pfunc, ins, to_mask)[0]  # type: ignore[no-any-return]
-
-    def revealTo(self, obj: MPObject, to_rank: Rank) -> MPObject:
-        return self.reveal(obj, to_mask=Mask.from_ranks(to_rank))
-
-
-class SEE(Enum):
-    """Secure Execution Environment."""
-
-    MOCK = 0
-    SPU = 1
-    TEE = 2
-
-
-# TODO(jint): move me to options.py
-mode: SEE = SEE.SPU
-
-
-def _get_sapi() -> SecureAPI:
-    """Get the current secure API based on the mode."""
-    if mode == SEE.MOCK:
-        return Delegation()
-    elif mode == SEE.SPU:
-        return SPU()
-    elif mode == SEE.TEE:
-        raise NotImplementedError("TEE is not implemented yet")
-    else:
-        raise ValueError(f"Unknown mode: {mode}")
-
-
-# seal :: m a -> [s a]
-def seal(obj: MPObject, frm_mask: Mask | None = None) -> list[MPObject]:
-    """Seal an simp object, result a list of sealed objects, with
-    the i'th element as the secret from the i'th party.
-    """
-    return _get_sapi().seal(obj, frm_mask=frm_mask)
-
-
-# sealFrom :: m a -> m Rank -> s a
-def sealFrom(obj: MPObject, root: Rank) -> MPObject:
-    """Seal an simp object from a specific root party."""
-    return _get_sapi().sealFrom(obj, root)
-
-
-# reveal :: s a -> m a
-def reveal(obj: MPObject, to_mask: Mask | None = None) -> MPObject:
-    """Reveal a sealed object to pmask'ed parties."""
-    to_mask = to_mask or Mask.all(psize())
-    return _get_sapi().reveal(obj, to_mask)
-
-
-# revealTo :: s a -> m Rank -> m a
-def revealTo(obj: MPObject, to_rank: Rank) -> MPObject:
-    return _get_sapi().revealTo(obj, to_rank)
-
-
-# srun :: (a -> a) -> s a -> s a
-def srun(pyfn: Callable, *, fe_type: str = "jax") -> Callable:
-    @wraps(pyfn)
-    def wrapped(*args: Any, **kwargs: Any) -> Any:
-        return _get_sapi().seval(fe_type, pyfn, *args, **kwargs)
-
-    return wrapped