mplang-nightly 0.1.dev163__py3-none-any.whl → 0.1.dev165__py3-none-any.whl

mplang/kernels/value.py

@@ -0,0 +1,626 @@
+# 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 abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, ClassVar, TypeVar
+
+from mplang.protos.v1alpha1 import value_pb2 as _value_pb2
+
+if TYPE_CHECKING:
+    import numpy as np
+
+__all__ = [
+    "BytesBlob",
+    "TableValue",
+    "TensorValue",
+    "Value",
+    "ValueDecodeError",
+    "ValueError",
+    "ValueProtoBuilder",
+    "ValueProtoReader",
+    "decode_value",
+    "encode_value",
+    "is_value_envelope",
+    "list_value_kinds",
+    "register_value",
+]
+
+
+class ValueError(Exception):  # shadow built-in intentionally local
+    """Base exception for backend Value related errors."""
+
+
+class ValueDecodeError(ValueError):
+    """Raised when decoding a Value envelope or payload fails."""
+
+
+class ValueProtoBuilder:
+    """Builder for creating ValueProto messages with a fluent API.
+
+    Provides a cleaner, more ergonomic interface than directly working with pb2 messages.
+
+    Example:
+        proto = ValueProtoBuilder("my.custom.Type", version=1) \
+            .set_attr("shape", [1, 2, 3]) \
+            .set_attr("dtype", "float32") \
+            .set_payload(b"some bytes") \
+            .build()
+    """
+
+    def __init__(self, kind: str, version: int):
+        """Initialize a new proto builder.
+
+        Args:
+            kind: The globally unique KIND identifier for this Value type
+            version: The WIRE_VERSION for this Value type
+        """
+        self._proto = _value_pb2.ValueProto()
+        self._proto.kind = kind
+        self._proto.value_version = version
+
+    def set_payload(self, payload: bytes) -> ValueProtoBuilder:
+        """Set the payload bytes.
+
+        Args:
+            payload: Raw bytes to store in the proto
+
+        Returns:
+            Self for method chaining
+        """
+        self._proto.payload = payload
+        return self
+
+    def set_attr(self, key: str, value: Any) -> ValueProtoBuilder:
+        """Set a single runtime attribute.
+
+        Args:
+            key: Attribute key
+            value: Attribute value. Supported types: bool, int, float, str, bytes,
+                   list[int/float/str]
+
+        Returns:
+            Self for method chaining
+        """
+        self._proto.runtime_attrs[key].CopyFrom(_python_to_attr_proto(value))
+        return self
+
+    def build(self) -> _value_pb2.ValueProto:
+        """Return the built proto.
+
+        Returns:
+            The fully constructed ValueProto
+        """
+        return self._proto
+
+
+class ValueProtoReader:
+    """Reader for extracting data from ValueProto messages.
+
+    Provides a convenient interface for reading proto fields and attributes.
+
+    Example:
+        reader = ValueProtoReader(proto)
+        shape = reader.get_attr("shape")
+        dtype = reader.get_attr("dtype")
+        payload = reader.payload
+    """
+
+    def __init__(self, proto: _value_pb2.ValueProto):
+        """Initialize a reader for an existing proto.
+
+        Args:
+            proto: The ValueProto to read from
+        """
+        self._proto = proto
+
+    @property
+    def kind(self) -> str:
+        """Get the KIND identifier."""
+        return self._proto.kind
+
+    @property
+    def version(self) -> int:
+        """Get the WIRE_VERSION."""
+        return self._proto.value_version
+
+    @property
+    def payload(self) -> bytes:
+        """Get the payload bytes."""
+        return self._proto.payload
+
+    def get_attr(self, key: str, default: Any = ...) -> Any:
+        """Get a single attribute with optional default.
+
+        Args:
+            key: Attribute key to retrieve
+            default: Default value if key is missing. If not provided (default),
+                     raises ValueDecodeError when key is missing.
+
+        Returns:
+            The decoded attribute value or default
+
+        Raises:
+            ValueDecodeError: If key is missing and no default provided
+        """
+        if key not in self._proto.runtime_attrs:
+            if default is ...:
+                raise ValueDecodeError(f"Missing required runtime_attr: {key}")
+            return default
+        return _attr_proto_to_python(self._proto.runtime_attrs[key])
+
+
+class Value(ABC):
+    """Abstract base for backend-level transferable values.
+
+    Subclasses MUST define:
+        KIND (ClassVar[str]): globally unique stable identifier
+        WIRE_VERSION (ClassVar[int]): per-kind payload version integer >=1
+
+    Use ValueProtoBuilder and ValueProtoReader for proto serialization.
+    """
+
+    KIND: ClassVar[str]
+    WIRE_VERSION: ClassVar[int] = 1
+
+    def estimated_wire_size(self) -> int | None:  # optional hint
+        return None
+
+    @abstractmethod
+    def to_proto(self) -> _value_pb2.ValueProto:
+        """Return fully-populated ValueProto for this value."""
+
+    @classmethod
+    @abstractmethod
+    def from_proto(cls, proto: _value_pb2.ValueProto) -> Value:
+        """Construct instance from a parsed ValueProto."""
+
+    def to_bool(self) -> bool:
+        """Convert value to bool (for predicates in control flow).
+
+        Default implementation raises NotImplementedError.
+        Subclasses should override to provide appropriate conversion.
+
+        Raises:
+            NotImplementedError: If the value type cannot be converted to bool
+        """
+        raise NotImplementedError(
+            f"{type(self).__name__} does not support conversion to bool"
+        )
+
+
+T = TypeVar("T", bound=Value)
+
+_VALUE_REGISTRY: dict[str, type[Value]] = {}
+
+
+def _python_to_attr_proto(value: Any) -> _value_pb2.ValueAttrProto:
+    attr = _value_pb2.ValueAttrProto()
+    if isinstance(value, bool):
+        attr.type = _value_pb2.ValueAttrProto.BOOL
+        attr.b = value
+    elif isinstance(value, int) and not isinstance(value, bool):
+        attr.type = _value_pb2.ValueAttrProto.INT
+        attr.i = value
+    elif isinstance(value, float):
+        attr.type = _value_pb2.ValueAttrProto.FLOAT
+        attr.f = value
+    elif isinstance(value, str):
+        attr.type = _value_pb2.ValueAttrProto.STRING
+        attr.s = value
+    elif isinstance(value, (bytes, bytearray, memoryview)):
+        attr.type = _value_pb2.ValueAttrProto.BYTES
+        attr.raw_bytes = bytes(value)
+    elif isinstance(value, (list, tuple)):
+        if not value:
+            # Represent empty list explicitly
+            attr.type = _value_pb2.ValueAttrProto.EMPTY
+            return attr
+        if all(isinstance(v, float) for v in value):
+            attr.type = _value_pb2.ValueAttrProto.FLOATS
+            attr.floats.extend(float(v) for v in value)
+        elif all(isinstance(v, int) and not isinstance(v, bool) for v in value):
+            attr.type = _value_pb2.ValueAttrProto.INTS
+            attr.ints.extend(int(v) for v in value)
+        elif all(isinstance(v, str) for v in value):
+            attr.type = _value_pb2.ValueAttrProto.STRINGS
+            attr.strs.extend(value)
+        else:
+            raise TypeError(
+                "Unsupported iterable element type for AttrProto: "
+                f"{type(value[0]).__name__}"
+            )
+    elif value is None:
+        attr.type = _value_pb2.ValueAttrProto.UNDEFINED
+    else:
+        raise TypeError(
+            "Unsupported runtime attr type for Value serialization: "
+            f"{type(value).__name__}"
+        )
+    return attr
+
+
+def _attr_proto_to_python(attr: _value_pb2.ValueAttrProto) -> Any:
+    if attr.type == _value_pb2.ValueAttrProto.FLOAT:
+        return attr.f
+    if attr.type == _value_pb2.ValueAttrProto.INT:
+        return attr.i
+    if attr.type == _value_pb2.ValueAttrProto.STRING:
+        return attr.s
+    if attr.type == _value_pb2.ValueAttrProto.BOOL:
+        return attr.b
+    if attr.type == _value_pb2.ValueAttrProto.BYTES:
+        return attr.raw_bytes
+    if attr.type == _value_pb2.ValueAttrProto.FLOATS:
+        return list(attr.floats)
+    if attr.type == _value_pb2.ValueAttrProto.INTS:
+        return list(attr.ints)
+    if attr.type == _value_pb2.ValueAttrProto.STRINGS:
+        return list(attr.strs)
+    if attr.type == _value_pb2.ValueAttrProto.EMPTY:
+        return []
+    if attr.type == _value_pb2.ValueAttrProto.UNDEFINED:
+        return None
+    raise ValueDecodeError(f"Unsupported AttrProto type {attr.type}")
+
+
+def _looks_like_pyarrow_table(obj: Any) -> bool:
+    if obj is None:
+        return False
+    module = getattr(obj.__class__, "__module__", "")
+    return (
+        module.startswith("pyarrow.")
+        and hasattr(obj, "schema")
+        and hasattr(obj, "column_names")
+        and hasattr(obj, "num_rows")
+    )
+
+
+def _looks_like_pandas_df(obj: Any) -> bool:
+    if obj is None:
+        return False
+    module = getattr(obj.__class__, "__module__", "")
+    return (
+        module.startswith("pandas.")
+        and hasattr(obj, "columns")
+        and hasattr(obj, "dtypes")
+    )
+
+
+def register_value(cls: type[T]) -> type[T]:
+    kind = getattr(cls, "KIND", None)
+    if not kind or not isinstance(kind, str):
+        raise ValueError(f"Value subclass {cls.__name__} missing KIND str")
+    if kind in _VALUE_REGISTRY:
+        raise ValueError(f"Duplicate Value KIND '{kind}'")
+    if getattr(cls, "WIRE_VERSION", None) is None:
+        raise ValueError(f"Value subclass {cls.__name__} missing WIRE_VERSION")
+    _VALUE_REGISTRY[kind] = cls
+    return cls
+
+
+def list_value_kinds() -> list[str]:
+    return sorted(_VALUE_REGISTRY.keys())
+
+
+def encode_value(val: Value) -> bytes:
+    """Encode using protobuf envelope.
+
+    Raises:
+        ValueError if protobuf module not available.
+    """
+    if _value_pb2 is None:  # pragma: no cover
+        raise ValueError("protobuf value_pb2 not generated yet")
+    proto = val.to_proto()
+    if not isinstance(proto, _value_pb2.ValueProto):
+        raise ValueError("Value.to_proto must return ValueProto")
+    if not proto.kind:
+        proto.kind = val.KIND
+    elif proto.kind != val.KIND:
+        raise ValueError(
+            f"ValueProto.kind mismatch: expected '{val.KIND}', got '{proto.kind}'"
+        )
+    if proto.value_version == 0:
+        proto.value_version = val.WIRE_VERSION
+    elif proto.value_version != val.WIRE_VERSION:
+        raise ValueError(
+            f"ValueProto.value_version mismatch: expected {val.WIRE_VERSION}, got {proto.value_version}"
+        )
+    return proto.SerializeToString()  # type: ignore[no-any-return]
+
+
+def is_value_envelope(data: bytes) -> bool:
+    if _value_pb2 is None:
+        return False
+    env = _value_pb2.ValueProto()
+    try:
+        env.ParseFromString(data)
+        return bool(env.kind)
+    except Exception:
+        return False
+
+
+def decode_value(data: bytes) -> Value:
+    if _value_pb2 is None:
+        raise ValueDecodeError("protobuf value_pb2 not available for decode")
+    env = _value_pb2.ValueProto()
+    try:
+        env.ParseFromString(data)
+    except Exception as e:  # pragma: no cover
+        raise ValueDecodeError(f"Failed parsing ValueProto: {e}") from e
+    if not env.kind:
+        raise ValueDecodeError("Envelope missing kind")
+    cls = _VALUE_REGISTRY.get(env.kind)
+    if cls is None:
+        raise ValueDecodeError(f"Unknown Value kind '{env.kind}'")
+    if env.value_version and env.value_version != cls.WIRE_VERSION:
+        raise ValueDecodeError(
+            f"Unsupported {cls.__name__} version {env.value_version}"
+        )
+    if env.value_version == 0:
+        env.value_version = cls.WIRE_VERSION
+    return cls.from_proto(env)
+
+
+@register_value
+class BytesBlob(Value):  # demo subclass
+    KIND = "mplang.demo.BytesBlob"
+    WIRE_VERSION = 1
+
+    def __init__(self, data: bytes):
+        self._data = data
+
+    def to_proto(self) -> _value_pb2.ValueProto:
+        return (
+            ValueProtoBuilder(self.KIND, self.WIRE_VERSION)
+            .set_payload(self._data)
+            .build()
+        )
+
+    @classmethod
+    def from_proto(cls, proto: _value_pb2.ValueProto) -> BytesBlob:
+        reader = ValueProtoReader(proto)
+        if reader.version != cls.WIRE_VERSION:
+            raise ValueDecodeError(f"Unsupported BytesBlob version {reader.version}")
+        if proto.runtime_attrs:
+            raise ValueDecodeError("BytesBlob does not expect runtime attributes")
+        return cls(reader.payload)
+
+    def __repr__(self) -> str:  # pragma: no cover
+        return f"BytesBlob(len={len(self._data)})"
+
+
+@register_value
+class TensorValue(Value):  # well-known tensor (ndarray) Value
+    """Numpy ndarray serialization via raw buffer + runtime metadata."""
+
+    KIND = "mplang.ndarray"
+    WIRE_VERSION = 1
+
+    def __init__(self, array):  # type: ignore[no-untyped-def]
+        import numpy as np
+
+        if not isinstance(array, np.ndarray):
+            raise TypeError("TensorValue expects a numpy.ndarray")
+        if not array.flags.c_contiguous:
+            array = np.ascontiguousarray(array)
+        self._arr = array
+
+    def to_proto(self) -> _value_pb2.ValueProto:
+        import numpy as np
+
+        arr = self._arr
+        if not arr.flags.c_contiguous:
+            arr = np.ascontiguousarray(arr)
+        return (
+            ValueProtoBuilder(self.KIND, self.WIRE_VERSION)
+            .set_attr("dtype", arr.dtype.str)
+            .set_attr("shape", [int(dim) for dim in arr.shape])
+            .set_payload(arr.tobytes())
+            .build()
+        )
+
+    @classmethod
+    def from_proto(cls, proto: _value_pb2.ValueProto) -> TensorValue:
+        import numpy as np
+
+        reader = ValueProtoReader(proto)
+        if reader.version != cls.WIRE_VERSION:
+            raise ValueDecodeError(f"Unsupported TensorValue version {reader.version}")
+        dtype_val = reader.get_attr("dtype")
+        if not isinstance(dtype_val, str):
+            raise ValueDecodeError("TensorValue runtime attr 'dtype' must be str")
+        shape_val = reader.get_attr("shape")
+        if not isinstance(shape_val, list):
+            raise ValueDecodeError("TensorValue runtime attr 'shape' must be list")
+        shape = tuple(int(dim) for dim in shape_val)
+        try:
+            arr = np.frombuffer(reader.payload, dtype=np.dtype(dtype_val)).reshape(
+                shape
+            )
+        except Exception as e:  # pragma: no cover
+            raise ValueDecodeError(f"Failed reconstruct ndarray: {e}") from e
+        return cls(np.array(arr, copy=True))
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        return tuple(self._arr.shape)
+
+    @property
+    def dtype(self) -> np.dtype[Any]:  # pragma: no cover - simple accessor
+        return self._arr.dtype  # type: ignore[no-any-return]
+
+    @property
+    def ndim(self) -> int:  # pragma: no cover - simple accessor
+        return int(self._arr.ndim)
+
+    def to_numpy(
+        self, *, copy: bool = False
+    ) -> np.ndarray[Any, Any]:  # pragma: no cover - simple accessor
+        if copy:
+            import numpy as np
+
+            return np.array(self._arr, copy=True)
+        return self._arr  # type: ignore[no-any-return]
+
+    def __array__(
+        self, dtype: np.dtype[Any] | None = None
+    ) -> np.ndarray[Any, Any]:  # pragma: no cover - numpy bridge
+        import numpy as np
+
+        return np.asarray(self._arr, dtype=dtype)
+
+    def to_bool(self) -> bool:
+        """Convert tensor to bool (for scalar predicates).
+
+        Returns:
+            bool value if tensor is scalar
+
+        Raises:
+            ValueError: If tensor is not a scalar
+        """
+        if self._arr.size != 1:
+            raise ValueError(
+                f"Cannot convert non-scalar tensor (shape={self._arr.shape}) to bool"
+            )
+        return bool(self._arr.item())
+
+    def __repr__(self) -> str:  # pragma: no cover
+        return f"TensorValue(shape={self._arr.shape}, dtype={self._arr.dtype})"
+
+
+@register_value
+class TableValue(Value):  # well-known table (Arrow IPC) Value
+    """Table value backed by PyArrow, serialized via Arrow IPC stream.
+
+    KIND: mplang.dataframe.arrow
+    WIRE_VERSION: increments if wire semantics (not DataFrame contents) change.
+
+    Internal representation: Always pyarrow.Table for consistency and performance.
+    Wire format: Arrow IPC stream.
+
+    Accepts pandas DataFrame or pyarrow.Table as input, but internally converts
+    everything to pyarrow.Table for unified handling.
+    """
+
+    KIND = "mplang.dataframe.arrow"
+    WIRE_VERSION = 1
+
+    def __init__(self, data):  # type: ignore[no-untyped-def]
+        """Initialize TableValue from pandas DataFrame or pyarrow.Table.
+
+        Args:
+            data: pandas.DataFrame, pyarrow.Table, or dict-like object
+        """
+        try:
+            import pyarrow as pa  # type: ignore
+        except ImportError as e:
+            raise ValueError("pyarrow is required for TableValue") from e
+
+        if _looks_like_pyarrow_table(data):
+            self._table = data
+        elif _looks_like_pandas_df(data):
+            try:
+                self._table = pa.Table.from_pandas(data, preserve_index=False)
+            except Exception as e:
+                raise ValueError(
+                    f"Cannot convert pandas DataFrame to Arrow: {e}"
+                ) from e
+        else:
+            # Try to convert dict-like or other structures
+            try:
+                self._table = pa.table(data)
+            except Exception as e:
+                raise TypeError(
+                    f"TableValue requires pandas.DataFrame or pyarrow.Table, got {type(data).__name__}"
+                ) from e
+
+    def to_proto(self) -> _value_pb2.ValueProto:
+        """Serialize to Arrow IPC stream format."""
+        import pyarrow as pa  # type: ignore
+        import pyarrow.ipc as pa_ipc  # type: ignore
+
+        sink = pa.BufferOutputStream()
+        with pa_ipc.new_stream(sink, self._table.schema) as writer:  # type: ignore[arg-type]
+            writer.write_table(self._table)  # type: ignore[arg-type]
+        return (
+            ValueProtoBuilder(self.KIND, self.WIRE_VERSION)
+            .set_payload(sink.getvalue().to_pybytes())
+            .build()
+        )
+
+    @classmethod
+    def from_proto(cls, proto: _value_pb2.ValueProto) -> TableValue:
+        """Deserialize from Arrow IPC stream format."""
+        reader = ValueProtoReader(proto)
+        if reader.version != cls.WIRE_VERSION:
+            raise ValueDecodeError(f"Unsupported TableValue version {reader.version}")
+        if proto.runtime_attrs:
+            raise ValueDecodeError("TableValue does not expect runtime attributes")
+
+        import pyarrow as pa  # type: ignore
+        import pyarrow.ipc as pa_ipc  # type: ignore
+
+        buf = pa.py_buffer(reader.payload)
+        ipc_reader = pa_ipc.open_stream(buf)
+        table = ipc_reader.read_all()
+        return cls(table)
+
+    def to_arrow(self) -> Any:  # pyarrow.Table
+        """Get the underlying pyarrow.Table (primary interface).
+
+        Returns:
+            pyarrow.Table: The table data
+        """
+        return self._table
+
+    def to_pandas(self) -> Any:  # pandas.DataFrame
+        """Convert to pandas DataFrame (compatibility interface).
+
+        Note: This creates a copy and converts from Arrow to pandas format.
+        For better performance, consider using to_arrow() and working with
+        Arrow-native APIs (DuckDB, Ibis, etc.) directly.
+
+        Returns:
+            pandas.DataFrame: Converted dataframe
+        """
+        return self._table.to_pandas()  # type: ignore[attr-defined]
+
+    @property
+    def columns(self) -> list[str]:
+        """Return column names (TableLike protocol compatibility)."""
+        return [str(name) for name in self._table.column_names]
+
+    @property
+    def dtypes(self) -> Any:  # pyarrow.Schema
+        """Return column dtypes as Arrow schema (TableLike protocol compatibility)."""
+        return self._table.schema
+
+    def num_rows(self) -> int:
+        """Get number of rows in the table.
+
+        Returns:
+            Number of rows
+        """
+        return self._table.num_rows  # type: ignore[attr-defined,return-value,no-any-return]
+
+    def __repr__(self) -> str:
+        """String representation of TableValue."""
+        try:
+            rows = self.num_rows()
+            cols = self.columns
+            return f"TableValue(rows={rows}, cols={cols})"
+        except Exception:
+            return "TableValue()"

mplang/ops/tee.py

@@ -14,11 +14,7 @@
 
 from __future__ import annotations
 
-from jax.tree_util import PyTreeDef, tree_flatten
-
 from mplang.core.dtype import UINT8
-from mplang.core.mpobject import MPObject
-from mplang.core.pfunc import PFunction
 from mplang.core.tensor import TensorType
 from mplang.ops.base import stateless_mod
 
@@ -32,20 +28,10 @@ def quote_gen(pk: TensorType) -> TensorType:
     return TensorType(UINT8, (-1,))
 
 
-@_TEE_MOD.op_def()
-def attest(
-    quote: MPObject, platform: str
-) -> tuple[PFunction, list[MPObject], PyTreeDef]:
-    """TEE quote verification returning the attested TEE public key."""
-
-    ins_info = [TensorType.from_obj(quote)]
-    outs_info = [TensorType(UINT8, (32,))]  # pk is always 32 bytes for x25519
-    pfunc = PFunction(
-        fn_type="tee.attest",
-        ins_info=ins_info,
-        outs_info=outs_info,
-        platform=platform,
-    )
-    _, treedef = tree_flatten(outs_info[0])
-
-    return pfunc, [quote], treedef
+@_TEE_MOD.simple_op()
+def attest(quote: TensorType) -> TensorType:
+    """TEE quote verification returning the attested TEE public key.
+    API (mock): attest(quote: u8[33]) -> tee_pk: u8[32]
+    """
+    _ = quote  # Mark as used for the decorator
+    return TensorType(UINT8, (32,))