Add partitioning for distributed training

[ZhengHongming888]

This code belongs to the part of the whole distributed training for PyG.

This class (partitioner.py) will implement

1. the partition algorithm based on pyg's clusterData
2. in each partition LocalGraphStore/LocalFeatureStore will be used to initialize the graph & node/edge feature data
3. each partition also contains the partition information book/map between node/edge ids and partition id
4. each of which above will be further saved as .pt file folders include graph/node_feat/edge_feat/labels/node_map/edge_map.

The partition folders as below-

• homo graph
  output_dir/
  |-- META.json
  |-- node_map.pt
  |-- edge_map.pt
  |-- part0/

             |-- graph.pt
             |-- node_feats.pt
             |-- edge_feats.pt

  |-- part1/
             |-- graph.pt
             |-- node_feats.pt
             |-- edge_feats.pt

* hetero graph
  output_dir/
  |-- META.json
  |-- node_map/
             |-- ntype1.pt
             |-- ntype2.pt
  |-- edge_map/
            |-- etype1.pt
            |-- etype2.pt
  |-- part0/
            |-- graph.pt
            |-- node_feats.pt
            |-- edge_feats.pt
  |-- part1/
           |-- graph.pt
           |-- node_feats.pt
           |-- edge_feats.pt

We also provide two example codes to help generate the homo/hetero graph partition based on ogbn-products/ogbn-mags under example/distributed folder.

One unit test code under /test folder is used to verify this partition algorithm based on FakeDataset/FakeHeteroDataset.

Any comments please let us know. thanks

[codecov]

Codecov Report

Merging #7502 (f8c98e4) into master (5c72f33) will decrease coverage by 0.29%.
The diff coverage is 99.10%.

❗ Current head f8c98e4 differs from pull request most recent head 89f8eb8. Consider uploading reports for the commit 89f8eb8 to get more accurate results

@@            Coverage Diff             @@
##           master    #7502      +/-   ##
==========================================
- Coverage   91.74%   91.45%   -0.29%     
==========================================
  Files         450      451       +1     
  Lines       25161    25270     +109     
==========================================
+ Hits        23084    23111      +27     
- Misses       2077     2159      +82     

Impacted Files	Coverage Δ
torch_geometric/distributed/partition.py	99.09% <99.09%> (ø)
torch_geometric/distributed/__init__.py	100.00% <100.00%> (ø)

... and 17 files with indirect coverage changes

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[kaixuanliu]

Hi @rusty1s , I added a unit test for graph partitioning, while it returns error ImportError: 'ClusterData' requires either 'pyg-lib' or 'torch-sparse' What should I do to avoid this?

[akihironitta]

@kaixuanliu You can place the WithPackage decorator to skip the test case when those optional packages are not present in the environment:

pytorch_geometric/test/utils/test_scatter.py

Lines 25 to 27 in 0f0e0da

	@withPackage('torch_scatter')
	@pytest.mark.parametrize('reduce', ['sum', 'add', 'mean', 'min', 'max'])
	def test_scatter(reduce, device):

[rusty1s]

Thank you. I cleaned this up a bit. Especially, Partitioner no longer stores LocalFeatureStore and LocalGraphStore instances, as it is bad practice to pickle arbitrary Python objects. Instead, it saves Python dictionaries now, so please construct LocalFeatureStore and LocalGraphStore instances from this when loading the data from disk.

Otherwise, looks good. I would be in favor of just adding the homogeneous code path though. I am not really 100% confident the heterogeneous code path is correct.

[kaixuanliu]

Thank you. I cleaned this up a bit. Especially, Partitioner no longer stores LocalFeatureStore and LocalGraphStore instances, as it is bad practice to pickle arbitrary Python objects. Instead, it saves Python dictionaries now, so please construct LocalFeatureStore and LocalGraphStore instances from this when loading the data from disk.

Otherwise, looks good. I would be in favor of just adding the homogeneous code path though. I am not really 100% confident the heterogeneous code path is correct.

Thanks Matthias! It is a better practice to replace LocalFeatureStore/LocalGraphStore with python dictionary, as it does not need customized data structure. For heterogeneous graph partition, we have checked the output of partition for ogbn-mags dataset, and we will further validate its correctness in later development for hetero graph distributed training.