mplang-nightly 0.1.dev150__tar.gz → 0.1.dev152__tar.gz

{mplang_nightly-0.1.dev150 → mplang_nightly-0.1.dev152}/PKG-INFO

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

{mplang_nightly-0.1.dev150 → mplang_nightly-0.1.dev152}/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_nightly-0.1.dev150/mplang/kernels/tee.py → mplang_nightly-0.1.dev152/mplang/kernels/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_nightly-0.1.dev150 → mplang_nightly-0.1.dev152}/mplang/kernels/sql_duckdb.py

@@ -25,6 +25,11 @@ def _duckdb_sql(pfunc: PFunction, *args: Any) -> Any:
     import duckdb
     import pandas as pd
 
+    # TODO: maybe we could translate the sql to duckdb dialect
+    # instead of raising an exception
+    if pfunc.attrs.get("dialect") != "duckdb":
+        raise ValueError("duckdb.run_sql must have dialect=duckdb attr")
+
     conn = duckdb.connect(":memory:")
     if args:
         in_names = pfunc.attrs.get("in_names")

{mplang_nightly-0.1.dev150 → mplang_nightly-0.1.dev152}/mplang/ops/base.py

@@ -30,7 +30,7 @@ from mplang.core.tensor import TensorType
 # The standard return contract for frontend operations (FeOperation.trace).
 #
 # Triad := (PFunction, list[MPObject], PyTreeDef)
-# - PFunction: Captures fn_type (routing key, e.g., "mlir.stablehlo", "sql.duckdb"),
+# - PFunction: Captures fn_type (routing key, e.g., "mlir.stablehlo", "sql.run"),
 #              input/output MPTypes and optional attributes.
 # - list[MPObject]: The flat positional MPObjects captured under the current
 #                   context (Trace/Interp). Order matches pfunc.ins_info.
@@ -422,6 +422,6 @@ def list_ops(module: str | None = None) -> dict[tuple[str, str], FeOperation]:
 # - For type-only kernels, use @module.typed_op(pfunc_name)(kernel). The op name is derived from the kernel function name.
 # - For complex ops (with Python callables/closures), subclass FeOperation and register
 #   using get_registry().register_op(module, name, op_instance) or use @module.feop with InlineFeOperation.
-# - Ensure PFunction.fn_type is set as the routing key (e.g., "mlir.stablehlo", "sql.duckdb").
+# - Ensure PFunction.fn_type is set as the routing key (e.g., "mlir.stablehlo", "sql.run").
 # - Keep device selection/routing out of frontend code; only set fn_type and attributes.
 # - Avoid moving MPObjects across contexts directly; capture within current ctx in trace().

{mplang_nightly-0.1.dev150 → mplang_nightly-0.1.dev152}/mplang/ops/ibis_cc.py

@@ -65,6 +65,7 @@ def ibis2sql(
         ins_info=tuple(ins_info),
         outs_info=tuple(outs_info),
         in_names=tuple(in_names),
+        dialect="duckdb",
     )
     return pfn
 

{mplang_nightly-0.1.dev150 → mplang_nightly-0.1.dev152}/mplang/ops/sql.py

@@ -46,12 +46,13 @@ class SqlFE(FeOperation):
                 in_vars.append(tbl)
 
         pfn = PFunction(
-            fn_type=f"sql[{self._dialect}]",
+            fn_type="sql.run",
             fn_name="",
             fn_text=sql,
             ins_info=tuple(ins_info),
             outs_info=(out_type,),
             in_names=tuple(in_names),
+            dialect=self._dialect,
         )
         _, treedef = tree_flatten(out_type)
         return pfn, in_vars, treedef

{mplang_nightly-0.1.dev150 → mplang_nightly-0.1.dev152}/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 → mplang_nightly-0.1.dev152}/tests/ops/test_crypto_tee.py

@@ -74,7 +74,9 @@ def _demo_flow():
 
 
 def test_crypto_enc_dec_and_tee_quote_attest_roundtrip():
-    sim = mplang.Simulator.simple(3)
+    # Create simulator with TEE bindings using the new initial_bindings parameter
+    tee_bindings = {"tee.quote": "mock_tee.quote", "tee.attest": "mock_tee.attest"}
+    sim = mplang.Simulator.simple(3, op_bindings=tee_bindings)
     p0, p1 = mplang.evaluate(sim, _demo_flow)
     a = mplang.fetch(sim, p0)
     b = mplang.fetch(sim, p1)

{mplang_nightly-0.1.dev150 → mplang_nightly-0.1.dev152}/tests/runtime/test_simulation.py

@@ -107,10 +107,10 @@ class TestSimulator:
     def test_simulator_creation(self):
         """Test Simulator creation."""
         sim = Simulator.simple(world_size=3)
-
         assert sim.world_size() == 3
         assert len(sim._comms) == 3
-        assert len(sim._evaluators) == 3
+        # persistent runtimes after refactor
+        assert len(sim._runtimes) == 3
 
         # Check that communicators are properly connected
         for i, comm in enumerate(sim._comms):

{mplang_nightly-0.1.dev150 → mplang_nightly-0.1.dev152}/tutorials/3_device.py

@@ -100,7 +100,9 @@ def run_spu():
 def run_tee():
     print("-" * 10, "millionaire (TEE)", "-" * 10)
 
-    sim = mplang.Simulator(cluster_spec)
+    # TEE operations need explicit binding for security
+    tee_bindings = {"tee.quote": "mock_tee.quote", "tee.attest": "mock_tee.attest"}
+    sim = mplang.Simulator(cluster_spec, op_bindings=tee_bindings)
     x_p0, y_p1, z_t, r_p0 = mplang.evaluate(sim, millionaire, "TEE0")
     print("x_p0:", x_p0, mpd.fetch(sim, x_p0))
     print("y_p1:", y_p1, mpd.fetch(sim, y_p1))

{mplang_nightly-0.1.dev150 → mplang_nightly-0.1.dev152}/tutorials/9_tee.py

@@ -116,7 +116,9 @@ def millionaire_manual():
 
 def main():
     print("-" * 10, "TEE millionaire: device vs manual (end-to-end IR)", "-" * 10)
-    sim = Simulator(cluster_spec)
+    # Create simulator with TEE bindings
+    tee_bindings = {"tee.quote": "mock_tee.quote", "tee.attest": "mock_tee.attest"}
+    sim = Simulator(cluster_spec, op_bindings=tee_bindings)
 
     compiled_dev = mplang.compile(sim, millionaire_device)
     compiled_man = mplang.compile(sim, millionaire_manual)