mplang-nightly 0.1.dev331__tar.gz → 0.1.dev333__tar.gz

{mplang_nightly-0.1.dev331 → mplang_nightly-0.1.dev333}/PKG-INFO

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

{mplang_nightly-0.1.dev331 → mplang_nightly-0.1.dev333}/mplang/backends/simp_driver/ops.py

@@ -45,8 +45,8 @@ def _wrap_op_as_graph(op: Operation) -> Graph:
     # Create graph inputs
     graph_inputs = [g.add_input(f"in_{i}", any_type) for i in range(len(op.inputs))]
 
-    # Determine output types
-    output_types = [out.type for out in op.outputs] if op.outputs else [any_type]
+    # Determine output types (faithfully mirror the original op)
+    output_types = [out.type for out in op.outputs]
 
     # Add the operation (this handles outputs and value registration)
     g.add_op(
@@ -77,7 +77,7 @@ def _collect_to_hostvars(results: list[Any], num_outputs: int, world_size: int)
         Single DriverVar if num_outputs == 1, else list of DriverVars
     """
     if num_outputs == 0:
-        return None
+        return []
 
     # Transpose [worker][output] -> [output][worker]
     # results[worker_idx] is a list of URIs for that worker's outputs
@@ -117,7 +117,7 @@ def _generic_simp_dispatch(interpreter: Any, op: Operation, *args: Any) -> Any:
     results = driver.collect(futures)
 
     # 4. Assemble into DriverVar(s)
-    num_outputs = len(op.outputs) if op.outputs else 1
+    num_outputs = len(op.outputs)
     return _collect_to_hostvars(results, num_outputs, world_size)
 
 

{mplang_nightly-0.1.dev331 → mplang_nightly-0.1.dev333}/mplang/backends/simp_worker/ops.py

@@ -53,13 +53,19 @@ def _pcall_static_worker_impl(
 
         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
+            # Use inner graph's actual output count for safety
+            n_out = len(fn_graph.outputs)
+            if n_out == 0:
+                return []
+            return result[0] if n_out == 1 else result
         finally:
             worker.current_parties = prev_parties
     else:
         # No data for this rank
-        return None if len(op.outputs) == 1 else [None] * len(op.outputs)
+        n_out = len(op.outputs)
+        if n_out == 0:
+            return []
+        return None if n_out == 1 else [None] * n_out
 
 
 def _pcall_dynamic_worker_impl(
@@ -68,7 +74,10 @@ def _pcall_dynamic_worker_impl(
     """Worker implementation of pcall_dynamic."""
     fn_graph = op.regions[0]
     result = interpreter.evaluate_graph(fn_graph, list(args))
-    return result[0] if len(op.outputs) == 1 else result
+    n_out = len(fn_graph.outputs)
+    if n_out == 0:
+        return []
+    return result[0] if n_out == 1 else result
 
 
 def _shuffle_static_worker_impl(

{mplang_nightly-0.1.dev331 → mplang_nightly-0.1.dev333}/mplang/backends/spu_state.py

@@ -51,10 +51,11 @@ class SPUState(DialectState):
     def __init__(self, infra: WorkerInfra | None = None) -> None:
         # Optional shared infrastructure (for per-request isolation via link.spawn)
         self._infra = infra
-        # Key: (local_rank, world_size, protocol, field, link_mode)
+        # Key: (local_rank, world_size, protocol, field, link_mode, spu_endpoints)
         # Value: (Runtime, Io)
         self._runtimes: dict[
-            tuple[int, int, str, str, str], tuple[spu_api.Runtime, spu_api.Io]
+            tuple[int, int, str, str, str, tuple[str, ...] | None],
+            tuple[spu_api.Runtime, spu_api.Io],
         ] = {}
         # Local template link cache (used when no WorkerInfra is provided)
         self._template_links: dict[tuple, libspu.link.Context] = {}
@@ -75,14 +76,14 @@ class SPUState(DialectState):
         """
 
         def _create() -> libspu.link.Context:
-            if communicator is not None:
+            if spu_endpoints:
+                return self._create_brpc_link(local_rank, spu_endpoints)
+            elif communicator is not None:
                 if parties is None:
                     raise ValueError("parties required when using communicator")
                 return self._create_channels_link(
                     local_rank, spu_world_size, communicator, parties
                 )
-            elif spu_endpoints:
-                return self._create_brpc_link(local_rank, spu_endpoints)
             else:
                 return self._create_mem_link(local_rank, spu_world_size)
 
@@ -104,13 +105,18 @@ class SPUState(DialectState):
     ) -> tuple[spu_api.Runtime, spu_api.Io]:
         """Get or create SPU Runtime and Io for the given configuration.
 
+        Link mode priority: spu_endpoints (BRPC) > communicator (Channels) > mem.
+        When ``spu_endpoints`` is provided it always takes precedence, even if
+        a ``communicator`` is also supplied.
+
         Args:
             local_rank: The local rank within the SPU device (0-indexed).
             spu_world_size: The number of parties in the SPU device.
             config: SPU configuration including protocol settings.
-            spu_endpoints: Optional list of BRPC endpoints. If None, use mem link.
+            spu_endpoints: Optional list of BRPC endpoints. Takes highest
+                priority when provided.
             communicator: Optional v2 communicator (ThreadCommunicator/HttpCommunicator).
-                If provided, use Channels mode to reuse existing communication.
+                Used only when ``spu_endpoints`` is not provided.
             parties: Optional list of global ranks for SPU parties.
                 Required when communicator is provided.
 
@@ -120,10 +126,10 @@ class SPUState(DialectState):
         from mplang.backends.spu_impl import to_runtime_config
 
         # Determine link mode
-        if communicator is not None:
-            link_mode = "channels"
-        elif spu_endpoints:
+        if spu_endpoints:
             link_mode = "brpc"
+        elif communicator is not None:
+            link_mode = "channels"
         else:
             link_mode = "mem"
 
@@ -133,6 +139,7 @@ class SPUState(DialectState):
             config.protocol,
             config.field,
             link_mode,
+            tuple(spu_endpoints) if spu_endpoints else None,
         )
 
         if cache_key in self._runtimes:

{mplang_nightly-0.1.dev331 → mplang_nightly-0.1.dev333}/mplang/runtime/interpreter.py

@@ -1082,7 +1082,9 @@ class Interpreter(AbstractInterpreter):
                     return
 
                 # Store results
-                if len(op.outputs) == 1:
+                if len(op.outputs) == 0:
+                    pass  # Void operation
+                elif len(op.outputs) == 1:
                     env[op.outputs[0]] = result
                 else:
                     for out_val, res in zip(op.outputs, result, strict=True):

{mplang_nightly-0.1.dev331 → mplang_nightly-0.1.dev333}/tests/backends/test_simp_integration.py

@@ -194,6 +194,50 @@ def test_nested_pcall():
         assert values == [20, 40]
 
 
+def test_pcall_static_zero_output():
+    """Regression: pcall_static body returning None must not crash (IndexError).
+
+    When a function returns None, its traced graph has 0 outputs.
+    The driver/worker dispatch must handle this correctly.
+    """
+    sim = simp.make_simulator(world_size=2)
+    with sim:
+
+        def side_effect_only(x):
+            # Produces a traced op but returns nothing
+            run_jax(jnp.abs, x)
+            return None
+
+        x0 = simp.constant((0,), 1)
+        x1 = simp.constant((1,), 2)
+        x_obj = simp.converge(x0, x1)
+
+        # Should not raise IndexError
+        result = pcall_static((0, 1), side_effect_only, x_obj)
+        assert result is None
+
+
+def test_pcall_static_zero_output_compile_evaluate():
+    """Regression: compile+evaluate path with zero-output pcall_static."""
+    sim = simp.make_simulator(world_size=2)
+    with sim:
+
+        def workflow():
+            x0 = simp.constant((0,), 42)
+            x1 = simp.constant((1,), 99)
+            x_obj = simp.converge(x0, x1)
+
+            def body(x):
+                run_jax(jnp.abs, x)
+                return None
+
+            pcall_static((0, 1), body, x_obj)
+
+        traced = mp.compile(workflow)
+        result = mp.evaluate(traced)
+        assert result is None
+
+
 def test_mp_function_decorator():
     """Test @mp.function decorator which implies pcall_static(ALL, ...)."""
     sim = simp.make_simulator(world_size=3)