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

mplang/v1/protos/v1alpha1/value_pb2.pyi

@@ -0,0 +1,169 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+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.
+"""
+import builtins
+import collections.abc
+import google.protobuf.descriptor
+import google.protobuf.internal.containers
+import google.protobuf.internal.enum_type_wrapper
+import google.protobuf.message
+import sys
+import typing
+
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
+
+@typing_extensions.final
+class ValueAttrProto(google.protobuf.message.Message):
+    """Lightweight attribute proto used solely for ValueProto runtime metadata."""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    class _AttrType:
+        ValueType = typing.NewType("ValueType", builtins.int)
+        V: typing_extensions.TypeAlias = ValueType
+
+    class _AttrTypeEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[ValueAttrProto._AttrType.ValueType], builtins.type):
+        DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+        UNDEFINED: ValueAttrProto._AttrType.ValueType  # 0
+        FLOAT: ValueAttrProto._AttrType.ValueType  # 1
+        INT: ValueAttrProto._AttrType.ValueType  # 2
+        STRING: ValueAttrProto._AttrType.ValueType  # 3
+        BOOL: ValueAttrProto._AttrType.ValueType  # 4
+        BYTES: ValueAttrProto._AttrType.ValueType  # 5
+        FLOATS: ValueAttrProto._AttrType.ValueType  # 6
+        INTS: ValueAttrProto._AttrType.ValueType  # 7
+        STRINGS: ValueAttrProto._AttrType.ValueType  # 8
+        EMPTY: ValueAttrProto._AttrType.ValueType  # 9
+        """Represents an explicitly empty attribute value (e.g., empty list)"""
+
+    class AttrType(_AttrType, metaclass=_AttrTypeEnumTypeWrapper): ...
+    UNDEFINED: ValueAttrProto.AttrType.ValueType  # 0
+    FLOAT: ValueAttrProto.AttrType.ValueType  # 1
+    INT: ValueAttrProto.AttrType.ValueType  # 2
+    STRING: ValueAttrProto.AttrType.ValueType  # 3
+    BOOL: ValueAttrProto.AttrType.ValueType  # 4
+    BYTES: ValueAttrProto.AttrType.ValueType  # 5
+    FLOATS: ValueAttrProto.AttrType.ValueType  # 6
+    INTS: ValueAttrProto.AttrType.ValueType  # 7
+    STRINGS: ValueAttrProto.AttrType.ValueType  # 8
+    EMPTY: ValueAttrProto.AttrType.ValueType  # 9
+    """Represents an explicitly empty attribute value (e.g., empty list)"""
+
+    TYPE_FIELD_NUMBER: builtins.int
+    F_FIELD_NUMBER: builtins.int
+    I_FIELD_NUMBER: builtins.int
+    S_FIELD_NUMBER: builtins.int
+    B_FIELD_NUMBER: builtins.int
+    RAW_BYTES_FIELD_NUMBER: builtins.int
+    FLOATS_FIELD_NUMBER: builtins.int
+    INTS_FIELD_NUMBER: builtins.int
+    STRS_FIELD_NUMBER: builtins.int
+    type: global___ValueAttrProto.AttrType.ValueType
+    f: builtins.float
+    i: builtins.int
+    s: builtins.str
+    b: builtins.bool
+    raw_bytes: builtins.bytes
+    @property
+    def floats(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.float]: ...
+    @property
+    def ints(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ...
+    @property
+    def strs(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: ...
+    def __init__(
+        self,
+        *,
+        type: global___ValueAttrProto.AttrType.ValueType = ...,
+        f: builtins.float = ...,
+        i: builtins.int = ...,
+        s: builtins.str = ...,
+        b: builtins.bool = ...,
+        raw_bytes: builtins.bytes = ...,
+        floats: collections.abc.Iterable[builtins.float] | None = ...,
+        ints: collections.abc.Iterable[builtins.int] | None = ...,
+        strs: collections.abc.Iterable[builtins.str] | None = ...,
+    ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["b", b"b", "f", b"f", "floats", b"floats", "i", b"i", "ints", b"ints", "raw_bytes", b"raw_bytes", "s", b"s", "strs", b"strs", "type", b"type"]) -> None: ...
+
+global___ValueAttrProto = ValueAttrProto
+
+@typing_extensions.final
+class ValueProto(google.protobuf.message.Message):
+    """Generic envelope for kernel-level transferable values.
+
+    DESIGN PRINCIPLES
+    * Small, stable schema: only descriptors needed for dynamic dispatch.
+    * Payload is opaque to the envelope; per-kind versioning lives in
+    value_version.
+    * Backward-compatible evolution: only append new optional fields.
+
+    Versioning Guidelines:
+    - value_version: per-kind semantic payload version (maintained by KernelValue
+    subclass).
+    - Adding fields: assign new unique field numbers; never reuse old numbers.
+    - Removing fields: reserve the field number & (optionally) name.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    @typing_extensions.final
+    class RuntimeAttrsEntry(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        KEY_FIELD_NUMBER: builtins.int
+        VALUE_FIELD_NUMBER: builtins.int
+        key: builtins.str
+        @property
+        def value(self) -> global___ValueAttrProto: ...
+        def __init__(
+            self,
+            *,
+            key: builtins.str = ...,
+            value: global___ValueAttrProto | None = ...,
+        ) -> None: ...
+        def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ...
+        def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ...
+
+    KIND_FIELD_NUMBER: builtins.int
+    VALUE_VERSION_FIELD_NUMBER: builtins.int
+    PAYLOAD_FIELD_NUMBER: builtins.int
+    RUNTIME_ATTRS_FIELD_NUMBER: builtins.int
+    kind: builtins.str
+    """Globally unique identifier for Value subclass, e.g. "mplang.ndarray"."""
+    value_version: builtins.int
+    """Per-kind payload schema version (>=1)."""
+    payload: builtins.bytes
+    """Primary data payload bytes (layout defined by each Value subclass)."""
+    @property
+    def runtime_attrs(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, global___ValueAttrProto]:
+        """Additional runtime metadata required to recreate the Value instance."""
+    def __init__(
+        self,
+        *,
+        kind: builtins.str = ...,
+        value_version: builtins.int = ...,
+        payload: builtins.bytes = ...,
+        runtime_attrs: collections.abc.Mapping[builtins.str, global___ValueAttrProto] | None = ...,
+    ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["kind", b"kind", "payload", b"payload", "runtime_attrs", b"runtime_attrs", "value_version", b"value_version"]) -> None: ...
+
+global___ValueProto = ValueProto

mplang/{runtime → v1/runtime}/__init__.py

@@ -20,8 +20,8 @@ This module contains runtime implementations including:
 - Driver for distributed execution
 """
 
-from mplang.runtime.driver import Driver, DriverVar
-from mplang.runtime.simulation import Simulator
+from mplang.v1.runtime.driver import Driver, DriverVar
+from mplang.v1.runtime.simulation import Simulator
 
 __all__ = [
     "Driver",

mplang/v1/runtime/channel.py

@@ -0,0 +1,230 @@
+# 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.
+
+"""SPU IChannel implementation that bridges to MPLang CommunicatorBase.
+
+This module provides BaseChannel, which allows SPU to reuse MPLang's
+existing communication layer (ThreadCommunicator/HttpCommunicator) instead
+of creating separate BRPC connections.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import spu.libspu as libspu
+
+if TYPE_CHECKING:
+    from mplang.v1.core.comm import CommunicatorBase
+
+
+class BaseChannel(libspu.link.IChannel):
+    """Bridge MPLang CommunicatorBase to SPU IChannel interface.
+
+    This adapter allows SPU to use MPLang's existing communication layer
+    (ThreadCommunicator or HttpCommunicator) instead of creating separate
+    BRPC connections.
+
+    Each BaseChannel represents a channel to ONE peer rank.
+
+    Communication Protocol:
+        - SPU calls send(tag, bytes_data) -> MPLang comm.send(peer, key, bytes_data)
+        - SPU calls recv(tag) -> bytes_data <- MPLang comm.recv(peer, key)
+
+    Tag Namespace:
+        All tags are prefixed with "spu:" to avoid collision with other
+        MPLang traffic on the same communicator.
+    """
+
+    def __init__(
+        self,
+        comm: CommunicatorBase,
+        local_rank: int,
+        peer_rank: int,
+        tag_prefix: str = "spu",
+    ):
+        """Initialize channel to a specific peer.
+
+        Args:
+            comm: MPLang communicator instance (Thread/Http)
+            local_rank: Global rank of this party (for logging/debugging)
+            peer_rank: Global rank of the peer party
+            tag_prefix: Prefix for all tags to avoid collision (default: "spu")
+        """
+        super().__init__()
+        self._comm = comm
+        self._local_rank = local_rank
+        self._peer_rank = peer_rank
+        self._tag_prefix = tag_prefix
+
+        logging.debug(
+            f"BaseChannel initialized: local_rank={local_rank}, "
+            f"peer_rank={peer_rank}, tag_prefix={tag_prefix}"
+        )
+
+    def _make_key(self, tag: str) -> str:
+        """Create unique key for MPLang comm.
+
+        Prefixes the tag to avoid collision with non-SPU traffic.
+
+        Args:
+            tag: SPU-provided tag (e.g., "send_0", "recv_0")
+
+        Returns:
+            Prefixed key (e.g., "spu:send_0")
+        """
+        return f"{self._tag_prefix}:{tag}"
+
+    def Send(self, tag: str, data: bytes) -> None:
+        """Send bytes to peer (synchronous in SPU semantics).
+
+        Args:
+            tag: Message tag for matching send/recv pairs
+            data: Raw bytes to send
+        """
+        key = self._make_key(tag)
+        logging.debug(
+            f"BaseChannel.Send: {self._local_rank} -> {self._peer_rank}, "
+            f"tag={tag}, key={key}, size={len(data)}"
+        )
+
+        # Send raw bytes directly
+        # Note: CommunicatorBase.send expects Any type, bytes is acceptable
+        self._comm.send(self._peer_rank, key, data)
+
+    def Recv(self, tag: str) -> bytes:
+        """Receive bytes from peer (blocking).
+
+        Args:
+            tag: Message tag for matching send/recv pairs
+
+        Returns:
+            Raw bytes received
+        """
+        key = self._make_key(tag)
+        logging.debug(
+            f"BaseChannel.Recv: {self._local_rank} <- {self._peer_rank}, "
+            f"tag={tag}, key={key}"
+        )
+
+        # Receive data (should be bytes)
+        data = self._comm.recv(self._peer_rank, key)
+
+        # Validate data type
+        if not isinstance(data, bytes):
+            raise TypeError(
+                f"Expected bytes from communicator, got {type(data).__name__}. "
+                f"Communicator must support raw bytes transmission for SPU channels."
+            )
+
+        logging.debug(
+            f"BaseChannel.Recv complete: {self._local_rank} <- {self._peer_rank}, "
+            f"tag={tag}, size={len(data)}"
+        )
+        return data
+
+    def SendAsync(self, tag: str, data: bytes) -> None:
+        """Async send (MPLang's send is already async at network layer).
+
+        For HttpCommunicator, the underlying httpx.put() is non-blocking
+        at the HTTP client level. For ThreadCommunicator, send is instant
+        (memory transfer).
+
+        Args:
+            tag: Message tag
+            data: Raw bytes to send
+        """
+        # Reuse synchronous send - it's already async underneath
+        self.Send(tag, data)
+
+    def SendAsyncThrottled(self, tag: str, data: bytes) -> None:
+        """Throttled async send.
+
+        Currently maps to regular SendAsync. Future optimization could
+        implement rate limiting if needed.
+
+        Args:
+            tag: Message tag
+            data: Raw bytes to send
+        """
+        self.SendAsync(tag, data)
+
+    def TestSend(self, timeout: int) -> None:
+        """Test if this channel can send a dummy msg to peer.
+
+        Uses fixed 0 seq_id as dummy msg's id to make this function reentrant.
+        ConnectToMesh will retry on this multiple times.
+
+        Args:
+            timeout: Timeout in milliseconds
+        """
+        # Send a handshake message to test connectivity
+        # Use fixed tag "__test__" to make this reentrant (idempotent)
+        test_data = b"\x00"  # Minimal 1-byte message with seq_id=0
+        self.Send("__test__", test_data)
+
+    def TestRecv(self) -> None:
+        """Wait for dummy msg from peer.
+
+        Timeout is controlled by recv_timeout_ms in link descriptor.
+        """
+        # Receive the handshake message from peer
+        # This blocks until message arrives (timeout from desc.recv_timeout_ms)
+        test_data = self.Recv("__test__")
+        # Validate it's the expected handshake message
+        if test_data != b"\x00":
+            logging.warning(
+                f"TestRecv: unexpected handshake data from {self._peer_rank}, "
+                f"expected b'\\x00', got {test_data!r}"
+            )
+
+    def WaitLinkTaskFinish(self) -> None:
+        """Wait for all pending async tasks.
+
+        For MPLang communicators:
+        - ThreadCommunicator: No-op (instant memory transfer)
+        - HttpCommunicator: No explicit wait needed (httpx handles it)
+
+        This is a no-op in current implementation.
+        """
+
+    def Abort(self) -> None:
+        """Abort communication (cleanup resources).
+
+        This could be extended to notify the communicator to drop pending
+        messages for this channel, but currently is a no-op.
+        """
+        logging.warning(
+            f"BaseChannel.Abort called: {self._local_rank} <-> {self._peer_rank}"
+        )
+        # Future: Could call comm.abort_session() if implemented
+
+    def SetThrottleWindowSize(self, size: int) -> None:
+        """Set throttle window size.
+
+        Not applicable to MPLang communicators. No-op.
+
+        Args:
+            size: Window size (ignored)
+        """
+
+    def SetChunkParallelSendSize(self, size: int) -> None:
+        """Set chunk parallel send size.
+
+        Not applicable to MPLang communicators. No-op.
+
+        Args:
+            size: Chunk size (ignored)
+        """

mplang/{runtime → v1/runtime}/cli.py

@@ -26,9 +26,9 @@ from typing import Any
 import uvicorn
 import yaml
 
-from mplang.core.cluster import ClusterSpec
-from mplang.runtime.client import HttpExecutorClient
-from mplang.runtime.server import app
+from mplang.v1.core import ClusterSpec
+from mplang.v1.runtime.client import HttpExecutorClient
+from mplang.v1.runtime.server import app
 
 
 def load_config(config_path: str) -> ClusterSpec:
@@ -195,14 +195,14 @@ def status_command(args: argparse.Namespace) -> int:
     """
 
     async def _get_node_status(
-        node_id: str, endpoint: str, details: bool = False, timeout: int = 60
+        node_id: str, endpoint: str, details: int = 0, timeout: int = 60
     ) -> dict[str, Any]:
         """Get status information for a single node.
 
         Args:
             node_id: Identifier for the node
             endpoint: HTTP endpoint of the node
-            details: Whether to include detailed session information
+            details: Verbosity level (0=basic, 1=-v, 2=-vv)
             timeout: HTTP request timeout in seconds (default: 60)
         """
 
@@ -224,21 +224,26 @@ def status_command(args: argparse.Namespace) -> int:
                 sessions = await client.list_sessions()
                 status["sessions"] = sessions
 
-                # Get detailed session info if requested
-                if details:
+                # Get detailed session info based on verbosity level
+                # details=1 (-v): show session names and basic counts
+                # details=2 (-vv): show full computation and symbol lists
+                if details >= 1:
                     session_details = []
                     for session_name in sessions:
                         try:
                             # Get computations and symbols for each session
                             computations = await client.list_computations(session_name)
                             symbols = await client.list_symbols(session_name)
-                            session_details.append({
+                            session_info = {
                                 "name": session_name,
                                 "computations": len(computations),
                                 "symbols": len(symbols),
-                                "computation_list": computations,
-                                "symbol_list": symbols,
-                            })
+                            }
+                            # Include full lists only at -vv level
+                            if details >= 2:
+                                session_info["computation_list"] = computations
+                                session_info["symbol_list"] = symbols
+                            session_details.append(session_info)
                         except Exception as e:
                             session_details.append({
                                 "name": session_name,
@@ -255,13 +260,13 @@ def status_command(args: argparse.Namespace) -> int:
         return status
 
     async def _collect_cluster_status(
-        nodes: dict[str, str], details: bool = False
+        nodes: dict[str, str], details: int = 0
     ) -> list[dict[str, Any] | BaseException]:
         """Collect status from all nodes concurrently.
 
         Args:
             nodes: Dictionary mapping node IDs to their HTTP endpoints
-            details: Whether to include detailed session information for each node
+            details: Verbosity level (0=basic, 1=-v, 2=-vv)
 
         Returns:
             List of status dictionaries or exceptions for each node
@@ -284,7 +289,8 @@ def status_command(args: argparse.Namespace) -> int:
         node_addrs = {node_id: node.endpoint for node_id, node in nodes.items()}
 
         # Collect status from all nodes
-        cluster_status = asyncio.run(_collect_cluster_status(node_addrs, args.details))
+        verbosity = getattr(args, "verbose", 0)
+        cluster_status = asyncio.run(_collect_cluster_status(node_addrs, verbosity))
 
         # Basic node health check
         print("Node Status:")
@@ -314,8 +320,8 @@ def status_command(args: argparse.Namespace) -> int:
                 print(f"{node_id:<15} {endpoint:<20} UNHEALTHY")
                 all_healthy = False
 
-        # If detailed status is requested, show detailed information
-        if args.details and valid_statuses:
+        # If verbose mode is enabled, show detailed information
+        if verbosity >= 1 and valid_statuses:
             print("\nDetailed Runtime Status:")
             print("-" * 50)
 
@@ -351,6 +357,14 @@ def status_command(args: argparse.Namespace) -> int:
                                 print(
                                     f"  - Session '{session_name}': {computations} computations, {symbols} symbols"
                                 )
+                                # At -vv level, show the actual lists
+                                if verbosity >= 2:
+                                    comp_list = session.get("computation_list", [])
+                                    symbol_list = session.get("symbol_list", [])
+                                    if comp_list:
+                                        print(f"    Computations: {comp_list}")
+                                    if symbol_list:
+                                        print(f"    Symbols: {symbol_list}")
                 elif not sessions:
                     print("  - No active sessions")
 
@@ -381,10 +395,11 @@ def main() -> int:
         "--config", "-c", required=True, help="Path to the YAML configuration file"
     )
     status_parser.add_argument(
-        "--details",
-        "-d",
-        action="store_true",
-        help="Show detailed runtime status including sessions, computations, and symbols",
+        "--verbose",
+        "-v",
+        action="count",
+        default=0,
+        help="Increase verbosity: -v for session details, -vv for full lists",
     )
     status_parser.set_defaults(func=status_command)
 

mplang/{runtime → v1/runtime}/client.py

@@ -25,9 +25,10 @@ from __future__ import annotations
 import base64
 from typing import Any
 
-import cloudpickle as pickle
 import httpx
 
+from mplang.v1.kernels.value import Value, decode_value, encode_value
+
 
 class ExecutionStatus:
     """Status of a computation execution."""
@@ -253,8 +254,10 @@ class HttpExecutorClient:
         """
         url = f"/sessions/{session_name}/symbols/{symbol_name}"
 
-        # Serialize data
-        data_bytes = pickle.dumps(data)
+        # Serialize data using Value envelope
+        if not isinstance(data, Value):
+            raise TypeError(f"Data must be a Value instance, got {type(data)}")
+        data_bytes = encode_value(data)
         data_b64 = base64.b64encode(data_bytes).decode("utf-8")
 
         payload = {"data": data_b64, "mptype": mptype or {}}
@@ -286,11 +289,15 @@ class HttpExecutorClient:
             response.raise_for_status()
             symbol_data = response.json()
 
-            # Deserialize data
+            # Deserialize data using Value envelope
             data_bytes = base64.b64decode(symbol_data["data"])
-            return pickle.loads(data_bytes)
+            return decode_value(data_bytes)
 
-        except (httpx.HTTPStatusError, httpx.RequestError) as e:
+        except httpx.HTTPStatusError as e:
+            if e.response is not None and e.response.status_code == 404:
+                return None
+            raise self._raise_http_error(f"get symbol {symbol_name}", e) from e
+        except httpx.RequestError as e:
             raise self._raise_http_error(f"get symbol {symbol_name}", e) from e
 
     async def delete_symbol(self, session_name: str, symbol_name: str) -> None:
@@ -403,8 +410,12 @@ class HttpExecutorClient:
         """
         url = f"/api/v1/symbols/{symbol_name}"
         try:
+            # Serialize using Value envelope
+            if not isinstance(data, Value):
+                raise TypeError(f"Data must be a Value instance, got {type(data)}")
+            data_bytes = encode_value(data)
             payload = {
-                "data": base64.b64encode(pickle.dumps(data)).decode("utf-8"),
+                "data": base64.b64encode(data_bytes).decode("utf-8"),
                 "mptype": mptype or {},
             }
             resp = await self._client.put(url, json=payload)
@@ -421,7 +432,7 @@ class HttpExecutorClient:
             resp.raise_for_status()
             payload = resp.json()
             data_bytes = base64.b64decode(payload["data"])
-            return pickle.loads(data_bytes)
+            return decode_value(data_bytes)
         except (httpx.HTTPStatusError, httpx.RequestError) as e:
             raise self._raise_http_error(f"get global symbol {symbol_name}", e) from e