mplang-nightly 0.1.dev151__py3-none-any.whl → 0.1.dev153__py3-none-any.whl

mplang/runtime/session.py

@@ -0,0 +1,285 @@
+# 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.
+
+"""Core Session model (pure, no global registries).
+
+Contents:
+    * SessionState dataclass
+    * LinkCommFactory (SPU link reuse cache)
+    * Session (topology derivation, runtime init, SPU env seeding, local symbol/computation storage)
+
+Process-wide registries (sessions, global symbols) live in the server layer
+(`server.py`) so this module remains portable and easy to unit test.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from dataclasses import dataclass, field
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, cast
+from urllib.parse import urlparse
+
+import spu.libspu as libspu
+
+from mplang.core.expr.ast import Expr
+from mplang.core.expr.evaluator import IEvaluator, create_evaluator
+from mplang.core.mask import Mask
+from mplang.kernels.context import RuntimeContext
+from mplang.kernels.spu import PFunction  # type: ignore
+from mplang.runtime.communicator import HttpCommunicator
+from mplang.runtime.exceptions import ResourceNotFound
+from mplang.runtime.link_comm import LinkCommunicator
+from mplang.utils.spu_utils import parse_field, parse_protocol
+
+if TYPE_CHECKING:  # pragma: no cover - import only for type checking
+    from mplang.core.cluster import ClusterSpec, Node, RuntimeInfo
+
+
+class LinkCommFactory:
+    """Factory for creating and caching link communicators."""
+
+    def __init__(self) -> None:
+        self._cache: dict[tuple[int, tuple[str, ...]], LinkCommunicator] = {}
+
+    def create_link(self, rel_rank: int, addrs: list[str]) -> LinkCommunicator:
+        key = (rel_rank, tuple(addrs))
+        link = self._cache.get(key)
+        if link is not None:
+            return link
+        logging.info(f"LinkCommunicator created: rel_rank={rel_rank} addrs={addrs}")
+        link = LinkCommunicator(rel_rank, addrs)
+        self._cache[key] = link
+        return link
+
+
+# Shared link factory (module-local, not global registry of sessions)
+g_link_factory = LinkCommFactory()
+
+
+@dataclass
+class Symbol:
+    name: str
+    mptype: Any
+    data: Any
+
+
+@dataclass
+class Computation:
+    name: str
+    expr: Expr
+
+
+@dataclass
+class SessionState:
+    runtime: RuntimeContext | None = None
+    computations: dict[str, Computation] = field(default_factory=dict)
+    symbols: dict[str, Symbol] = field(default_factory=dict)
+    spu_seeded: bool = False
+    created_ts: float = field(default_factory=time.time)
+    last_access_ts: float = field(default_factory=time.time)
+
+
+class Session:
+    """Represents the per-rank execution context.
+
+    Immutable config: name, rank, cluster_spec.
+    Derived: node, runtime_info, endpoints, spu_device, spu_mask, protocol/field, is_spu_party.
+    Mutable: state (runtime object, symbols, computations, seeded flag).
+    """
+
+    def __init__(self, name: str, rank: int, cluster_spec: ClusterSpec):
+        self.name = name
+        self.rank = rank
+        self.cluster_spec = cluster_spec
+        self.state = SessionState()
+        self.communicator = HttpCommunicator(
+            session_name=name, rank=rank, endpoints=self.endpoints
+        )
+
+    # --- Derived topology ---
+    @cached_property
+    def node(self) -> Node:
+        return self.cluster_spec.get_node_by_rank(self.rank)
+
+    @property
+    def runtime_info(self) -> RuntimeInfo:
+        return self.node.runtime_info
+
+    @cached_property
+    def endpoints(self) -> list[str]:
+        eps: list[str] = []
+        for n in sorted(
+            self.cluster_spec.nodes.values(),
+            key=lambda x: x.rank,  # type: ignore[attr-defined]
+        ):
+            ep = n.endpoint
+            if not ep.startswith(("http://", "https://")):
+                ep = f"http://{ep}"
+            eps.append(ep)
+        return eps
+
+    @cached_property
+    def spu_device(self):  # type: ignore
+        devs = self.cluster_spec.get_devices_by_kind("SPU")
+        if len(devs) != 1:
+            raise RuntimeError(
+                f"Expected exactly one SPU device, got {len(devs)} (session={self.name})"
+            )
+        return devs[0]
+
+    @cached_property
+    def spu_mask(self) -> Mask:
+        return Mask.from_ranks([m.rank for m in self.spu_device.members])
+
+    @property
+    def spu_protocol(self) -> str:
+        return cast(str, self.spu_device.config.get("protocol", "SEMI2K"))
+
+    @property
+    def spu_field(self) -> str:
+        return cast(str, self.spu_device.config.get("field", "FM64"))
+
+    @property
+    def is_spu_party(self) -> bool:
+        return self.rank in self.spu_mask
+
+    # --- Runtime helpers ---
+    def ensure_runtime(self) -> RuntimeContext:
+        if self.state.runtime is None:
+            self.state.runtime = RuntimeContext(
+                rank=self.rank,
+                world_size=len(self.cluster_spec.nodes),  # type: ignore[attr-defined]
+                initial_bindings=(
+                    self.runtime_info.op_bindings if self.runtime_info else {}
+                ),
+            )
+        return self.state.runtime
+
+    def ensure_spu_env(self) -> None:
+        """Ensure SPU kernel env (config/world[/link]) registered on this runtime.
+
+        Previous logic only seeded SPU parties; non-participating ranks then raised
+        a hard error when the evaluator encountered SPU ops in the global program,
+        because the kernel pocket lacked config/world. For now we register the
+        config/world on ALL parties (idempotent) and only attach a link context for
+        participating SPU ranks. Non-parties will still error later if they try to
+        execute a link-dependent SPU kernel (which should be guarded by masks in the
+        IR), but they will no longer fail early with a misleading
+        "SPU kernel state not initialized" message.
+        """
+        if self.state.spu_seeded:
+            return
+
+        link_ctx = None
+        # Fixed port offset for SPU runtime link services (legacy value retained).
+        # TODO: make configurable if future deployments require dynamic offset.
+        SPU_PORT_OFFSET = 100
+
+        if self.is_spu_party:
+            # Build SPU address list across all endpoints for ranks in mask
+            spu_addrs: list[str] = []
+            for r, addr in enumerate(self.communicator.endpoints):
+                if r in self.spu_mask:
+                    if "//" not in addr:
+                        addr = f"//{addr}"
+                    parsed = urlparse(addr)
+                    assert isinstance(parsed.port, int)
+                    new_addr = f"{parsed.hostname}:{parsed.port + SPU_PORT_OFFSET}"
+                    spu_addrs.append(new_addr)
+            rel_index = sum(1 for r in range(self.rank) if r in self.spu_mask)
+            link_ctx = g_link_factory.create_link(rel_index, spu_addrs)
+
+        spu_config = libspu.RuntimeConfig(
+            protocol=parse_protocol(self.spu_protocol),
+            field=parse_field(self.spu_field),
+            fxp_fraction_bits=18,
+        )
+        seed_pfunc = PFunction(
+            fn_type="spu.seed_env",
+            ins_info=(),
+            outs_info=(),
+            config=spu_config,
+            world=self.spu_mask.num_parties(),
+            link=link_ctx,
+        )
+        self.ensure_runtime().run_kernel(seed_pfunc, [])
+        self.state.spu_seeded = True
+
+    # --- Computations & Symbols (instance-local) ---
+    def add_computation(self, computation: Computation) -> None:
+        self.state.computations[computation.name] = computation
+
+    def get_computation(self, name: str) -> Computation | None:
+        return self.state.computations.get(name)
+
+    def add_symbol(self, symbol: Symbol) -> None:
+        self.state.symbols[symbol.name] = symbol
+
+    def get_symbol(self, name: str) -> Symbol | None:
+        return self.state.symbols.get(name)
+
+    def list_symbols(self) -> list[str]:  # pragma: no cover - trivial
+        return list(self.state.symbols.keys())
+
+    def delete_symbol(self, name: str) -> bool:
+        if name in self.state.symbols:
+            del self.state.symbols[name]
+            return True
+        return False
+
+    def list_computations(self) -> list[str]:  # pragma: no cover - trivial
+        return list(self.state.computations.keys())
+
+    def delete_computation(self, name: str) -> bool:
+        if name in self.state.computations:
+            del self.state.computations[name]
+            return True
+        return False
+
+    # --- Execution ---
+    def execute(
+        self, computation: Computation, input_names: list[str], output_names: list[str]
+    ) -> None:
+        env: dict[str, Any] = {}
+        for in_name in input_names:
+            sym = self.get_symbol(in_name)
+            if sym is None:
+                raise ResourceNotFound(
+                    f"Input symbol '{in_name}' not found in session '{self.name}'"
+                )
+            env[in_name] = sym.data
+        rt = self.ensure_runtime()
+        self.ensure_spu_env()
+        evaluator: IEvaluator = create_evaluator(
+            rank=self.rank, env=env, comm=self.communicator, runtime=rt
+        )
+        results = evaluator.evaluate(computation.expr)
+        if results and len(results) != len(output_names):
+            raise RuntimeError(
+                f"Expected {len(output_names)} results, got {len(results)}"
+            )
+        for name, val in zip(output_names, results, strict=True):
+            self.add_symbol(Symbol(name=name, mptype={}, data=val))
+
+    # --- Convenience constructor ---
+    @classmethod
+    def from_cluster_spec_dict(cls, name: str, rank: int, spec_dict: dict) -> Session:
+        from mplang.core.cluster import ClusterSpec  # local import to avoid cycles
+
+        spec = ClusterSpec.from_dict(spec_dict)
+        if len(spec.get_devices_by_kind("SPU")) == 0:
+            raise RuntimeError("No SPU device found in cluster_spec")
+        return cls(name=name, rank=rank, cluster_spec=spec)

mplang/runtime/simulation.py

@@ -86,20 +86,17 @@ class Simulator(InterpContext):
         cluster_spec: ClusterSpec,
         *,
         trace_ranks: list[int] | None = None,
-        op_bindings: dict[str, str] | None = None,
     ) -> None:
         """Initialize a simulator with the given cluster specification.
 
         Args:
             cluster_spec: The cluster specification defining the simulation environment.
             trace_ranks: List of ranks to trace execution for debugging.
-            op_bindings: Optional op->kernel binding template applied to all
-                RuntimeContexts. These are static dispatch overrides (merged
-                with project defaults) and are orthogonal to the per-evaluate
-                variable ``bindings`` dict passed into ``evaluate``.
+                Per-node op binding overrides should now be provided via
+                each node's `runtime_info.op_bindings` in the supplied
+                `cluster_spec`.
         """
         super().__init__(cluster_spec)
-        self._op_bindings_template = op_bindings or {}
         self._trace_ranks = trace_ranks or []
 
         spu_devices = cluster_spec.get_devices_by_kind("SPU")
@@ -145,21 +142,22 @@ class Simulator(InterpContext):
 
         # Persistent per-rank RuntimeContext instances (reused across evaluates).
         # We no longer pre-create evaluators since each evaluate has different env bindings.
-        self._runtimes: list[RuntimeContext] = [
-            RuntimeContext(
+        # Build per-rank runtime contexts.
+        self._runtimes: list[RuntimeContext] = []
+        for rank in range(self.world_size()):
+            node = self.cluster_spec.get_node_by_rank(rank)
+            rt = RuntimeContext(
                 rank=rank,
                 world_size=self.world_size(),
-                # Static op bindings template cloned into each runtime. These are kernel
-                # dispatch mappings, not per-evaluate variable bindings.
-                initial_bindings=self._op_bindings_template,
+                initial_bindings=node.runtime_info.op_bindings,
             )
-            for rank in range(self.world_size())
-        ]
+            self._runtimes.append(rt)
 
     @classmethod
     def simple(
         cls,
         world_size: int,
+        op_bindings: dict[str, str] | None = None,
         **kwargs: Any,
     ) -> Simulator:
         """Create a simple simulator with the given number of parties.
@@ -175,6 +173,10 @@ class Simulator(InterpContext):
             A Simulator instance with a simple cluster configuration.
         """
         cluster_spec = ClusterSpec.simple(world_size)
+        if op_bindings:
+            # Apply the same op_bindings to every node's runtime_info for convenience
+            for node in cluster_spec.nodes.values():
+                node.runtime_info.op_bindings.update(op_bindings)
         return cls(cluster_spec, **kwargs)
 
     def _do_evaluate(self, expr: Expr, evaluator_engine: IEvaluator) -> Any:

{mplang_nightly-0.1.dev151.dist-info → mplang_nightly-0.1.dev153.dist-info}/METADATA

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

{mplang_nightly-0.1.dev151.dist-info → mplang_nightly-0.1.dev153.dist-info}/RECORD

@@ -4,7 +4,7 @@ mplang/device.py,sha256=RmjnhzHxJkkNmtBKtYMEbpQYBZpuC43qlllkCOp-QD8,12548
 mplang/analysis/__init__.py,sha256=CTHFvRsi-nFngojqjn08UaR3RY9i7CJ7T2UdR95kCrk,1056
 mplang/analysis/diagram.py,sha256=ffwgD12gL1_KH1uJ_EYkjmIlDrfxYJJkWj-wHl09_Xk,19520
 mplang/core/__init__.py,sha256=lWxlEKfRwX7FNDzgyKZ1fiDMaCiqkyg0j5mKlZD_v7g,2244
-mplang/core/cluster.py,sha256=gqMJenvXUfHhE181Dd5JiUkD4nT07RLoicBnvsGmRkE,8598
+mplang/core/cluster.py,sha256=IqXHLogetegUEEAzmD8cWRash-UID06Wo3OBeZFwatg,11800
 mplang/core/comm.py,sha256=MByyu3etlQh_TkP1vKCFLIAPPuJOpl9Kjs6hOj6m4Yc,8843
 mplang/core/context_mgr.py,sha256=R0QJAod-1nYduVoOknLfAsxZiy-RtmuQcp-07HABYZU,1541
 mplang/core/dtype.py,sha256=0rZqFaFikFu9RxtdO36JLEgFL-E-lo3hH10whwkTVVY,10213
@@ -34,33 +34,33 @@ mplang/kernels/crypto.py,sha256=TWixli1uRQ_7OjA49qQXUXa2ldHDEwaCMXXPSHdAPi8,3812
 mplang/kernels/mock_tee.py,sha256=ifVqb03FgQ5EDK8r14PMpZKMMaSeqMDorcpGLiy00mM,2233
 mplang/kernels/phe.py,sha256=8-_1IFPOaGECGj9mbYja8XoqbMYnYqfpDNVyMJd8J1Y,65247
 mplang/kernels/spu.py,sha256=Kkg1ZQbmklXl7YkIeBfxqs3o4wX7ygBE8hXLpx90L9Y,9307
-mplang/kernels/sql_duckdb.py,sha256=C2XdNLoE2Apma-Fs7OYzDzkBAXAVuShuROxhCWCHDG4,1502
+mplang/kernels/sql_duckdb.py,sha256=UN1Ev6-MxF_-65zMExUsLScC9PlmEIEcN8YziIoX_rY,1724
 mplang/kernels/stablehlo.py,sha256=jDsu-lIHRAm94FcUcxZgK02c6BhFJpbO8cf2hP2DFgk,2937
 mplang/ops/__init__.py,sha256=dpe7WWiYapOFzJeGoKFYBr5mnd6P5SdOyvdYaM2Nhm0,1408
-mplang/ops/base.py,sha256=rGtfBejcDh9mTRxOdJK5VUlG5vYiVJSir8X72X0Huvc,18264
+mplang/ops/base.py,sha256=h67_SHWNZGUuTCuMll-9kDgGvlPhlFov7WAQCHTmUvw,18258
 mplang/ops/builtin.py,sha256=D7T8rRF9g05VIw9T72lsncF5cDQqaT37eapBieRKvRI,9363
 mplang/ops/crypto.py,sha256=9CeFJrYmvjmgx-3WQl6jHXh8VafRpT4QBunbzsPF8Uc,3646
-mplang/ops/ibis_cc.py,sha256=bWKN1dL8Nluwvu5TLi8iUwytcnpXtWakZDCL793zBRA,4230
+mplang/ops/ibis_cc.py,sha256=G8TobBlkyT_IxOWXmMm4APPUnj6sH9bJ-1STfw4kIcw,4256
 mplang/ops/jax_cc.py,sha256=42czYg3hNQbI_nUebXnshlU8ULwM-oBDe_TQoApLNVA,7802
 mplang/ops/phe.py,sha256=SatswExjZWPed8y3qA33BCwIWbvsgHCuCAz_pv2RLLw,6790
 mplang/ops/spu.py,sha256=UHr5DSoqG08xDYER_11OsMVjGGNXXxsvkFoVvXU8uik,4989
-mplang/ops/sql.py,sha256=p-u0wQPk9KlgveltYvQcF1UefScJoqBCqhzYPeLBB5Y,1994
+mplang/ops/sql.py,sha256=HyY2i5aGC5W7r62JryFSjQCUDXH3kQz82YADwn4z5uc,2015
 mplang/ops/tee.py,sha256=gwzP81y2idH-d-Du84H6oNZpLaGD-3fEgm8G1uxWpUA,1388
 mplang/protos/v1alpha1/mpir_pb2.py,sha256=Bros37t-4LMJbuUYVSM65rImUYTtZDhNTIADGbZCKp0,7522
 mplang/protos/v1alpha1/mpir_pb2.pyi,sha256=GwXR4wPB_kB_36iYS9x-cGI9KDKFMq89KhdLhW_xmvE,19342
 mplang/protos/v1alpha1/mpir_pb2_grpc.py,sha256=xYOs94SXiNYAlFodACnsXW5QovLsHY5tCk3p76RH5Zc,158
 mplang/runtime/__init__.py,sha256=IRPP3TtpFC4iSt7_uaq-S4dL7CwrXL0XBMeaBoEYLlg,948
 mplang/runtime/cli.py,sha256=WehDodeVB4AukSWx1LJxxtKUqGmLPY4qjayrPlOg3bE,14438
-mplang/runtime/client.py,sha256=w8sPuQzqaJI5uS_3JHu2mf0tLaFmZH3f6-SeUBfMLMY,15737
+mplang/runtime/client.py,sha256=vkJUFSDcKIdbKiGUM5AosCKTZygl9g8uZFEjw2xwKig,15249
 mplang/runtime/communicator.py,sha256=Lek6_h_Wmr_W-_JpT-vMxL3CHxcVZdtf7jdaLGuxPgQ,3199
 mplang/runtime/data_providers.py,sha256=hH2butEOYNGq2rRZjVBDfXLxe3YUin2ftAF6htbTfLA,8226
-mplang/runtime/driver.py,sha256=Ok1jY301ctN1_KTb4jwSxOdB0lI_xhx9AwhtEGJ-VLQ,11300
+mplang/runtime/driver.py,sha256=pq2EQFZK9tH90Idops_yeF6fj0cfFVD_5mFcmy4Hzco,11089
 mplang/runtime/exceptions.py,sha256=c18U0xK20dRmgZo0ogTf5vXlkix9y3VAFuzkHxaXPEk,981
 mplang/runtime/http_api.md,sha256=-re1DhEqMplAkv_wnqEU-PSs8tTzf4-Ml0Gq0f3Go6s,4883
 mplang/runtime/link_comm.py,sha256=uNqTCGZVwWeuHAb7yXXQf0DUsMXLa8leHCkrcZdzYMU,4559
-mplang/runtime/resource.py,sha256=xNke4UpNDjsjWcr09oXWNBXsMfSZFOwsKD7FWdCVPbc,11688
-mplang/runtime/server.py,sha256=LQ5uJi95tYrKmgHwZaxUQi-aiqwSsT3W4z7pZ9dQaUQ,14716
-mplang/runtime/simulation.py,sha256=_cmUsYL58mvc6msHZ2fDjFAEHHLdJ-TRzJV8BxOP_WA,11473
+mplang/runtime/server.py,sha256=vYjuWTWhhSLHUpsO8FDnOQ8kFzPhE-fXDDyL8GHVPj4,16673
+mplang/runtime/session.py,sha256=4TQ_RPRmriv0H0S6rl_GSabxS7XrwMkdZIdcnyE8bHw,10374
+mplang/runtime/simulation.py,sha256=WyIs8ta3ZM5o3RB0Bcb0MUu6Yh88Iujr27KvZFqGxig,11497
 mplang/simp/__init__.py,sha256=xNXnA8-jZAANa2A1W39b3lYO7D02zdCXl0TpivkTGS4,11579
 mplang/simp/mpi.py,sha256=Wv_Q16TQ3rdLam6OzqXiefIGSMmagGkso09ycyOkHEs,4774
 mplang/simp/random.py,sha256=7PVgWNL1j7Sf3MqT5PRiWplUu-0dyhF3Ub566iqX86M,3898
@@ -70,8 +70,8 @@ mplang/utils/crypto.py,sha256=rvPomBFtznRHc3RPi6Aip9lsU8zW2oxBqGv1K3vn7Rs,1052
 mplang/utils/func_utils.py,sha256=vCJcZmu0bEbqhOQKdpttV2_MBllIcPSN0b8U4WjNGGo,5164
 mplang/utils/spu_utils.py,sha256=S3L9RBkBe2AvSuMSQQ12cBY5Y1NPthubvErSX_7nj1A,4158
 mplang/utils/table_utils.py,sha256=aC-IZOKkSmFkpr3NZchLM0Wt0GOn-rg_xHBHREWBwAU,2202
-mplang_nightly-0.1.dev151.dist-info/METADATA,sha256=sCQECTJOQoKyr3XXAf8Kma7lrB5KEt6toJbQA9a5nEA,16547
-mplang_nightly-0.1.dev151.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-mplang_nightly-0.1.dev151.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
-mplang_nightly-0.1.dev151.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mplang_nightly-0.1.dev151.dist-info/RECORD,,
+mplang_nightly-0.1.dev153.dist-info/METADATA,sha256=4dEwwbuB0n0oRHxO09vuMY2Al57Ol8O8KdXGlDpEZqo,16547
+mplang_nightly-0.1.dev153.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+mplang_nightly-0.1.dev153.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
+mplang_nightly-0.1.dev153.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mplang_nightly-0.1.dev153.dist-info/RECORD,,