mplang-nightly 0.1.dev163__py3-none-any.whl → 0.1.dev164__py3-none-any.whl

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/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/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.dev163.dist-info → mplang_nightly-0.1.dev164.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev163
+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.dev163.dist-info → mplang_nightly-0.1.dev164.dist-info}/RECORD

@@ -1,6 +1,6 @@
 mplang/__init__.py,sha256=ofO-F-CNoVIxpMpTJtTJoQtKegJcHwcOJLzoVispiyc,1852
 mplang/api.py,sha256=ssmv0_CyZPFORhOUJ84Jo6NwRJSK7_Ono3n7ZjEg4sA,3058
-mplang/device.py,sha256=UxfsgaGrJaQyAOzrQBWYArqU6CHIvyONWxgAJHXVQ_0,13041
+mplang/device.py,sha256=oFdu_XGaJ7AaSePCg1IEQmt6pVoJ3YWHzFyFX4UQBGM,12552
 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
@@ -31,7 +31,7 @@ mplang/kernels/base.py,sha256=-YV4Aj5fs6GT4ehS6Tyi8WQ-amxn5edHTFJRQzyjHXY,3826
 mplang/kernels/builtin.py,sha256=nSuM79cn7M6M27A6Y8ycilXT_qAlB1ktkwkRX6dv_VQ,7052
 mplang/kernels/context.py,sha256=n-Z7fz7HjHb3UY380iZcasmn2sK-OQUGEIWJk2-fT18,13602
 mplang/kernels/crypto.py,sha256=s7R0yd4Fk5cI2Qd3LpLc-kmbVuk8fFsbKbfKi43R0aE,3892
-mplang/kernels/mock_tee.py,sha256=wzXgiCpDWqQFS1HJHa5z_GxRJM7-r5l6T7KBpNcIZlg,2457
+mplang/kernels/mock_tee.py,sha256=N8Ft50sD0t1V3-NfkGfQxYa6N5nJcz7uuwXEVj1F_mc,2310
 mplang/kernels/phe.py,sha256=8-_1IFPOaGECGj9mbYja8XoqbMYnYqfpDNVyMJd8J1Y,65247
 mplang/kernels/spu.py,sha256=LkM8tNzhwTa8lufNgClHfnI4LNu25cdWLQZdJsMDEO8,9301
 mplang/kernels/sql_duckdb.py,sha256=UN1Ev6-MxF_-65zMExUsLScC9PlmEIEcN8YziIoX_rY,1724
@@ -45,7 +45,7 @@ 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=HyY2i5aGC5W7r62JryFSjQCUDXH3kQz82YADwn4z5uc,2015
-mplang/ops/tee.py,sha256=1yoaFFF5NI9gJge1-bhbT2lsphjBXErfDsleYlEMoWs,1664
+mplang/ops/tee.py,sha256=bOpS_BXG12D6bONikzdF2yt0oVZj9Jyd0g_3IXP8VgE,1281
 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
@@ -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.dev163.dist-info/METADATA,sha256=2JegY_nYHjbyCsl7Tl-0afJb7ufKl7SoMfvz_9E3v50,16547
-mplang_nightly-0.1.dev163.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-mplang_nightly-0.1.dev163.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
-mplang_nightly-0.1.dev163.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mplang_nightly-0.1.dev163.dist-info/RECORD,,
+mplang_nightly-0.1.dev164.dist-info/METADATA,sha256=7TEhvNRPHlSW4___A1ctr8Jo4fM36TiDKx-9JSQsjWE,16547
+mplang_nightly-0.1.dev164.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+mplang_nightly-0.1.dev164.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
+mplang_nightly-0.1.dev164.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mplang_nightly-0.1.dev164.dist-info/RECORD,,