diff --git a/examples/distributed.py b/examples/distributed.py index 733b41f74..119ffe462 100644 --- a/examples/distributed.py +++ b/examples/distributed.py @@ -30,6 +30,8 @@ def main(): # Find the partition outputs = pt.DictOfNamedArrays({"out": y}) + + pt.verify_distributed_dag_pre_partition(comm, outputs) distributed_parts = find_distributed_partition(outputs) distributed_parts, _ = number_distributed_tags( comm, distributed_parts, base_tag=42) diff --git a/pytato/__init__.py b/pytato/__init__.py index 7f7a61066..c9bee29c3 100644 --- a/pytato/__init__.py +++ b/pytato/__init__.py @@ -100,7 +100,8 @@ def set_debug_enabled(flag: bool) -> None: from pytato.distributed.partition import ( find_distributed_partition, DistributedGraphPart, DistributedGraphPartition) from pytato.distributed.tags import number_distributed_tags -from pytato.distributed.verify import verify_distributed_partition +from pytato.distributed.verify import (verify_distributed_partition, + verify_distributed_dag_pre_partition) from pytato.distributed.execute import execute_distributed_partition from pytato.transform.lower_to_index_lambda import to_index_lambda @@ -161,6 +162,7 @@ def set_debug_enabled(flag: bool) -> None: "number_distributed_tags", "execute_distributed_partition", "verify_distributed_partition", + "verify_distributed_dag_pre_partition", "generate_code_for_partition", diff --git a/pytato/distributed/verify.py b/pytato/distributed/verify.py index e37d8a83a..8354e0865 100644 --- a/pytato/distributed/verify.py +++ b/pytato/distributed/verify.py @@ -30,14 +30,17 @@ """ -from typing import Any, FrozenSet, Dict, Set, Optional, Sequence, TYPE_CHECKING +from typing import FrozenSet, Dict, Set, Optional, Sequence, TYPE_CHECKING, Union +from immutables import Map -import numpy as np - -from pytato.distributed.nodes import CommTagType, DistributedRecv +from pytato.distributed.nodes import (CommTagType, DistributedRecv, + DistributedSendRefHolder, DistributedSend) from pytato.partition import PartId from pytato.distributed.partition import DistributedGraphPartition -from pytato.array import ShapeType +from pytato.transform import UsersCollector, ArrayOrNames +from pytato import DictOfNamedArrays + +from pytools import UniqueNameGenerator import attrs @@ -58,9 +61,6 @@ class _SummarizedDistributedSend: dest_rank: int comm_tag: CommTagType - shape: ShapeType - dtype: np.dtype[Any] - @attrs.define(frozen=True) class _DistributedPartId: @@ -74,6 +74,12 @@ class _DistributedName: name: str +@attrs.define(frozen=True) +class _DistributedNode: + rank: int + node: Union[DistributedSend, ArrayOrNames, _SummarizedDistributedSend] + + @attrs.define(frozen=True) class _SummarizedDistributedGraphPart: pid: _DistributedPartId @@ -89,6 +95,20 @@ def rank(self) -> int: return self.pid.rank +@attrs.define(frozen=True) +class _SummarizedDistributedGraph: + rank: int + node_to_users: Dict[_DistributedNode, Set[_DistributedNode]] + input_name_to_recv_node: Dict[_DistributedName, DistributedRecv] + output_name_to_send_node: Dict[_DistributedName, _SummarizedDistributedSend] + + def __hash__(self) -> int: + return (hash(self.rank) + ^ hash(Map(self.node_to_users)) + ^ hash(Map(self.input_name_to_recv_node)) + ^ hash(Map(self.output_name_to_send_node))) + + @attrs.define(frozen=True) class _CommIdentifier: src_rank: int @@ -122,6 +142,152 @@ class MissingRecvError(DistributedPartitionVerificationError): # }}} +# {{{ verify_distributed_dag_pre_partition + +class _DistributedDAGGatherer(UsersCollector): + def __init__(self, dist_name_generator: UniqueNameGenerator, my_rank: int) \ + -> None: + super().__init__() + + self.name_generator = dist_name_generator + self.my_rank = my_rank + + self.input_name_to_recv_node: Dict[str, DistributedRecv] = {} + self.output_name_to_send_node: Dict[str, _SummarizedDistributedSend] = {} + + def map_distributed_recv(self, expr: DistributedRecv) -> None: + super().map_distributed_recv(expr) + new_name = self.name_generator() + self.input_name_to_recv_node[new_name] = expr + + def map_distributed_send_ref_holder( + self, expr: DistributedSendRefHolder) -> None: + super().map_distributed_send_ref_holder(expr) + s = _SummarizedDistributedSend( + src_rank=self.my_rank, + dest_rank=expr.send.dest_rank, + comm_tag=expr.send.comm_tag) + + new_name = self.name_generator() + self.output_name_to_send_node[new_name] = s + + +def verify_distributed_dag_pre_partition(mpi_communicator: mpi4py.MPI.Comm, + outputs: DictOfNamedArrays) -> None: + """ + .. warning:: + + This is an MPI-collective operation. + """ + my_rank = mpi_communicator.rank + root_rank = 0 + + ung = UniqueNameGenerator(forced_prefix="_pt_verify_dist_") + + dg = _DistributedDAGGatherer(ung, my_rank) + dg(outputs) + + def dist_send_to_summarized_dist_send(node: + Union[ArrayOrNames, _SummarizedDistributedSend, DistributedSend]) \ + -> Union[ArrayOrNames, _SummarizedDistributedSend]: + if (not isinstance(node, DistributedSend) + and not isinstance(node, DistributedSendRefHolder)): + return node + + if isinstance(node, DistributedSend): + return _SummarizedDistributedSend( + src_rank=my_rank, + dest_rank=node.dest_rank, + comm_tag=node.comm_tag,) + elif isinstance(node, DistributedSendRefHolder): + return _SummarizedDistributedSend( + src_rank=my_rank, + dest_rank=node.send.dest_rank, + comm_tag=node.send.comm_tag,) + + summarized_dag = _SummarizedDistributedGraph( + rank=my_rank, + node_to_users={_DistributedNode(my_rank, k): + set((_DistributedNode(my_rank, dist_send_to_summarized_dist_send(n)) + for n in v)) for k, v in dg.node_to_users.items()}, + input_name_to_recv_node={_DistributedName(my_rank, name): recv + for name, recv in dg.input_name_to_recv_node.items()}, + output_name_to_send_node={ + _DistributedName(my_rank, name): + _SummarizedDistributedSend( + src_rank=my_rank, + dest_rank=send.dest_rank, + comm_tag=send.comm_tag,) + for name, send in dg.output_name_to_send_node.items()}) + + all_outputs = mpi_communicator.gather(summarized_dag, root=root_rank) + + if my_rank == root_rank: + assert all_outputs + + all_summarized_outputs = { + rank: rank_outputs + for rank, rank_outputs in enumerate(all_outputs)} + + all_recvs: Set[_CommIdentifier] = set() + + send_recv_deps: \ + Dict[_DistributedNode, Set[_DistributedNode]] = {} + + def add_send_recv_dep(recv: _DistributedNode, + send: _DistributedNode) -> None: + send_recv_deps.setdefault(recv, set()).add(send) + + # {{{ gather information on senders + + comm_id_to_sending_node = {} + + for sumdag in all_summarized_outputs.values(): + for sumsend in sumdag.output_name_to_send_node.values(): + comm_id = _CommIdentifier( + src_rank=sumsend.src_rank, + dest_rank=sumsend.dest_rank, + comm_tag=sumsend.comm_tag) + + if comm_id in comm_id_to_sending_node: + raise DuplicateSendError( + f"duplicate send for comm id: '{comm_id}'") + comm_id_to_sending_node[comm_id] = sumsend + + # }}} + + for sumdag in all_summarized_outputs.values(): + send_recv_deps.update(sumdag.node_to_users) + + for dname, dist_recv in sumdag.input_name_to_recv_node.items(): + comm_id = _CommIdentifier( + src_rank=dist_recv.src_rank, + dest_rank=dname.rank, + comm_tag=dist_recv.comm_tag) + + if comm_id in all_recvs: + raise DuplicateRecvError(f"Duplicate recv: '{comm_id}'") + + all_recvs.add(comm_id) + + # Add edges between sends and receives (cross-rank) + try: + sending_node = comm_id_to_sending_node[comm_id] + except KeyError: + raise MissingSendError( + f"no matching send for recv on '{comm_id}'") + + add_send_recv_dep(_DistributedNode(comm_id.dest_rank, dist_recv), + _DistributedNode(comm_id.src_rank, sending_node)) + + from pytools.graph import compute_topological_order + compute_topological_order(send_recv_deps) + + logger.info("verify_distributed_dag_pre_partition completed successfully.") + +# }}} + + # {{{ verify_distributed_partition def verify_distributed_partition(mpi_communicator: mpi4py.MPI.Comm, @@ -169,9 +335,7 @@ def verify_distributed_partition(mpi_communicator: mpi4py.MPI.Comm, _SummarizedDistributedSend( src_rank=my_rank, dest_rank=send.dest_rank, - comm_tag=send.comm_tag, - shape=send.data.shape, - dtype=send.data.dtype) + comm_tag=send.comm_tag) for name, send in part.output_name_to_send_node.items()}) # Gather the _SummarizedDistributedGraphPart's to rank 0 diff --git a/test/test_distributed.py b/test/test_distributed.py index 95f220b7e..b9dbf7497 100644 --- a/test/test_distributed.py +++ b/test/test_distributed.py @@ -308,7 +308,7 @@ def test_deterministic_partitioning(): # }}} -# {{{ test verify_distributed_partition +# {{{ test verify_distributed_partition and verify_distributed_dag_pre_partition def test_verify_distributed_partition(): run_test_with_mpi(2, _do_verify_distributed_partition) @@ -331,6 +331,13 @@ def _do_verify_distributed_partition(ctx_factory): src_rank=(rank+1) % size, comm_tag=42, shape=(4, 4), dtype=int) outputs = pt.make_dict_of_named_arrays({"out": y}) + + if rank == 0: + with pytest.raises(MissingSendError): + pt.verify_distributed_dag_pre_partition(comm, outputs) + else: + pt.verify_distributed_dag_pre_partition(comm, outputs) + distributed_parts = pt.find_distributed_partition(outputs) if rank == 0: @@ -350,6 +357,13 @@ def _do_verify_distributed_partition(ctx_factory): outputs = pt.make_dict_of_named_arrays({"out": send}) distributed_parts = pt.find_distributed_partition(outputs) + if rank == 0: + # FIXME: this should raise + with pytest.raises(MissingRecvError): + pt.verify_distributed_dag_pre_partition(comm, outputs) + else: + pt.verify_distributed_dag_pre_partition(comm, outputs) + if rank == 0: with pytest.raises(MissingRecvError): pt.verify_distributed_partition(comm, distributed_parts) @@ -368,6 +382,13 @@ def _do_verify_distributed_partition(ctx_factory): src_rank=(rank+1) % size, comm_tag=42, shape=(4, 4), dtype=int)) outputs = pt.make_dict_of_named_arrays({"out": x+send}) + + if rank == 0: + with pytest.raises(MissingSendError): + pt.verify_distributed_dag_pre_partition(comm, outputs) + else: + pt.verify_distributed_dag_pre_partition(comm, outputs) + distributed_parts = pt.find_distributed_partition(outputs) if rank == 0: @@ -388,6 +409,13 @@ def _do_verify_distributed_partition(ctx_factory): dest_rank=(rank-1) % size, comm_tag=42, stapled_to=x) outputs = pt.make_dict_of_named_arrays({"out": send+send2}) + + if rank == 0: + with pytest.raises(DuplicateSendError): + pt.verify_distributed_dag_pre_partition(comm, outputs) + else: + pt.verify_distributed_dag_pre_partition(comm, outputs) + distributed_parts = pt.find_distributed_partition(outputs) if rank == 0: @@ -413,6 +441,13 @@ def _do_verify_distributed_partition(ctx_factory): stapled_to=recv) outputs = pt.make_dict_of_named_arrays({"out": send+send2}) + + if rank == 0: + with pytest.raises(MissingSendError): + pt.verify_distributed_dag_pre_partition(comm, outputs) + else: + pt.verify_distributed_dag_pre_partition(comm, outputs) + distributed_parts = pt.find_distributed_partition(outputs) if rank == 0: