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

mplang/{simp → v1/simp}/random.py

@@ -21,26 +21,25 @@ import jax.numpy as jnp
 import jax.random as jr
 from jax.typing import ArrayLike
 
-import mplang.core.primitive as prim
-from mplang import simp
-from mplang.core import MPObject, Shape
+from mplang.v1.core import MPObject, Shape, function, pmask, psize
+from mplang.v1.simp.api import prand, prank, run_jax
 
 
-@prim.function
+@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 = simp.run(jr.split)(key) does not work.
+        # 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 simp.run(kernel)(key)  # type: ignore[no-any-return]
+    return run_jax(kernel, key)  # type: ignore[no-any-return]
 
 
-@prim.function
+@function
 def ukey(seed: int | ArrayLike) -> MPObject:
     """Party uniformly generate a random key."""
 
@@ -49,10 +48,10 @@ def ukey(seed: int | ArrayLike) -> MPObject:
         # Note: key.dtype is jax._src.prng.KeyTy, which could not be handled by MPObject.
         return jax.random.key_data(key)
 
-    return simp.run(kernel)()  # type: ignore[no-any-return]
+    return run_jax(kernel)  # type: ignore[no-any-return]
 
 
-@prim.function
+@function
 def urandint(
     key: MPObject | ArrayLike,
     low: int,
@@ -61,13 +60,13 @@ def urandint(
 ) -> MPObject:
     """Party uniformly generate a random integer in the range [low, high) with the given shape."""
 
-    return simp.run(partial(jr.randint, minval=low, maxval=high, shape=shape))(key)  # type: ignore[no-any-return]
+    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.
 
 
-@prim.function
+@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."""
 
@@ -80,11 +79,11 @@ def prandint(low: int, high: int, shape: Shape = ()) -> MPObject:
         result = low + remainder.astype(jnp.int64)
         return result
 
-    rand_u64 = prim.prand(shape)
-    return simp.run(kernel)(rand_u64)  # type: ignore[no-any-return]
+    rand_u64 = prand(shape)
+    return run_jax(kernel, rand_u64)  # type: ignore[no-any-return]
 
 
-@prim.function
+@function
 def pperm(key: MPObject) -> MPObject:
     """Party jointly generate a random permutation.
 
@@ -97,25 +96,25 @@ def pperm(key: MPObject) -> MPObject:
     if key.pmask is None:
         raise ValueError("dynamic pmask is not supported for pperm")
 
-    full_mask = (1 << prim.psize()) - 1
+    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 {prim.psize()}"
+            f"got {key.pmask} with world size {psize()}"
         )
 
-    if prim.pmask() is None or prim.pmask() != full_mask:
+    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 {prim.psize()}"
+            f"got {key.pmask} with world size {psize()}"
         )
 
-    size = prim.psize()
+    size = psize()
 
     def kernel(key: jax.Array) -> jax.Array:
         return jr.permutation(key, size)
 
-    perm = simp.run(kernel)(key)
-    rank = prim.prank()
-    return simp.run(lambda perm, rank: perm[rank])(perm, rank)  # type: ignore[no-any-return]
+    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

@@ -0,0 +1,238 @@
+# 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/table_utils.py

@@ -0,0 +1,185 @@
+# 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 io
+from typing import Any
+
+import pyarrow as pa
+import pyarrow.csv as pa_csv
+import pyarrow.orc as pa_orc
+import pyarrow.parquet as pa_pq
+
+from mplang.v1.core.table import TableLike
+
+__all__ = ["decode_table", "encode_table", "read_table", "write_table"]
+
+
+def _parse_kwargs(kwargs: dict[str, Any], keys: list[str]) -> dict[str, Any] | None:
+    if not kwargs:
+        return None
+
+    return {key: kwargs[key] for key in keys if key in kwargs}
+
+
+_csv_read_option_keys = [
+    "skip_rows",
+    "skip_rows_after_names",
+    "column_names",
+    "autogenerate_column_names",
+    "encoding",
+]
+_csv_parse_option_keys = [
+    "delimiter",
+    "quote_char",
+    "double_quote",
+    "escape_char",
+    "newlines_in_values",
+    "ignore_empty_lines",
+]
+_csv_convert_option_keys = [
+    "check_utf8",
+    "column_types",
+    "null_values",
+    "true_values",
+    "false_values",
+    "decimal_point",
+    "strings_can_be_null",
+    "quoted_strings_can_be_null",
+    "include_columns",
+    "include_missing_columns",
+    "auto_dict_encode",
+    "auto_dict_max_cardinality",
+    "timestamp_parsers",
+]
+
+
+def read_table(
+    source: Any,
+    format: str = "parquet",
+    columns: list[str] | None = None,
+    **kwargs: Any,
+) -> pa.Table:
+    """Read data from a file and return a PyArrow table.
+
+    Args:
+        source: The source to read data from (file path, file-like object, etc.)
+        format: The format of the data source ("parquet", "csv", or "orc")
+        columns: List of column names to read (None means all columns)
+        **kwargs: Additional keyword arguments passed to the underlying reader
+
+    Returns:
+        A PyArrow Table containing the data from the source
+
+    Raises:
+        ValueError: If an unsupported format is specified
+    """
+    match format:
+        case "csv":
+            if columns:
+                kwargs["include_columns"] = columns
+            read_args = _parse_kwargs(kwargs, _csv_read_option_keys)
+            parse_args = _parse_kwargs(kwargs, _csv_parse_option_keys)
+            convert_args = _parse_kwargs(kwargs, _csv_convert_option_keys)
+
+            read_opts = pa_csv.ReadOptions(**read_args) if read_args else None
+            parse_opts = pa_csv.ParseOptions(**parse_args) if parse_args else None
+            conv_opts = pa_csv.ConvertOptions(**convert_args) if convert_args else None
+            return pa_csv.read_csv(
+                source,
+                read_options=read_opts,
+                parse_options=parse_opts,
+                convert_options=conv_opts,
+            )
+        case "orc":
+            return pa_orc.read_table(source, columns=columns, **kwargs)
+        case "parquet":
+            return pa_pq.read_table(source, columns=columns, **kwargs)
+        case _:
+            raise ValueError(f"unsupported data format. {format}")
+
+
+def write_table(
+    data: TableLike,
+    where: Any,
+    format: str = "parquet",
+    **kwargs: Any,
+) -> None:
+    """Write a table-like object to a file in the specified format.
+
+    Args:
+        data: The table-like object to write (PyArrow Table or other compatible format)
+        where: The destination to write to (file path, file-like object, etc.)
+        format: The format to write the data in ("parquet", "csv", or "orc")
+        **kwargs: Additional keyword arguments passed to the underlying writer
+
+    Returns:
+        None
+
+    Raises:
+        ValueError: If the table has no columns or an unsupported format is specified
+    """
+    # Convert data to PyArrow Table if needed
+    table = data if isinstance(data, pa.Table) else pa.table(data)
+    if len(table.column_names) == 0:
+        raise ValueError("Cannot convert Table with no columns.")
+
+    match format:
+        case "csv":
+            options = pa_csv.WriteOptions(**kwargs) if kwargs else None
+            pa_csv.write_csv(table, where, write_options=options)
+        case "orc":
+            pa_orc.write_table(table, where, **kwargs)
+        case "parquet":
+            pa_pq.write_table(table, where, **kwargs)
+        case _:
+            raise ValueError(f"unsupported data format. {format}")
+
+
+def decode_table(
+    data: bytes,
+    format: str = "parquet",
+    columns: list[str] | None = None,
+    **kwargs: Any,
+) -> pa.Table:
+    """Decode a bytes object into a PyArrow table.
+
+    Args:
+        data: The bytes object containing the encoded table data
+        format: The format of the encoded data ("parquet", "csv", or "orc")
+        columns: List of column names to decode (None means all columns)
+        **kwargs: Additional keyword arguments passed to the underlying reader
+
+    Returns:
+        A PyArrow Table decoded from the bytes data
+    """
+    buffer = io.BytesIO(data)
+    return read_table(buffer, format=format, columns=columns, **kwargs)
+
+
+def encode_table(data: TableLike, format: str = "parquet", **kwargs: Any) -> bytes:
+    """Encode a table-like object into bytes.
+
+    Args:
+        data: The table-like object to encode (PyArrow Table or other compatible format)
+        format: The format to encode the data in ("parquet", "csv", or "orc")
+        **kwargs: Additional keyword arguments passed to the underlying writer
+
+    Returns:
+        Bytes object containing the encoded table data
+    """
+    buffer = io.BytesIO()
+    write_table(data, buffer, format, **kwargs)
+    return buffer.getvalue()