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

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

@@ -17,7 +17,7 @@ from __future__ import annotations
 from collections.abc import Callable
 from typing import Any, cast
 
-from mplang.core import (
+from mplang.v1.core import (
     Mask,
     MPObject,
     Rank,
@@ -28,8 +28,8 @@ from mplang.core import (
     builtin_function,
     peval,
 )
-from mplang.ops import basic, ibis_cc, jax_cc, sql_cc
-from mplang.ops.base import FeOperation
+from mplang.v1.ops import basic, jax_cc, nnx_cc, sql_cc
+from mplang.v1.ops.base import FeOperation
 
 
 def run(
@@ -273,23 +273,81 @@ def run_jax_at(rank: Rank, jax_fn: Callable, *args: Any, **kwargs: Any) -> Any:
     return run_at(rank, jax_cc.run_jax, jax_fn, *args, **kwargs)
 
 
-def run_ibis(ibis_expr: Any, *args: Any, **kwargs: Any) -> Any:
-    # TODO(jint): add docstring, add type hints, describe args and kwargs constraints.
-    return run(None, ibis_cc.run_ibis, ibis_expr, *args, **kwargs)
-
-
-def run_ibis_at(rank: Rank, ibis_fn: Any, *args: Any, **kwargs: Any) -> Any:
-    return run_at(rank, ibis_cc.run_ibis, ibis_fn, *args, **kwargs)
-
-
 def run_sql(
     query: str, out_type: Any, in_tables: dict[str, MPObject] | None = None
 ) -> Any:
     # TODO(jint): add docstring, drop out_type.
-    return run(None, sql_cc.run_sql, query, out_type, in_tables)
+    return run(None, sql_cc.run_sql_raw, query, out_type, in_tables)
 
 
 def run_sql_at(
     rank: Rank, query: str, out_type: Any, in_tables: dict[str, MPObject] | None = None
 ) -> Any:
-    return run_at(rank, sql_cc.run_sql, query, out_type, in_tables)
+    return run_at(rank, sql_cc.run_sql_raw, query, out_type, in_tables)
+
+
+def run_nnx(nnx_fn: Callable, *args: Any, **kwargs: Any) -> Any:
+    """Run an NNX function.
+
+    Args:
+        nnx_fn: The NNX function to be executed.
+        *args: Positional arguments to pass to the NNX function.
+        **kwargs: Keyword arguments to pass to the NNX function.
+
+    Returns:
+        The result of evaluating the NNX function through the mplang system.
+
+    Raises:
+        TypeError: If the function compilation or evaluation fails.
+        RuntimeError: If the underlying peval execution encounters errors.
+
+    Notes:
+        Argument binding semantics with respect to NNX static arguments:
+
+        - If an argument (or any leaf within a PyTree argument) is an
+          :class:`~mplana.v1.core.mpobject.MPObject`, it is captured as a runtime
+          variable (dynamic value) in the traced program and is not treated as a
+          NNX static argument.
+        - If an argument contains no :class:`MPObject` leaves, it is treated as a
+          constant configuration with respect to NNX; effectively it behaves
+          like a static argument and may contribute to NNX compilation cache
+          keys (similar to ``static_argnums`` semantics). Changing such constant
+          arguments can lead to different compiled variants/cached entries.
+
+    Examples:
+        Defining and running a simple NNX function:
+
+        >>> from flax import nnx
+        >>> import jax.numpy as jnp
+        >>> def nnx_linear(inputs, weights, bias):
+        ...     return jnp.dot(inputs, weights) + bias
+        >>> result = run_nnx(nnx_linear, inputs, weights, bias)
+
+        Running an NNX model:
+
+        >>> class LinearModel(nnx.Module):
+        ...     def __init__(self, features: int, rngs: nnx.Rngs):
+        ...         self.linear = nnx.Linear(features, features, rngs=rngs)
+        ...
+        ...     def __call__(self, x):
+        ...         return self.linear(x)
+        >>> def forward_pass(model, x):
+        ...     return model(x)
+        >>> output = run_nnx(forward_pass, model, input_data)
+    """
+    return run(None, nnx_cc.run_nnx, nnx_fn, *args, **kwargs)
+
+
+def run_nnx_at(rank: Rank, nnx_fn: Callable, *args: Any, **kwargs: Any) -> Any:
+    """Run an NNX function at a specific rank.
+
+    Args:
+        rank: The rank where the NNX function should be executed.
+        nnx_fn: The NNX function to be executed.
+        *args: Positional arguments to pass to the NNX function.
+        **kwargs: Keyword arguments to pass to the NNX function.
+
+    Returns:
+        The result of evaluating the NNX function at the specified rank.
+    """
+    return run_at(rank, nnx_cc.run_nnx, nnx_fn, *args, **kwargs)

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

@@ -16,7 +16,7 @@ from __future__ import annotations
 
 import logging
 
-from mplang.core import Mask, MPObject, Rank, function, pconv, pshfl_s
+from mplang.v1.core import Mask, MPObject, Rank, function, pconv, pshfl_s
 
 
 # scatter :: [m a] -> m Rank -> m a

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

@@ -22,9 +22,9 @@ from functools import wraps
 from types import ModuleType
 from typing import Any
 
-from mplang.ops.base import FeOperation
-from mplang.simp.api import run_at, run_jax_at
-from mplang.simp.mpi import p2p
+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:
@@ -163,7 +163,7 @@ def _load_prelude_modules() -> None:
     unwieldy we can switch to an allowlist.
     """
     try:
-        import mplang.ops as _fe  # type: ignore
+        import mplang.v1.ops as _fe  # type: ignore
     except (ImportError, ModuleNotFoundError):  # pragma: no cover
         # Frontend package not present (minimal install); safe to skip.
         return
@@ -173,7 +173,7 @@ def _load_prelude_modules() -> None:
         if m.name.startswith("_"):
             continue
         if m.name not in _NAMESPACE_REGISTRY:
-            _NAMESPACE_REGISTRY[m.name] = f"mplang.ops.{m.name}"
+            _NAMESPACE_REGISTRY[m.name] = f"mplang.v1.ops.{m.name}"
 
 
 def load_module(module: str, alias: str | None = None) -> None:

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

@@ -21,8 +21,8 @@ import jax.numpy as jnp
 import jax.random as jr
 from jax.typing import ArrayLike
 
-from mplang.core import MPObject, Shape, function, pmask, psize
-from mplang.simp.api import prand, prank, run_jax
+from mplang.v1.core import MPObject, Shape, function, pmask, psize
+from mplang.v1.simp.api import prand, prank, run_jax
 
 
 @function

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()