Repository: pyg-team/pytorch_geometric
Pull request #7451 (2023-05-27 18:17:44)

Add `LocalGraphStore` for distributed training
This code belongs to the part of the whole distributed training for PyG. 

This class extends from GraphStore and use the store dict to save the graph topology and also use edge_ids dict to save the global eids. 

We also include the unit test under /test folder to show how the graph edge_index/ edge_ids will be saving/getting and verified. 

Any comments please let us know. thanks

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7451?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7451](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7451?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (f782b91) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/f2957584f04aa3a631baa79a303426e359f00310?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (f295758) will **increase** coverage by `0.00%`.
> The diff coverage is `96.00%`.

> :exclamation: Current head f782b91 differs from pull request most recent head 537928a. Consider uploading reports for the commit 537928a to get more accurate results

```diff
@@           Coverage Diff           @@
##           master    #7451   +/-   ##
=======================================
  Coverage   91.42%   91.43%           
=======================================
  Files         442      444    +2     
  Lines       24658    24681   +23     
=======================================
+ Hits        22543    22566   +23     
  Misses       2115     2115           
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7451?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/data/local\_graph\_store.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7451?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvbG9jYWxfZ3JhcGhfc3RvcmUucHk=) | `96.00% <96.00%> (ø)` | |

... and [3 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7451/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index fef2b8554..15ebe1b9c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added a `LocalGraphStore` implementation for distributed training ([#7451](https://github.com/pyg-team/pytorch_geometric/pull/7451))
 - Added the `GDELTLite` dataset ([#7442](https://github.com/pyg-team/pytorch_geometric/pull/7442))
 - Added the `approx_knn` function for approximated nearest neighbor search ([#7421](https://github.com/pyg-team/pytorch_geometric/pull/7421))
 - Added the `IGMCDataset` ([#7441](https://github.com/pyg-team/pytorch_geometric/pull/7441))
diff --git a/torch_geometric/distributed/__init__.py b/torch_geometric/distributed/__init__.py
new file mode 100644
index 000000000..489e0063d
--- /dev/null
+++ b/torch_geometric/distributed/__init__.py
@@ -0,0 +1,5 @@
+from .local_graph_store import LocalGraphStore
+
+__all__ = classes = [
+    'LocalGraphStore',
+]
diff --git a/torch_geometric/distributed/local_graph_store.py b/torch_geometric/distributed/local_graph_store.py
new file mode 100644
index 000000000..aa5281bd3
--- /dev/null
+++ b/torch_geometric/distributed/local_graph_store.py
@@ -0,0 +1,47 @@
+from typing import Dict, List, Optional, Tuple
+
+from torch import Tensor
+
+from torch_geometric.data import EdgeAttr, GraphStore
+from torch_geometric.typing import EdgeTensorType
+
+
+class LocalGraphStore(GraphStore):
+    r"""This class implements the :class:`torch_geometric.data.GraphStore`
+    interface to act as a local graph store for distributed training.
+    """
+    def __init__(self):
+        super().__init__()
+        self._edge_index: Dict[Tuple, EdgeTensorType] = {}
+        self._edge_id: Dict[Tuple, Tensor] = {}
+
+    @staticmethod
+    def key(attr: EdgeAttr) -> Tuple:
+        return (attr.edge_type, attr.layout.value, attr.is_sorted, attr.size)
+
+    def put_edge_id(self, edge_id: Tensor, *args, **kwargs) -> bool:
+        edge_attr = self._edge_attr_cls.cast(*args, **kwargs)
+        self._edge_id[self.key(edge_attr)] = edge_id
+        return True
+
+    def get_edge_id(self, *args, **kwargs) -> Optional[EdgeTensorType]:
+        edge_attr = self._edge_attr_cls.cast(*args, **kwargs)
+        return self._edge_id[self.key(edge_attr)]
+
+    def remove_edge_id(self, *args, **kwargs) -> bool:
+        edge_attr = self._edge_attr_cls.cast(*args, **kwargs)
+        return self._edge_id.pop(self.key(edge_attr), None) is not None
+
+    def _put_edge_index(self, edge_index: EdgeTensorType,
+                        edge_attr: EdgeAttr) -> bool:
+        self._edge_index[self.key(edge_attr)] = edge_index
+        return True
+
+    def _get_edge_index(self, edge_attr: EdgeAttr) -> Optional[EdgeTensorType]:
+        return self._edge_index.get(self.key(edge_attr), None)
+
+    def _remove_edge_index(self, edge_attr: EdgeAttr) -> bool:
+        return self._edge_index.pop(self.key(edge_attr), None) is not None
+
+    def get_all_edge_attrs(self) -> List[EdgeAttr]:
+        return [EdgeAttr(*key) for key in self._edge_index.keys()]


Test patch:
diff --git a/test/distributed/test_local_graph_store.py b/test/distributed/test_local_graph_store.py
new file mode 100644
index 000000000..2047db8ef
--- /dev/null
+++ b/test/distributed/test_local_graph_store.py
@@ -0,0 +1,26 @@
+import torch
+
+from torch_geometric.distributed import LocalGraphStore
+from torch_geometric.testing import get_random_edge_index
+
+
+def test_local_graph_store():
+    graph_store = LocalGraphStore()
+
+    edge_index = get_random_edge_index(100, 100, 300)
+    edge_id = torch.tensor([1, 2, 3, 5, 8, 4], dtype=torch.int64)
+
+    graph_store.put_edge_index(edge_index, edge_type=None, layout='coo',
+                               size=(100, 100))
+
+    graph_store.put_edge_id(edge_id, edge_type=None, layout='coo',
+                            size=(100, 100))
+
+    assert len(graph_store.get_all_edge_attrs()) == 1
+    edge_attr = graph_store.get_all_edge_attrs()[0]
+    assert torch.equal(graph_store.get_edge_index(edge_attr), edge_index)
+    assert torch.equal(graph_store.get_edge_id(edge_attr), edge_id)
+
+    graph_store.remove_edge_index(edge_attr)
+    graph_store.remove_edge_id(edge_attr)
+    assert len(graph_store.get_all_edge_attrs()) == 0