mplang-nightly 0.1.dev180__tar.gz → 0.1.dev182__tar.gz

{mplang_nightly-0.1.dev180 → mplang_nightly-0.1.dev182}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev180
+Version: 0.1.dev182
 Summary: Multi-Party Programming Language
 Author-email: SecretFlow Team <secretflow-contact@service.alipay.com>
 License:                                  Apache License
@@ -277,7 +277,7 @@ sim = mplang.Simulator.simple(2)
 result = mplang.evaluate(sim, millionaire)
 
 # Securely fetch the result (reveals SPU value)
-print("Is Alice poorer than Bob?", mpd.fetch(sim, result))
+print("Is Alice poorer than Bob?", mplang.fetch(sim, result))
 ```
 
 ## Learn More

{mplang_nightly-0.1.dev180 → mplang_nightly-0.1.dev182}/README.md

@@ -57,7 +57,7 @@ sim = mplang.Simulator.simple(2)
 result = mplang.evaluate(sim, millionaire)
 
 # Securely fetch the result (reveals SPU value)
-print("Is Alice poorer than Bob?", mpd.fetch(sim, result))
+print("Is Alice poorer than Bob?", mplang.fetch(sim, result))
 ```
 
 ## Learn More

{mplang_nightly-0.1.dev180 → mplang_nightly-0.1.dev182}/mplang/core/primitive.py

@@ -718,7 +718,7 @@ def pshfl(src: MPObject, index: MPObject) -> MPObject:
 
     Raises:
         ValueError: If the index tensor is not a scalar.
-        RuntimeError: If src[index[i]] is None for any valid index[i] (i.e/,
+        RuntimeError: If src[index[i]] is None for any valid index[i] (i.e.,
                      trying to fetch from a party that doesn't hold the data).
 
     Examples:

{mplang_nightly-0.1.dev180 → mplang_nightly-0.1.dev182}/mplang/host.py

@@ -38,6 +38,16 @@ def evaluate(
 
     This function accepts arbitrary types as it's designed to handle
     any multi-party computation function and arguments.
+
+    Args:
+        interp: The interpreter context for evaluating the multi-party function.
+        mpfn: The multi-party function to evaluate.
+        *args: Positional arguments to pass to the function.
+        **kwargs: Keyword arguments to pass to the function.
+
+    Returns:
+        Any: The result of evaluating the multi-party function, which can be
+             any type depending on the function's return type.
     """
     assert isinstance(interp, InterpContext), f"Expect InterpContext, got {interp}"
     with with_ctx(interp):
@@ -49,6 +59,16 @@ def fetch(interp: InterpContext | None, objs: Any) -> Any:  # type: ignore[misc]
 
     This function uses tree_map to handle arbitrary nested structures,
     so it needs to accept and return Any type.
+
+    Args:
+        interp: The interpreter context for fetching results. If None, uses the
+                current context from cur_ctx().
+        objs: The objects containing MPObject instances to fetch. Can be any
+              nested structure.
+
+    Returns:
+        Any: The fetched results with the same structure as the input objects,
+             but with MPObject instances replaced by their computed values.
     """
     ctx = interp or cur_ctx()
     assert isinstance(ctx, InterpContext), f"Expect MPExecutor, got {ctx}"
@@ -94,5 +114,17 @@ class CompileOptions(MPContext):
 def compile(
     mctx: MPContext, fn: Callable[..., Any], *args: Any, **kwargs: Any
 ) -> TracedFunction:
+    """Compile a multi-party function into a TracedFunction.
+
+    Args:
+        mctx: The multi-party context for compilation.
+        fn: The function to compile.
+        *args: Positional arguments to pass during compilation.
+        **kwargs: Keyword arguments to pass during compilation.
+
+    Returns:
+        TracedFunction: The compiled function representation that can be
+                       evaluated in multi-party contexts.
+    """
     trace_ctx = TraceContext(mctx.cluster_spec)
     return trace(trace_ctx, fn, *args, **kwargs)

{mplang_nightly-0.1.dev180 → mplang_nightly-0.1.dev182}/mplang/simp/mpi.py

@@ -57,9 +57,9 @@ def gather_m(src_mask: Mask, root: Rank, arg: MPObject) -> list[MPObject]:
     """Gather the object from pmask'ed parties to the root party.
 
     Args:
-        src_pmask: The mask of the parties that will gather the object.
+        src_mask: The mask of the parties that will gather the object.
         root: The rank of the root party.
-        arg: The object to be gathered, which must be the subset of pmask.
+        arg: The object to be gathered. It must be held by all parties specified in `src_mask`.
 
     Returns:
         A list of objects, with length equal to the number of parties in pmask.

{mplang_nightly-0.1.dev180 → mplang_nightly-0.1.dev182}/tests/runtime/test_communicator.py

@@ -150,7 +150,9 @@ def test_communicator_properties(http_servers):  # noqa: F811
     assert id1 != id2  # Should generate unique IDs
 
 
-def run_party_e2e_process(rank: int, return_dict: dict, assigned_ports: dict):
+def run_party_e2e_process(
+    rank: int, return_dict: dict, assigned_ports: dict, barrier: multiprocessing.Barrier
+):
     """
     Run a complete party process with server and communication logic.
     This is the proper single-process-per-party architecture.
@@ -247,6 +249,8 @@ def run_party_e2e_process(rank: int, return_dict: dict, assigned_ports: dict):
         if not server_ready:
             raise RuntimeError("Server failed to start within timeout")
 
+        barrier.wait(timeout=30)  # Synchronize with other parties
+
         # Run party-specific communication logic
         if rank == 0:
             # Party 0: Send message to Party 1
@@ -309,15 +313,21 @@ def test_end_to_end_communication():
         assigned_ports[0] = p0
         assigned_ports[1] = p1
 
+        worlds = [0, 1]
+        barrier = mp_ctx.Barrier(len(worlds) + 1)  # size(parties) + main process
+
         # Start both party processes
         processes = []
-        for rank in [0, 1]:
+        for rank in worlds:
             process = mp_ctx.Process(
-                target=run_party_e2e_process, args=(rank, return_dict, assigned_ports)
+                target=run_party_e2e_process,
+                args=(rank, return_dict, assigned_ports, barrier),
             )
             process.start()
             processes.append(process)
 
+        barrier.wait(timeout=30)  # Wait for both parties to be ready
+
         # Wait for both processes to complete
         for process in processes:
             process.join(timeout=30)
@@ -328,7 +338,7 @@ def test_end_to_end_communication():
                     process.kill()
 
         # Validate results
-        for rank in [0, 1]:
+        for rank in worlds:
             assert rank in return_dict, f"Party {rank} did not complete"
             result = return_dict[rank]
             assert result.get("status") == "success", f"Party {rank} failed: {result}"