mplang-nightly 0.1.dev162__tar.gz → 0.1.dev164__tar.gz

{mplang_nightly-0.1.dev162 → mplang_nightly-0.1.dev164}/PKG-INFO

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

{mplang_nightly-0.1.dev162 → mplang_nightly-0.1.dev164}/mplang/device.py

@@ -207,15 +207,8 @@ 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, platform
-        )
+        sess_p, sess_t = _ensure_tee_session(frm_dev_id, to_dev_id, frm_rank, tee_rank)
         # 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)
@@ -229,15 +222,8 @@ 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 = frm_dev.config.get("platform")
-        if not platform:
-            raise ValueError(
-                f"TEE device '{frm_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, platform
-        )
+        sess_p, sess_t = _ensure_tee_session(to_dev_id, frm_dev_id, ppu_rank, tee_rank)
         obj_ty = TensorType.from_obj(obj)
         b = simp.runAt(tee_rank, builtin.pack)(obj)
         ct = simp.runAt(tee_rank, crypto.enc)(b, sess_t)
@@ -259,7 +245,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, platform: str
+    frm_dev_id: str, to_dev_id: str, frm_rank: int, tee_rank: int
 ) -> tuple[MPObject, MPObject]:
     """Ensure a TEE session (sess_p at sender, sess_t at TEE) exists.
 
@@ -281,7 +267,7 @@ def _ensure_tee_session(
 
     # 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, platform)
+    tee_pk_at_sender = simp.runAt(frm_rank, tee.attest)(quote_at_sender)
 
     # 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.dev162 → mplang_nightly-0.1.dev164}/mplang/kernels/mock_tee.py

@@ -64,9 +64,6 @@ 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)")

{mplang_nightly-0.1.dev162 → mplang_nightly-0.1.dev164}/mplang/ops/tee.py

@@ -14,11 +14,7 @@
 
 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
 
@@ -32,20 +28,10 @@ def quote_gen(pk: TensorType) -> TensorType:
     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
+@_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,))

{mplang_nightly-0.1.dev162 → mplang_nightly-0.1.dev164}/tests/ops/test_crypto_tee.py

@@ -31,8 +31,8 @@ def _demo_flow():
     # 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, "TDX")
-    t_pk1_for_p1 = simp.runAt(P1, tee.attest)(q1_for_p1, "TDX")
+    t_pk0_for_p0 = simp.runAt(P0, tee.attest)(q0_for_p0)
+    t_pk1_for_p1 = simp.runAt(P1, tee.attest)(q1_for_p1)
 
     # Each party generates its own ephemeral keypair and shares pk with TEE
     v_sk0, v_pk0 = simp.runAt(P0, crypto.kem_keygen)("x25519")

{mplang_nightly-0.1.dev162 → mplang_nightly-0.1.dev164}/tutorials/9_tee.py

@@ -49,7 +49,6 @@ cluster_spec = ClusterSpec.from_dict({
         "TEE0": {
             "kind": "TEE",
             "members": ["node_2"],
-            "config": {"platform": "TDX"},
         },
     },
 })
@@ -77,7 +76,7 @@ def millionaire_manual():
     # P0 <-> TEE handshake and transfer x (using sugar)
     tee_sk0, tee_pk0 = P2.crypto.kem_keygen("x25519")
     quote0 = P2.tee.quote_gen(tee_pk0)
-    tee_pk0_at_p0 = P0.tee.attest(P2P(P2, P0, quote0), "TDX")
+    tee_pk0_at_p0 = P0.tee.attest(P2P(P2, P0, quote0))
     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")
@@ -93,7 +92,7 @@ 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_gen(tee_pk1)
-    tee_pk1_at_p1 = P1.tee.attest(P2P(P2, P1, quote1), "TDX")
+    tee_pk1_at_p1 = P1.tee.attest(P2P(P2, P1, quote1))
     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")