Memory Robustness code refactoring to increase capacity and reduce memory footprint #5449
Conversation
|
To trigger regression tests:
|
|
Is there a way to verify the following statement in the unit test?
|
| offset to be used when assigning edge global ids in the current partition | ||
| return_orig_ids : bool, optional | ||
| Indicates whether to return original node/edge IDs. | ||
| schema : json dictionary |
There was a problem hiding this comment.
Can you follow the DGL convention, e.g.
dgl/python/dgl/data/graph_serialize.py
Line 100 in f5ddb11
There was a problem hiding this comment.
and other types, such as numpy array -> numpy.array
| schema : json dictionary | ||
| dictionary object created by reading the metadata.json file for the | ||
| current dataset | ||
| part_id : integer |
| return_orig_ids : bool, optional | ||
| Indicates whether to return original node/edge IDs. | ||
| schema : json dictionary | ||
| dictionary object created by reading the metadata.json file for the |
There was a problem hiding this comment.
Capitalize the first char.
Same for others.
| numpy array : | ||
| shuffle_global_nids, assigned after the data-shuffling phase, for the | ||
| nodes in the current partition | ||
| tuple : |
There was a problem hiding this comment.
tuple of what? follow the convention: tuple[blablabla]
| gc.collect() | ||
| logging.info( | ||
| f"There are {len(shuffle_global_src_id)} edges in partition {part_id}" | ||
| f"[Rank: {part_id}] There are {len(shuffle_global_src_id)} edges in partition {part_id}" |
There was a problem hiding this comment.
too long, over 80 chars.
| shuffle_global_dst_id = None | ||
| global_src_id = None | ||
| global_dst_id = None | ||
| # global_edge_id = None |
| return_orig_eids=False, | ||
| ): | ||
| """ | ||
There was a problem hiding this comment.
unnecessary blank line.
| Returns: | ||
| -------- | ||
| dgl object |
There was a problem hiding this comment.
please fix the types above and below according to the comment.
| dictionary | ||
| map between edge type(string) and edge_type_id(int) | ||
| dict of tensors | ||
| If `return_orig_nids=True`, return a dict of 1D tensors whose key is the node type |
| dstids = dstids - nids[0] + offset | ||
|
|
||
| # return the values | ||
| return uniques, idxes, srcids, dstids |
There was a problem hiding this comment.
[Optional] Since you changed almost 80% of the file, I'd suggest you to also update the other piece to follow the convention.
| ) | ||
| memory_snapshot("ShuffleGlobalID_Lookup_Complete: ", rank) | ||
|
|
||
| def prepare_local_data(src_data, local_part_id): |
There was a problem hiding this comment.
Add unit test for this? make sure the src_data are properly poped?
| REV_DATA_TYPE_ID, | ||
| ) | ||
|
|
||
| DATA_TYPE_ID = { |
There was a problem hiding this comment.
what's the difference better this DATA_TYPE_ID and utils.DATA_TYPE_ID, why do you need to create a new one?
... as part of improving the unit test coverage of the data pre processing pipeline.
Description
With benchmarking felarge/webgraph/searchctr graphs, it is observed that the memory bottleneck is because of dgl graph creation, in the convert_partition.py module. Currently we can only partition felarge graph no less than 12 partitions and webgraph no less than 32 partitions.
To overcome this issue, so as to increase the size of the graph partition per node following code changes are made using this PR:
No additional (unit) test cases are instrumented to this PR because all the existing end-to-end test cases in the unit test framework will be sufficient to test this functional refactoring of the code.
Checklist
Please feel free to remove inapplicable items for your PR.
Changes