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

mplang/v1/host.py

@@ -1,130 +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 collections.abc import Callable
-from typing import Any
-
-from jax.tree_util import tree_map
-
-from mplang.v1.core import (
-    ClusterSpec,
-    InterpContext,
-    MPContext,
-    MPObject,
-    TraceContext,
-    TracedFunction,
-    trace,
-)
-from mplang.v1.core.context_mgr import cur_ctx, with_ctx
-
-
-def evaluate(
-    interp: InterpContext, mpfn: Callable[..., Any], *args: Any, **kwargs: Any
-) -> Any:  # type: ignore[misc]
-    """Evaluate a multi-party function with the given interpreter context.
-
-    This function accepts arbitrary types as it's designed to handle
-    any multi-party computation function and arguments.
-
-    Args:
-        interp: The interpreter context for evaluating the multi-party function.
-        mpfn: The multi-party function to evaluate.
-        *args: Positional arguments to pass to the function.
-        **kwargs: Keyword arguments to pass to the function.
-
-    Returns:
-        Any: The result of evaluating the multi-party function, which can be
-             any type depending on the function's return type.
-    """
-    assert isinstance(interp, InterpContext), f"Expect InterpContext, got {interp}"
-    with with_ctx(interp):
-        return mpfn(*args, **kwargs)
-
-
-def fetch(interp: InterpContext | None, objs: Any) -> Any:  # type: ignore[misc]
-    """Fetch computed results from MPObject instances in nested data structures.
-
-    This function uses tree_map to handle arbitrary nested structures,
-    so it needs to accept and return Any type.
-
-    Args:
-        interp: The interpreter context for fetching results. If None, uses the
-                current context from cur_ctx().
-        objs: The objects containing MPObject instances to fetch. Can be any
-              nested structure.
-
-    Returns:
-        Any: The fetched results with the same structure as the input objects,
-             but with MPObject instances replaced by their computed values.
-    """
-    ctx = interp or cur_ctx()
-    assert isinstance(ctx, InterpContext), f"Expect MPExecutor, got {ctx}"
-
-    evaluated = evaluate(ctx, lambda x: x, objs)
-
-    def fetch_impl(arg: MPObject | Any) -> Any:
-        if not isinstance(arg, MPObject):
-            return arg
-
-        return ctx.fetch(arg)
-
-    return tree_map(fetch_impl, evaluated)
-
-
-class CompileOptions(MPContext):
-    """
-    Lightweight ``MPContext`` used for ahead-of-time (AOT) compilation.
-
-    Args:
-        psize: Number of participating parties.
-        spu_mask: Bitmask indicating which parties own an SPU device. Defaults
-            to a mask that enables all parties.
-    """
-
-    def __init__(self, cluster_spec: Any) -> None:
-        super().__init__(cluster_spec)
-
-    @classmethod
-    def simple(cls, world_size: int) -> CompileOptions:
-        """Create a simple CompileOptions with the given party size and SPU mask.
-
-        Args:
-            world_size: Number of participating parties.
-
-        Returns:
-            A CompileOptions instance.
-        """
-        cluster_spec = ClusterSpec.simple(world_size)
-        return cls(cluster_spec)
-
-
-def compile(
-    mctx: MPContext, fn: Callable[..., Any], *args: Any, **kwargs: Any
-) -> TracedFunction:
-    """Compile a multi-party function into a TracedFunction.
-
-    Args:
-        mctx: The multi-party context for compilation.
-        fn: The function to compile.
-        *args: Positional arguments to pass during compilation.
-        **kwargs: Keyword arguments to pass during compilation.
-
-    Returns:
-        TracedFunction: The compiled function representation that can be
-                       evaluated in multi-party contexts.
-    """
-    trace_ctx = TraceContext(mctx.cluster_spec)
-    return trace(trace_ctx, fn, *args, **kwargs)

mplang/v1/kernels/__init__.py

@@ -1,41 +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 mplang.v1.kernels.value import (
-    BytesBlob,
-    TableValue,
-    TensorValue,
-    Value,
-    ValueDecodeError,
-    ValueError,
-    decode_value,
-    encode_value,
-    is_value_envelope,
-    list_value_kinds,
-    register_value,
-)
-
-__all__ = [
-    "BytesBlob",
-    "TableValue",
-    "TensorValue",
-    "Value",
-    "ValueDecodeError",
-    "ValueError",
-    "decode_value",
-    "encode_value",
-    "is_value_envelope",
-    "list_value_kinds",
-    "register_value",
-]

mplang/v1/kernels/base.py

@@ -1,125 +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.
-
-"""Backend kernel registry: mapping kernel_id -> implementation.
-
-This module provides a lightweight registry for backend kernel implementations.
-It does not track or decide which kernel handles a given semantic operation;
-that policy (op -> kernel_id) is managed externally by each ``RuntimeContext``.
-
-Exposed primitives:
-* ``@kernel_def(kernel_id)``: decorator to register a kernel implementation.
-* ``get_kernel_spec(kernel_id)``: look up a previously registered kernel.
-* ``cur_kctx()`` / ``KernelContext``: execution context available only
-    inside a kernel body (rank, world_size, per-backend state pockets, and a
-    runtime-wide cache shared by kernels of the same runtime instance).
-
-No global op binding table exists here; callers resolve an op to a kernel_id
-before invoking the kernel function.
-"""
-
-from __future__ import annotations
-
-import contextvars
-from collections.abc import Callable
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
-
-if TYPE_CHECKING:
-    from mplang.v1.kernels.context import RuntimeContext
-
-__all__ = [
-    "KernelContext",
-    "KernelSpec",
-    "cur_kctx",
-    "get_kernel_spec",
-    "kernel_exists",
-    "list_kernels",
-]
-
-
-@dataclass
-class KernelContext:
-    """Ephemeral per-kernel invocation context.
-
-    Cross-kernel persistent state (RNGs, compiled artifacts, environment handles)
-    should be stored in RuntimeContext.
-    """
-
-    rank: int
-    world_size: int
-    runtime: RuntimeContext
-
-
-_CTX_VAR: contextvars.ContextVar[KernelContext | None] = contextvars.ContextVar(
-    "_flat_backend_ctx", default=None
-)
-
-
-def cur_kctx() -> KernelContext:
-    """Return current kernel execution context (only valid inside kernel)."""
-    ctx = _CTX_VAR.get()
-    if ctx is None:
-        raise RuntimeError("cur_kctx() called outside backend kernel execution")
-    return ctx
-
-
-# ---------------- Registry ----------------
-
-# Kernel callable signature: (pfunc, *args) -> Any | sequence (no **kwargs)
-KernelFn = Callable[..., Any]
-
-
-@dataclass
-class KernelSpec:
-    kernel_id: str
-    fn: KernelFn
-    meta: dict[str, Any]
-
-
-# All registered kernel implementations: kernel_id -> spec
-_KERNELS: dict[str, KernelSpec] = {}
-
-
-def kernel_def(kernel_id: str, /, **meta: Any) -> Callable[[KernelFn], KernelFn]:
-    """Decorator to register a concrete kernel implementation.
-
-    This ONLY registers the implementation (kernel_id -> fn). It does NOT bind
-    any op. Higher layer must call ``bind_op(op_type, kernel_id)`` explicitly.
-    """
-
-    def _decorator(fn: KernelFn) -> KernelFn:
-        if kernel_id in _KERNELS:
-            raise ValueError(f"duplicate kernel_id={kernel_id}")
-        _KERNELS[kernel_id] = KernelSpec(kernel_id=kernel_id, fn=fn, meta=dict(meta))
-        return fn
-
-    return _decorator
-
-
-def get_kernel_spec(kernel_id: str) -> KernelSpec:
-    """Return KernelSpec for a registered kernel_id (no op binding lookup)."""
-    spec = _KERNELS.get(kernel_id)
-    if spec is None:
-        raise KeyError(f"kernel_id {kernel_id} not registered")
-    return spec
-
-
-def list_kernels() -> list[str]:
-    return sorted(_KERNELS.keys())
-
-
-def kernel_exists(kernel_id: str) -> bool:
-    """Return True if a kernel_id has been registered."""
-    return kernel_id in _KERNELS

mplang/v1/kernels/basic.py

@@ -1,240 +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 numpy as np
-
-from mplang.v1.core import PFunction, TableType, TensorType
-from mplang.v1.kernels.base import cur_kctx, kernel_def
-from mplang.v1.kernels.value import TableValue, TensorValue, Value
-from mplang.v1.runtime.data_providers import get_provider, resolve_uri
-from mplang.v1.utils import table_utils
-
-
-@kernel_def("basic.identity")
-def _identity(pfunc: PFunction, value: Value) -> Value:
-    # Runtime guarantees exactly one argument; no extra arity checks here.
-    return value
-
-
-@kernel_def("basic.read")
-def _read(pfunc: PFunction) -> Value:
-    path = pfunc.attrs.get("path")
-    if path is None:
-        raise ValueError("missing path attr for basic.read")
-    out_t = pfunc.outs_info[0]
-    uri = resolve_uri(str(path))
-    prov = get_provider(uri.scheme)
-    if prov is None:
-        raise NotImplementedError(f"no resource provider for scheme: {uri.scheme}")
-    ctx = cur_kctx()
-    try:
-        data = prov.read(uri, out_t, ctx=ctx)
-    except Exception as e:  # pragma: no cover - provider errors
-        raise RuntimeError(f"basic.read failed: {e}") from e
-
-    if isinstance(data, Value):
-        return data
-
-    if isinstance(out_t, TableType):
-        return TableValue(data)
-    elif isinstance(out_t, TensorType):
-        return TensorValue(np.asarray(data))
-    else:
-        raise TypeError(
-            f"basic.read only supports TableType/TensorType outputs, got {type(out_t).__name__}"
-        )
-
-
-@kernel_def("basic.write")
-def _write(pfunc: PFunction, obj: Value) -> Value:
-    path = pfunc.attrs.get("path")
-    if path is None:
-        raise ValueError("missing path attr for basic.write")
-    uri = resolve_uri(str(path))
-    prov = get_provider(uri.scheme)
-    if prov is None:
-        raise NotImplementedError(f"no resource provider for scheme: {uri.scheme}")
-    # Pass Value object directly to provider - let provider decide how to handle it
-    ctx = cur_kctx()
-    try:
-        prov.write(uri, obj, ctx=ctx)
-    except Exception as e:  # pragma: no cover
-        raise RuntimeError(f"basic.write failed: {e}") from e
-    return obj
-
-
-@kernel_def("basic.constant")
-def _constant(pfunc: PFunction) -> Value:
-    """Return constants as Value types (TensorValue or TableValue)."""
-    data_bytes = pfunc.attrs.get("data_bytes")
-    if data_bytes is None:
-        raise ValueError("missing data_bytes attr for basic.constant")
-    out_t = pfunc.outs_info[0]
-    fmt = pfunc.attrs.get("data_format")
-    if isinstance(out_t, TableType):
-        if fmt != "bytes[parquet]":
-            raise ValueError(f"unsupported table constant format {fmt}")
-        df = table_utils.decode_table(data_bytes, format="parquet")
-        return TableValue(df)
-    # tensor path
-    shape = out_t.shape  # type: ignore[attr-defined,union-attr]
-    dtype = out_t.dtype.numpy_dtype()  # type: ignore[attr-defined,union-attr]
-    arr = np.frombuffer(data_bytes, dtype=dtype).reshape(shape)
-    return TensorValue(arr)
-
-
-@kernel_def("basic.rank")
-def _rank(pfunc: PFunction) -> TensorValue:
-    """Return rank as TensorValue."""
-    ctx = cur_kctx()
-    arr = np.array(ctx.rank, dtype=np.uint64)
-    return TensorValue(arr)
-
-
-@kernel_def("basic.prand")
-def _prand(pfunc: PFunction) -> TensorValue:
-    """Return random data as TensorValue."""
-    shape = pfunc.attrs.get("shape", ())
-    rng = np.random.default_rng()
-    info = np.iinfo(np.uint64)
-    data = rng.integers(
-        low=info.min, high=info.max, size=shape, dtype=np.uint64, endpoint=True
-    )
-    return TensorValue(data)
-
-
-@kernel_def("basic.table_to_tensor")
-def _table_to_tensor(pfunc: PFunction, table: TableValue) -> TensorValue:
-    """Convert table to tensor, return as TensorValue."""
-    arrow_table = table.to_arrow()
-    if arrow_table.num_columns == 0:
-        raise ValueError("cannot pack empty table")
-    # Convert Arrow columns to numpy arrays and stack
-    mat = np.column_stack([
-        arrow_table.column(i).to_numpy() for i in range(arrow_table.num_columns)
-    ])
-    return TensorValue(mat)
-
-
-@kernel_def("basic.tensor_to_table")
-def _tensor_to_table(pfunc: PFunction, tensor: TensorValue) -> TableValue:
-    """Convert tensor to table, return as TableValue."""
-    import pyarrow as pa  # type: ignore
-
-    arr = tensor.to_numpy()
-    if arr.ndim != 2:
-        raise ValueError("tensor_to_table expects rank-2 array")
-    col_names = pfunc.attrs.get("column_names")
-    if col_names is None:
-        raise ValueError("missing column_names attr")
-    # Create Arrow table directly from numpy array columns
-    arrays = [pa.array(arr[:, i]) for i in range(arr.shape[1])]
-    arrow_table = pa.table(dict(zip(col_names, arrays, strict=True)))
-    return TableValue(arrow_table)
-
-
-def _summ(v: Value) -> str:
-    try:
-        if isinstance(v, TableValue):
-            # Use Arrow's native string representation (more efficient)
-            arrow_table = v.to_arrow()
-            # Show first 8 rows
-            preview = arrow_table.slice(0, min(8, arrow_table.num_rows))
-            return str(preview)
-        if isinstance(v, TensorValue):
-            arr = v.to_numpy()
-            return str(
-                np.array2string(
-                    arr, threshold=64, edgeitems=3, precision=6, suppress_small=True
-                )
-            )
-        return repr(v)
-    except Exception as e:  # pragma: no cover
-        return f"<unprintable {type(v).__name__}: {e}>"
-
-
-@kernel_def("basic.debug_print")
-def _debug_print(pfunc: PFunction, val: Value) -> Value:
-    prefix = pfunc.attrs.get("prefix", "")
-    ctx = cur_kctx()
-    print(f"[debug_print][rank={ctx.rank}] {prefix}{_summ(val)}")
-    return val
-
-
-@kernel_def("basic.pack")
-def _pack(pfunc: PFunction, value: Value) -> TensorValue:
-    outs_info = pfunc.outs_info
-    if len(outs_info) != 1:
-        raise ValueError("basic.pack expects single output type")
-    out_ty = outs_info[0]
-    if not isinstance(out_ty, TensorType):
-        raise TypeError("basic.pack must return TensorType")
-    if out_ty.dtype.numpy_dtype() != np.uint8:
-        raise TypeError("basic.pack output dtype must be uint8")
-
-    if isinstance(value, TableValue):
-        # Serialize Arrow table using IPC stream for consistency with Value serde
-        import pyarrow as pa  # type: ignore
-        import pyarrow.ipc as pa_ipc  # type: ignore
-
-        arrow_table = value.to_arrow()
-        sink = pa.BufferOutputStream()
-        with pa_ipc.new_stream(sink, arrow_table.schema) as writer:  # type: ignore[arg-type]
-            writer.write_table(arrow_table)  # type: ignore[arg-type]
-        ipc_bytes = sink.getvalue().to_pybytes()
-        return TensorValue(np.frombuffer(ipc_bytes, dtype=np.uint8))
-
-    if isinstance(value, TensorValue):
-        arr = value.to_numpy()
-        return TensorValue(np.frombuffer(arr.tobytes(order="C"), dtype=np.uint8))
-
-    raise TypeError(f"basic.pack does not support Value type {type(value).__name__}")
-
-
-@kernel_def("basic.unpack")
-def _unpack(pfunc: PFunction, packed: TensorValue) -> Value:
-    outs_info = pfunc.outs_info
-    if len(outs_info) != 1:
-        raise ValueError("basic.unpack expects single output type")
-    out_ty = outs_info[0]
-
-    b = packed.to_numpy().astype(np.uint8, copy=False).reshape(-1)
-
-    if isinstance(out_ty, TensorType):
-        np_dtype = out_ty.dtype.numpy_dtype()
-        shape = tuple(out_ty.shape)
-        if any(dim < 0 for dim in shape):
-            raise ValueError("basic.unpack does not support dynamic tensor shapes")
-        elem_count = int(np.prod(shape))
-        expected = elem_count * np.dtype(np_dtype).itemsize
-        if b.size != expected:
-            raise ValueError(
-                f"unpack size mismatch: got {b.size} bytes, expect {expected} for {np_dtype} {shape}"
-            )
-        arr = np.frombuffer(b.tobytes(), dtype=np_dtype)
-        return TensorValue(arr.reshape(shape))
-
-    if isinstance(out_ty, TableType):
-        # Deserialize Arrow IPC stream back to TableValue
-        import pyarrow as pa  # type: ignore
-        import pyarrow.ipc as pa_ipc  # type: ignore
-
-        buf = pa.py_buffer(b.tobytes())
-        reader = pa_ipc.open_stream(buf)
-        table = reader.read_all()
-        return TableValue(table)
-
-    raise TypeError("basic.unpack output type must be TensorType or TableType")