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

mplang/v2/backends/spu_state.py

@@ -0,0 +1,187 @@
+# 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 Dialect State.
+
+Manages SPU Runtime lifecycle as a dialect state, enabling reuse across
+multiple executions while binding to the Interpreter's lifecycle.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import spu.api as spu_api
+import spu.libspu as libspu
+
+from mplang.v2.runtime.dialect_state import DialectState
+
+if TYPE_CHECKING:
+    from mplang.v2.dialects import spu
+
+
+class SPUState(DialectState):
+    """SPU Runtime cache as dialect state.
+
+    Caches SPU Runtime and Io objects per (local_rank, world_size, config, link_mode)
+    to enable reuse across multiple SPU kernel executions.
+
+    This replaces the previous global `_SPU_RUNTIMES` cache with a properly
+    lifecycle-managed dialect state.
+    """
+
+    dialect_name: str = "spu"
+
+    def __init__(self) -> None:
+        # Key: (local_rank, world_size, protocol, field, link_mode)
+        # Value: (Runtime, Io)
+        self._runtimes: dict[
+            tuple[int, int, str, str, str], tuple[spu_api.Runtime, spu_api.Io]
+        ] = {}
+
+    def get_or_create(
+        self,
+        local_rank: int,
+        spu_world_size: int,
+        config: spu.SPUConfig,
+        spu_endpoints: list[str] | None = None,
+        communicator: object | None = None,
+        parties: list[int] | None = None,
+    ) -> tuple[spu_api.Runtime, spu_api.Io]:
+        """Get or create SPU Runtime and Io for the given configuration.
+
+        Args:
+            local_rank: The local rank within the SPU device (0-indexed).
+            spu_world_size: The number of parties in the SPU device.
+            config: SPU configuration including protocol settings.
+            spu_endpoints: Optional list of BRPC endpoints. If None, use mem link.
+            communicator: Optional v2 communicator (ThreadCommunicator/HttpCommunicator).
+                If provided, use Channels mode to reuse existing communication.
+            parties: Optional list of global ranks for SPU parties.
+                Required when communicator is provided.
+
+        Returns:
+            A tuple of (Runtime, Io) for this party.
+        """
+        from mplang.v2.backends.spu_impl import to_runtime_config
+
+        # Determine link mode
+        if communicator is not None:
+            link_mode = "channels"
+        elif spu_endpoints:
+            link_mode = "brpc"
+        else:
+            link_mode = "mem"
+
+        cache_key = (
+            local_rank,
+            spu_world_size,
+            config.protocol,
+            config.field,
+            link_mode,
+        )
+
+        if cache_key in self._runtimes:
+            return self._runtimes[cache_key]
+
+        # Create Link
+        if communicator is not None:
+            if parties is None:
+                raise ValueError("parties required when using communicator")
+            link = self._create_channels_link(
+                local_rank, spu_world_size, communicator, parties
+            )
+        elif spu_endpoints:
+            link = self._create_brpc_link(local_rank, spu_endpoints)
+        else:
+            link = self._create_mem_link(local_rank, spu_world_size)
+
+        # Create Runtime and Io
+        runtime_config = to_runtime_config(config)
+        runtime = spu_api.Runtime(link, runtime_config)
+        io = spu_api.Io(spu_world_size, runtime_config)
+
+        self._runtimes[cache_key] = (runtime, io)
+        return runtime, io
+
+    def _create_mem_link(
+        self, local_rank: int, spu_world_size: int
+    ) -> libspu.link.Context:
+        """Create in-memory link for simulation."""
+        desc = libspu.link.Desc()  # type: ignore
+        desc.recv_timeout_ms = 30 * 1000
+        for i in range(spu_world_size):
+            desc.add_party(f"P{i}", f"mem:{i}")
+        return libspu.link.create_mem(desc, local_rank)
+
+    def _create_channels_link(
+        self,
+        local_rank: int,
+        spu_world_size: int,
+        communicator: Any,
+        parties: list[int],
+    ) -> libspu.link.Context:
+        """Create link using custom channels (reuse v2 communicator).
+
+        Args:
+            local_rank: SPU local rank (0-indexed, already converted from global)
+            spu_world_size: Number of SPU parties
+            communicator: v2 communicator (ThreadCommunicator/HttpCommunicator)
+            parties: List of global ranks for SPU parties (ordered by local rank)
+
+        Returns:
+            libspu link context using BaseChannel adapters
+        """
+        from mplang.v2.backends.channel import BaseChannel
+
+        # Get this worker's global rank
+        global_rank = parties[local_rank]
+
+        # Create channels list (world_size elements, self = None)
+        channels = []
+        for idx, peer_global_rank in enumerate(parties):
+            if idx == local_rank:
+                # Self channel must be None
+                channel = None
+            else:
+                # Create channel to peer
+                channel = BaseChannel(communicator, global_rank, peer_global_rank)
+            channels.append(channel)
+
+        # Create link descriptor
+        desc = libspu.link.Desc()  # type: ignore
+        desc.recv_timeout_ms = 100 * 1000  # 100 seconds
+
+        # Add party info (required for world_size inference)
+        for idx in range(spu_world_size):
+            desc.add_party(f"P{idx}", f"dummy_{parties[idx]}")
+
+        return libspu.link.create_with_channels(desc, local_rank, channels)
+
+    def _create_brpc_link(
+        self, local_rank: int, spu_endpoints: list[str]
+    ) -> libspu.link.Context:
+        """Create BRPC link for distributed execution."""
+        desc = libspu.link.Desc()  # type: ignore
+        desc.recv_timeout_ms = 100 * 1000  # 100 seconds
+        desc.http_max_payload_size = 32 * 1024 * 1024  # 32MB
+
+        for i, endpoint in enumerate(spu_endpoints):
+            desc.add_party(f"P{i}", endpoint)
+
+        return libspu.link.create_brpc(desc, local_rank)
+
+    def shutdown(self) -> None:
+        """Clear all cached runtimes."""
+        self._runtimes.clear()

mplang/v2/backends/store_impl.py

@@ -0,0 +1,62 @@
+# 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.
+
+"""Store Runtime Implementation."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from mplang.v2.dialects import store
+from mplang.v2.edsl.graph import Operation
+from mplang.v2.runtime.interpreter import Interpreter
+
+
+def _get_uri(uri_base: str) -> str:
+    """Generate URI: {uri_base}."""
+    # Handle different schemes if necessary, for now assume simple path joining
+    # or scheme preservation.
+    if "://" in uri_base:
+        scheme, _, path = uri_base.partition("://")
+        # Ensure we don't double slash if path is absolute
+        return f"{scheme}://{path}"
+    else:
+        # Default to fs:// for absolute paths (sandboxed)
+        return f"fs://{uri_base}"
+
+
+@store.save_p.def_impl
+def save_impl(interpreter: Interpreter, op: Operation, obj_val: Any) -> Any:
+    """Save implementation."""
+    uri_base: str = op.attrs["uri_base"]
+
+    # Use ObjectStore to put the value
+    # Note: obj_val is the runtime value (e.g. TensorValue, TableValue, or raw)
+    # We store it as is (pickle).
+    if interpreter.store is None:
+        raise RuntimeError("Interpreter has no ObjectStore configured. Cannot save.")
+    interpreter.store.put(obj_val, uri=_get_uri(uri_base))
+
+    return obj_val
+
+
+@store.load_p.def_impl
+def load_impl(interpreter: Interpreter, op: Operation) -> Any:
+    """Load implementation."""
+    uri_base: str = op.attrs["uri_base"]
+    # expected_type is in attrs but not needed for runtime loading (pickle handles it)
+
+    if interpreter.store is None:
+        raise RuntimeError("Interpreter has no ObjectStore configured. Cannot load.")
+    return interpreter.store.get(_get_uri(uri_base))