mplang-nightly 0.1.dev332__tar.gz → 0.1.dev334__tar.gz

{mplang_nightly-0.1.dev332 → mplang_nightly-0.1.dev334}/PKG-INFO

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

{mplang_nightly-0.1.dev332 → mplang_nightly-0.1.dev334}/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.dev332 → mplang_nightly-0.1.dev334}/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.dev332 → mplang_nightly-0.1.dev334}/mplang/backends/spu_state.py

@@ -20,6 +20,9 @@ multiple executions while binding to the Interpreter's lifecycle.
 
 from __future__ import annotations
 
+import logging
+import os
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
 import spu.api as spu_api
@@ -31,6 +34,64 @@ if TYPE_CHECKING:
     from mplang.backends.simp_worker.infra import WorkerInfra
     from mplang.dialects import spu
 
+logger = logging.getLogger(__name__)
+
+
+def _opt_int_env(name: str) -> int | None:
+    """Read an optional int env var; return None when unset/empty."""
+    val = os.getenv(name)
+    return int(val) if val else None
+
+
+@dataclass(frozen=True)
+class BrpcLinkConfig:
+    """Tunables for the brpc-backed SPU link.
+
+    Defaults come from environment variables so deployments can switch
+    protocol/timeouts without code changes; explicit constructor args
+    still win for tests and programmatic overrides.
+
+    Set ``protocol`` to ``"http"`` (HTTP/1.1) or ``"h2"`` (HTTP/2) when
+    traffic must traverse an L7 gateway that does not understand the
+    default ``baidu_std`` binary protocol. HTTP/1.1 requires
+    ``connection_type="pooled"`` to avoid head-of-line blocking.
+
+    ``http_timeout_ms`` defaults to ``None`` (use brpc's own default) —
+    long-running MPC jobs traversing a gateway with idle timeout should
+    set this explicitly to cover the longest single RPC.
+
+    Frozen so instances are hashable and safe to embed in cache keys.
+    """
+
+    protocol: str = field(
+        default_factory=lambda: os.getenv("MPLANG_BRPC_PROTOCOL") or "http"
+    )
+    connection_type: str = field(
+        default_factory=lambda: os.getenv("MPLANG_BRPC_CONN_TYPE") or "pooled"
+    )
+    recv_timeout_ms: int = field(
+        default_factory=lambda: int(os.getenv("MPLANG_BRPC_RECV_TIMEOUT_MS") or 100000)
+    )
+    http_max_payload_size: int = field(
+        default_factory=lambda: int(
+            os.getenv("MPLANG_BRPC_HTTP_MAX_PAYLOAD") or 32 * 1024 * 1024
+        )
+    )
+    # Per-RPC HTTP timeout. None = inherit brpc default. Set when going
+    # through a gateway with its own idle/request timeout.
+    http_timeout_ms: int | None = field(
+        default_factory=lambda: _opt_int_env("MPLANG_BRPC_HTTP_TIMEOUT_MS")
+    )
+    # Initial connect retry — tolerates peers that come up later.
+    connect_retry_times: int = field(
+        default_factory=lambda: int(os.getenv("MPLANG_BRPC_CONNECT_RETRY_TIMES") or 60)
+    )
+    connect_retry_interval_ms: int = field(
+        default_factory=lambda: int(
+            os.getenv("MPLANG_BRPC_CONNECT_RETRY_INTERVAL_MS") or 1000
+        )
+    )
+
 
 class SPUState(DialectState):
     """SPU Runtime cache as dialect state.
@@ -48,13 +109,22 @@ class SPUState(DialectState):
 
     dialect_name: str = "spu"
 
-    def __init__(self, infra: WorkerInfra | None = None) -> None:
+    def __init__(
+        self,
+        infra: WorkerInfra | None = None,
+        brpc_config: BrpcLinkConfig | 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, spu_endpoints)
+        self._brpc_config = brpc_config or BrpcLinkConfig()
+        # Key: (local_rank, world_size, protocol, field, link_mode, spu_endpoints, brpc_config)
+        # ``brpc_config`` participates only when link_mode == "brpc" (else None) so that
+        # callers passing custom configs don't silently reuse the first config's link.
         # Value: (Runtime, Io)
         self._runtimes: dict[
-            tuple[int, int, str, str, str, tuple[str, ...] | None],
+            tuple[
+                int, int, str, str, str, tuple[str, ...] | None, BrpcLinkConfig | None
+            ],
             tuple[spu_api.Runtime, spu_api.Io],
         ] = {}
         # Local template link cache (used when no WorkerInfra is provided)
@@ -140,6 +210,7 @@ class SPUState(DialectState):
             config.field,
             link_mode,
             tuple(spu_endpoints) if spu_endpoints else None,
+            self._brpc_config if link_mode == "brpc" else None,
         )
 
         if cache_key in self._runtimes:
@@ -222,13 +293,36 @@ class SPUState(DialectState):
         self, local_rank: int, spu_endpoints: list[str]
     ) -> libspu.link.Context:
         """Create BRPC link for distributed execution."""
+        cfg = self._brpc_config
         desc = libspu.link.Desc()  # type: ignore
-        desc.recv_timeout_ms = 100 * 1000  # 100 seconds
-        desc.http_max_payload_size = 32 * 1024 * 1024  # 32MB
+        desc.recv_timeout_ms = cfg.recv_timeout_ms
+        desc.http_max_payload_size = cfg.http_max_payload_size
+        desc.brpc_channel_protocol = cfg.protocol
+        desc.brpc_channel_connection_type = cfg.connection_type
+        desc.connect_retry_times = cfg.connect_retry_times
+        desc.connect_retry_interval_ms = cfg.connect_retry_interval_ms
+        if cfg.http_timeout_ms is not None:
+            desc.http_timeout_ms = cfg.http_timeout_ms
 
         for i, endpoint in enumerate(spu_endpoints):
             desc.add_party(f"P{i}", endpoint)
 
+        logger.info(
+            "Creating SPU brpc link: rank=%d, endpoints=%s, protocol=%s, "
+            "connection_type=%s, recv_timeout_ms=%d, http_timeout_ms=%s, "
+            "http_max_payload_size=%d, connect_retry_times=%d, "
+            "connect_retry_interval_ms=%d",
+            local_rank,
+            spu_endpoints,
+            cfg.protocol,
+            cfg.connection_type,
+            cfg.recv_timeout_ms,
+            cfg.http_timeout_ms,
+            cfg.http_max_payload_size,
+            cfg.connect_retry_times,
+            cfg.connect_retry_interval_ms,
+        )
+
         return libspu.link.create_brpc(desc, local_rank)
 
     def shutdown(self) -> None:

{mplang_nightly-0.1.dev332 → mplang_nightly-0.1.dev334}/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.dev332 → mplang_nightly-0.1.dev334}/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)