From 87bed91cbb20191c14f8fcb4838df3fba8a4397b Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 18 Nov 2022 17:58:50 -0600 Subject: [PATCH 1/9] Add distributed prepartitioning check --- examples/distributed.py | 2 + pytato/distributed.py | 136 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 138 insertions(+) diff --git a/examples/distributed.py b/examples/distributed.py index 1f12839f0..e3bf11a5f 100644 --- a/examples/distributed.py +++ b/examples/distributed.py @@ -30,6 +30,8 @@ def main(): # Find the partition outputs = pt.DictOfNamedArrays({"out": y}) + + pt.distributed.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/distributed.py b/pytato/distributed.py index d52d2df91..9ff77cd3e 100644 --- a/pytato/distributed.py +++ b/pytato/distributed.py @@ -1078,6 +1078,13 @@ def rank(self) -> int: return self.pid.rank +@attrs.define(frozen=True) +class _SummarizedDistributedGraph: + rank: int + input_name_to_recv_node: Dict[_DistributedName, DistributedRecv] + output_name_to_send_node: Dict[_DistributedName, _SummarizedDistributedSend] + + @attrs.define(frozen=True) class _CommIdentifier: src_rank: int @@ -1105,6 +1112,135 @@ class MissingRecvError(DistributedPartitionVerificationError): pass +@optimize_mapper(drop_args=True, drop_kwargs=True, inline_get_cache_key=True) +class _DistributedDAGGatherer(CachedWalkMapper): + 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] = {} + + # type-ignore-reason: dropped the extra `*args, **kwargs`. + def get_cache_key(self, expr: ArrayOrNames) -> int: # type: ignore[override] + return id(expr) + + def map_distributed_recv( # type: ignore[override] + self, expr: DistributedRecv) -> None: + new_name = self.name_generator() + self.input_name_to_recv_node[new_name] = expr + + def map_distributed_send_ref_holder( # type: ignore[override] + self, expr: DistributedSendRefHolder) -> None: + s = _SummarizedDistributedSend( + src_rank=self.my_rank, + dest_rank=expr.send.dest_rank, + comm_tag=expr.send.comm_tag, + shape=expr.send.data.shape, + dtype=expr.send.data.dtype) + + new_name = self.name_generator() + self.output_name_to_send_node[new_name] = s + + def map_distributed_send(self, expr: DistributedSend) -> None: + raise ValueError("Unpartitioned DAG should not have DistributedSend nodes") + + +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) + + summarized_dag = _SummarizedDistributedGraph( + rank=my_rank, + 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, + shape=send.shape, + dtype=send.dtype) + 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)} + + print(f"{all_summarized_outputs=}") + + all_recvs: Set[_CommIdentifier] = set() + comm_id_to_sending_dag: Dict[_CommIdentifier, _SummarizedDistributedGraph] \ + = {} + + # Every node in the graph is a _SummarizedDistributedGraph + dags_to_needed_dags: \ + Dict[_SummarizedDistributedGraph, Set[_SummarizedDistributedGraph]] = {} + + def add_needed_dag(dag: _SummarizedDistributedGraph, + needed_dag: _SummarizedDistributedGraph) -> None: + dags_to_needed_dags.setdefault(dag, set()).add(needed_dag) + + 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_dag: + raise DuplicateSendError( + f"duplicate send for comm id: '{comm_id}'") + comm_id_to_sending_dag[comm_id] = sumdag + + 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_dag = comm_id_to_sending_dag[comm_id] + except KeyError: + raise MissingSendError( + f"no matching send for recv on '{comm_id}'") + + add_needed_dag(sumdag, sending_dag) + + print(dags_to_needed_dags) + from pytools.graph import compute_topological_order + compute_topological_order(dags_to_needed_dags) + + logger.info("verify_distributed_dag_pre_partition completed successfully.") + + def verify_distributed_partition(mpi_communicator: mpi4py.MPI.Comm, partition: DistributedGraphPartition) -> None: """ From 824ecb81134bbf119a1c1a665b1d21aff52a5921 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Mon, 21 Nov 2022 22:39:54 -0600 Subject: [PATCH 2/9] derive from UsersCollector --- pytato/distributed.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/pytato/distributed.py b/pytato/distributed.py index 9ff77cd3e..a561a4601 100644 --- a/pytato/distributed.py +++ b/pytato/distributed.py @@ -40,7 +40,7 @@ NamedArray) from pytato.transform import (ArrayOrNames, CopyMapper, Mapper, CachedWalkMapper, CopyMapperWithExtraArgs, - CombineMapper) + CombineMapper, UsersCollector) from pytato.partition import GraphPart, GraphPartition, PartId, GraphPartitioner from pytato.target import BoundProgram @@ -48,6 +48,7 @@ from functools import cached_property from pytato.scalar_expr import SCALAR_CLASSES, INT_CLASSES from pymbolic.mapper.optimize import optimize_mapper +from immutables import Map import numpy as np @@ -1084,6 +1085,11 @@ class _SummarizedDistributedGraph: 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.input_name_to_recv_node)) + ^ hash(Map(self.output_name_to_send_node))) + @attrs.define(frozen=True) class _CommIdentifier: @@ -1112,8 +1118,7 @@ class MissingRecvError(DistributedPartitionVerificationError): pass -@optimize_mapper(drop_args=True, drop_kwargs=True, inline_get_cache_key=True) -class _DistributedDAGGatherer(CachedWalkMapper): +class _DistributedDAGGatherer(UsersCollector): def __init__(self, dist_name_generator: UniqueNameGenerator, my_rank: int) \ -> None: super().__init__() @@ -1124,17 +1129,15 @@ def __init__(self, dist_name_generator: UniqueNameGenerator, my_rank: int) \ self.input_name_to_recv_node: Dict[str, DistributedRecv] = {} self.output_name_to_send_node: Dict[str, _SummarizedDistributedSend] = {} - # type-ignore-reason: dropped the extra `*args, **kwargs`. - def get_cache_key(self, expr: ArrayOrNames) -> int: # type: ignore[override] - return id(expr) - def map_distributed_recv( # type: ignore[override] 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( # type: ignore[override] self, expr: DistributedSendRefHolder) -> None: + super().map_distributed_send_ref_holder(expr) s = _SummarizedDistributedSend( src_rank=self.my_rank, dest_rank=expr.send.dest_rank, @@ -1146,7 +1149,8 @@ def map_distributed_send_ref_holder( # type: ignore[override] self.output_name_to_send_node[new_name] = s def map_distributed_send(self, expr: DistributedSend) -> None: - raise ValueError("Unpartitioned DAG should not have DistributedSend nodes") + raise ValueError("Unpartitioned DAG should not have DistributedSend nodes. " + "outside of DistributedSendRefHolder nodes.") def verify_distributed_dag_pre_partition(mpi_communicator: mpi4py.MPI.Comm, @@ -1234,7 +1238,7 @@ def add_needed_dag(dag: _SummarizedDistributedGraph, add_needed_dag(sumdag, sending_dag) - print(dags_to_needed_dags) + print(f"{dags_to_needed_dags=}") from pytools.graph import compute_topological_order compute_topological_order(dags_to_needed_dags) From cc4642796197ed59381fe91868fcdf0af996a9ce Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Mon, 28 Nov 2022 13:48:59 -0600 Subject: [PATCH 3/9] expose verify_distributed_dag_pre_partition globally --- pytato/__init__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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", From ba7160d8daca85147442d41226e74fd3e5f2bd23 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Mon, 28 Nov 2022 13:55:39 -0600 Subject: [PATCH 4/9] add to tests --- test/test_distributed.py | 37 ++++++++++++++++++++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) 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: From 451c6431f9e7afebedbf5d5e1eb73b8dbe1cff91 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Thu, 1 Dec 2022 23:56:24 -0600 Subject: [PATCH 5/9] use node_to_users --- pytato/distributed/verify.py | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/pytato/distributed/verify.py b/pytato/distributed/verify.py index fd58219e9..1f9259a6d 100644 --- a/pytato/distributed/verify.py +++ b/pytato/distributed/verify.py @@ -80,6 +80,12 @@ class _DistributedName: name: str +@attrs.define(frozen=True) +class _DistributedNode: + rank: int + node: str + + @attrs.define(frozen=True) class _SummarizedDistributedGraphPart: pid: _DistributedPartId @@ -98,11 +104,13 @@ def rank(self) -> int: @attrs.define(frozen=True) class _SummarizedDistributedGraph: rank: int + node_to_users: Dict[ArrayOrNames, Set[ArrayOrNames]] 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))) @@ -194,6 +202,7 @@ def verify_distributed_dag_pre_partition(mpi_communicator: mpi4py.MPI.Comm, summarized_dag = _SummarizedDistributedGraph( rank=my_rank, + node_to_users={_DistributedNode(my_rank, k): frozenset((_DistributedNode(my_rank, 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={ @@ -219,8 +228,6 @@ def verify_distributed_dag_pre_partition(mpi_communicator: mpi4py.MPI.Comm, print(f"{all_summarized_outputs=}") all_recvs: Set[_CommIdentifier] = set() - comm_id_to_sending_dag: Dict[_CommIdentifier, _SummarizedDistributedGraph] \ - = {} # Every node in the graph is a _SummarizedDistributedGraph dags_to_needed_dags: \ @@ -230,6 +237,8 @@ def add_needed_dag(dag: _SummarizedDistributedGraph, needed_dag: _SummarizedDistributedGraph) -> None: dags_to_needed_dags.setdefault(dag, set()).add(needed_dag) + + comm_id_to_sending_node: Dict[_CommIdentifier, _SummarizedDistributedGraph] = {} for sumdag in all_summarized_outputs.values(): for sumsend in sumdag.output_name_to_send_node.values(): comm_id = _CommIdentifier( @@ -237,11 +246,14 @@ def add_needed_dag(dag: _SummarizedDistributedGraph, dest_rank=sumsend.dest_rank, comm_tag=sumsend.comm_tag) - if comm_id in comm_id_to_sending_dag: + if comm_id in comm_id_to_sending_node: raise DuplicateSendError( f"duplicate send for comm id: '{comm_id}'") - comm_id_to_sending_dag[comm_id] = sumdag + comm_id_to_sending_node[comm_id] = sumsend + + for sumdag in all_summarized_outputs.values(): + dags_to_needed_dags.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, @@ -255,14 +267,13 @@ def add_needed_dag(dag: _SummarizedDistributedGraph, # Add edges between sends and receives (cross-rank) try: - sending_dag = comm_id_to_sending_dag[comm_id] + sending_node = comm_id_to_sending_node[comm_id] except KeyError: raise MissingSendError( f"no matching send for recv on '{comm_id}'") + add_needed_dag(_DistributedNode(comm_id.dest_rank, dist_recv), _DistributedNode(comm_id.dest_rank, sending_node)) - add_needed_dag(sumdag, sending_dag) - print(f"{dags_to_needed_dags=}") from pytools.graph import compute_topological_order compute_topological_order(dags_to_needed_dags) From e5ca9013b088b3630bb1b7d42925b226e1aad0b6 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Thu, 1 Dec 2022 23:57:00 -0600 Subject: [PATCH 6/9] remove map_distributed_send --- pytato/distributed/verify.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/pytato/distributed/verify.py b/pytato/distributed/verify.py index 1f9259a6d..02f2fa49d 100644 --- a/pytato/distributed/verify.py +++ b/pytato/distributed/verify.py @@ -180,10 +180,6 @@ def map_distributed_send_ref_holder( # type: ignore[override] new_name = self.name_generator() self.output_name_to_send_node[new_name] = s - def map_distributed_send(self, expr: DistributedSend) -> None: - raise ValueError("Unpartitioned DAG should not have DistributedSend nodes. " - "outside of DistributedSendRefHolder nodes.") - def verify_distributed_dag_pre_partition(mpi_communicator: mpi4py.MPI.Comm, outputs: DictOfNamedArrays) -> None: From 8a5df37ecebb8bb19d5411518bd1673479844581 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Mon, 5 Dec 2022 14:07:41 -0600 Subject: [PATCH 7/9] remove shape, dtype from summarizeddistsend --- pytato/distributed/verify.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/pytato/distributed/verify.py b/pytato/distributed/verify.py index 02f2fa49d..1739a68fe 100644 --- a/pytato/distributed/verify.py +++ b/pytato/distributed/verify.py @@ -30,17 +30,14 @@ """ -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, DistributedSendRefHolder, DistributedSend) from pytato.partition import PartId from pytato.distributed.partition import DistributedGraphPartition -from pytato.array import ShapeType -from pytato.transform import UsersCollector +from pytato.transform import UsersCollector, ArrayOrNames from pytato import DictOfNamedArrays from pytools import UniqueNameGenerator @@ -64,9 +61,6 @@ class _SummarizedDistributedSend: dest_rank: int comm_tag: CommTagType - shape: ShapeType - dtype: np.dtype[Any] - @attrs.define(frozen=True) class _DistributedPartId: @@ -325,9 +319,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 From 2ca2810b5ed34ab8c32cfd0d234ab54e075a1443 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Mon, 5 Dec 2022 14:16:29 -0600 Subject: [PATCH 8/9] finish impl --- pytato/distributed/verify.py | 66 ++++++++++++++++++++++-------------- 1 file changed, 41 insertions(+), 25 deletions(-) diff --git a/pytato/distributed/verify.py b/pytato/distributed/verify.py index 1739a68fe..4001a789d 100644 --- a/pytato/distributed/verify.py +++ b/pytato/distributed/verify.py @@ -77,7 +77,7 @@ class _DistributedName: @attrs.define(frozen=True) class _DistributedNode: rank: int - node: str + node: Union[DistributedSend, ArrayOrNames, _SummarizedDistributedSend] @attrs.define(frozen=True) @@ -98,7 +98,7 @@ def rank(self) -> int: @attrs.define(frozen=True) class _SummarizedDistributedGraph: rank: int - node_to_users: Dict[ArrayOrNames, Set[ArrayOrNames]] + node_to_users: Dict[_DistributedNode, Set[_DistributedNode]] input_name_to_recv_node: Dict[_DistributedName, DistributedRecv] output_name_to_send_node: Dict[_DistributedName, _SummarizedDistributedSend] @@ -155,21 +155,18 @@ def __init__(self, dist_name_generator: UniqueNameGenerator, my_rank: int) \ self.input_name_to_recv_node: Dict[str, DistributedRecv] = {} self.output_name_to_send_node: Dict[str, _SummarizedDistributedSend] = {} - def map_distributed_recv( # type: ignore[override] - self, expr: DistributedRecv) -> None: + 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( # type: ignore[override] + 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, - shape=expr.send.data.shape, - dtype=expr.send.data.dtype) + comm_tag=expr.send.comm_tag) new_name = self.name_generator() self.output_name_to_send_node[new_name] = s @@ -190,9 +187,29 @@ def verify_distributed_dag_pre_partition(mpi_communicator: mpi4py.MPI.Comm, 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): frozenset((_DistributedNode(my_rank, n) for n in v)) for k, v in dg.node_to_users.items()}, + 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={ @@ -200,9 +217,7 @@ def verify_distributed_dag_pre_partition(mpi_communicator: mpi4py.MPI.Comm, _SummarizedDistributedSend( src_rank=my_rank, dest_rank=send.dest_rank, - comm_tag=send.comm_tag, - shape=send.shape, - dtype=send.dtype) + comm_tag=send.comm_tag,) for name, send in dg.output_name_to_send_node.items()}) all_outputs = \ @@ -215,20 +230,19 @@ def verify_distributed_dag_pre_partition(mpi_communicator: mpi4py.MPI.Comm, rank: rank_outputs for rank, rank_outputs in enumerate(all_outputs)} - print(f"{all_summarized_outputs=}") - all_recvs: Set[_CommIdentifier] = set() - # Every node in the graph is a _SummarizedDistributedGraph - dags_to_needed_dags: \ - Dict[_SummarizedDistributedGraph, Set[_SummarizedDistributedGraph]] = {} + send_recv_deps: \ + Dict[_DistributedNode, Set[_DistributedNode]] = {} - def add_needed_dag(dag: _SummarizedDistributedGraph, - needed_dag: _SummarizedDistributedGraph) -> None: - dags_to_needed_dags.setdefault(dag, set()).add(needed_dag) + 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 = {} - comm_id_to_sending_node: Dict[_CommIdentifier, _SummarizedDistributedGraph] = {} for sumdag in all_summarized_outputs.values(): for sumsend in sumdag.output_name_to_send_node.values(): comm_id = _CommIdentifier( @@ -241,9 +255,11 @@ def add_needed_dag(dag: _SummarizedDistributedGraph, f"duplicate send for comm id: '{comm_id}'") comm_id_to_sending_node[comm_id] = sumsend + # }}} for sumdag in all_summarized_outputs.values(): - dags_to_needed_dags.update(sumdag.node_to_users) + 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, @@ -261,11 +277,11 @@ def add_needed_dag(dag: _SummarizedDistributedGraph, except KeyError: raise MissingSendError( f"no matching send for recv on '{comm_id}'") - add_needed_dag(_DistributedNode(comm_id.dest_rank, dist_recv), _DistributedNode(comm_id.dest_rank, sending_node)) - + 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(dags_to_needed_dags) + compute_topological_order(send_recv_deps) logger.info("verify_distributed_dag_pre_partition completed successfully.") From 41faa7c860e9be833c5ca0b3c73876f3aa22e2f2 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 9 Dec 2022 20:40:32 -0600 Subject: [PATCH 9/9] minor fixes --- examples/distributed.py | 2 +- pytato/distributed/verify.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/distributed.py b/examples/distributed.py index e3bf11a5f..e1eba67aa 100644 --- a/examples/distributed.py +++ b/examples/distributed.py @@ -31,7 +31,7 @@ def main(): # Find the partition outputs = pt.DictOfNamedArrays({"out": y}) - pt.distributed.verify_distributed_dag_pre_partition(comm, outputs) + 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/distributed/verify.py b/pytato/distributed/verify.py index 4001a789d..8354e0865 100644 --- a/pytato/distributed/verify.py +++ b/pytato/distributed/verify.py @@ -220,8 +220,7 @@ def dist_send_to_summarized_dist_send(node: 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) + all_outputs = mpi_communicator.gather(summarized_dag, root=root_rank) if my_rank == root_rank: assert all_outputs @@ -236,7 +235,7 @@ def dist_send_to_summarized_dist_send(node: Dict[_DistributedNode, Set[_DistributedNode]] = {} def add_send_recv_dep(recv: _DistributedNode, - send: _DistributedNode) -> None: + send: _DistributedNode) -> None: send_recv_deps.setdefault(recv, set()).add(send) # {{{ gather information on senders @@ -277,6 +276,7 @@ def add_send_recv_dep(recv: _DistributedNode, 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))