cugraph_pyg.data.GraphStore#

class cugraph_pyg.data.GraphStore(location='cpu')[source]#

cuGraph-backed PyG GraphStore implementation that distributes the graph across workers. Users can repeatedly call put_edge_index to store edge indices in WholeGraph memory. The final cuGraph graph is always constructed in GPU device memory, regardless of the WholeGraph memory location. Construction occurs lazily when the graph is needed by a loader, or eagerly when the user calls finalize().

Supports single-node/single-GPU, single-node/multi-GPU, and multi-node/multi-GPU graph storage.

Each worker should have a slice of the graph locally, and call put_edge_index with its slice.

When a COO DistMatrix is passed to put_edge_index, it is stored by reference and its data is not copied. Likewise, when a pair of DistTensor objects is passed, the resulting matrix retains references to those tensors. Subsequent changes to these distributed objects are therefore visible to the store.

Attributes:

is_homogeneous
is_multi_gpu

Methods

`finalize`([weight_attr, time_attr])	Finalizes the graph store, constructing the cuGraph graph object on device, and deleting the precusor edge index tensors.
`get_all_edge_attrs`()	Returns all registered edge attributes.

__init__(location='cpu')[source]#

Constructs a new, empty GraphStore object. This object represents one slice of a graph on particular worker.

Parameters:

location: str(optional, default=’cpu’): The WholeGraph memory location (‘cpu’ or ‘cuda’) used to store edge indices before the cuGraph graph is constructed. This does not affect the final graph location; the cuGraph graph is always constructed in GPU device memory when it’s needed.

Methods

`__init__`([location])	Constructs a new, empty GraphStore object.
`coo`([edge_types, store])	Returns the edge indices in the `GraphStore` in COO format.
`csc`([edge_types, store])	Returns the edge indices in the `GraphStore` in CSC format.
`csr`([edge_types, store])	Returns the edge indices in the `GraphStore` in CSR format.
`finalize`([weight_attr, time_attr])	Finalizes the graph store, constructing the cuGraph graph object on device, and deleting the precusor edge index tensors.
`get_all_edge_attrs`()	Returns all registered edge attributes.
`get_edge_index`(args, *kwargs)	Synchronously obtains an `edge_index` tuple from the `GraphStore`.
`put_edge_index`(edge_index, args, *kwargs)	Synchronously adds an `edge_index` tuple to the `GraphStore`.
`remove_edge_index`(args, *kwargs)	Synchronously deletes an `edge_index` tuple from the `GraphStore`.

Attributes

`is_homogeneous`
`is_multi_gpu`