Repository: pyg-team/pytorch_geometric
Pull request #7452 (2023-05-27 23:35:28)

Add `LocalFeatureStore` for distributed training
This code belongs to the part of the whole distributed training for PyG.

This class extends from FeatureStore and use the store dict to save the node/edge features and also use global_idx & id2index dict to save the global node ids and the id mapping between node ids and local feature index, these of which will be used by sampling and feature lookup stage.

We also include the unit test under /test/distributed folder to show how the node/edge feature will be saving/getting and verified by two kinds of method.

Any comments please let us know. thanks

ZhengHongming888: @rusty1s anything i need change here? :-) Thanks.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7452?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7452](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7452?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (4fcba1d) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/af4360a2ed5490e17b12d06f46bca13e6bbf46e1?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (af4360a) will **decrease** coverage by `0.25%`.
> The diff coverage is `83.90%`.

```diff
@@            Coverage Diff             @@
##           master    #7452      +/-   ##
==========================================
- Coverage   91.68%   91.44%   -0.25%     
==========================================
  Files         445      446       +1     
  Lines       24699    24775      +76     
==========================================
+ Hits        22645    22655      +10     
- Misses       2054     2120      +66     
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7452?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/distributed/local\_feature\_store.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7452?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2Rpc3RyaWJ1dGVkL2xvY2FsX2ZlYXR1cmVfc3RvcmUucHk=) | `76.66% <76.66%> (ø)` | |
| [torch\_geometric/distributed/\_\_init\_\_.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7452?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2Rpc3RyaWJ1dGVkL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/sampler/neighbor\_sampler.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7452?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3NhbXBsZXIvbmVpZ2hib3Jfc2FtcGxlci5weQ==) | `94.79% <100.00%> (ø)` | |
| [torch\_geometric/utils/cross\_entropy.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7452?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL2Nyb3NzX2VudHJvcHkucHk=) | `100.00% <100.00%> (ø)` | |

... and [16 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7452/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 72aace80f..7ae81a384 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,7 +8,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ### Added
 
 - Added hierarichial heterogeneous GraphSAGE example on OGB-MAG ([#7425](https://github.com/pyg-team/pytorch_geometric/pull/7425))
-- Added a `LocalGraphStore` implementation for distributed training ([#7451](https://github.com/pyg-team/pytorch_geometric/pull/7451))
+- Added the `torch_geometric.distributed` package ([#7451](https://github.com/pyg-team/pytorch_geometric/pull/7451), [#7452](https://github.com/pyg-team/pytorch_geometric/pull/7452))
 - 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
index 489e0063d..1412f174a 100644
--- a/torch_geometric/distributed/__init__.py
+++ b/torch_geometric/distributed/__init__.py
@@ -1,5 +1,7 @@
+from .local_feature_store import LocalFeatureStore
 from .local_graph_store import LocalGraphStore
 
 __all__ = classes = [
+    'LocalFeatureStore',
     'LocalGraphStore',
 ]
diff --git a/torch_geometric/distributed/local_feature_store.py b/torch_geometric/distributed/local_feature_store.py
new file mode 100644
index 000000000..5abbfa2f1
--- /dev/null
+++ b/torch_geometric/distributed/local_feature_store.py
@@ -0,0 +1,102 @@
+import copy
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+import torch
+from torch import Tensor
+
+from torch_geometric.data import FeatureStore, TensorAttr
+from torch_geometric.data.feature_store import _field_status
+
+
+@dataclass
+class LocalTensorAttr(TensorAttr):
+    r"""Tensor attribute for storing features without :obj:`index`."""
+    def __init__(
+        self,
+        group_name: Optional[str] = _field_status.UNSET,
+        attr_name: Optional[str] = _field_status.UNSET,
+        index=None,
+    ):
+        super().__init__(group_name, attr_name, index)
+
+
+class LocalFeatureStore(FeatureStore):
+    r"""This class implements the :class:`torch_geometric.data.FeatureStore`
+    interface to act as a local feature store for distributed training."""
+    def __init__(self):
+        super().__init__(tensor_attr_cls=LocalTensorAttr)
+
+        self._feat: Dict[Tuple[str, str], Tensor] = {}
+
+        # Save the global node/edge IDs:
+        self._global_id: Dict[Tuple[str, str], Tensor] = {}
+
+        # Save the mapping from global node/edge IDs to indices in `_feat`:
+        self._global_id_to_index: Dict[Tuple[str, str], Tensor] = {}
+
+    @staticmethod
+    def key(attr: TensorAttr) -> Tuple[str, str]:
+        return (attr.group_name, attr.attr_name)
+
+    def put_global_id(self, global_id: Tensor, *args, **kwargs) -> bool:
+        attr = self._tensor_attr_cls.cast(*args, **kwargs)
+        self._global_id[self.key(attr)] = global_id
+        self._set_global_id_to_index(attr)
+        return True
+
+    def get_global_id(self, *args, **kwargs) -> Optional[Tensor]:
+        attr = self._tensor_attr_cls.cast(*args, **kwargs)
+        return self._global_id.get(self.key(attr))
+
+    def remove_global_id(self, *args, **kwargs) -> bool:
+        attr = self._tensor_attr_cls.cast(*args, **kwargs)
+        return self._global_id.pop(self.key(attr), None) is not None
+
+    def _set_global_id_to_index(self, *args, **kwargs):
+        attr = self._tensor_attr_cls.cast(*args, **kwargs)
+        global_id = self.get_global_id(attr)
+
+        if global_id is None:
+            return
+
+        # TODO Compute this mapping without materializing a full-sized tensor:
+        global_id_to_index = global_id.new_full((int(global_id.max()) + 1, ),
+                                                fill_value=-1)
+        global_id_to_index[global_id] = torch.arange(global_id.numel())
+        self._global_id_to_index[self.key(attr)] = global_id_to_index
+
+    def _put_tensor(self, tensor: Tensor, attr: TensorAttr) -> bool:
+        assert attr.index is None
+        self._feat[self.key(attr)] = tensor
+        return True
+
+    def _get_tensor(self, attr: TensorAttr) -> Optional[Tensor]:
+        tensor = self._feat.get(self.key(attr))
+
+        if tensor is None:
+            return None
+
+        if attr.index is None:  # Empty indices return the full tensor:
+            return tensor
+
+        return tensor[attr.index]
+
+    def _remove_tensor(self, attr: TensorAttr) -> bool:
+        assert attr.index is None
+        return self._feat.pop(self.key(attr), None) is not None
+
+    def get_tensor_from_global_id(self, *args, **kwargs) -> Optional[Tensor]:
+        attr = self._tensor_attr_cls.cast(*args, **kwargs)
+        assert attr.index is not None
+
+        attr = copy.copy(attr)
+        attr.index = self._global_id_to_index[self.key(attr)][attr.index]
+
+        return self.get_tensor(attr)
+
+    def _get_tensor_size(self, attr: TensorAttr) -> Tuple[int, ...]:
+        return self._get_tensor(attr).size()
+
+    def get_all_tensor_attrs(self) -> List[LocalTensorAttr]:
+        return [self._tensor_attr_cls.cast(*key) for key in self._feat.keys()]
diff --git a/torch_geometric/distributed/local_graph_store.py b/torch_geometric/distributed/local_graph_store.py
index aa5281bd3..d83e301a1 100644
--- a/torch_geometric/distributed/local_graph_store.py
+++ b/torch_geometric/distributed/local_graph_store.py
@@ -8,8 +8,7 @@ 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.
-    """
+    interface to act as a local graph store for distributed training."""
     def __init__(self):
         super().__init__()
         self._edge_index: Dict[Tuple, EdgeTensorType] = {}
@@ -26,7 +25,7 @@ class LocalGraphStore(GraphStore):
 
     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)]
+        return self._edge_id.get(self.key(edge_attr))
 
     def remove_edge_id(self, *args, **kwargs) -> bool:
         edge_attr = self._edge_attr_cls.cast(*args, **kwargs)


Test patch:
diff --git a/test/distributed/test_local_feature_store.py b/test/distributed/test_local_feature_store.py
new file mode 100644
index 000000000..1ccb15b3e
--- /dev/null
+++ b/test/distributed/test_local_feature_store.py
@@ -0,0 +1,59 @@
+import torch
+
+from torch_geometric.distributed import LocalFeatureStore
+
+
+def test_local_feature_store_global_id():
+    store = LocalFeatureStore()
+
+    feat = torch.Tensor([
+        [0, 0, 0],
+        [1, 1, 1],
+        [2, 2, 2],
+        [3, 3, 3],
+        [4, 4, 4],
+        [5, 5, 5],
+        [6, 6, 6],
+        [7, 7, 7],
+        [8, 8, 8],
+    ])
+
+    kwargs = dict(group_name='part1', attr_name='feat')
+    part1_global_id = torch.tensor([1, 2, 3, 5, 8, 4])
+    part1_feat = feat[part1_global_id]
+
+    store.put_global_id(part1_global_id, **kwargs)
+    store.put_tensor(part1_feat, **kwargs)
+
+    out = store.get_tensor_from_global_id(index=torch.tensor([3, 8, 4]),
+                                          **kwargs)
+    assert torch.equal(out, feat[torch.tensor([3, 8, 4])])
+
+
+def test_local_feature_store_utils():
+    store = LocalFeatureStore()
+
+    feat = torch.Tensor([
+        [0, 0, 0],
+        [1, 1, 1],
+        [2, 2, 2],
+        [3, 3, 3],
+        [4, 4, 4],
+        [5, 5, 5],
+        [6, 6, 6],
+        [7, 7, 7],
+        [8, 8, 8],
+    ])
+
+    kwargs = dict(group_name='part1', attr_name='feat')
+    part1_global_id = torch.tensor([1, 2, 3, 5, 8, 4])
+    part1_feat = feat[part1_global_id]
+
+    store.put_tensor(part1_feat, **kwargs)
+
+    assert len(store.get_all_tensor_attrs()) == 1
+    attr = store.get_all_tensor_attrs()[0]
+    assert attr.group_name == 'part1'
+    assert attr.attr_name == 'feat'
+    assert attr.index is None
+    assert store.get_tensor_size(attr) == (6, 3)
diff --git a/torch_geometric/testing/feature_store.py b/torch_geometric/testing/feature_store.py
index 55fa633e1..fb74ba982 100644
--- a/torch_geometric/testing/feature_store.py
+++ b/torch_geometric/testing/feature_store.py
@@ -24,12 +24,12 @@ class MyFeatureStore(FeatureStore):
             index = torch.arange(0, tensor.shape[0])
 
         # Store the index:
-        self.store[MyFeatureStore.key(attr)] = (index, tensor)
+        self.store[self.key(attr)] = (index, tensor)
 
         return True
 
     def _get_tensor(self, attr: TensorAttr) -> Optional[FeatureTensorType]:
-        index, tensor = self.store.get(MyFeatureStore.key(attr), (None, None))
+        index, tensor = self.store.get(self.key(attr), (None, None))
         if tensor is None:
             return None
 
@@ -47,14 +47,13 @@ class MyFeatureStore(FeatureStore):
         return tensor[idx]
 
     def _remove_tensor(self, attr: TensorAttr) -> bool:
-        del self.store[MyFeatureStore.key(attr)]
-        return True
+        return self.store.pop(self.key(attr), None) is not None
 
-    def _get_tensor_size(self, attr: TensorAttr) -> Tuple:
+    def _get_tensor_size(self, attr: TensorAttr) -> Tuple[int, ...]:
         return self._get_tensor(attr).size()
 
-    def get_all_tensor_attrs(self) -> List[str]:
-        return [TensorAttr(*key) for key in self.store.keys()]
+    def get_all_tensor_attrs(self) -> List[TensorAttr]:
+        return [self._tensor_attr_cls.cast(*key) for key in self.store.keys()]
 
     def __len__(self):
         raise NotImplementedError