[Roadmap] Remote Backend Support and Integration 🚀

Motivation

PyG currently requires users to store graphs (and associated node + edge features) in Data and HeteroData objects, which are accepted by loaders to run forward/backward passes on an accelerator of choice. This abstraction, however, does not scale to large graphs (or large feature tensors), which can quickly oversubscribe CPU DRAM (despite the GPU VRAM requirements only being the memory consumption of each sampled subgraph and its associated node and edge features). Indeed, one can imagine storing graph features (and the graph itself) in “remote backends”, which provide fixed operators that can be used to integrate cleanly with downstream PyG samplers and loaders.

The goal of this roadmap is to track the integration of native remote backend support into PyG. At a high level, this will be accomplished through the feature store, graph store, and sampler abstractions into PyG. For more freeform discussion, please visit the #scalability channel in the PyG Slack community.

Implementation

Abstractions: FeatureStore, GraphStore, Sampler

• Let Data and HeteroData implement the FeatureStore abstraction (#4807)
• Define a GraphStore abstraction that is intended to hold an edge_index in memory (#4816)
• Let Data and HeteroData implement the GraphStore abstraction (#4816)
• Modify NeighborLoader to call FeatureStore and GraphStore methods instead of their Data/HeteroData counterparts. Note that this will require moving filtering of data into the feature store. The new interface will look like data: Union[Union[Data, HeteroData], Tuple[FeatureStore, GraphStore]] (#4817, #4883)
• Implement BaseSampler and refactor existing samplers behind a common interface (#5312, #5365, #5402)
• Introduce NodeLoader and LinkLoader, refactor existing loaders behind loader + sampler interface (#5404, #5418)
• Support (optional) methods to obtain a TensorAttr or EdgeAttr from a FeatureStore/GraphStore from their first dataclass attribute, and refactor existing computations that get all (tensors, edges) and subsequently filter to use these methods.
• Support variable samplers in LightningNodeData and LightningLinkData

Implementations

• Implement a concrete FeatureStore, GraphStore, and Sampler with a popular backend to provide example usage. Some thoughts here include a Ray RandomAccessDataset for a feature store and a Neo4j graph for a graph store.
• Implement a validation class that operates on Tuple[FeatureStore, GraphStore] to perform basic sanity checks (in a similar way that Data and HeteroData do today)
• Implement sampling from edges in the HGTSampler
• Implement (to the extent possible) samplers in torch_geometric/loader (e.g. GraphSAINT, ShaDow) behind the sampler interface, enabling (a) easy extension to sampling from edges and (b) ease of extension to reote backedns in the future.

Code Health

• Implement a remote backend utility class to consolidate common methods across feature and graph stores (#5307)
• Consolidate conditionals for Data, HeteroData, and Tuple[FeatureStore, GraphStore] throughout the PyG codebase into a single conditional. This should be possible since both Data and HeteroData are FeatureStore and GraphStores

[mananshah99]

Hi folks, this roadmap has been updated a bit to describe latest changes and a few potential further directions (cc @Padarn, I hope this helps address some of your questions as well). Feel free to add on, or let me know if you have any questions/comments/concerns!

[rusty1s]

Yes, @wsad1. I think these are good points. One thing we could do to showcase is to have a short example/tutorial on how to connect to a neo4j graph database or similar.