mplang-nightly 0.1.dev286__tar.gz → 0.1.dev288__tar.gz

{mplang_nightly-0.1.dev286 → mplang_nightly-0.1.dev288}/PKG-INFO

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

mplang_nightly-0.1.dev288/mplang/backends/simp_worker/collective_algorithms.py

@@ -0,0 +1,228 @@
+# Copyright 2026 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Collective communication algorithms (communicator-only).
+
+This module contains *pure* collective algorithms implemented only in terms of
+(a) a communicator and (b) an explicit participant set.
+
+It intentionally does NOT depend on:
+- Interpreter execution IDs / graph keys
+- SimpWorker current_parties
+- Operation objects
+
+Callers are expected to provide a collision-free key prefix for each collective
+instance.
+"""
+
+from __future__ import annotations
+
+import operator
+from collections.abc import Callable, Sequence
+from typing import Any, Protocol
+
+
+class Communicator(Protocol):
+    """Minimal communicator interface required by the algorithms."""
+
+    rank: int
+    world_size: int
+
+    def send(
+        self, to: int, key: str, data: Any, *, is_raw_bytes: bool = False
+    ) -> None: ...
+
+    def recv(self, frm: int, key: str) -> Any: ...
+
+
+def normalize_participants(
+    comm: Communicator, participants: Sequence[int]
+) -> tuple[int, ...]:
+    ps = tuple(sorted({int(r) for r in participants}))
+    if not ps:
+        raise ValueError("participants must be non-empty")
+    if any(r < 0 or r >= comm.world_size for r in ps):
+        raise ValueError(
+            f"participants out of range: {ps}, world_size={comm.world_size}"
+        )
+    if comm.rank not in ps:
+        raise ValueError(f"rank {comm.rank} is not in participants {ps}")
+    return ps
+
+
+def barrier(
+    comm: Communicator, *, participants: Sequence[int], key_prefix: str
+) -> None:
+    """Barrier using root gather + root release."""
+
+    ps = normalize_participants(comm, participants)
+    root = ps[0]
+
+    arrive_key = f"{key_prefix}_arrive"
+    release_key = f"{key_prefix}_release"
+
+    if comm.rank != root:
+        comm.send(root, arrive_key, True)
+        comm.recv(root, release_key)
+        return
+
+    for r in ps:
+        if r == root:
+            continue
+        _ = comm.recv(r, arrive_key)
+
+    for r in ps:
+        if r == root:
+            continue
+        comm.send(r, release_key, True)
+
+
+def broadcast(
+    comm: Communicator,
+    value: Any,
+    *,
+    root: int,
+    participants: Sequence[int],
+    key_prefix: str,
+) -> Any:
+    """Broadcast a value from root to all participants."""
+
+    ps = normalize_participants(comm, participants)
+    if root not in ps:
+        raise ValueError(f"root {root} must be in participants {ps}")
+
+    bcast_key = f"{key_prefix}_bcast"
+
+    if comm.rank == root:
+        for r in ps:
+            if r == root:
+                continue
+            comm.send(r, bcast_key, value)
+        return value
+
+    return comm.recv(root, bcast_key)
+
+
+def allgather(
+    comm: Communicator, value: Any, *, participants: Sequence[int], key_prefix: str
+) -> list[Any]:
+    """Allgather implemented as gather-to-root then root broadcast."""
+
+    ps = normalize_participants(comm, participants)
+    root = ps[0]
+
+    gather_key = f"{key_prefix}_gather"
+    bcast_key = f"{key_prefix}_bcast"
+
+    if comm.rank != root:
+        comm.send(root, gather_key, value)
+        gathered = comm.recv(root, bcast_key)
+        if not isinstance(gathered, list):
+            raise TypeError(f"expected list from root broadcast, got {type(gathered)}")
+        return gathered
+
+    values_by_rank: dict[int, Any] = {root: value}
+    for r in ps:
+        if r == root:
+            continue
+        values_by_rank[r] = comm.recv(r, gather_key)
+
+    gathered = [values_by_rank[r] for r in ps]
+
+    for r in ps:
+        if r == root:
+            continue
+        comm.send(r, bcast_key, gathered)
+
+    return gathered
+
+
+def allreduce_bool_and(
+    comm: Communicator,
+    value: bool,
+    *,
+    participants: Sequence[int],
+    key_prefix: str,
+) -> bool:
+    return _allreduce_bool(
+        comm,
+        value,
+        participants=participants,
+        key_prefix=key_prefix,
+        combine=operator.and_,
+    )
+
+
+def allreduce_bool_or(
+    comm: Communicator,
+    value: bool,
+    *,
+    participants: Sequence[int],
+    key_prefix: str,
+) -> bool:
+    return _allreduce_bool(
+        comm,
+        value,
+        participants=participants,
+        key_prefix=key_prefix,
+        combine=operator.or_,
+    )
+
+
+def allreduce_bool_xor(
+    comm: Communicator,
+    value: bool,
+    *,
+    participants: Sequence[int],
+    key_prefix: str,
+) -> bool:
+    return _allreduce_bool(
+        comm,
+        value,
+        participants=participants,
+        key_prefix=key_prefix,
+        combine=operator.xor,
+    )
+
+
+def _allreduce_bool(
+    comm: Communicator,
+    value: bool,
+    *,
+    participants: Sequence[int],
+    key_prefix: str,
+    combine: Callable[[bool, bool], bool],
+) -> bool:
+    ps = normalize_participants(comm, participants)
+    root = ps[0]
+
+    gather_key = f"{key_prefix}_gather"
+    bcast_key = f"{key_prefix}_bcast"
+
+    if comm.rank != root:
+        comm.send(root, gather_key, bool(value))
+        return bool(comm.recv(root, bcast_key))
+
+    acc = bool(value)
+    for r in ps:
+        if r == root:
+            continue
+        acc = combine(acc, bool(comm.recv(r, gather_key)))
+
+    for r in ps:
+        if r == root:
+            continue
+        comm.send(r, bcast_key, acc)
+
+    return acc

mplang_nightly-0.1.dev288/mplang/backends/simp_worker/collectives.py

@@ -0,0 +1,275 @@
+# Copyright 2026 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Simp worker-side collectives (wrapper layer).
+
+This module is the *context-aware wrapper* on top of
+`mplang.backends.simp_worker.collective_algorithms`.
+
+Responsibilities here:
+- Resolve "participants" from (explicit arg / op.attrs["parties"] /
+  worker.current_parties / world).
+- Build collision-free `key_prefix` using interpreter execution IDs.
+
+The underlying algorithms only depend on the communicator interface.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any, Protocol
+
+from mplang.backends.simp_worker import collective_algorithms as algo
+from mplang.edsl.graph import Operation
+
+
+class _ExecContext(Protocol):
+    def current_op_exec_id(self) -> int: ...
+
+    def current_graph_exec_key(self) -> str: ...
+
+
+class _Worker(Protocol):
+    rank: int
+    world_size: int
+    communicator: algo.Communicator
+    current_parties: tuple[int, ...] | None
+
+
+def resolve_participants(
+    worker: _Worker,
+    *,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+) -> Sequence[int]:
+    """Resolve participant ranks.
+
+    Priority:
+    1) explicit participants argument
+    2) op.attrs["parties"] if present
+    3) worker.current_parties if set (pcall_static dynamic scope)
+    4) all ranks [0, world_size)
+
+    Note:
+        Normalization/validation (sorting, emptiness, range checks, rank
+        inclusion) is intentionally delegated to the lower-level algorithms in
+        `collective_algorithms`.
+    """
+
+    if participants is not None:
+        return participants
+
+    if op is not None:
+        parties = op.attrs.get("parties")
+        if parties is not None:
+            if not isinstance(parties, Sequence):
+                raise TypeError(
+                    "op.attrs['parties'] must be a sequence of rank integers"
+                )
+            return tuple(int(r) for r in parties)
+
+    if worker.current_parties is not None:
+        return worker.current_parties
+
+    return tuple(range(worker.world_size))
+
+
+def _collective_prefix(
+    interpreter: _ExecContext, *, op: Operation | None, name: str
+) -> str:
+    exec_id = interpreter.current_op_exec_id()
+    graph_key = interpreter.current_graph_exec_key()
+    op_name = op.name if op is not None else "_"
+    return f"coll_{graph_key}_{op_name}_{exec_id}_{name}"
+
+
+def barrier(
+    interpreter: _ExecContext,
+    worker: _Worker,
+    *,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+    name: str = "barrier",
+) -> None:
+    ps = resolve_participants(worker, op=op, participants=participants)
+    prefix = _collective_prefix(interpreter, op=op, name=name)
+    algo.barrier(worker.communicator, participants=ps, key_prefix=prefix)
+
+
+def broadcast(
+    interpreter: _ExecContext,
+    worker: _Worker,
+    value: Any,
+    *,
+    root: int,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+    name: str = "broadcast",
+) -> Any:
+    ps = resolve_participants(worker, op=op, participants=participants)
+    prefix = _collective_prefix(interpreter, op=op, name=name)
+    return algo.broadcast(
+        worker.communicator,
+        value,
+        root=int(root),
+        participants=ps,
+        key_prefix=prefix,
+    )
+
+
+def allgather_obj(
+    interpreter: _ExecContext,
+    worker: _Worker,
+    value: Any,
+    *,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+    name: str = "allgather_obj",
+) -> list[Any]:
+    ps = resolve_participants(worker, op=op, participants=participants)
+    prefix = _collective_prefix(interpreter, op=op, name=name)
+    return algo.allgather(
+        worker.communicator, value, participants=ps, key_prefix=prefix
+    )
+
+
+def allgather_bool(
+    interpreter: _ExecContext,
+    worker: _Worker,
+    value: bool,
+    *,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+    name: str = "allgather_bool",
+) -> list[bool]:
+    gathered = allgather_obj(
+        interpreter,
+        worker,
+        bool(value),
+        op=op,
+        participants=participants,
+        name=name,
+    )
+    return [bool(v) for v in gathered]
+
+
+def allreduce_bool_and(
+    interpreter: _ExecContext,
+    worker: _Worker,
+    value: bool,
+    *,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+    name: str = "allreduce_bool_and",
+) -> bool:
+    ps = resolve_participants(worker, op=op, participants=participants)
+    prefix = _collective_prefix(interpreter, op=op, name=name)
+    return algo.allreduce_bool_and(
+        worker.communicator,
+        bool(value),
+        participants=ps,
+        key_prefix=prefix,
+    )
+
+
+def allreduce_bool_or(
+    interpreter: _ExecContext,
+    worker: _Worker,
+    value: bool,
+    *,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+    name: str = "allreduce_bool_or",
+) -> bool:
+    ps = resolve_participants(worker, op=op, participants=participants)
+    prefix = _collective_prefix(interpreter, op=op, name=name)
+    return algo.allreduce_bool_or(
+        worker.communicator,
+        bool(value),
+        participants=ps,
+        key_prefix=prefix,
+    )
+
+
+def allreduce_bool_xor(
+    interpreter: _ExecContext,
+    worker: _Worker,
+    value: bool,
+    *,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+    name: str = "allreduce_bool_xor",
+) -> bool:
+    ps = resolve_participants(worker, op=op, participants=participants)
+    prefix = _collective_prefix(interpreter, op=op, name=name)
+    return algo.allreduce_bool_xor(
+        worker.communicator,
+        bool(value),
+        participants=ps,
+        key_prefix=prefix,
+    )
+
+
+def verify_uniform_predicate(
+    interpreter: _ExecContext,
+    worker: _Worker,
+    pred: bool,
+    *,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+    name: str = "uniform_predicate",
+) -> bool:
+    """Verify that `pred` is uniform across participants.
+
+    Uses AND/OR all-reduce to detect mismatch. If mismatch is detected, runs an
+    allgather to provide a helpful error message. All participants execute the
+    same comm steps to avoid deadlocks.
+    """
+
+    ps = resolve_participants(worker, op=op, participants=participants)
+
+    all_and = allreduce_bool_and(
+        interpreter,
+        worker,
+        bool(pred),
+        op=op,
+        participants=ps,
+        name=f"{name}_and",
+    )
+    all_or = allreduce_bool_or(
+        interpreter,
+        worker,
+        bool(pred),
+        op=op,
+        participants=ps,
+        name=f"{name}_or",
+    )
+
+    if all_and != all_or:
+        gathered = allgather_bool(
+            interpreter,
+            worker,
+            bool(pred),
+            op=op,
+            participants=ps,
+            name=f"{name}_gather",
+        )
+        ps_norm = algo.normalize_participants(worker.communicator, ps)
+        dist = dict(zip(ps_norm, gathered, strict=True))
+        raise RuntimeError(
+            "simp.uniform_cond predicate is not uniform across participants: "
+            f"participants={ps_norm}, values={dist}"
+        )
+
+    return bool(pred)

{mplang_nightly-0.1.dev286 → mplang_nightly-0.1.dev288}/mplang/backends/simp_worker/ops.py

@@ -22,6 +22,7 @@ from __future__ import annotations
 
 from typing import Any
 
+from mplang.backends.simp_worker.collectives import verify_uniform_predicate
 from mplang.dialects import simp
 from mplang.edsl.graph import Operation
 from mplang.runtime.interpreter import Interpreter
@@ -117,12 +118,14 @@ def _uniform_cond_worker_impl(
     """Worker implementation of simp.uniform_cond."""
     from mplang.backends.tensor_impl import TensorValue
 
-    if op.attrs.get("verify_uniform", True):
-        pass  # TODO: Implement AllReduce verification
+    worker = _ensure_worker_context(interpreter, "uniform_cond_impl")
 
     if isinstance(pred, TensorValue):
         pred = bool(pred.unwrap())
 
+    if op.attrs.get("verify_uniform", True):
+        pred = verify_uniform_predicate(interpreter, worker, bool(pred), op=op)
+
     if pred:
         result = interpreter.evaluate_graph(op.regions[0], list(args))
     else:

{mplang_nightly-0.1.dev286 → mplang_nightly-0.1.dev288}/mplang/backends/table_impl.py

@@ -346,7 +346,7 @@ class ParquetReader(pa.RecordBatchReader):
             pass
         if batches:
             return pa.Table.from_batches(batches)
-        return pa.Table.from_batches([])
+        return pa.Table.from_batches([], schema=self._schema)
 
     def read_next_batch(self) -> pa.RecordBatch:
         batch = next(self._iter)

mplang_nightly-0.1.dev288/tests/backends/simp_worker/test_collectives.py

@@ -0,0 +1,141 @@
+# Copyright 2026 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import concurrent.futures
+from dataclasses import dataclass
+from typing import Any
+
+from mplang.backends.simp_worker.collectives import (
+    allgather_obj,
+    allreduce_bool_and,
+    allreduce_bool_or,
+    barrier,
+    broadcast,
+)
+from mplang.backends.simp_worker.mem import LocalMesh
+
+
+class _FakeInterpreter:
+    def current_op_exec_id(self) -> int:
+        return 1
+
+    def current_graph_exec_key(self) -> str:
+        return "test_graph"
+
+
+@dataclass
+class _Worker:
+    rank: int
+    world_size: int
+    communicator: Any
+    current_parties: tuple[int, ...] | None = None
+
+
+def test_broadcast_roundtrip() -> None:
+    mesh = LocalMesh(world_size=3)
+    interp: Any = _FakeInterpreter()
+
+    def run_rank(rank: int) -> Any:
+        worker = _Worker(rank, 3, mesh.comms[rank])
+        return broadcast(
+            interp,
+            worker,
+            {"x": 123},
+            root=0,
+            participants=(0, 1, 2),
+            name="test_broadcast",
+        )
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as ex:
+        results = list(ex.map(run_rank, range(3)))
+
+    assert results == [{"x": 123}, {"x": 123}, {"x": 123}]
+    mesh.shutdown()
+
+
+def test_allgather_obj_order() -> None:
+    mesh = LocalMesh(world_size=3)
+    interp: Any = _FakeInterpreter()
+
+    def run_rank(rank: int) -> Any:
+        worker = _Worker(rank, 3, mesh.comms[rank])
+        gathered = allgather_obj(
+            interp,
+            worker,
+            rank,
+            participants=(0, 1, 2),
+            name="test_allgather",
+        )
+        return tuple(gathered)
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as ex:
+        results = list(ex.map(run_rank, range(3)))
+
+    assert results == [(0, 1, 2), (0, 1, 2), (0, 1, 2)]
+    mesh.shutdown()
+
+
+def test_allreduce_bool_and_or() -> None:
+    mesh = LocalMesh(world_size=3)
+    interp: Any = _FakeInterpreter()
+
+    inputs = {0: True, 1: True, 2: False}
+
+    def run_rank(rank: int) -> tuple[bool, bool]:
+        worker = _Worker(rank, 3, mesh.comms[rank])
+        v = inputs[rank]
+        r_and = allreduce_bool_and(
+            interp,
+            worker,
+            v,
+            participants=(0, 1, 2),
+            name="test_allreduce_and",
+        )
+        r_or = allreduce_bool_or(
+            interp,
+            worker,
+            v,
+            participants=(0, 1, 2),
+            name="test_allreduce_or",
+        )
+        return r_and, r_or
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as ex:
+        results = list(ex.map(run_rank, range(3)))
+
+    assert results == [(False, True), (False, True), (False, True)]
+    mesh.shutdown()
+
+
+def test_barrier_completes() -> None:
+    mesh = LocalMesh(world_size=3)
+    interp: Any = _FakeInterpreter()
+
+    def run_rank(rank: int) -> int:
+        worker = _Worker(rank, 3, mesh.comms[rank])
+        barrier(
+            interp,
+            worker,
+            participants=(0, 1, 2),
+            name="test_barrier",
+        )
+        return rank
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as ex:
+        results = list(ex.map(run_rank, range(3)))
+
+    assert sorted(results) == [0, 1, 2]
+    mesh.shutdown()