[Roadmap] Remote Backend Support and Integration 🚀
See original GitHub issueMotivation
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
andHeteroData
implement theFeatureStore
abstraction (#4807) - Define a
GraphStore
abstraction that is intended to hold anedge_index
in memory (#4816) - Let
Data
andHeteroData
implement theGraphStore
abstraction (#4816) - Modify
NeighborLoader
to callFeatureStore
andGraphStore
methods instead of theirData/HeteroData
counterparts. Note that this will require moving filtering of data into the feature store. The new interface will look likedata: Union[Union[Data, HeteroData], Tuple[FeatureStore, GraphStore]]
(#4817, #4883) - Implement
BaseSampler
and refactor existing samplers behind a common interface (#5312, #5365, #5402) - Introduce
NodeLoader
andLinkLoader
, refactor existing loaders behind loader + sampler interface (#5404, #5418) - Support (optional) methods to obtain a
TensorAttr
orEdgeAttr
from aFeatureStore
/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
andLightningLinkData
Implementations
- Implement a concrete
FeatureStore
,GraphStore
, andSampler
with a popular backend to provide example usage. Some thoughts here include a RayRandomAccessDataset
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 thatData
andHeteroData
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
, andTuple[FeatureStore, GraphStore]
throughout the PyG codebase into a single conditional. This should be possible since bothData
andHeteroData
areFeatureStore
andGraphStore
s
Issue Analytics
- State:
- Created a year ago
- Reactions:4
- Comments:8 (8 by maintainers)
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!
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.