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

mplang/kernels/context.py

@@ -37,11 +37,11 @@ def _ensure_impl_imported() -> None:
         return
     from mplang.kernels import builtin as _impl_builtin  # noqa: F401
     from mplang.kernels import crypto as _impl_crypto  # noqa: F401
+    from mplang.kernels import mock_tee as _impl_tee  # noqa: F401
     from mplang.kernels import phe as _impl_phe  # noqa: F401
     from mplang.kernels import spu as _impl_spu  # noqa: F401
     from mplang.kernels import sql_duckdb as _impl_sql_duckdb  # noqa: F401
     from mplang.kernels import stablehlo as _impl_stablehlo  # noqa: F401
-    from mplang.kernels import tee as _impl_tee  # noqa: F401
 
     _IMPL_IMPORTED = True
 
@@ -91,8 +91,8 @@ _DEFAULT_BINDINGS: dict[str, str] = {
     # generic SQL op; backend-specific kernel id for duckdb
     "sql.run": "duckdb.run_sql",
     # tee
-    "tee.quote": "tee.quote",
-    "tee.attest": "tee.attest",
+    # "tee.quote": "mock_tee.quote",
+    # "tee.attest": "mock_tee.attest",
 }
 
 
@@ -102,6 +102,10 @@ _DEFAULT_BINDINGS: dict[str, str] = {
 class RuntimeContext:
     """Per-runtime execution context with isolated op->kernel bindings.
 
+    This object owns ONLY static dispatch metadata ("op bindings") and mutable
+    per-rank kernel side state/cache/stats. It does NOT store per-evaluation
+    variable bindings (those are provided to the evaluator at evaluation time).
+
     Parameters
     ----------
     rank : int
@@ -110,9 +114,10 @@ class RuntimeContext:
         Total number of participants.
     initial_bindings : Mapping[str, str] | None, optional
         Optional partial overrides applied on top of the default binding table
-        during construction (override semantics, not replace). After
+        during construction (override semantics, not replace). These map
+        op_type -> kernel_id and form a *template* for dispatch. After
         initialization, all (re)binding must go through ``bind_op`` /
-        ``rebind_op``.
+        ``rebind_op`` on this context (scoped to THIS runtime only).
     state / cache / stats : dict, optional
         Mutable pockets reused across kernel invocations. If omitted, new
         dictionaries are created.

mplang/kernels/{tee.py → mock_tee.py}

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import os
+import warnings
 
 import numpy as np
 from numpy.typing import NDArray
@@ -43,16 +44,24 @@ def _quote_from_pk(pk: np.ndarray) -> NDArray[np.uint8]:
     return out
 
 
-@kernel_def("tee.quote")
+@kernel_def("mock_tee.quote")
 def _tee_quote(pfunc: PFunction, pk: object) -> NDArray[np.uint8]:
+    warnings.warn(
+        "Insecure mock TEE kernel 'mock_tee.quote' in use. NOT secure; for local testing only.",
+        stacklevel=3,
+    )
     pk = np.asarray(pk, dtype=np.uint8)
     # rng access ensures deterministic seeding per rank even if unused now
     _rng()
     return _quote_from_pk(pk)
 
 
-@kernel_def("tee.attest")
+@kernel_def("mock_tee.attest")
 def _tee_attest(pfunc: PFunction, quote: object) -> NDArray[np.uint8]:
+    warnings.warn(
+        "Insecure mock TEE kernel 'mock_tee.attest' in use. NOT secure; for local testing only.",
+        stacklevel=3,
+    )
     quote = np.asarray(quote, dtype=np.uint8)
     if quote.size != 33:
         raise ValueError("mock quote must be 33 bytes (1 header + 32 pk)")

mplang/runtime/simulation.py

@@ -34,10 +34,6 @@ from mplang.core.mpir import Reader, Writer
 from mplang.core.mpobject import MPObject
 from mplang.core.mptype import MPType, TensorLike
 from mplang.core.pfunc import PFunction  # for spu.seed_env kernel seeding
-
-# New explicit binding model: we only need RuntimeContext which ensures
-# bindings via bind_all_ops() on creation; per-module side-effect imports
-# are no longer required here.
 from mplang.kernels.context import RuntimeContext
 from mplang.runtime.link_comm import LinkCommunicator
 from mplang.utils.spu_utils import parse_field, parse_protocol
@@ -90,14 +86,20 @@ 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``.
         """
         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")
@@ -141,20 +143,18 @@ class Simulator(InterpContext):
         self._spu_world = spu_mask.num_parties()
         self._spu_mask = spu_mask
 
-        # No per-backend handlers needed anymore (all flat kernels)
-        self._handlers: list[list[Any]] = [[] for _ in range(self.world_size())]
-
-        self._evaluators: list[IEvaluator] = []
-        for rank in range(self.world_size()):
-            runtime = RuntimeContext(rank=rank, world_size=self.world_size())
-            ev = create_evaluator(
-                rank,
-                {},  # the global environment for this rank
-                self._comms[rank],
-                runtime,
-                None,
+        # 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(
+                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,
             )
-            self._evaluators.append(ev)
+            for rank in range(self.world_size())
+        ]
 
     @classmethod
     def simple(
@@ -215,10 +215,10 @@ class Simulator(InterpContext):
             for rank in range(self.world_size())
         ]
 
-        # Build per-rank evaluators with the per-party environment
+        # Build per-rank evaluators with the per-party environment (runtime reused)
         pts_evaluators: list[IEvaluator] = []
         for rank in range(self.world_size()):
-            runtime = RuntimeContext(rank=rank, world_size=self.world_size())
+            runtime = self._runtimes[rank]
             ev = create_evaluator(
                 rank,
                 pts_env[rank],
@@ -226,17 +226,21 @@ class Simulator(InterpContext):
                 runtime,
                 None,
             )
-            link_ctx = self._spu_link_ctxs[rank]
-            seed_fn = PFunction(
-                fn_type="spu.seed_env",
-                ins_info=(),
-                outs_info=(),
-                config=self._spu_runtime_cfg,
-                world=self._spu_world,
-                link=link_ctx,
-            )
-            # Seed SPU backend environment explicitly via runtime (no evaluator fast-path)
-            ev.runtime.run_kernel(seed_fn, [])  # type: ignore[arg-type]
+            # Seed SPU once per runtime (idempotent logical requirement)
+            # Use setdefault to both retrieve and create metadata dict in one step.
+            spu_meta = runtime.state.setdefault("_spu", {})
+            if not spu_meta.get("inited", False):
+                link_ctx = self._spu_link_ctxs[rank]
+                seed_fn = PFunction(
+                    fn_type="spu.seed_env",
+                    ins_info=(),
+                    outs_info=(),
+                    config=self._spu_runtime_cfg,
+                    world=self._spu_world,
+                    link=link_ctx,
+                )
+                ev.runtime.run_kernel(seed_fn, [])  # type: ignore[arg-type]
+                spu_meta["inited"] = True
             pts_evaluators.append(ev)
 
         # Collect evaluation results from all parties

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

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

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

@@ -29,13 +29,13 @@ mplang/core/expr/walk.py,sha256=lXkGJEEuvKGDqQihbxXPxfz2RfR1Q1zYUlt11iooQW0,1188
 mplang/kernels/__init__.py,sha256=2WE4cmW96Xkzyq2yRRYNww4kZ5o6u6NbPV0BxqZG698,581
 mplang/kernels/base.py,sha256=eizxj16sWkUvBvXWS0Zl-S9uuqalJmNUzB1xLhBg8S8,4920
 mplang/kernels/builtin.py,sha256=nSuM79cn7M6M27A6Y8ycilXT_qAlB1ktkwkRX6dv_VQ,7052
-mplang/kernels/context.py,sha256=VcVcRUaNISFgy9ATc4WFsX3Ng2mB9YlPdPEeHDQ1tA0,10952
+mplang/kernels/context.py,sha256=EiqsYq9Cw4Z_2lLvRNc8xNsu5EHkG_738QKFJ3Wf0zE,11329
 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/stablehlo.py,sha256=jDsu-lIHRAm94FcUcxZgK02c6BhFJpbO8cf2hP2DFgk,2937
-mplang/kernels/tee.py,sha256=v4_njXjb7cGQuE-D50V5byp2JDQJ79SwzyKSHUElbI0,1918
 mplang/ops/__init__.py,sha256=dpe7WWiYapOFzJeGoKFYBr5mnd6P5SdOyvdYaM2Nhm0,1408
 mplang/ops/base.py,sha256=rGtfBejcDh9mTRxOdJK5VUlG5vYiVJSir8X72X0Huvc,18264
 mplang/ops/builtin.py,sha256=D7T8rRF9g05VIw9T72lsncF5cDQqaT37eapBieRKvRI,9363
@@ -60,7 +60,7 @@ mplang/runtime/http_api.md,sha256=-re1DhEqMplAkv_wnqEU-PSs8tTzf4-Ml0Gq0f3Go6s,48
 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=xnK07RTcXwzHZzGVQw0w6JNNTMvj43q8OyTZW8_Ohmc,10958
+mplang/runtime/simulation.py,sha256=_cmUsYL58mvc6msHZ2fDjFAEHHLdJ-TRzJV8BxOP_WA,11473
 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.dev150.dist-info/METADATA,sha256=1oeiSoNdCsznO3MokEH5aCFsnafkjiN4dllVN9UiSpA,16547
-mplang_nightly-0.1.dev150.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-mplang_nightly-0.1.dev150.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
-mplang_nightly-0.1.dev150.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mplang_nightly-0.1.dev150.dist-info/RECORD,,
+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,,