From adbefc6c77a6d3deac16fef0cfe2f2022ec56852 Mon Sep 17 00:00:00 2001 From: Wenxuan Cao <90617523+CfromBU@users.noreply.github.com> Date: Wed, 25 Sep 2024 10:40:07 +0800 Subject: [PATCH] [DistGB]Restore the gb_convert_single_dgl_partition of the public function in partition.py (#7803) Co-authored-by: Ubuntu Co-authored-by: Ubuntu Co-authored-by: Rhett Ying <85214957+Rhett-Ying@users.noreply.github.com> --- python/dgl/distributed/__init__.py | 1 + python/dgl/distributed/partition.py | 26 +++++++++++++------------- 2 files changed, 14 insertions(+), 13 deletions(-) diff --git a/python/dgl/distributed/__init__.py b/python/dgl/distributed/__init__.py index 947fadebc253..667dbc410598 100644 --- a/python/dgl/distributed/__init__.py +++ b/python/dgl/distributed/__init__.py @@ -17,6 +17,7 @@ from .nn import * from .partition import ( dgl_partition_to_graphbolt, + gb_convert_single_dgl_partition, load_partition, load_partition_book, load_partition_feats, diff --git a/python/dgl/distributed/partition.py b/python/dgl/distributed/partition.py index 48005ffb4d27..de782a5f144d 100644 --- a/python/dgl/distributed/partition.py +++ b/python/dgl/distributed/partition.py @@ -687,7 +687,7 @@ def _partition_to_graphbolt( part_config=part_config, part_id=part_i, part_metadata=part_metadata ) graph = parts[part_i] - csc_graph = gb_convert_single_dgl_partition( + csc_graph = _convert_dgl_partition_to_gb( ntypes=ntypes, etypes=etypes, gpb=gpb, @@ -1702,7 +1702,7 @@ def _create_attributes_gb( return node_attributes, edge_attributes, type_per_edge -def gb_convert_single_dgl_partition( +def _convert_dgl_partition_to_gb( ntypes, etypes, gpb, @@ -1820,30 +1820,30 @@ def gb_convert_single_dgl_partition( return csc_graph -def _convert_partition_to_graphbolt( - part_config, +def gb_convert_single_dgl_partition( part_id, - graph_formats=None, - store_eids=False, - store_inner_node=False, - store_inner_edge=False, + graph_formats, + part_config, + store_eids=True, + store_inner_node=True, + store_inner_edge=True, ): """ The pipeline converting signle partition to graphbolt. Parameters ---------- - part_config : str - The path of the partition config file. part_id : int The partition ID. - graph_formats : str or list[str], optional + graph_formats : str or list[str] Save partitions in specified formats. It could be any combination of `coo`, `csc`. As `csc` format is mandatory for `FusedCSCSamplingGraph`, it is not necessary to specify this argument. It's mainly for specifying `coo` format to save edge ID mapping and destination node IDs. If not specified, whether to save `coo` format is determined by the availability of the format in DGL partitions. Default: None. + part_config : str + The path of the partition config file. store_eids : bool, optional Whether to store edge IDs in the new graph. Default: True. store_inner_node : bool, optional @@ -1861,7 +1861,7 @@ def _convert_partition_to_graphbolt( ) part = _load_part(part_config, part_id) part_meta = copy.deepcopy(_load_part_config(part_config)) - csc_graph = gb_convert_single_dgl_partition( + csc_graph = _convert_dgl_partition_to_gb( graph=part, ntypes=ntypes, etypes=etypes, @@ -1896,7 +1896,7 @@ def _convert_partition_to_graphbolt_wrapper( # Iterate over partitions. convert_with_format = partial( - _convert_partition_to_graphbolt, + gb_convert_single_dgl_partition, part_config=part_config, graph_formats=graph_formats, store_eids=store_eids,