Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[GraphBolt] Allow using multiple processes for GraphBolt partition co… #7497

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[GraphBolt] Allow using multiple processes for GraphBolt partition co… #7497

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 17 additions & 3 deletions python/dgl/distributed/partition.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,8 @@
import os
import time

import numpy as np

Check warning on line 9 in python/dgl/distributed/partition.py

View workflow job for this annotation

GitHub Actions / lintrunner

UFMT format 

Run `lintrunner -a` to apply this patch.
from joblib import Parallel, delayed
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could we achieve parallelism with standard module from python such as https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ProcessPoolExecutor instead of adding new dependency?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure, I'll switch to using multiprocessing. I think joblib is pulled in with other DGL deps, but it's not explicitly added yeah.


import torch

Expand Down Expand Up @@ -1384,6 +1385,7 @@
store_inner_node=False,
store_inner_edge=False,
graph_formats=None,
n_jobs=1,
):
"""Convert partitions of dgl to FusedCSCSamplingGraph of GraphBolt.

Expand Down Expand Up @@ -1411,6 +1413,9 @@
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.
n_jobs: int
Number of parallel jobs to run during partition conversion. Max parallelism
is determined by the partition count.
"""
debug_mode = "DGL_DIST_DEBUG" in os.environ
if debug_mode:
Expand Down Expand Up @@ -1439,8 +1444,7 @@
# But this is not a problem since such information is not used in sampling.
# We can simply pass None to it.

# Iterate over partitions.
for part_id in range(num_parts):
def convert_partition(part_id, graph_formats):
graph, _, _, gpb, _, _, _ = load_partition(
part_config, part_id, load_feats=False
)
Expand Down Expand Up @@ -1564,10 +1568,20 @@
)
torch.save(csc_graph, csc_graph_path)


return os.path.relpath(csc_graph_path, os.path.dirname(part_config))
# Update graph path.

# Iterate over partitions.
partition_paths = Parallel(n_jobs=min(num_parts, n_jobs))(
delayed(convert_partition)(part_id, graph_formats) for part_id in range(num_parts)
)

for part_id, part_path in enumerate(partition_paths):
new_part_meta[f"part-{part_id}"][
"part_graph_graphbolt"
] = os.path.relpath(csc_graph_path, os.path.dirname(part_config))
] = part_path


# Save dtype info into partition config.
# [TODO][Rui] Always use int64_t for node/edge IDs in GraphBolt. See more
Expand Down
Loading