mplang-nightly 0.1.dev256__tar.gz → 0.1.dev258__tar.gz

{mplang_nightly-0.1.dev256 → mplang_nightly-0.1.dev258}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev256
+Version: 0.1.dev258
 Summary: Multi-Party Programming Language
 Author-email: SecretFlow Team <secretflow-contact@service.alipay.com>
 License:                                  Apache License
@@ -219,7 +219,7 @@ Requires-Dist: pandas>=2.0.0
 Requires-Dist: protobuf<6.0,>=5.0
 Requires-Dist: pyarrow>=14.0.0
 Requires-Dist: pyyaml>=6.0
-Requires-Dist: spu>=0.10.0.dev20251211
+Requires-Dist: spu>=0.10.0.dev20251208
 Requires-Dist: sqlglot>=23.0.0
 Requires-Dist: tenseal==0.3.16
 Requires-Dist: typing-extensions

{mplang_nightly-0.1.dev256 → mplang_nightly-0.1.dev258}/mplang/v2/backends/simp_worker/ops.py

@@ -47,18 +47,15 @@ def _pcall_static_worker_impl(
 
     if worker.rank in parties:
         fn_graph = op.regions[0]
-        prev_parties = getattr(interpreter, "current_parties", None)
-        interpreter.current_parties = parties  # type: ignore[attr-defined]
+        prev_parties = worker.current_parties
+        worker.current_parties = parties
 
         try:
             result = interpreter.evaluate_graph(fn_graph, list(args))
             # Return single value for single output (interpreter expects this)
             return result[0] if len(op.outputs) == 1 else result
         finally:
-            if prev_parties is None:
-                del interpreter.current_parties  # type: ignore[attr-defined]
-            else:
-                interpreter.current_parties = prev_parties  # type: ignore[attr-defined]
+            worker.current_parties = prev_parties
     else:
         # No data for this rank
         return None if len(op.outputs) == 1 else [None] * len(op.outputs)

{mplang_nightly-0.1.dev256 → mplang_nightly-0.1.dev258}/mplang/v2/backends/spu_impl.py

@@ -164,26 +164,19 @@ def exec_impl(interpreter: Interpreter, op: Operation, *args: Any) -> Any:
     # Get SPU config from attrs (passed through from run_jax)
     config: spu.SPUConfig = op.attrs["config"]
 
-    # Get parties from interpreter context (injected by pcall_static_impl)
-    parties = getattr(interpreter, "current_parties", None)
+    # Get parties from SimpWorker state (injected by pcall_static_impl)
+    context = interpreter.get_dialect_state("simp")
+    if not isinstance(context, SimpWorker):
+        raise RuntimeError(f"spu.exec requires SimpWorker, got {type(context)}")
+
+    parties = context.current_parties
     if parties is None:
         raise RuntimeError(
-            "spu.exec requires 'current_parties' in interpreter context. "
+            "spu.exec requires 'current_parties' in SimpWorker state. "
             "Ensure it is called within a pcall_static block."
         )
 
-    # Get global rank from interpreter or its context
-    # Use SimpWorker if available
-    context = interpreter.get_dialect_state("simp")
-    if isinstance(context, SimpWorker):
-        global_rank = context.rank
-    else:
-        # Fallback for other contexts or direct interpreter usage?
-        # User said: "directly ensure simp_context is there"
-        # If not SimpWorker, we can't run spu.exec?
-        # But maybe integration tests run differently?
-        # Let's trust user: "ensure simp_context is there"
-        raise RuntimeError(f"spu.exec requires SimpWorker, got {type(context)}")
+    global_rank = context.rank
 
     if global_rank not in parties:
         raise RuntimeError(

{mplang_nightly-0.1.dev256 → mplang_nightly-0.1.dev258}/pyproject.toml

@@ -24,7 +24,7 @@ dependencies = [
     "protobuf>=5.0,<6.0",
     "typing_extensions",
     # both FE & BE
-    "spu>=0.10.0.dev20251211",
+    "spu>=0.10.0.dev20251208",
     "numpy>=2.0.0",            # Updated for JAX 0.8.0 compatibility
     "jax[cpu]==0.8.0",         # Align with SPU JAX version
     "flax>=0.12.0",            # NNX is part of Flax 0.12+