mplang-nightly 0.1.dev156__py3-none-any.whl → 0.1.dev158__py3-none-any.whl

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 = 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)
+        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/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",
 }
 

mplang/kernels/mock_tee.py

@@ -45,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)
@@ -64,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/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.dev156.dist-info → mplang_nightly-0.1.dev158.dist-info}/METADATA

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

{mplang_nightly-0.1.dev156.dist-info → mplang_nightly-0.1.dev158.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=RmjnhzHxJkkNmtBKtYMEbpQYBZpuC43qlllkCOp-QD8,12548
+mplang/device.py,sha256=UxfsgaGrJaQyAOzrQBWYArqU6CHIvyONWxgAJHXVQ_0,13041
 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
@@ -29,9 +29,9 @@ mplang/core/expr/walk.py,sha256=lXkGJEEuvKGDqQihbxXPxfz2RfR1Q1zYUlt11iooQW0,1188
 mplang/kernels/__init__.py,sha256=2WE4cmW96Xkzyq2yRRYNww4kZ5o6u6NbPV0BxqZG698,581
 mplang/kernels/base.py,sha256=-YV4Aj5fs6GT4ehS6Tyi8WQ-amxn5edHTFJRQzyjHXY,3826
 mplang/kernels/builtin.py,sha256=nSuM79cn7M6M27A6Y8ycilXT_qAlB1ktkwkRX6dv_VQ,7052
-mplang/kernels/context.py,sha256=yJjQUyQmBBl6btLb7KlmA-Ejf4-cgK3KmtC-m0sBbb8,13594
+mplang/kernels/context.py,sha256=n-Z7fz7HjHb3UY380iZcasmn2sK-OQUGEIWJk2-fT18,13602
 mplang/kernels/crypto.py,sha256=s7R0yd4Fk5cI2Qd3LpLc-kmbVuk8fFsbKbfKi43R0aE,3892
-mplang/kernels/mock_tee.py,sha256=173QSzPgkrLo0zn0jsx6nNmq1WvfUhIM67FM4Dn30aA,2297
+mplang/kernels/mock_tee.py,sha256=wzXgiCpDWqQFS1HJHa5z_GxRJM7-r5l6T7KBpNcIZlg,2457
 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=gwzP81y2idH-d-Du84H6oNZpLaGD-3fEgm8G1uxWpUA,1388
+mplang/ops/tee.py,sha256=1yoaFFF5NI9gJge1-bhbT2lsphjBXErfDsleYlEMoWs,1664
 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.dev156.dist-info/METADATA,sha256=iaTA5SM0ALRdvM5h40B44CV8WMxsMQXP6DqsPrqeCY4,16547
-mplang_nightly-0.1.dev156.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-mplang_nightly-0.1.dev156.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
-mplang_nightly-0.1.dev156.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mplang_nightly-0.1.dev156.dist-info/RECORD,,
+mplang_nightly-0.1.dev158.dist-info/METADATA,sha256=LCrLu9AfQJfLUeapTmp0dAVGjYUjl92YfGLJh4WXKXQ,16547
+mplang_nightly-0.1.dev158.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+mplang_nightly-0.1.dev158.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
+mplang_nightly-0.1.dev158.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mplang_nightly-0.1.dev158.dist-info/RECORD,,