mplang-nightly 0.1.dev254__tar.gz → 0.1.dev255__tar.gz

{mplang_nightly-0.1.dev254 → mplang_nightly-0.1.dev255}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev254
+Version: 0.1.dev255
 Summary: Multi-Party Programming Language
 Author-email: SecretFlow Team <secretflow-contact@service.alipay.com>
 License:                                  Apache License

{mplang_nightly-0.1.dev254 → mplang_nightly-0.1.dev255}/mplang/v2/backends/spu_impl.py

@@ -27,6 +27,7 @@ import spu.api as spu_api
 import spu.libspu as libspu
 
 from mplang.v2.backends.simp_worker import SimpWorker
+from mplang.v2.backends.spu_state import SPUState
 from mplang.v2.backends.tensor_impl import TensorValue
 from mplang.v2.dialects import spu
 from mplang.v2.edsl import serde
@@ -107,84 +108,6 @@ def to_runtime_config(config: spu.SPUConfig) -> libspu.RuntimeConfig:
     return runtime_config
 
 
-# Global cache for SPU runtimes per (local_rank, world_size) pair
-# Key: (local_rank, spu_world_size, protocol, field, link_mode), Value: (Runtime, Io)
-_SPU_RUNTIMES: dict[
-    tuple[int, int, str, str, str], tuple[spu_api.Runtime, spu_api.Io]
-] = {}
-
-
-def _create_mem_link(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_brpc_link(local_rank: int, spu_endpoints: list[str]) -> libspu.link.Context:
-    """Create BRPC link for distributed execution.
-
-    Args:
-        local_rank: The local rank within the SPU device (0-indexed).
-        spu_endpoints: List of BRPC endpoints for all SPU parties.
-
-    Returns:
-        A libspu.link.Context for BRPC communication.
-    """
-    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 _get_spu_ctx(
-    local_rank: int,
-    spu_world_size: int,
-    config: spu.SPUConfig,
-    spu_endpoints: list[str] | None = None,
-) -> tuple[spu_api.Runtime, spu_api.Io]:
-    """Get or create SPU runtime and IO for the given local rank within SPU.
-
-    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.
-
-    Returns:
-        A tuple of (Runtime, Io) for this party.
-    """
-    # Determine link mode
-    link_mode = "brpc" if spu_endpoints else "mem"
-
-    # Include protocol, field, and link_mode in cache key
-    cache_key = (local_rank, spu_world_size, config.protocol, config.field, link_mode)
-    if cache_key in _SPU_RUNTIMES:
-        return _SPU_RUNTIMES[cache_key]
-
-    # Create Link
-    if spu_endpoints:
-        link = _create_brpc_link(local_rank, spu_endpoints)
-    else:
-        link = _create_mem_link(local_rank, spu_world_size)
-
-    # Use config from SPUConfig
-    runtime_config = to_runtime_config(config)
-
-    # Create Runtime and Io
-    runtime = spu_api.Runtime(link, runtime_config)
-    io = spu_api.Io(spu_world_size, runtime_config)
-
-    _SPU_RUNTIMES[cache_key] = (runtime, io)
-    return runtime, io
-
-
 @spu.makeshares_p.def_impl
 def makeshares_impl(
     interpreter: Interpreter, op: Operation, data: TensorValue
@@ -293,7 +216,15 @@ def exec_impl(interpreter: Interpreter, op: Operation, *args: Any) -> Any:
                 )
             spu_endpoints.append(spu_endpoints_map[party_rank])
 
-    runtime, io = _get_spu_ctx(local_rank, spu_world_size, config, spu_endpoints)
+    # Get or create SPUState for caching Runtime/Io
+    spu_state = interpreter.get_dialect_state(SPUState.dialect_name)
+    if not isinstance(spu_state, SPUState):
+        spu_state = SPUState()
+        interpreter.set_dialect_state(SPUState.dialect_name, spu_state)
+
+    runtime, io = spu_state.get_or_create(
+        local_rank, spu_world_size, config, spu_endpoints
+    )
 
     executable_code = op.attrs["executable"]
     input_names = op.attrs["input_names"]

mplang_nightly-0.1.dev255/mplang/v2/backends/spu_state.py

@@ -0,0 +1,124 @@
+# 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
+
+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,
+    ) -> 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.
+
+        Returns:
+            A tuple of (Runtime, Io) for this party.
+        """
+        from mplang.v2.backends.spu_impl import to_runtime_config
+
+        link_mode = "brpc" if spu_endpoints else "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 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_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()