mplang-nightly 0.1.dev269__py3-none-any.whl → 0.1.dev270__py3-none-any.whl

mplang/v1/simp/party.py

@@ -1,225 +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 importlib
-import pathlib
-import pkgutil
-from collections.abc import Callable
-from functools import wraps
-from types import ModuleType
-from typing import Any
-
-from mplang.v1.ops.base import FeOperation
-from mplang.v1.simp.api import run_at, run_jax_at
-from mplang.v1.simp.mpi import p2p
-
-
-def P2P(src: Party, dst: Party, value: Any) -> Any:
-    """Point-to-point transfer using Party objects instead of raw ranks.
-
-    Equivalent to ``p2p(src.rank, dst.rank, value)`` but improves readability
-    and reduces magic numbers in user code / tutorials.
-
-    Parameters
-    ----------
-    src : Party
-        Source party object.
-    dst : Party
-        Destination party object.
-    value : Any
-        Value to transfer.
-
-    Returns
-    -------
-    Any
-        The same value representation at destination context (as defined by
-        underlying ``p2p`` primitive semantics).
-    """
-    if not isinstance(src, Party) or not isinstance(dst, Party):  # defensive
-        raise TypeError("P2P expects Party objects, e.g. P2P(P0, P2, value)")
-    return p2p(src.rank, dst.rank, value)
-
-
-"""Party-scoped module registration & dispatch.
-
-This module provides a light-weight mechanism to expose *module-like* groups
-of callable operations bound to a specific party (rank) via attribute access:
-
-    load_module("mplang.ops.crypto", alias="crypto")
-    P0.crypto.encrypt(x)  # executes encrypt() with pmask = {rank 0}
-
-Core concepts:
-* Registry (``_NAMESPACE_REGISTRY``): maps alias -> importable module path.
-* Lazy import: underlying module is imported on first attribute access.
-* Wrapping: fetched callables are wrapped so that invocation automatically
-    routes through ``run_impl`` with that party's mask.
-
-Only *callable* attributes are exposed; non-callable attributes raise
-``AttributeError`` to avoid surprising divergent local vs. distributed
-semantics.
-
-The public API surface intentionally stays small (`Party`, `P`, `run`,
-`runAt`, and `load_module`). Internal details (proxy class / registry) are
-not part of the stability guarantee.
-"""
-
-_NAMESPACE_REGISTRY: dict[str, str] = {}
-
-
-class _PartyModuleProxy:
-    """Lazy module proxy bound to a specific party.
-
-    Attribute access resolves a callable inside the registered module and
-    returns a wrapped function that executes with the party's mask.
-    Non-callable attributes are rejected explicitly to keep semantics clear.
-    """
-
-    def __init__(self, party: Party, name: str):
-        self._party: Party = party
-        self._name: str = name
-        self._module: ModuleType | None = None  # loaded lazily
-
-    def _ensure(self) -> None:
-        if self._module is None:
-            self._module = importlib.import_module(_NAMESPACE_REGISTRY[self._name])
-
-    def __getattr__(self, item: str) -> Callable[..., Any]:
-        self._ensure()
-        op = getattr(self._module, item)
-        if not callable(op):
-            raise AttributeError(
-                f"Attribute '{item}' of party module '{self._name}' is not callable (got {type(op).__name__})"
-            )
-
-        @wraps(op)
-        def _wrapped(*args: Any, **kw: Any) -> Any:
-            # Inline runAt to reduce an extra partial layer while preserving semantics.
-            return run_at(self._party.rank, op, *args, **kw)
-
-        # Provide a party-qualified name for debugging / logs without losing original metadata.
-        base_name = getattr(op, "__name__", None)
-        if base_name is None:
-            # Frontend FeOperation or object without __name__; try .name attribute (FeOperation contract) or fallback to repr
-            base_name = getattr(op, "name", None) or type(op).__name__
-        try:
-            _wrapped.__name__ = f"{base_name}@P{self._party.rank}"
-        except Exception:  # pragma: no cover - assignment may fail for exotic wrappers
-            pass
-        return _wrapped
-
-
-class Party:
-    def __init__(self, rank: int) -> None:
-        self.rank: int = int(rank)
-
-    def __repr__(self) -> str:  # pragma: no cover
-        return f"Party(rank={self.rank})"
-
-    def __call__(self, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
-        if not callable(fn):
-            raise TypeError(
-                f"First argument to Party({self.rank}) must be callable, got {fn!r}"
-            )
-        # Use run_op_at for FeOperation, run_jax_at for plain callables
-        if isinstance(fn, FeOperation):
-            return run_at(self.rank, fn, *args, **kwargs)
-        else:
-            # TODO(jint): implicitly assume non-FeOperation as JAX function is a bit too magical?
-            return run_jax_at(self.rank, fn, *args, **kwargs)
-
-    def __getattr__(self, name: str) -> _PartyModuleProxy:
-        if name in _NAMESPACE_REGISTRY:
-            return _PartyModuleProxy(self, name)
-        raise AttributeError(
-            f"Party has no attribute '{name}'. Registered: {list(_NAMESPACE_REGISTRY)}"
-        )
-
-
-class _PartyIndex:
-    def __getitem__(self, rank: int) -> Party:
-        return Party(rank)
-
-
-def _load_prelude_modules() -> None:
-    """Auto-register public frontend submodules for party namespace access.
-
-    Implementation detail: we treat every non-underscore immediate child of
-    ``mplang.ops`` as public and make it available as ``P0.<name>``.
-    This keeps user ergonomics high (no manual load_module calls for core
-    frontends) but slightly increases implicit surface area. If this grows
-    unwieldy we can switch to an allowlist.
-    """
-    try:
-        import mplang.v1.ops as _fe  # type: ignore
-    except (ImportError, ModuleNotFoundError):  # pragma: no cover
-        # Frontend package not present (minimal install); safe to skip.
-        return
-
-    pkg_path = pathlib.Path(_fe.__file__).parent
-    for m in pkgutil.iter_modules([str(pkg_path)]):
-        if m.name.startswith("_"):
-            continue
-        if m.name not in _NAMESPACE_REGISTRY:
-            _NAMESPACE_REGISTRY[m.name] = f"mplang.v1.ops.{m.name}"
-
-
-def load_module(module: str, alias: str | None = None) -> None:
-    """Register a module for party-scoped (per-rank) callable access.
-
-    After registration, each party object (e.g. ``P0``) can access callable
-    attributes of the target module through the chosen alias and have them
-    executed under that party's mask automatically. Non-callable attributes
-    are intentionally not exposed to avoid ambiguity between local data and
-    distributed execution semantics.
-
-    Parameters
-    ----------
-    module : str
-        The fully-qualified import path of the module to expose. It must be
-        importable via ``importlib.import_module``.
-    alias : str | None, optional
-        The short name used as an attribute on ``Party``/``P0``/``P1``/... .
-        If omitted, the last path segment of ``module`` is used.
-
-    Raises
-    ------
-    ValueError
-        If the alias is already registered to a *different* module path.
-
-    Notes
-    -----
-    Registration is idempotent when the alias maps to the same module. The
-    actual module object is imported lazily on first attribute access, so
-    calling ``load_module`` has negligible upfront cost.
-
-    Examples
-    --------
-    >>> load_module("mplang.ops.crypto", alias="crypto")
-    >>> # Now call an op on party 0
-    >>> P0.crypto.encrypt(data)
-    """
-    if alias is None:
-        alias = module.rsplit(".", 1)[-1]
-    prev = _NAMESPACE_REGISTRY.get(alias)
-    if prev and prev != module:
-        raise ValueError(f"Alias '{alias}' already registered for '{prev}'")
-    _NAMESPACE_REGISTRY[alias] = module
-
-
-P = _PartyIndex()
-P0, P1, P2 = Party(0), Party(1), Party(2)
-
-_load_prelude_modules()

mplang/v1/simp/random.py

@@ -1,120 +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
-
-from functools import partial
-
-import jax
-import jax.numpy as jnp
-import jax.random as jr
-from jax.typing import ArrayLike
-
-from mplang.v1.core import MPObject, Shape, function, pmask, psize
-from mplang.v1.simp.api import prand, prank, run_jax
-
-
-@function
-def key_split(key: MPObject) -> tuple[MPObject, MPObject]:
-    """Split the key into two keys."""
-
-    def kernel(key: jax.Array) -> tuple[jax.Array, jax.Array]:
-        # TODO: since MPObject tensor does not implement slicing yet.
-        # subkey, key = run_jax(jr.split, key) does not work.
-        # we workaround it by splitting inside tracer.
-        subkey, key = jr.split(key)
-        return subkey, key
-
-    return run_jax(kernel, key)  # type: ignore[no-any-return]
-
-
-@function
-def ukey(seed: int | ArrayLike) -> MPObject:
-    """Party uniformly generate a random key."""
-
-    def kernel() -> jax.Array:
-        key = jax.random.key(seed)
-        # Note: key.dtype is jax._src.prng.KeyTy, which could not be handled by MPObject.
-        return jax.random.key_data(key)
-
-    return run_jax(kernel)  # type: ignore[no-any-return]
-
-
-@function
-def urandint(
-    key: MPObject | ArrayLike,
-    low: int,
-    high: int,
-    shape: Shape = (),
-) -> MPObject:
-    """Party uniformly generate a random integer in the range [low, high) with the given shape."""
-
-    return run_jax(partial(jr.randint, minval=low, maxval=high, shape=shape), key)  # type: ignore[no-any-return]
-
-
-# Private(different per-party) related functions begin.
-
-
-@function
-def prandint(low: int, high: int, shape: Shape = ()) -> MPObject:
-    """Party privately generate a random integer in the range [low, high) with the given shape."""
-
-    def kernel(rand_u64: jnp.ndarray) -> jnp.ndarray:
-        range_size = high - low
-        if range_size <= 0:
-            raise ValueError("'high' must be greater than 'low'")
-
-        remainder = jax.lax.rem(rand_u64, jnp.uint64(range_size))
-        result = low + remainder.astype(jnp.int64)
-        return result
-
-    rand_u64 = prand(shape)
-    return run_jax(kernel, rand_u64)  # type: ignore[no-any-return]
-
-
-@function
-def pperm(key: MPObject) -> MPObject:
-    """Party jointly generate a random permutation.
-
-    That is, each party holds a random number in range(size), and all parties as a whole
-    hold a random permutation of integers from 0 to size-1.
-
-    Note: this function is NOT 'secure', that is, all parties know the permutation result.
-    """
-
-    if key.pmask is None:
-        raise ValueError("dynamic pmask is not supported for pperm")
-
-    full_mask = (1 << psize()) - 1
-
-    if key.pmask != full_mask:
-        raise ValueError(
-            "key must be a MPObject with mask covering all parties, "
-            f"got {key.pmask} with world size {psize()}"
-        )
-
-    if pmask() is None or pmask() != full_mask:
-        raise ValueError(
-            "pperm must be run with a mask covering all parties, "
-            f"got {key.pmask} with world size {psize()}"
-        )
-
-    size = psize()
-
-    def kernel(key: jax.Array) -> jax.Array:
-        return jr.permutation(key, size)
-
-    perm = run_jax(kernel, key)
-    rank = prank()
-    return run_jax(lambda perm, rank: perm[rank], perm, rank)  # type: ignore[no-any-return]

mplang/v1/simp/smpc.py

@@ -1,238 +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.
-
-"""
-SMPC on simp: conventions and object semantics
-
-Overview
-- simp is party-centric. Objects produced purely by simp code carry only an execution
-    mask ("pmask") and have no security device semantics by default.
-- Secure semantics (secret sharing, protected execution, declassification) are introduced
-    only when using the device API or the helpers in this module: "seal", "srun", "reveal".
-
-Definitions
-- "__device__" attribute is attached by the device API to indicate the concrete device
-    an object is bound to (e.g., an SPU/TEE/PPU name). See mplang.device.DEVICE_ATTR_NAME.
-- pmask describes which parties currently hold/execute the value under the simp model.
-
-Conventions
-1) If an object has NO "__device__" attribute (i.e., it has not gone through the device API):
-     - It is a simp object, privately owned on the parties indicated by its pmask.
-     - When sealed via "seal(obj)", we infer target PPU device(s) from pmask:
-         • one-hot pmask {pi} → route to PPU(pi).
-         • multi-party pmask → fan out per party and seal independently to each party's PPU.
-     - Such objects CANNOT be passed to "srun"/"reveal" directly; seal first.
-
-2) If an object HAS a "__device__" attribute:
-     - Its behavior follows the bound device (e.g., SPU/TEE/PPU) and its membership.
-     - "srun" executes on that device; "reveal" declassifies from that device to the requested parties.
-     - pmask must be consistent with the device membership during transitions; inconsistencies raise errors.
-
-Notes
-- "seal"/"seal_from" construct secret shares on the chosen secure device and attach the
-    "__device__" attribute to outputs. "srun"/"reveal" assume inputs are already sealed
-    (device-bound) and validate pmask ↔ device-membership consistency.
-- These rules align with "design/simp_vs_device.md" and keep routing unambiguous.
-
-Examples (obj state → interpretation/behavior)
-- {pmask={A}, dev_attr=None}: simp plaintext on party A. "seal" routes to PPU(A);
-    must "seal" before "srun"/"reveal".
-- {pmask={A,B}, dev_attr=None}: simp plaintext held by A and B. "seal" produces two
-    per-party sealed objects via PPU(A) and PPU(B), respectively.
-- {pmask={A,B}, dev_attr="spu:spu0"}: device object on SPU(spu0) whose members are {A,B};
-    "srun" runs on spu0; "reveal(to={A})" reveals result to party A.
-- {pmask={A}, dev_attr="ppu:A"}: device object on PPU(A); "reveal(to={A})" returns A's plaintext.
-- {pmask=None, dev_attr=None}: dynamic pmask; "seal" is unsupported and will error.
-- {pmask={A}, dev_attr="spu:spu0"} where A ∉ members(spu0): inconsistent; operations will error.
-"""
-
-from collections.abc import Callable
-from typing import Any
-
-from mplang.v1 import _device
-from mplang.v1.core import Mask, MPObject, Rank, psize
-from mplang.v1.core.cluster import Device
-from mplang.v1.core.context_mgr import cur_ctx
-from mplang.v1.core.primitive import pconv
-from mplang.v1.simp.api import set_mask
-from mplang.v1.utils.func_utils import normalize_fn
-
-
-def _determine_secure_device(*args: MPObject) -> Device:
-    """Determine secure device from args, or find any available if no args."""
-    if not args:
-        # Find an available secure device (fallback when no args provided).
-        devices = cur_ctx().cluster_spec.get_devices_by_kind("SPU")
-        if devices:
-            return devices[0]
-
-        devices = cur_ctx().cluster_spec.get_devices_by_kind("TEE")
-        if devices:
-            return devices[0]
-
-        raise ValueError(
-            "No secure device (SPU or TEE) found in the cluster specification"
-        )
-
-    dev_names: list[str] = []
-    for arg in args:
-        if not _device.is_device_obj(arg):
-            raise ValueError(
-                "srun/reveal expect sealed inputs with a device attribute; "
-                f"got an unsealed object: {arg}. Please call seal()/seal_from() first."
-            )
-        dev_names.append(_device.get_dev_attr(arg))
-
-    if len(set(dev_names)) != 1:
-        raise ValueError(f"Ambiguous secure devices among arguments: {dev_names}")
-
-    dev_name = dev_names[0]
-
-    cluster_spec = cur_ctx().cluster_spec
-    assert dev_name in cluster_spec.devices
-    return cluster_spec.devices[dev_name]
-
-
-def _get_ppu_from_rank(rank: Rank) -> Device:
-    """Get the PPU device for a specific rank."""
-    for dev in cur_ctx().cluster_spec.get_devices_by_kind("PPU"):
-        assert len(dev.members) == 1, "Expected single member PPU devices."
-        if dev.members[0].rank == rank:
-            return dev
-    raise ValueError(f"No PPU device found for rank {rank}.")
-
-
-def seal(obj: MPObject) -> list[MPObject] | MPObject:
-    """Seal a simp object to a secure device.
-
-    Args:
-        obj: The simp object to seal.
-
-    Returns:
-        The sealed object(s). If the input is a plaintext simp object with a multi-party
-        mask, a list of sealed objects (one per party) is returned. Otherwise, a
-        single sealed object is returned.
-    """
-
-    if obj.pmask is None:
-        raise ValueError("Seal does not support dynamic masks.")
-
-    if _device.is_device_obj(obj):
-        sdev = _determine_secure_device()
-        return _device._d2d(sdev.name, obj)
-    else:
-        # it's a normal plaintext simp object, treat as a list of PPU objects
-        rets: list[MPObject] = []
-        for rank in obj.pmask:
-            ppu_obj = set_mask(obj, Mask.from_ranks([rank]))
-            _device.set_dev_attr(ppu_obj, _get_ppu_from_rank(rank).name)
-            sealed = seal(ppu_obj)
-            assert isinstance(sealed, MPObject), (
-                "Expected single sealed object per rank"
-            )
-            rets.append(sealed)
-        return rets
-
-
-def seal_from(from_rank: Rank, obj: MPObject) -> MPObject:
-    """Seal a simp object from a specific party to its PPU.
-
-    Args:
-        from_rank: The party rank from which to seal the object.
-        obj: The simp object to seal.
-
-    Returns:
-        The sealed object.
-    """
-    obj = set_mask(obj, Mask.from_ranks([from_rank]))
-    out = seal(obj)
-    assert isinstance(out, list), "seal_from should return a list of sealed objects."
-    assert len(out) == 1, "seal_from should return a single sealed object."
-    return out[0]
-
-
-# reveal :: s a -> m a
-def reveal(obj: MPObject, to_mask: Mask | None = None) -> MPObject:
-    """Reveal a sealed object to pmask'ed parties."""
-    assert isinstance(obj, MPObject), "reveal expects an MPObject."
-
-    if not _device.is_device_obj(obj):
-        raise ValueError(f"reveal does not support non-device object={obj}.")
-
-    if to_mask is None:
-        ranks = []
-        for rank in range(psize()):
-            try:
-                _get_ppu_from_rank(rank)
-            except ValueError:
-                continue
-            ranks.append(rank)
-        to_mask = Mask.from_ranks(ranks)
-    rets = [reveal_to(rank, obj) for rank in to_mask]
-    return pconv(rets)
-
-
-def reveal_to(to_rank: Rank, obj: MPObject) -> MPObject:
-    """Reveal a sealed object to a specific party."""
-    if not _device.is_device_obj(obj):
-        raise ValueError("reveal_to expects a device object (sealed value).")
-
-    to_dev = _get_ppu_from_rank(to_rank)
-    return _device._d2d(to_dev.name, obj)
-
-
-def srun(fe_type: str, pyfn: Callable, *args: Any, **kwargs: Any) -> Any:
-    """Run a function on sealed values securely.
-
-    This function executes a computation on sealed (secret-shared) values
-    using secure multi-party computation (MPC).
-
-    Args:
-        fe_type: The front-end type, e.g., "jax"
-        pyfn: A function to run on sealed values
-        *args: Positional arguments (sealed values)
-        **kwargs: Keyword arguments (sealed values)
-
-    Returns:
-        The result of the computation, still in sealed form
-    """
-
-    fn_flat, args_flat = normalize_fn(
-        pyfn, args, kwargs, lambda x: isinstance(x, MPObject)
-    )
-
-    dev_info = _determine_secure_device(*args_flat)
-
-    dev_kind = dev_info.kind.upper()
-    if dev_kind in {"SPU", "TEE"}:
-        return _device.device(dev_info.name, fe_type=fe_type)(fn_flat)(args_flat)
-    else:
-        raise ValueError(f"Unsupported secure device kind: {dev_kind}")
-
-
-def srun_jax(jax_fn: Callable, *args: Any, **kwargs: Any) -> Any:
-    """Run a jax function on sealed values securely.
-
-    This function executes a JAX computation on sealed (secret-shared) values
-    using secure multi-party computation (MPC).
-
-    Args:
-        jax_fn: A JAX function to run on sealed values
-        *args: Positional arguments (sealed values)
-        **kwargs: Keyword arguments (sealed values)
-
-    Returns:
-        The result of the computation, still in sealed form
-    """
-    return srun("jax", jax_fn, *args, **kwargs)

mplang/v1/utils/__init__.py

@@ -1,13 +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.

mplang/v1/utils/crypto.py

@@ -1,32 +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.
-
-"""Lightweight cryptographic utilities.
-
-Currently only exposes `blake2b` which is used by mock crypto backends.
-Previously contained mock signing/key utilities which were unused in code
-paths and have been removed to avoid confusion.
-"""
-
-from __future__ import annotations
-
-import hashlib
-
-
-def blake2b(data: bytes) -> bytes:
-    """Return 32-byte BLAKE2b digest for the given data."""
-    return hashlib.blake2b(data, digest_size=32).digest()
-
-
-__all__ = ["blake2b"]