mplang-nightly 0.1.dev155__tar.gz → 0.1.dev157__tar.gz

{mplang_nightly-0.1.dev155 → mplang_nightly-0.1.dev157}/PKG-INFO

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

{mplang_nightly-0.1.dev155 → mplang_nightly-0.1.dev157}/mplang/device.py

@@ -207,8 +207,15 @@ def _d2d(to_dev_id: str, obj: MPObject) -> MPObject:
         assert len(frm_dev.members) == 1 and len(to_dev.members) == 1
         frm_rank = frm_dev.members[0].rank
         tee_rank = to_dev.members[0].rank
+        platform = to_dev.config.get("platform")
+        if not platform:
+            raise ValueError(
+                f"TEE device '{to_dev_id}' is missing 'platform' in its config."
+            )
         # Ensure sessions (both directions) exist for this PPU<->TEE pair
-        sess_p, sess_t = _ensure_tee_session(frm_dev_id, to_dev_id, frm_rank, tee_rank)
+        sess_p, sess_t = _ensure_tee_session(
+            frm_dev_id, to_dev_id, frm_rank, tee_rank, platform
+        )
         # Bytes-only path: pack -> enc -> p2p -> dec -> unpack (with static out type)
         obj_ty = TensorType.from_obj(obj)
         b = simp.runAt(frm_rank, builtin.pack)(obj)
@@ -222,8 +229,15 @@ def _d2d(to_dev_id: str, obj: MPObject) -> MPObject:
         assert len(frm_dev.members) == 1 and len(to_dev.members) == 1
         tee_rank = frm_dev.members[0].rank
         ppu_rank = to_dev.members[0].rank
+        platform = to_dev.config.get("platform")
+        if not platform:
+            raise ValueError(
+                f"TEE device '{to_dev_id}' is missing 'platform' in its config."
+            )
         # Ensure bidirectional session established for this pair
-        sess_p, sess_t = _ensure_tee_session(to_dev_id, frm_dev_id, ppu_rank, tee_rank)
+        sess_p, sess_t = _ensure_tee_session(
+            to_dev_id, frm_dev_id, ppu_rank, tee_rank, platform
+        )
         obj_ty = TensorType.from_obj(obj)
         b = simp.runAt(tee_rank, builtin.pack)(obj)
         ct = simp.runAt(tee_rank, crypto.enc)(b, sess_t)
@@ -245,7 +259,7 @@ def _d2d(to_dev_id: str, obj: MPObject) -> MPObject:
 
 
 def _ensure_tee_session(
-    frm_dev_id: str, to_dev_id: str, frm_rank: int, tee_rank: int
+    frm_dev_id: str, to_dev_id: str, frm_rank: int, tee_rank: int, platform: str
 ) -> tuple[MPObject, MPObject]:
     """Ensure a TEE session (sess_p at sender, sess_t at TEE) exists.
 
@@ -263,11 +277,11 @@ def _ensure_tee_session(
     # 1) TEE generates (sk, pk) and quote(pk)
     # KEM suite currently constant; future: read from tee device config (e.g. cluster_spec.devices[dev_id].config)
     tee_sk, tee_pk = simp.runAt(tee_rank, crypto.kem_keygen)(_TEE_KEM_SUITE)
-    quote = simp.runAt(tee_rank, tee.quote)(tee_pk)
+    quote = simp.runAt(tee_rank, tee.quote_gen)(tee_pk)
 
     # 2) Send quote to sender and attest to obtain TEE pk
     quote_at_sender = mpi.p2p(tee_rank, frm_rank, quote)
-    tee_pk_at_sender = simp.runAt(frm_rank, tee.attest)(quote_at_sender)
+    tee_pk_at_sender = simp.runAt(frm_rank, tee.attest)(quote_at_sender, platform)
 
     # 3) Sender generates its ephemeral keypair and sends its pk to TEE
     v_sk, v_pk = simp.runAt(frm_rank, crypto.kem_keygen)(_TEE_KEM_SUITE)

{mplang_nightly-0.1.dev155 → mplang_nightly-0.1.dev157}/mplang/kernels/base.py

@@ -34,7 +34,10 @@ from __future__ import annotations
 import contextvars
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from mplang.kernels.context import RuntimeContext
 
 __all__ = [
     "KernelContext",
@@ -48,12 +51,15 @@ __all__ = [
 
 @dataclass
 class KernelContext:
-    """Ephemeral call context set via contextvar while a kernel runs."""
+    """Ephemeral per-kernel invocation context.
+
+    Cross-kernel persistent state (RNGs, compiled artifacts, environment handles)
+    should be stored in RuntimeContext.
+    """
 
     rank: int
     world_size: int
-    state: dict[str, dict[str, Any]]  # backend namespace -> pocket
-    cache: dict[str, Any]  # runtime-level shared cache (per BackendRuntime)
+    runtime: RuntimeContext
 
 
 _CTX_VAR: contextvars.ContextVar[KernelContext | None] = contextvars.ContextVar(
@@ -62,37 +68,7 @@ _CTX_VAR: contextvars.ContextVar[KernelContext | None] = contextvars.ContextVar(
 
 
 def cur_kctx() -> KernelContext:
-    """Return the current kernel execution context (only valid inside a kernel).
-
-    Two storages:
-      - state: namespaced pockets (dict[str, dict]) for backend-local mutable helpers
-      - cache: global (per runtime) shared dict; prefer state unless truly cross-backend
-
-    Examples:
-      1) Compile cache::
-            @kernel_def("mlir.stablehlo")
-            def _exec(pfunc, args):
-                ctx = cur_kctx()
-                pocket = ctx.state.setdefault("stablehlo", {})
-                cache = pocket.setdefault("compile_cache", {})
-                text = pfunc.fn_text
-                mod = cache.get(text)
-                if mod is None:
-                    mod = compile_mlir(text)
-                    cache[text] = mod
-                return run(mod, args)
-
-      2) Deterministic RNG::
-            @kernel_def("crypto.keygen")
-            def _keygen(pfunc, args):
-                ctx = cur_kctx()
-                pocket = ctx.state.setdefault("crypto", {})
-                rng = pocket.get("rng")
-                if rng is None:
-                    rng = np.random.default_rng(1234 + ctx.rank * 7919)
-                    pocket["rng"] = rng
-                return (rng.integers(0, 256, size=(32,), dtype=np.uint8),)
-    """
+    """Return current kernel execution context (only valid inside kernel)."""
     ctx = _CTX_VAR.get()
     if ctx is None:
         raise RuntimeError("cur_kctx() called outside backend kernel execution")

{mplang_nightly-0.1.dev155 → mplang_nightly-0.1.dev157}/mplang/kernels/context.py

@@ -91,7 +91,7 @@ _DEFAULT_BINDINGS: dict[str, str] = {
     # generic SQL op; backend-specific kernel id for duckdb
     "sql.run": "duckdb.run_sql",
     # tee
-    # "tee.quote": "mock_tee.quote",
+    # "tee.quote_gen": "mock_tee.quote_gen",
     # "tee.attest": "mock_tee.attest",
 }
 
@@ -118,12 +118,21 @@ class RuntimeContext:
         op_type -> kernel_id and form a *template* for dispatch. After
         initialization, all (re)binding must go through ``bind_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.
+    state : dict, optional
+        Mutable per-runtime key/value storage for kernels. Flat key space;
+        callers SHOULD use dotted prefixes (e.g. "stablehlo.compile_cache").
+        Kernels own their *state* (functional correctness data, caches,
+        handles, compiled objects, RNGs, etc.). Runtime does not interpret
+        structure—values may themselves be dicts if a kernel wants its own
+        pocket. Created empty when omitted.
+    stats : dict, optional
+        Mutable statistics/telemetry owned by the runtime (usage counters,
+        timings, profiling aids). Kernels may increment counters but should
+        avoid storing functional state here. A default "op_calls" mapping is
+        ensured. Created empty when omitted.
     """
 
-    __slots__ = ("_ibindings", "cache", "rank", "state", "stats", "world_size")
+    __slots__ = ("_ibindings", "rank", "state", "stats", "world_size")
 
     def __init__(
         self,
@@ -131,8 +140,7 @@ class RuntimeContext:
         world_size: int,
         initial_bindings: Mapping[str, str] | None = None,
         *,
-        state: dict[str, dict[str, Any]] | None = None,
-        cache: dict[str, Any] | None = None,
+        state: dict[str, Any] | None = None,
         stats: dict[str, Any] | None = None,
     ) -> None:
         _ensure_impl_imported()
@@ -144,7 +152,6 @@ class RuntimeContext:
             **(initial_bindings or {}),
         }
         self.state = state if state is not None else {}
-        self.cache = cache if cache is not None else {}
         self.stats = stats if stats is not None else {}
         self.stats.setdefault("op_calls", {})
 
@@ -168,19 +175,15 @@ class RuntimeContext:
             if isinstance(ins_spec, TensorType):
                 _validate_tensor_arg(fn_type, idx, ins_spec, val)
                 continue
+
         # install kernel context
-        kctx = KernelContext(
-            rank=self.rank,
-            world_size=self.world_size,
-            state=self.state,
-            cache=self.cache,
-        )
-        token = base._CTX_VAR.set(kctx)  # type: ignore[attr-defined]
+        kctx = KernelContext(rank=self.rank, world_size=self.world_size, runtime=self)
+        token = base._CTX_VAR.set(kctx)
         try:
             raw = fn(pfunc, *arg_list)
         finally:
-            base._CTX_VAR.reset(token)  # type: ignore[attr-defined]
-        # Stats (best effort)
+            base._CTX_VAR.reset(token)
+
         try:
             op_calls = self.stats.setdefault("op_calls", {})
             op_calls[fn_type] = op_calls.get(fn_type, 0) + 1
@@ -213,7 +216,57 @@ class RuntimeContext:
 
     def reset(self) -> None:
         self.state.clear()
-        self.cache.clear()
+
+    # ---- runtime state API (flat key space) ----
+    # Keys are treated atomically; convention encourages dotted prefixes
+    # (e.g. 'stablehlo.compile_cache.hash', 'crypto.rng'). Implementation
+    # does NOT parse or create hierarchical dicts—any grouping is purely
+    # by string prefix. Values themselves MAY be dicts if callers want a
+    # manual pocket. This keeps semantics simple and predictable.
+
+    def ensure_state(self, key: str, factory: type | Any = dict) -> Any:
+        """Return value for key; if absent create via factory and store.
+
+        Key is not parsed; dotted forms are allowed but treated as a single
+        map key. Use consistent prefixes for grouping (e.g. 'spu.config').
+        """
+        if not key:
+            raise ValueError("empty state key")
+        val = self.state.get(key)
+        if val is None:
+            val = factory()
+            self.state[key] = val
+        return val
+
+    def get_state(self, key: str, default: Any | None = None) -> Any:
+        if not key:
+            raise ValueError("empty state key")
+        return self.state.get(key, default)
+
+    def set_state(self, key: str, value: Any) -> None:
+        if not key:
+            raise ValueError("empty state key")
+        self.state[key] = value
+
+    def del_state(self, key: str) -> None:
+        if not key:
+            raise ValueError("empty state key")
+        self.state.pop(key, None)
+
+    def list_state(self, prefix: str = "") -> dict[str, Any]:
+        """Return mapping of key -> value; optional prefix filter.
+
+        Prefix match is string-based; if prefix is non-empty include keys
+        where key == prefix or key starts with prefix + '.'.
+        """
+        if not prefix:
+            return dict(self.state)
+        pref = prefix if prefix.endswith(".") else prefix + "."
+        out: dict[str, Any] = {}
+        for k, v in self.state.items():
+            if k == prefix or k.startswith(pref):
+                out[k] = v
+        return out
 
     # ---- explicit (re)binding API ----
     def bind_op(self, op_type: str, kernel_id: str, *, force: bool = False) -> None:

{mplang_nightly-0.1.dev155 → mplang_nightly-0.1.dev157}/mplang/kernels/crypto.py

@@ -27,15 +27,19 @@ __all__: list[str] = []  # flat kernels only
 
 
 def _get_rng() -> np.random.Generator:
-    """Get (and lazily create) per-rank RNG for crypto kernels."""
+    """Get (and lazily create) per-rank RNG for crypto kernels.
+
+    Runtime state is untyped, so we narrow the type explicitly for mypy.
+    """
     kctx = cur_kctx()
-    pocket = kctx.state.setdefault("crypto", {})
-    rng = pocket.get("rng")
-    if rng is None:
+    rt = kctx.runtime
+    rng_obj = rt.get_state("crypto.rng")
+    if rng_obj is None:
         seed = int(os.environ.get("MPLANG_CRYPTO_SEED", "0")) + kctx.rank * 7919
-        rng = np.random.default_rng(seed)
-        pocket["rng"] = rng
-    return rng
+        rng_obj = np.random.default_rng(seed)
+        rt.set_state("crypto.rng", rng_obj)
+    assert isinstance(rng_obj, np.random.Generator)  # narrow
+    return rng_obj
 
 
 def _keystream(key: bytes, nonce: bytes, length: int) -> bytes:

{mplang_nightly-0.1.dev155 → mplang_nightly-0.1.dev157}/mplang/kernels/mock_tee.py

@@ -28,12 +28,13 @@ __all__: list[str] = []
 
 def _rng() -> np.random.Generator:
     kctx = cur_kctx()
-    pocket = kctx.state.setdefault("tee", {})
-    r = pocket.get("rng")
+    rt = kctx.runtime
+    r = rt.get_state("tee.rng")
     if r is None:
         seed = int(os.environ.get("MPLANG_TEE_SEED", "0")) + kctx.rank * 10007
         r = np.random.default_rng(seed)
-        pocket["rng"] = r
+        rt.set_state("tee.rng", r)
+    assert isinstance(r, np.random.Generator)  # type narrowing for mypy
     return r
 
 
@@ -44,10 +45,10 @@ def _quote_from_pk(pk: np.ndarray) -> NDArray[np.uint8]:
     return out
 
 
-@kernel_def("mock_tee.quote")
-def _tee_quote(pfunc: PFunction, pk: object) -> NDArray[np.uint8]:
+@kernel_def("mock_tee.quote_gen")
+def _tee_quote_gen(pfunc: PFunction, pk: object) -> NDArray[np.uint8]:
     warnings.warn(
-        "Insecure mock TEE kernel 'mock_tee.quote' in use. NOT secure; for local testing only.",
+        "Insecure mock TEE kernel 'mock_tee.quote_gen' in use. NOT secure; for local testing only.",
         stacklevel=3,
     )
     pk = np.asarray(pk, dtype=np.uint8)
@@ -63,6 +64,10 @@ def _tee_attest(pfunc: PFunction, quote: object) -> NDArray[np.uint8]:
         stacklevel=3,
     )
     quote = np.asarray(quote, dtype=np.uint8)
+    platform = pfunc.attrs.get("platform")
+    if platform is None:
+        raise ValueError("missing required 'platform' attribute in PFunction")
+
     if quote.size != 33:
         raise ValueError("mock quote must be 33 bytes (1 header + 32 pk)")
     return quote[1:33].astype(np.uint8)

{mplang_nightly-0.1.dev155 → mplang_nightly-0.1.dev157}/mplang/kernels/spu.py

@@ -63,14 +63,10 @@ class SpuValue:
         return f"SpuValue({self.shape},{self.dtype},{self.vtype})"
 
 
-def _get_spu_pocket() -> dict[str, Any]:
-    return cur_kctx().state.setdefault("spu", {})
-
-
 def _get_spu_config_and_world() -> tuple[libspu.RuntimeConfig, int]:
-    pocket = _get_spu_pocket()
-    cfg = pocket.get("config")
-    world = pocket.get("world")
+    kctx = cur_kctx()
+    cfg = kctx.runtime.get_state("spu.config")
+    world = kctx.runtime.get_state("spu.world")
     if cfg is None or world is None:
         raise RuntimeError("SPU kernel state not initialized (config/world)")
     return cfg, int(world)
@@ -84,12 +80,12 @@ def _register_spu_env(
     Idempotent: if config/world already set, they must match; link is recorded per rank.
     This replaces previous global fallback seeding logic.
     """
-    pocket = _get_spu_pocket()
-    prev_cfg = pocket.get("config")
-    prev_world = pocket.get("world")
+    kctx = cur_kctx()
+    prev_cfg = kctx.runtime.get_state("spu.config")
+    prev_world = kctx.runtime.get_state("spu.world")
     if prev_cfg is None:
-        pocket["config"] = config
-        pocket["world"] = world_size
+        kctx.runtime.set_state("spu.config", config)
+        kctx.runtime.set_state("spu.world", world_size)
     else:
         # libspu RuntimeConfig may not implement __eq__; compare serialized repr
         same_cfg = (
@@ -102,7 +98,7 @@ def _register_spu_env(
             raise RuntimeError("Conflicting SPU env registration")
     # Store single link per runtime (one runtime per rank)
     if link_ctx is not None:
-        pocket["link"] = link_ctx
+        kctx.runtime.set_state("spu.link", link_ctx)
 
 
 @kernel_def("spu.seed_env")
@@ -197,16 +193,16 @@ def _spu_run_mlir(pfunc: PFunction, *args: Any) -> Any:
         )
 
     cfg, _ = _get_spu_config_and_world()
-    pocket = _get_spu_pocket()
-    link_ctx: LinkCommunicator | None = pocket.get("link")
+    kctx = cur_kctx()
+    link_ctx = kctx.runtime.get_state("spu.link")
     if link_ctx is None:
         raise RuntimeError("Rank not participating in SPU; no link set via seed_env")
 
-    # Lazy runtime cache
-    spu_rt = pocket.get("runtime")
+    # Lazy runtime cache under key spu.runtime
+    spu_rt = kctx.runtime.get_state("spu.runtime")
     if spu_rt is None:
         spu_rt = spu_api.Runtime(link_ctx.get_lctx(), cfg)
-        pocket["runtime"] = spu_rt
+        kctx.runtime.set_state("spu.runtime", spu_rt)
 
     # Validate that all inputs are SpuValue objects
     for i, arg in enumerate(args):

{mplang_nightly-0.1.dev155 → mplang_nightly-0.1.dev157}/mplang/kernels/stablehlo.py

@@ -36,11 +36,14 @@ def _stablehlo_exec(pfunc: PFunction, *args: Any) -> Any:
     if isinstance(mlir_text, bytes):
         mlir_text = mlir_text.decode("utf-8")
 
-    # Simple compile cache per runtime (state pocket per backend namespace)
+    # Flat-key compile cache: stablehlo.compile_cache.<hash>
     ctx = cur_kctx()
-    pocket = ctx.state.setdefault("stablehlo", {})
-    cache = pocket.setdefault("compile_cache", {})
-    compiled = cache.get(mlir_text)
+    rt = ctx.runtime
+    import hashlib
+
+    h = hashlib.sha256(mlir_text.encode("utf-8")).hexdigest()[:16]
+    key = f"stablehlo.compile_cache.{h}"
+    compiled = rt.get_state(key)
     if compiled is None:
         backend = jax.default_backend()
         client = xla_bridge.get_backend(backend)
@@ -49,7 +52,7 @@ def _stablehlo_exec(pfunc: PFunction, *args: Any) -> Any:
             compiled = client.compile(mlir_text, compile_options)
         except Exception as e:  # pragma: no cover
             raise RuntimeError(f"StableHLO compile failed: {e}") from e
-        cache[mlir_text] = compiled
+        rt.set_state(key, compiled)
 
     # Handle JAX's unused parameter elimination via arg_keep_map
     runtime_args = args

{mplang_nightly-0.1.dev155 → mplang_nightly-0.1.dev157}/mplang/ops/tee.py

@@ -14,7 +14,11 @@
 
 from __future__ import annotations
 
+from jax.tree_util import PyTreeDef, tree_flatten
+
 from mplang.core.dtype import UINT8
+from mplang.core.mpobject import MPObject
+from mplang.core.pfunc import PFunction
 from mplang.core.tensor import TensorType
 from mplang.ops.base import stateless_mod
 
@@ -22,21 +26,26 @@ _TEE_MOD = stateless_mod("tee")
 
 
 @_TEE_MOD.simple_op()
-def quote(pk: TensorType) -> TensorType:
-    """TEE quote generation binding the provided ephemeral public key.
-
-    API (mock): quote(pk: u8[32]) -> (quote: u8[33])
-    The mock encodes a 1-byte header + 32-byte pk.
-    """
+def quote_gen(pk: TensorType) -> TensorType:
+    """TEE quote generation binding the provided ephemeral public key."""
     _ = pk  # Mark as used for the decorator
-    return TensorType(UINT8, (33,))
-
-
-@_TEE_MOD.simple_op()
-def attest(quote: TensorType) -> TensorType:
-    """TEE quote verification returning the attested TEE public key.
-
-    API (mock): attest(quote: u8[33]) -> tee_pk: u8[32]
-    """
-    _ = quote  # Mark as used for the decorator
-    return TensorType(UINT8, (32,))
+    return TensorType(UINT8, (-1,))
+
+
+@_TEE_MOD.op_def()
+def attest(
+    quote: MPObject, platform: str
+) -> tuple[PFunction, list[MPObject], PyTreeDef]:
+    """TEE quote verification returning the attested TEE public key."""
+
+    ins_info = [TensorType.from_obj(quote)]
+    outs_info = [TensorType(UINT8, (32,))]  # pk is always 32 bytes for x25519
+    pfunc = PFunction(
+        fn_type="tee.attest",
+        ins_info=ins_info,
+        outs_info=outs_info,
+        platform=platform,
+    )
+    _, treedef = tree_flatten(outs_info[0])
+
+    return pfunc, [quote], treedef

{mplang_nightly-0.1.dev155 → mplang_nightly-0.1.dev157}/mplang/runtime/data_providers.py

@@ -173,38 +173,32 @@ class FileProvider(DataProvider):
         np.save(path, np.asarray(value))
 
 
-class _KeyedPocket:
-    """Small helper to keep a dict in KernelContext.state under a namespaced key."""
-
-    def __init__(self, ns: str):
-        self.ns = ns
-
-    def get_map(self, ctx: KernelContext) -> dict[str, Any]:
-        pocket = ctx.state.setdefault("resource.providers", {})
-        store = pocket.get(self.ns)
-        if store is None:
-            store = {}
-            pocket[self.ns] = store
-        return store  # type: ignore[return-value]
-
-
 class MemProvider(DataProvider):
     """In-memory per-runtime KV provider (per rank, per session/runtime)."""
 
-    def __init__(self) -> None:
-        self._pocket = _KeyedPocket("mem")
+    STATE_KEY = "resource.providers.mem"
+
+    @staticmethod
+    def _store(ctx: KernelContext) -> dict[str, Any]:
+        # Use ensure_state so creation is atomic & centralized; enforce dict.
+        store = ctx.runtime.ensure_state(MemProvider.STATE_KEY, dict)
+        if not isinstance(store, dict):  # pragma: no cover - defensive
+            raise TypeError(
+                f"runtime state key '{MemProvider.STATE_KEY}' expected dict, got {type(store).__name__}"
+            )
+        return store  # type: ignore[return-value]
 
     def read(
         self, uri: ResolvedURI, out_spec: TensorType | TableType, *, ctx: KernelContext
     ) -> Any:
-        store = self._pocket.get_map(ctx)
+        store = self._store(ctx)
         key = uri.raw
         if key not in store:
             raise FileNotFoundError(f"mem resource not found: {key}")
         return store[key]
 
     def write(self, uri: ResolvedURI, value: Any, *, ctx: KernelContext) -> None:
-        store = self._pocket.get_map(ctx)
+        store = self._store(ctx)
         store[uri.raw] = value
 
 

{mplang_nightly-0.1.dev155 → mplang_nightly-0.1.dev157}/tests/ops/test_crypto_tee.py

@@ -25,14 +25,14 @@ def _demo_flow():
     # TEE generates two ephemeral keypairs and quotes binding their pk
     t_sk0, t_pk0 = simp.runAt(P2, crypto.kem_keygen)("x25519")
     t_sk1, t_pk1 = simp.runAt(P2, crypto.kem_keygen)("x25519")
-    q0 = simp.runAt(P2, tee.quote)(t_pk0)
-    q1 = simp.runAt(P2, tee.quote)(t_pk1)
+    q0 = simp.runAt(P2, tee.quote_gen)(t_pk0)
+    q1 = simp.runAt(P2, tee.quote_gen)(t_pk1)
 
     # Send quotes to P0/P1 and attest to obtain TEE public keys
     q0_for_p0 = simp.p2p(P2, P0, q0)
     q1_for_p1 = simp.p2p(P2, P1, q1)
-    t_pk0_for_p0 = simp.runAt(P0, tee.attest)(q0_for_p0)
-    t_pk1_for_p1 = simp.runAt(P1, tee.attest)(q1_for_p1)
+    t_pk0_for_p0 = simp.runAt(P0, tee.attest)(q0_for_p0, "TDX")
+    t_pk1_for_p1 = simp.runAt(P1, tee.attest)(q1_for_p1, "TDX")
 
     # Each party generates its own ephemeral keypair and shares pk with TEE
     v_sk0, v_pk0 = simp.runAt(P0, crypto.kem_keygen)("x25519")
@@ -75,7 +75,10 @@ def _demo_flow():
 
 def test_crypto_enc_dec_and_tee_quote_attest_roundtrip():
     # Create simulator with TEE bindings using the new initial_bindings parameter
-    tee_bindings = {"tee.quote": "mock_tee.quote", "tee.attest": "mock_tee.attest"}
+    tee_bindings = {
+        "tee.quote_gen": "mock_tee.quote_gen",
+        "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)

{mplang_nightly-0.1.dev155 → mplang_nightly-0.1.dev157}/tests/simp/test_sugar.py

@@ -35,7 +35,7 @@ def test_basic_callable_and_namespace():
         a, b = simp.P0(crypto.kem_keygen, "x25519")
         # namespace form (tee side key, then quote)
         t_sk, t_pk = simp.P[2].crypto.kem_keygen("x25519")
-        _ = simp.P[2].tee.quote(t_pk)
+        _ = simp.P[2].tee.quote_gen(t_pk)
         # derive something simple at party 0 to ensure run path works
         _ = simp.P0(lambda x: x + 1, 41)
         return a, b, t_sk, t_pk

{mplang_nightly-0.1.dev155 → mplang_nightly-0.1.dev157}/tutorials/3_device.py

@@ -101,7 +101,10 @@ def run_tee():
     print("-" * 10, "millionaire (TEE)", "-" * 10)
 
     # TEE operations need explicit binding for security
-    tee_bindings = {"tee.quote": "mock_tee.quote", "tee.attest": "mock_tee.attest"}
+    tee_bindings = {
+        "tee.quote_gen": "mock_tee.quote_gen",
+        "tee.attest": "mock_tee.attest",
+    }
     # Apply tee bindings across nodes before constructing simulator
     for n in cluster_spec.nodes.values():
         n.runtime_info.op_bindings.update(tee_bindings)

{mplang_nightly-0.1.dev155 → mplang_nightly-0.1.dev157}/tutorials/9_tee.py

@@ -46,7 +46,11 @@ cluster_spec = ClusterSpec.from_dict({
         },
         "P0": {"kind": "PPU", "members": ["node_0"], "config": {}},
         "P1": {"kind": "PPU", "members": ["node_1"], "config": {}},
-        "TEE0": {"kind": "TEE", "members": ["node_2"], "config": {}},
+        "TEE0": {
+            "kind": "TEE",
+            "members": ["node_2"],
+            "config": {"platform": "TDX"},
+        },
     },
 })
 
@@ -72,8 +76,8 @@ def millionaire_manual():
 
     # P0 <-> TEE handshake and transfer x (using sugar)
     tee_sk0, tee_pk0 = P2.crypto.kem_keygen("x25519")
-    quote0 = P2.tee.quote(tee_pk0)
-    tee_pk0_at_p0 = P0.tee.attest(P2P(P2, P0, quote0))
+    quote0 = P2.tee.quote_gen(tee_pk0)
+    tee_pk0_at_p0 = P0.tee.attest(P2P(P2, P0, quote0), "TDX")
     v_sk0, v_pk0 = P0.crypto.kem_keygen("x25519")
     shared0_p = P0.crypto.kem_derive(v_sk0, tee_pk0_at_p0, "x25519")
     shared0_t = P2.crypto.kem_derive(tee_sk0, P2P(P0, P2, v_pk0), "x25519")
@@ -88,8 +92,8 @@ def millionaire_manual():
 
     # P1 <-> TEE handshake and transfer y (still show original style for contrast)
     tee_sk1, tee_pk1 = P2.crypto.kem_keygen("x25519")
-    quote1 = P2.tee.quote(tee_pk1)
-    tee_pk1_at_p1 = P1.tee.attest(P2P(P2, P1, quote1))
+    quote1 = P2.tee.quote_gen(tee_pk1)
+    tee_pk1_at_p1 = P1.tee.attest(P2P(P2, P1, quote1), "TDX")
     v_sk1, v_pk1 = P1.crypto.kem_keygen("x25519")
     shared1_p = P1.crypto.kem_derive(v_sk1, tee_pk1_at_p1, "x25519")
     shared1_t = P2.crypto.kem_derive(tee_sk1, P2P(P1, P2, v_pk1), "x25519")
@@ -117,7 +121,10 @@ def millionaire_manual():
 def main():
     print("-" * 10, "TEE millionaire: device vs manual (end-to-end IR)", "-" * 10)
     # Create simulator with TEE bindings
-    tee_bindings = {"tee.quote": "mock_tee.quote", "tee.attest": "mock_tee.attest"}
+    tee_bindings = {
+        "tee.quote_gen": "mock_tee.quote_gen",
+        "tee.attest": "mock_tee.attest",
+    }
     # Apply tee_bindings per-node (preferred) then construct Simulator
     for n in cluster_spec.nodes.values():
         n.runtime_info.op_bindings.update(tee_bindings)