Skip to content

GitLab

Unverified Commit aa42aaeb authored by kylasa's avatar kylasa Committed by GitHub
Browse files

[DistDGL][Lintrunner]Lintrunner for tools directory (#5261) 

* lintrunner patch for gloo_wrapper.py

* lintrunner changes to the tools directory.
parent 8b47bad5
Hide whitespace changes
Inline Side-by-side
Showing with 1679 additions and 885 deletions
tools/change_etype_to_canonical_etype.py
View file @ aa42aaeb
...@@ -4,9 +4,9 @@ import logging ...@@ -4,9 +4,9 @@ import logging
import os import os
import time import time
import torch
import dgl import dgl
import torch
from dgl._ffi.base import DGLError from dgl._ffi.base import DGLError
from dgl.data.utils import load_graphs from dgl.data.utils import load_graphs
from dgl.utils import toindex from dgl.utils import toindex
......
tools/chunk_graph.py
View file @ aa42aaeb
...@@ -5,11 +5,12 @@ import os ...@@ -5,11 +5,12 @@ import os
import pathlib import pathlib
from contextlib import contextmanager from contextlib import contextmanager
import dgl
import torch import torch
from distpartitioning import array_readwriter from distpartitioning import array_readwriter
from files import setdir from files import setdir
import dgl
def chunk_numpy_array(arr, fmt_meta, chunk_sizes, path_fmt): def chunk_numpy_array(arr, fmt_meta, chunk_sizes, path_fmt):
paths = [] paths = []
...@@ -26,7 +27,9 @@ def chunk_numpy_array(arr, fmt_meta, chunk_sizes, path_fmt): ...@@ -26,7 +27,9 @@ def chunk_numpy_array(arr, fmt_meta, chunk_sizes, path_fmt):
return paths return paths
def _chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path, data_fmt): def _chunk_graph(
g, name, ndata_paths, edata_paths, num_chunks, output_path, data_fmt
):
# First deal with ndata and edata that are homogeneous (i.e. not a dict-of-dict) # First deal with ndata and edata that are homogeneous (i.e. not a dict-of-dict)
if len(g.ntypes) == 1 and not isinstance( if len(g.ntypes) == 1 and not isinstance(
next(iter(ndata_paths.values())), dict next(iter(ndata_paths.values())), dict
...@@ -96,7 +99,7 @@ def _chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path, dat ...@@ -96,7 +99,7 @@ def _chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path, dat
# Chunk node data # Chunk node data
reader_fmt_meta, writer_fmt_meta = {"name": "numpy"}, {"name": data_fmt} reader_fmt_meta, writer_fmt_meta = {"name": "numpy"}, {"name": data_fmt}
file_suffix = 'npy' if data_fmt == 'numpy' else 'parquet' file_suffix = "npy" if data_fmt == "numpy" else "parquet"
metadata["node_data"] = {} metadata["node_data"] = {}
with setdir("node_data"): with setdir("node_data"):
for ntype, ndata_per_type in ndata_paths.items(): for ntype, ndata_per_type in ndata_paths.items():
...@@ -154,7 +157,9 @@ def _chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path, dat ...@@ -154,7 +157,9 @@ def _chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path, dat
logging.info("Saved metadata in %s" % os.path.abspath(metadata_path)) logging.info("Saved metadata in %s" % os.path.abspath(metadata_path))
def chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path, data_fmt='numpy'): def chunk_graph(
g, name, ndata_paths, edata_paths, num_chunks, output_path, data_fmt="numpy"
):
""" """
Split the graph into multiple chunks. Split the graph into multiple chunks.
...@@ -185,7 +190,9 @@ def chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path, data ...@@ -185,7 +190,9 @@ def chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path, data
for key in edata.keys(): for key in edata.keys():
edata[key] = os.path.abspath(edata[key]) edata[key] = os.path.abspath(edata[key])
with setdir(output_path): with setdir(output_path):
_chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path, data_fmt) _chunk_graph(
g, name, ndata_paths, edata_paths, num_chunks, output_path, data_fmt
)
if __name__ == "__main__": if __name__ == "__main__":
......
tools/dispatch_data.py
View file @ aa42aaeb
...@@ -64,7 +64,9 @@ def submit_jobs(args) -> str: ...@@ -64,7 +64,9 @@ def submit_jobs(args) -> str:
argslist = "" argslist = ""
argslist += "--world-size {} ".format(num_ips) argslist += "--world-size {} ".format(num_ips)
argslist += "--partitions-dir {} ".format(os.path.abspath(args.partitions_dir)) argslist += "--partitions-dir {} ".format(
os.path.abspath(args.partitions_dir)
)
argslist += "--input-dir {} ".format(os.path.abspath(args.in_dir)) argslist += "--input-dir {} ".format(os.path.abspath(args.in_dir))
argslist += "--graph-name {} ".format(graph_name) argslist += "--graph-name {} ".format(graph_name)
argslist += "--schema {} ".format(schema_path) argslist += "--schema {} ".format(schema_path)
...@@ -74,7 +76,9 @@ def submit_jobs(args) -> str: ...@@ -74,7 +76,9 @@ def submit_jobs(args) -> str:
argslist += "--log-level {} ".format(args.log_level) argslist += "--log-level {} ".format(args.log_level)
argslist += "--save-orig-nids " if args.save_orig_nids else "" argslist += "--save-orig-nids " if args.save_orig_nids else ""
argslist += "--save-orig-eids " if args.save_orig_eids else "" argslist += "--save-orig-eids " if args.save_orig_eids else ""
argslist += f"--graph-formats {args.graph_formats} " if args.graph_formats else "" argslist += (
f"--graph-formats {args.graph_formats} " if args.graph_formats else ""
)
# (BarclayII) Is it safe to assume all the workers have the Python executable at the same path? # (BarclayII) Is it safe to assume all the workers have the Python executable at the same path?
pipeline_cmd = os.path.join(INSTALL_DIR, PIPELINE_SCRIPT) pipeline_cmd = os.path.join(INSTALL_DIR, PIPELINE_SCRIPT)
...@@ -153,9 +157,9 @@ def main(): ...@@ -153,9 +157,9 @@ def main():
type=str, type=str,
default=None, default=None,
help="Save partitions in specified formats. It could be any combination(joined with ``,``) " help="Save partitions in specified formats. It could be any combination(joined with ``,``) "
"of ``coo``, ``csc`` and ``csr``. If not specified, save one format only according to " "of ``coo``, ``csc`` and ``csr``. If not specified, save one format only according to "
"what format is available. If multiple formats are available, selection priority " "what format is available. If multiple formats are available, selection priority "
"from high to low is ``coo``, ``csc``, ``csr``.", "from high to low is ``coo``, ``csc``, ``csr``.",
) )
args, udf_command = parser.parse_known_args() args, udf_command = parser.parse_known_args()
......
tools/distgraphlaunch.py
View file @ aa42aaeb
"""Launching tool for DGL distributed training""" """Launching tool for DGL distributed training"""
import os
import stat
import sys
import subprocess
import argparse import argparse
import signal
import logging
import time
import json import json
import logging
import multiprocessing import multiprocessing
import os
import re import re
import signal
import stat
import subprocess
import sys
import time
from functools import partial from functools import partial
from threading import Thread from threading import Thread
from typing import Optional from typing import Optional
DEFAULT_PORT = 30050 DEFAULT_PORT = 30050
def cleanup_proc(get_all_remote_pids, conn): def cleanup_proc(get_all_remote_pids, conn):
'''This process tries to clean up the remote training tasks. """This process tries to clean up the remote training tasks."""
''' print("cleanupu process runs")
print('cleanupu process runs')
# This process should not handle SIGINT. # This process should not handle SIGINT.
signal.signal(signal.SIGINT, signal.SIG_IGN) signal.signal(signal.SIGINT, signal.SIG_IGN)
data = conn.recv() data = conn.recv()
# If the launch process exits normally, this process doesn't need to do anything. # If the launch process exits normally, this process doesn't need to do anything.
if data == 'exit': if data == "exit":
sys.exit(0) sys.exit(0)
else: else:
remote_pids = get_all_remote_pids() remote_pids = get_all_remote_pids()
# Otherwise, we need to ssh to each machine and kill the training jobs. # Otherwise, we need to ssh to each machine and kill the training jobs.
for (ip, port), pids in remote_pids.items(): for (ip, port), pids in remote_pids.items():
kill_process(ip, port, pids) kill_process(ip, port, pids)
print('cleanup process exits') print("cleanup process exits")
def kill_process(ip, port, pids): def kill_process(ip, port, pids):
'''ssh to a remote machine and kill the specified processes. """ssh to a remote machine and kill the specified processes."""
'''
curr_pid = os.getpid() curr_pid = os.getpid()
killed_pids = [] killed_pids = []
# If we kill child processes first, the parent process may create more again. This happens # If we kill child processes first, the parent process may create more again. This happens
...@@ -44,8 +44,14 @@ def kill_process(ip, port, pids): ...@@ -44,8 +44,14 @@ def kill_process(ip, port, pids):
pids.sort() pids.sort()
for pid in pids: for pid in pids:
assert curr_pid != pid assert curr_pid != pid
print('kill process {} on {}:{}'.format(pid, ip, port), flush=True) print("kill process {} on {}:{}".format(pid, ip, port), flush=True)
kill_cmd = 'ssh -o StrictHostKeyChecking=no -p ' + str(port) + ' ' + ip + ' \'kill {}\''.format(pid) kill_cmd = (
"ssh -o StrictHostKeyChecking=no -p "
+ str(port)
+ " "
+ ip
+ " 'kill {}'".format(pid)
)
subprocess.run(kill_cmd, shell=True) subprocess.run(kill_cmd, shell=True)
killed_pids.append(pid) killed_pids.append(pid)
# It's possible that some of the processes are not killed. Let's try again. # It's possible that some of the processes are not killed. Let's try again.
...@@ -56,29 +62,41 @@ def kill_process(ip, port, pids): ...@@ -56,29 +62,41 @@ def kill_process(ip, port, pids):
else: else:
killed_pids.sort() killed_pids.sort()
for pid in killed_pids: for pid in killed_pids:
print('kill process {} on {}:{}'.format(pid, ip, port), flush=True) print(
kill_cmd = 'ssh -o StrictHostKeyChecking=no -p ' + str(port) + ' ' + ip + ' \'kill -9 {}\''.format(pid) "kill process {} on {}:{}".format(pid, ip, port), flush=True
)
kill_cmd = (
"ssh -o StrictHostKeyChecking=no -p "
+ str(port)
+ " "
+ ip
+ " 'kill -9 {}'".format(pid)
)
subprocess.run(kill_cmd, shell=True) subprocess.run(kill_cmd, shell=True)
def get_killed_pids(ip, port, killed_pids): def get_killed_pids(ip, port, killed_pids):
'''Get the process IDs that we want to kill but are still alive. """Get the process IDs that we want to kill but are still alive."""
'''
killed_pids = [str(pid) for pid in killed_pids] killed_pids = [str(pid) for pid in killed_pids]
killed_pids = ','.join(killed_pids) killed_pids = ",".join(killed_pids)
ps_cmd = 'ssh -o StrictHostKeyChecking=no -p ' + str(port) + ' ' + ip + ' \'ps -p {} -h\''.format(killed_pids) ps_cmd = (
"ssh -o StrictHostKeyChecking=no -p "
+ str(port)
+ " "
+ ip
+ " 'ps -p {} -h'".format(killed_pids)
)
res = subprocess.run(ps_cmd, shell=True, stdout=subprocess.PIPE) res = subprocess.run(ps_cmd, shell=True, stdout=subprocess.PIPE)
pids = [] pids = []
for p in res.stdout.decode('utf-8').split('\n'): for p in res.stdout.decode("utf-8").split("\n"):
l = p.split() l = p.split()
if len(l) > 0: if len(l) > 0:
pids.append(int(l[0])) pids.append(int(l[0]))
return pids return pids
def execute_remote( def execute_remote(
cmd: str, cmd: str, ip: str, port: int, username: Optional[str] = ""
ip: str,
port: int,
username: Optional[str] = ""
) -> Thread: ) -> Thread:
"""Execute command line on remote machine via ssh. """Execute command line on remote machine via ssh.
...@@ -115,15 +133,21 @@ def execute_remote( ...@@ -115,15 +133,21 @@ def execute_remote(
thread.start() thread.start()
return thread return thread
def get_remote_pids(ip, port, cmd_regex): def get_remote_pids(ip, port, cmd_regex):
"""Get the process IDs that run the command in the remote machine. """Get the process IDs that run the command in the remote machine."""
"""
pids = [] pids = []
curr_pid = os.getpid() curr_pid = os.getpid()
# Here we want to get the python processes. We may get some ssh processes, so we should filter them out. # Here we want to get the python processes. We may get some ssh processes, so we should filter them out.
ps_cmd = 'ssh -o StrictHostKeyChecking=no -p ' + str(port) + ' ' + ip + ' \'ps -aux | grep python | grep -v StrictHostKeyChecking\'' ps_cmd = (
"ssh -o StrictHostKeyChecking=no -p "
+ str(port)
+ " "
+ ip
+ " 'ps -aux | grep python | grep -v StrictHostKeyChecking'"
)
res = subprocess.run(ps_cmd, shell=True, stdout=subprocess.PIPE) res = subprocess.run(ps_cmd, shell=True, stdout=subprocess.PIPE)
for p in res.stdout.decode('utf-8').split('\n'): for p in res.stdout.decode("utf-8").split("\n"):
l = p.split() l = p.split()
if len(l) < 2: if len(l) < 2:
continue continue
...@@ -132,28 +156,34 @@ def get_remote_pids(ip, port, cmd_regex): ...@@ -132,28 +156,34 @@ def get_remote_pids(ip, port, cmd_regex):
if res is not None and int(l[1]) != curr_pid: if res is not None and int(l[1]) != curr_pid:
pids.append(l[1]) pids.append(l[1])
pid_str = ','.join([str(pid) for pid in pids]) pid_str = ",".join([str(pid) for pid in pids])
ps_cmd = 'ssh -o StrictHostKeyChecking=no -p ' + str(port) + ' ' + ip + ' \'pgrep -P {}\''.format(pid_str) ps_cmd = (
"ssh -o StrictHostKeyChecking=no -p "
+ str(port)
+ " "
+ ip
+ " 'pgrep -P {}'".format(pid_str)
)
res = subprocess.run(ps_cmd, shell=True, stdout=subprocess.PIPE) res = subprocess.run(ps_cmd, shell=True, stdout=subprocess.PIPE)
pids1 = res.stdout.decode('utf-8').split('\n') pids1 = res.stdout.decode("utf-8").split("\n")
all_pids = [] all_pids = []
for pid in set(pids + pids1): for pid in set(pids + pids1):
if pid == '' or int(pid) == curr_pid: if pid == "" or int(pid) == curr_pid:
continue continue
all_pids.append(int(pid)) all_pids.append(int(pid))
all_pids.sort() all_pids.sort()
return all_pids return all_pids
def get_all_remote_pids(hosts, ssh_port, udf_command): def get_all_remote_pids(hosts, ssh_port, udf_command):
'''Get all remote processes. """Get all remote processes."""
'''
remote_pids = {} remote_pids = {}
for node_id, host in enumerate(hosts): for node_id, host in enumerate(hosts):
ip, _ = host ip, _ = host
# When creating training processes in remote machines, we may insert some arguments # When creating training processes in remote machines, we may insert some arguments
# in the commands. We need to use regular expressions to match the modified command. # in the commands. We need to use regular expressions to match the modified command.
cmds = udf_command.split() cmds = udf_command.split()
new_udf_command = ' .*'.join(cmds) new_udf_command = " .*".join(cmds)
pids = get_remote_pids(ip, ssh_port, new_udf_command) pids = get_remote_pids(ip, ssh_port, new_udf_command)
remote_pids[(ip, ssh_port)] = pids remote_pids[(ip, ssh_port)] = pids
return remote_pids return remote_pids
...@@ -164,7 +194,7 @@ def construct_torch_dist_launcher_cmd( ...@@ -164,7 +194,7 @@ def construct_torch_dist_launcher_cmd(
num_nodes: int, num_nodes: int,
node_rank: int, node_rank: int,
master_addr: str, master_addr: str,
master_port: int master_port: int,
) -> str: ) -> str:
"""Constructs the torch distributed launcher command. """Constructs the torch distributed launcher command.
Helper function. Helper function.
...@@ -179,18 +209,20 @@ def construct_torch_dist_launcher_cmd( ...@@ -179,18 +209,20 @@ def construct_torch_dist_launcher_cmd(
Returns: Returns:
cmd_str. cmd_str.
""" """
torch_cmd_template = "-m torch.distributed.launch " \ torch_cmd_template = (
"--nproc_per_node={nproc_per_node} " \ "-m torch.distributed.launch "
"--nnodes={nnodes} " \ "--nproc_per_node={nproc_per_node} "
"--node_rank={node_rank} " \ "--nnodes={nnodes} "
"--master_addr={master_addr} " \ "--node_rank={node_rank} "
"--master_port={master_port}" "--master_addr={master_addr} "
"--master_port={master_port}"
)
return torch_cmd_template.format( return torch_cmd_template.format(
nproc_per_node=num_trainers, nproc_per_node=num_trainers,
nnodes=num_nodes, nnodes=num_nodes,
node_rank=node_rank, node_rank=node_rank,
master_addr=master_addr, master_addr=master_addr,
master_port=master_port master_port=master_port,
) )
...@@ -233,7 +265,7 @@ def wrap_udf_in_torch_dist_launcher( ...@@ -233,7 +265,7 @@ def wrap_udf_in_torch_dist_launcher(
num_nodes=num_nodes, num_nodes=num_nodes,
node_rank=node_rank, node_rank=node_rank,
master_addr=master_addr, master_addr=master_addr,
master_port=master_port master_port=master_port,
) )
# Auto-detect the python binary that kicks off the distributed trainer code. # Auto-detect the python binary that kicks off the distributed trainer code.
# Note: This allowlist order matters, this will match with the FIRST matching entry. Thus, please add names to this # Note: This allowlist order matters, this will match with the FIRST matching entry. Thus, please add names to this
...@@ -241,9 +273,14 @@ def wrap_udf_in_torch_dist_launcher( ...@@ -241,9 +273,14 @@ def wrap_udf_in_torch_dist_launcher(
# (python3.7, python3.8) -> (python3) # (python3.7, python3.8) -> (python3)
# The allowed python versions are from this: https://www.dgl.ai/pages/start.html # The allowed python versions are from this: https://www.dgl.ai/pages/start.html
python_bin_allowlist = ( python_bin_allowlist = (
"python3.6", "python3.7", "python3.8", "python3.9", "python3", "python3.6",
"python3.7",
"python3.8",
"python3.9",
"python3",
# for backwards compatibility, accept python2 but technically DGL is a py3 library, so this is not recommended # for backwards compatibility, accept python2 but technically DGL is a py3 library, so this is not recommended
"python2.7", "python2", "python2.7",
"python2",
) )
# If none of the candidate python bins match, then we go with the default `python` # If none of the candidate python bins match, then we go with the default `python`
python_bin = "python" python_bin = "python"
...@@ -258,7 +295,9 @@ def wrap_udf_in_torch_dist_launcher( ...@@ -258,7 +295,9 @@ def wrap_udf_in_torch_dist_launcher(
# python -m torch.distributed.launch [DIST TORCH ARGS] path/to/dist_trainer.py arg0 arg1 # python -m torch.distributed.launch [DIST TORCH ARGS] path/to/dist_trainer.py arg0 arg1
# Note: if there are multiple python commands in `udf_command`, this may do the Wrong Thing, eg launch each # Note: if there are multiple python commands in `udf_command`, this may do the Wrong Thing, eg launch each
# python command within the torch distributed launcher. # python command within the torch distributed launcher.
new_udf_command = udf_command.replace(python_bin, f"{python_bin} {torch_dist_cmd}") new_udf_command = udf_command.replace(
python_bin, f"{python_bin} {torch_dist_cmd}"
)
return new_udf_command return new_udf_command
...@@ -322,6 +361,7 @@ def wrap_cmd_with_local_envvars(cmd: str, env_vars: str) -> str: ...@@ -322,6 +361,7 @@ def wrap_cmd_with_local_envvars(cmd: str, env_vars: str) -> str:
# https://stackoverflow.com/a/45993803 # https://stackoverflow.com/a/45993803
return f"(export {env_vars}; {cmd})" return f"(export {env_vars}; {cmd})"
def wrap_cmd_with_extra_envvars(cmd: str, env_vars: list) -> str: def wrap_cmd_with_extra_envvars(cmd: str, env_vars: list) -> str:
"""Wraps a CLI command with extra env vars """Wraps a CLI command with extra env vars
...@@ -341,6 +381,7 @@ def wrap_cmd_with_extra_envvars(cmd: str, env_vars: list) -> str: ...@@ -341,6 +381,7 @@ def wrap_cmd_with_extra_envvars(cmd: str, env_vars: list) -> str:
env_vars = " ".join(env_vars) env_vars = " ".join(env_vars)
return wrap_cmd_with_local_envvars(cmd, env_vars) return wrap_cmd_with_local_envvars(cmd, env_vars)
def submit_jobs(args, udf_command): def submit_jobs(args, udf_command):
"""Submit distributed jobs (server and client processes) via ssh""" """Submit distributed jobs (server and client processes) via ssh"""
hosts = [] hosts = []
...@@ -348,7 +389,7 @@ def submit_jobs(args, udf_command): ...@@ -348,7 +389,7 @@ def submit_jobs(args, udf_command):
server_count_per_machine = 0 server_count_per_machine = 0
# Get the IP addresses of the cluster. # Get the IP addresses of the cluster.
#ip_config = os.path.join(args.workspace, args.ip_config) # ip_config = os.path.join(args.workspace, args.ip_config)
ip_config = args.ip_config ip_config = args.ip_config
with open(ip_config) as f: with open(ip_config) as f:
for line in f: for line in f:
...@@ -376,58 +417,76 @@ def submit_jobs(args, udf_command): ...@@ -376,58 +417,76 @@ def submit_jobs(args, udf_command):
server_env_vars_cur = f"{server_env_vars} RANK={i} MASTER_ADDR={hosts[0][0]} MASTER_PORT={args.master_port}" server_env_vars_cur = f"{server_env_vars} RANK={i} MASTER_ADDR={hosts[0][0]} MASTER_PORT={args.master_port}"
cmd = wrap_cmd_with_local_envvars(udf_command, server_env_vars_cur) cmd = wrap_cmd_with_local_envvars(udf_command, server_env_vars_cur)
print(cmd) print(cmd)
thread_list.append(execute_remote(cmd, ip, args.ssh_port, username=args.ssh_username)) thread_list.append(
execute_remote(cmd, ip, args.ssh_port, username=args.ssh_username)
)
# Start a cleanup process dedicated for cleaning up remote training jobs. # Start a cleanup process dedicated for cleaning up remote training jobs.
conn1,conn2 = multiprocessing.Pipe() conn1, conn2 = multiprocessing.Pipe()
func = partial(get_all_remote_pids, hosts, args.ssh_port, udf_command) func = partial(get_all_remote_pids, hosts, args.ssh_port, udf_command)
process = multiprocessing.Process(target=cleanup_proc, args=(func, conn1)) process = multiprocessing.Process(target=cleanup_proc, args=(func, conn1))
process.start() process.start()
def signal_handler(signal, frame): def signal_handler(signal, frame):
logging.info('Stop launcher') logging.info("Stop launcher")
# We need to tell the cleanup process to kill remote training jobs. # We need to tell the cleanup process to kill remote training jobs.
conn2.send('cleanup') conn2.send("cleanup")
sys.exit(0) sys.exit(0)
signal.signal(signal.SIGINT, signal_handler) signal.signal(signal.SIGINT, signal_handler)
for thread in thread_list: for thread in thread_list:
thread.join() thread.join()
# The training processes complete. We should tell the cleanup process to exit. # The training processes complete. We should tell the cleanup process to exit.
conn2.send('exit') conn2.send("exit")
process.join() process.join()
def main(): def main():
parser = argparse.ArgumentParser(description='Launch a distributed job') parser = argparse.ArgumentParser(description="Launch a distributed job")
parser.add_argument('--ssh_port', type=int, default=22, help='SSH Port.') parser.add_argument("--ssh_port", type=int, default=22, help="SSH Port.")
parser.add_argument( parser.add_argument(
"--ssh_username", default="", "--ssh_username",
default="",
help="Optional. When issuing commands (via ssh) to cluster, use the provided username in the ssh cmd. " help="Optional. When issuing commands (via ssh) to cluster, use the provided username in the ssh cmd. "
"Example: If you provide --ssh_username=bob, then the ssh command will be like: 'ssh bob@1.2.3.4 CMD' " "Example: If you provide --ssh_username=bob, then the ssh command will be like: 'ssh bob@1.2.3.4 CMD' "
"instead of 'ssh 1.2.3.4 CMD'" "instead of 'ssh 1.2.3.4 CMD'",
)
parser.add_argument(
"--num_proc_per_machine",
type=int,
help="The number of server processes per machine",
)
parser.add_argument(
"--master_port",
type=int,
help="This port is used to form gloo group (randevouz server)",
)
parser.add_argument(
"--ip_config",
type=str,
help="The file (in workspace) of IP configuration for server processes",
) )
parser.add_argument('--num_proc_per_machine', type=int,
help='The number of server processes per machine')
parser.add_argument('--master_port', type=int,
help='This port is used to form gloo group (randevouz server)')
parser.add_argument('--ip_config', type=str,
help='The file (in workspace) of IP configuration for server processes')
args, udf_command = parser.parse_known_args() args, udf_command = parser.parse_known_args()
assert len(udf_command) == 1, 'Please provide user command line.' assert len(udf_command) == 1, "Please provide user command line."
assert args.num_proc_per_machine is not None and args.num_proc_per_machine > 0, \ assert (
'--num_proc_per_machine must be a positive number.' args.num_proc_per_machine is not None and args.num_proc_per_machine > 0
assert args.ip_config is not None, \ ), "--num_proc_per_machine must be a positive number."
'A user has to specify an IP configuration file with --ip_config.' assert (
args.ip_config is not None
), "A user has to specify an IP configuration file with --ip_config."
udf_command = str(udf_command[0]) udf_command = str(udf_command[0])
if 'python' not in udf_command: if "python" not in udf_command:
raise RuntimeError("DGL launching script can only support Python executable file.") raise RuntimeError(
"DGL launching script can only support Python executable file."
)
submit_jobs(args, udf_command) submit_jobs(args, udf_command)
if __name__ == '__main__':
fmt = '%(asctime)s %(levelname)s %(message)s' if __name__ == "__main__":
fmt = "%(asctime)s %(levelname)s %(message)s"
logging.basicConfig(format=fmt, level=logging.INFO) logging.basicConfig(format=fmt, level=logging.INFO)
main() main()
tools/distpartitioning/array_readwriter/parquet.py
View file @ aa42aaeb
...@@ -30,7 +30,9 @@ class ParquetArrayParser(object): ...@@ -30,7 +30,9 @@ class ParquetArrayParser(object):
# Spark ML feature processing produces single-column parquet files where each row is a vector object # Spark ML feature processing produces single-column parquet files where each row is a vector object
if len(data_types) == 1 and isinstance(data_types[0], pyarrow.ListType): if len(data_types) == 1 and isinstance(data_types[0], pyarrow.ListType):
arr = np.array(table.to_pandas().iloc[:, 0].to_list()) arr = np.array(table.to_pandas().iloc[:, 0].to_list())
logging.debug(f"Parquet data under {path} converted from single vector per row to ndarray") logging.debug(
f"Parquet data under {path} converted from single vector per row to ndarray"
)
else: else:
arr = table.to_pandas().to_numpy() arr = table.to_pandas().to_numpy()
if not shape: if not shape:
...@@ -49,8 +51,8 @@ class ParquetArrayParser(object): ...@@ -49,8 +51,8 @@ class ParquetArrayParser(object):
array = array.reshape(shape[0], -1) array = array.reshape(shape[0], -1)
if vector_rows: if vector_rows:
table = pyarrow.table( table = pyarrow.table(
[pyarrow.array(array.tolist())], [pyarrow.array(array.tolist())], names=["vector"]
names=["vector"]) )
logging.info("Writing to %s using single-vector rows..." % path) logging.info("Writing to %s using single-vector rows..." % path)
else: else:
table = pyarrow.Table.from_pandas(pd.DataFrame(array)) table = pyarrow.Table.from_pandas(pd.DataFrame(array))
......
tools/distpartitioning/constants.py
View file @ aa42aaeb
...@@ -37,4 +37,4 @@ STR_NAME = "name" ...@@ -37,4 +37,4 @@ STR_NAME = "name"
STR_GRAPH_NAME = "graph_name" STR_GRAPH_NAME = "graph_name"
STR_NODE_FEATURES = "node_features" STR_NODE_FEATURES = "node_features"
STR_EDGE_FEATURES = "edge_features" STR_EDGE_FEATURES = "edge_features"
\ No newline at end of file
tools/distpartitioning/convert_partition.py
View file @ aa42aaeb
...@@ -5,43 +5,50 @@ import logging ...@@ -5,43 +5,50 @@ import logging
import os import os
import time import time
import constants
import dgl import dgl
import numpy as np import numpy as np
import pandas as pd import pandas as pd
import pyarrow import pyarrow
import torch as th import torch as th
from pyarrow import csv
import constants
from utils import get_idranges, memory_snapshot, read_json
from dgl.distributed.partition import ( from dgl.distributed.partition import (
RESERVED_FIELD_DTYPE,
_etype_str_to_tuple, _etype_str_to_tuple,
_etype_tuple_to_str, _etype_tuple_to_str,
RESERVED_FIELD_DTYPE,
) )
from pyarrow import csv
from utils import get_idranges, memory_snapshot, read_json
def create_dgl_object(schema, part_id, node_data, edge_data, edgeid_offset, def create_dgl_object(
return_orig_nids=False, return_orig_eids=False): schema,
part_id,
node_data,
edge_data,
edgeid_offset,
return_orig_nids=False,
return_orig_eids=False,
):
""" """
This function creates dgl objects for a given graph partition, as in function This function creates dgl objects for a given graph partition, as in function
arguments. arguments.
The "schema" argument is a dictionary, which contains the metadata related to node ids The "schema" argument is a dictionary, which contains the metadata related to node ids
and edge ids. It contains two keys: "nid" and "eid", whose value is also a dictionary and edge ids. It contains two keys: "nid" and "eid", whose value is also a dictionary
with the following structure. with the following structure.
1. The key-value pairs in the "nid" dictionary has the following format. 1. The key-value pairs in the "nid" dictionary has the following format.
"ntype-name" is the user assigned name to this node type. "format" describes the "ntype-name" is the user assigned name to this node type. "format" describes the
format of the contents of the files. and "data" is a list of lists, each list has format of the contents of the files. and "data" is a list of lists, each list has
3 elements: file-name, start_id and end_id. File-name can be either absolute or 3 elements: file-name, start_id and end_id. File-name can be either absolute or
relative path to this file and starting and ending ids are type ids of the nodes relative path to this file and starting and ending ids are type ids of the nodes
which are contained in this file. These type ids are later used to compute global ids which are contained in this file. These type ids are later used to compute global ids
of these nodes which are used throughout the processing of this pipeline. of these nodes which are used throughout the processing of this pipeline.
"ntype-name" : { "ntype-name" : {
"format" : "csv", "format" : "csv",
"data" : [ "data" : [
[ <path-to-file>/ntype0-name-0.csv, start_id0, end_id0], [ <path-to-file>/ntype0-name-0.csv, start_id0, end_id0],
[ <path-to-file>/ntype0-name-1.csv, start_id1, end_id1], [ <path-to-file>/ntype0-name-1.csv, start_id1, end_id1],
... ...
[ <path-to-file>/ntype0-name-<p-1>.csv, start_id<p-1>, end_id<p-1>], [ <path-to-file>/ntype0-name-<p-1>.csv, start_id<p-1>, end_id<p-1>],
...@@ -50,11 +57,11 @@ def create_dgl_object(schema, part_id, node_data, edge_data, edgeid_offset, ...@@ -50,11 +57,11 @@ def create_dgl_object(schema, part_id, node_data, edge_data, edgeid_offset,
2. The key-value pairs in the "eid" dictionary has the following format. 2. The key-value pairs in the "eid" dictionary has the following format.
As described for the "nid" dictionary the "eid" dictionary is similarly structured As described for the "nid" dictionary the "eid" dictionary is similarly structured
except that these entries are for edges. except that these entries are for edges.
"etype-name" : { "etype-name" : {
"format" : "csv", "format" : "csv",
"data" : [ "data" : [
[ <path-to-file>/etype0-name-0, start_id0, end_id0], [ <path-to-file>/etype0-name-0, start_id0, end_id0],
[ <path-to-file>/etype0-name-1 start_id1, end_id1], [ <path-to-file>/etype0-name-1 start_id1, end_id1],
... ...
[ <path-to-file>/etype0-name-1 start_id<p-1>, end_id<p-1>] [ <path-to-file>/etype0-name-1 start_id<p-1>, end_id<p-1>]
...@@ -62,8 +69,8 @@ def create_dgl_object(schema, part_id, node_data, edge_data, edgeid_offset, ...@@ -62,8 +69,8 @@ def create_dgl_object(schema, part_id, node_data, edge_data, edgeid_offset,
} }
In "nid" dictionary, the type_nids are specified that In "nid" dictionary, the type_nids are specified that
should be assigned to nodes which are read from the corresponding nodes file. should be assigned to nodes which are read from the corresponding nodes file.
Along the same lines dictionary for the key "eid" is used for edges in the Along the same lines dictionary for the key "eid" is used for edges in the
input graph. input graph.
These type ids, for nodes and edges, are used to compute global ids for nodes These type ids, for nodes and edges, are used to compute global ids for nodes
...@@ -79,14 +86,14 @@ def create_dgl_object(schema, part_id, node_data, edge_data, edgeid_offset, ...@@ -79,14 +86,14 @@ def create_dgl_object(schema, part_id, node_data, edge_data, edgeid_offset,
node_data, where each row is of the following format: node_data, where each row is of the following format:
<global_nid> <ntype_id> <global_type_nid> <global_nid> <ntype_id> <global_type_nid>
edge_data : numpy ndarray edge_data : numpy ndarray
edge_data, where each row is of the following format: edge_data, where each row is of the following format:
<global_src_id> <global_dst_id> <etype_id> <global_type_eid> <global_src_id> <global_dst_id> <etype_id> <global_type_eid>
edgeid_offset : int edgeid_offset : int
offset to be used when assigning edge global ids in the current partition offset to be used when assigning edge global ids in the current partition
return_orig_ids : bool, optional return_orig_ids : bool, optional
Indicates whether to return original node/edge IDs. Indicates whether to return original node/edge IDs.
Returns: Returns:
-------- --------
dgl object dgl object
dgl object created for the current graph partition dgl object created for the current graph partition
...@@ -107,12 +114,16 @@ def create_dgl_object(schema, part_id, node_data, edge_data, edgeid_offset, ...@@ -107,12 +114,16 @@ def create_dgl_object(schema, part_id, node_data, edge_data, edgeid_offset,
and value is a 1D tensor mapping between shuffled edge IDs and the original edge and value is a 1D tensor mapping between shuffled edge IDs and the original edge
IDs for each edge type. Otherwise, ``None`` is returned. IDs for each edge type. Otherwise, ``None`` is returned.
""" """
#create auxiliary data structures from the schema object # create auxiliary data structures from the schema object
memory_snapshot("CreateDGLObj_Begin", part_id) memory_snapshot("CreateDGLObj_Begin", part_id)
_, global_nid_ranges = get_idranges(schema[constants.STR_NODE_TYPE], _, global_nid_ranges = get_idranges(
schema[constants.STR_NUM_NODES_PER_CHUNK]) schema[constants.STR_NODE_TYPE],
_, global_eid_ranges = get_idranges(schema[constants.STR_EDGE_TYPE], schema[constants.STR_NUM_NODES_PER_CHUNK],
schema[constants.STR_NUM_EDGES_PER_CHUNK]) )
_, global_eid_ranges = get_idranges(
schema[constants.STR_EDGE_TYPE],
schema[constants.STR_NUM_EDGES_PER_CHUNK],
)
id_map = dgl.distributed.id_map.IdMap(global_nid_ranges) id_map = dgl.distributed.id_map.IdMap(global_nid_ranges)
...@@ -147,15 +158,15 @@ def create_dgl_object(schema, part_id, node_data, edge_data, edgeid_offset, ...@@ -147,15 +158,15 @@ def create_dgl_object(schema, part_id, node_data, edge_data, edgeid_offset,
assert np.all(shuffle_global_nids[1:] - shuffle_global_nids[:-1] == 1) assert np.all(shuffle_global_nids[1:] - shuffle_global_nids[:-1] == 1)
shuffle_global_nid_range = (shuffle_global_nids[0], shuffle_global_nids[-1]) shuffle_global_nid_range = (shuffle_global_nids[0], shuffle_global_nids[-1])
# Determine the node ID ranges of different node types. # Determine the node ID ranges of different node types.
for ntype_name in global_nid_ranges: for ntype_name in global_nid_ranges:
ntype_id = ntypes_map[ntype_name] ntype_id = ntypes_map[ntype_name]
type_nids = shuffle_global_nids[ntype_ids == ntype_id] type_nids = shuffle_global_nids[ntype_ids == ntype_id]
node_map_val[ntype_name].append( node_map_val[ntype_name].append(
[int(type_nids[0]), int(type_nids[-1]) + 1]) [int(type_nids[0]), int(type_nids[-1]) + 1]
)
#process edges # process edges
memory_snapshot("CreateDGLObj_AssignEdgeData: ", part_id) memory_snapshot("CreateDGLObj_AssignEdgeData: ", part_id)
shuffle_global_src_id = edge_data[constants.SHUFFLE_GLOBAL_SRC_ID] shuffle_global_src_id = edge_data[constants.SHUFFLE_GLOBAL_SRC_ID]
edge_data.pop(constants.SHUFFLE_GLOBAL_SRC_ID) edge_data.pop(constants.SHUFFLE_GLOBAL_SRC_ID)
...@@ -180,28 +191,43 @@ def create_dgl_object(schema, part_id, node_data, edge_data, edgeid_offset, ...@@ -180,28 +191,43 @@ def create_dgl_object(schema, part_id, node_data, edge_data, edgeid_offset,
etype_ids = edge_data[constants.ETYPE_ID] etype_ids = edge_data[constants.ETYPE_ID]
edge_data.pop(constants.ETYPE_ID) edge_data.pop(constants.ETYPE_ID)
edge_data = None edge_data = None
gc.collect() gc.collect()
logging.info(f'There are {len(shuffle_global_src_id)} edges in partition {part_id}') logging.info(
f"There are {len(shuffle_global_src_id)} edges in partition {part_id}"
)
# It's not guaranteed that the edges are sorted based on edge type. # It's not guaranteed that the edges are sorted based on edge type.
# Let's sort edges and all attributes on the edges. # Let's sort edges and all attributes on the edges.
if not np.all(np.diff(etype_ids) >= 0): if not np.all(np.diff(etype_ids) >= 0):
sort_idx = np.argsort(etype_ids) sort_idx = np.argsort(etype_ids)
shuffle_global_src_id, shuffle_global_dst_id, global_src_id, global_dst_id, global_edge_id, etype_ids = \ (
shuffle_global_src_id[sort_idx], shuffle_global_dst_id[sort_idx], global_src_id[sort_idx], \ shuffle_global_src_id,
global_dst_id[sort_idx], global_edge_id[sort_idx], etype_ids[sort_idx] shuffle_global_dst_id,
global_src_id,
global_dst_id,
global_edge_id,
etype_ids,
) = (
shuffle_global_src_id[sort_idx],
shuffle_global_dst_id[sort_idx],
global_src_id[sort_idx],
global_dst_id[sort_idx],
global_edge_id[sort_idx],
etype_ids[sort_idx],
)
assert np.all(np.diff(etype_ids) >= 0) assert np.all(np.diff(etype_ids) >= 0)
else: else:
print(f'[Rank: {part_id} Edge data is already sorted !!!') print(f"[Rank: {part_id} Edge data is already sorted !!!")
# Determine the edge ID range of different edge types. # Determine the edge ID range of different edge types.
edge_id_start = edgeid_offset edge_id_start = edgeid_offset
for etype_name in global_eid_ranges: for etype_name in global_eid_ranges:
etype = _etype_str_to_tuple(etype_name) etype = _etype_str_to_tuple(etype_name)
assert len(etype) == 3 assert len(etype) == 3
etype_id = etypes_map[etype] etype_id = etypes_map[etype]
edge_map_val[etype].append([edge_id_start, edge_map_val[etype].append(
edge_id_start + np.sum(etype_ids == etype_id)]) [edge_id_start, edge_id_start + np.sum(etype_ids == etype_id)]
)
edge_id_start += np.sum(etype_ids == etype_id) edge_id_start += np.sum(etype_ids == etype_id)
memory_snapshot("CreateDGLObj_UniqueNodeIds: ", part_id) memory_snapshot("CreateDGLObj_UniqueNodeIds: ", part_id)
...@@ -209,25 +235,38 @@ def create_dgl_object(schema, part_id, node_data, edge_data, edgeid_offset, ...@@ -209,25 +235,38 @@ def create_dgl_object(schema, part_id, node_data, edge_data, edgeid_offset,
# Here the order of nodes is defined by the `np.unique` function # Here the order of nodes is defined by the `np.unique` function
# node order is as listed in the uniq_ids array # node order is as listed in the uniq_ids array
ids = np.concatenate( ids = np.concatenate(
[shuffle_global_src_id, shuffle_global_dst_id, [
np.arange(shuffle_global_nid_range[0], shuffle_global_nid_range[1] + 1)]) shuffle_global_src_id,
shuffle_global_dst_id,
np.arange(
shuffle_global_nid_range[0], shuffle_global_nid_range[1] + 1
),
]
)
uniq_ids, idx, inverse_idx = np.unique( uniq_ids, idx, inverse_idx = np.unique(
ids, return_index=True, return_inverse=True) ids, return_index=True, return_inverse=True
)
assert len(uniq_ids) == len(idx) assert len(uniq_ids) == len(idx)
# We get the edge list with their node IDs mapped to a contiguous ID range. # We get the edge list with their node IDs mapped to a contiguous ID range.
part_local_src_id, part_local_dst_id = np.split(inverse_idx[:len(shuffle_global_src_id) * 2], 2) part_local_src_id, part_local_dst_id = np.split(
inner_nodes = th.as_tensor(np.logical_and( inverse_idx[: len(shuffle_global_src_id) * 2], 2
)
inner_nodes = th.as_tensor(
np.logical_and(
uniq_ids >= shuffle_global_nid_range[0], uniq_ids >= shuffle_global_nid_range[0],
uniq_ids <= shuffle_global_nid_range[1])) uniq_ids <= shuffle_global_nid_range[1],
)
)
#get the list of indices, from inner_nodes, which will sort inner_nodes as [True, True, ...., False, False, ...] # get the list of indices, from inner_nodes, which will sort inner_nodes as [True, True, ...., False, False, ...]
#essentially local nodes will be placed before non-local nodes. # essentially local nodes will be placed before non-local nodes.
reshuffle_nodes = th.arange(len(uniq_ids)) reshuffle_nodes = th.arange(len(uniq_ids))
reshuffle_nodes = th.cat([reshuffle_nodes[inner_nodes.bool()], reshuffle_nodes = th.cat(
reshuffle_nodes[inner_nodes == 0]]) [reshuffle_nodes[inner_nodes.bool()], reshuffle_nodes[inner_nodes == 0]]
)
''' """
Following procedure is used to map the part_local_src_id, part_local_dst_id to account for Following procedure is used to map the part_local_src_id, part_local_dst_id to account for
reshuffling of nodes (to order localy owned nodes prior to non-local nodes in a partition) reshuffling of nodes (to order localy owned nodes prior to non-local nodes in a partition)
1. Form a node_map, in this case a numpy array, which will be used to map old node-ids (pre-reshuffling) 1. Form a node_map, in this case a numpy array, which will be used to map old node-ids (pre-reshuffling)
...@@ -261,59 +300,100 @@ def create_dgl_object(schema, part_id, node_data, edge_data, edgeid_offset, ...@@ -261,59 +300,100 @@ def create_dgl_object(schema, part_id, node_data, edge_data, edgeid_offset,
Since the edge are re-ordered in any way, there is no reordering required for edge related data Since the edge are re-ordered in any way, there is no reordering required for edge related data
during the DGL object creation. during the DGL object creation.
''' """
#create the mappings to generate mapped part_local_src_id and part_local_dst_id # create the mappings to generate mapped part_local_src_id and part_local_dst_id
#This map will map from unshuffled node-ids to reshuffled-node-ids (which are ordered to prioritize # This map will map from unshuffled node-ids to reshuffled-node-ids (which are ordered to prioritize
#locally owned nodes). # locally owned nodes).
nid_map = np.zeros((len(reshuffle_nodes,))) nid_map = np.zeros(
(
len(
reshuffle_nodes,
)
)
)
nid_map[reshuffle_nodes] = np.arange(len(reshuffle_nodes)) nid_map[reshuffle_nodes] = np.arange(len(reshuffle_nodes))
#Now map the edge end points to reshuffled_values. # Now map the edge end points to reshuffled_values.
part_local_src_id, part_local_dst_id = nid_map[part_local_src_id], nid_map[part_local_dst_id] part_local_src_id, part_local_dst_id = (
nid_map[part_local_src_id],
nid_map[part_local_dst_id],
)
#create the graph here now. # create the graph here now.
part_graph = dgl.graph(data=(part_local_src_id, part_local_dst_id), num_nodes=len(uniq_ids)) part_graph = dgl.graph(
data=(part_local_src_id, part_local_dst_id), num_nodes=len(uniq_ids)
)
part_graph.edata[dgl.EID] = th.arange( part_graph.edata[dgl.EID] = th.arange(
edgeid_offset, edgeid_offset + part_graph.number_of_edges(), dtype=th.int64) edgeid_offset,
part_graph.edata[dgl.ETYPE] = th.as_tensor(etype_ids, dtype=RESERVED_FIELD_DTYPE[dgl.ETYPE]) edgeid_offset + part_graph.number_of_edges(),
part_graph.edata['inner_edge'] = th.ones(part_graph.number_of_edges(), dtype=th.int64,
dtype=RESERVED_FIELD_DTYPE['inner_edge']) )
part_graph.edata[dgl.ETYPE] = th.as_tensor(
etype_ids, dtype=RESERVED_FIELD_DTYPE[dgl.ETYPE]
#compute per_type_ids and ntype for all the nodes in the graph. )
global_ids = np.concatenate( part_graph.edata["inner_edge"] = th.ones(
[global_src_id, global_dst_id, global_homo_nid]) part_graph.number_of_edges(), dtype=RESERVED_FIELD_DTYPE["inner_edge"]
)
# compute per_type_ids and ntype for all the nodes in the graph.
global_ids = np.concatenate([global_src_id, global_dst_id, global_homo_nid])
part_global_ids = global_ids[idx] part_global_ids = global_ids[idx]
part_global_ids = part_global_ids[reshuffle_nodes] part_global_ids = part_global_ids[reshuffle_nodes]
ntype, per_type_ids = id_map(part_global_ids) ntype, per_type_ids = id_map(part_global_ids)
#continue with the graph creation # continue with the graph creation
part_graph.ndata[dgl.NTYPE] = th.as_tensor(ntype, dtype=RESERVED_FIELD_DTYPE[dgl.NTYPE]) part_graph.ndata[dgl.NTYPE] = th.as_tensor(
ntype, dtype=RESERVED_FIELD_DTYPE[dgl.NTYPE]
)
part_graph.ndata[dgl.NID] = th.as_tensor(uniq_ids[reshuffle_nodes]) part_graph.ndata[dgl.NID] = th.as_tensor(uniq_ids[reshuffle_nodes])
part_graph.ndata['inner_node'] = th.as_tensor(inner_nodes[reshuffle_nodes], part_graph.ndata["inner_node"] = th.as_tensor(
dtype=RESERVED_FIELD_DTYPE['inner_node']) inner_nodes[reshuffle_nodes], dtype=RESERVED_FIELD_DTYPE["inner_node"]
)
orig_nids = None orig_nids = None
orig_eids = None orig_eids = None
if return_orig_nids: if return_orig_nids:
orig_nids = {} orig_nids = {}
for ntype, ntype_id in ntypes_map.items(): for ntype, ntype_id in ntypes_map.items():
mask = th.logical_and(part_graph.ndata[dgl.NTYPE] == ntype_id, mask = th.logical_and(
part_graph.ndata['inner_node']) part_graph.ndata[dgl.NTYPE] == ntype_id,
part_graph.ndata["inner_node"],
)
orig_nids[ntype] = th.as_tensor(per_type_ids[mask]) orig_nids[ntype] = th.as_tensor(per_type_ids[mask])
if return_orig_eids: if return_orig_eids:
orig_eids = {} orig_eids = {}
for etype, etype_id in etypes_map.items(): for etype, etype_id in etypes_map.items():
mask = th.logical_and(part_graph.edata[dgl.ETYPE] == etype_id, mask = th.logical_and(
part_graph.edata['inner_edge']) part_graph.edata[dgl.ETYPE] == etype_id,
orig_eids[_etype_tuple_to_str(etype)] = th.as_tensor(global_edge_id[mask]) part_graph.edata["inner_edge"],
)
orig_eids[_etype_tuple_to_str(etype)] = th.as_tensor(
return part_graph, node_map_val, edge_map_val, ntypes_map, etypes_map, \ global_edge_id[mask]
orig_nids, orig_eids )
def create_metadata_json(graph_name, num_nodes, num_edges, part_id, num_parts, node_map_val, \ return (
edge_map_val, ntypes_map, etypes_map, output_dir ): part_graph,
node_map_val,
edge_map_val,
ntypes_map,
etypes_map,
orig_nids,
orig_eids,
)
def create_metadata_json(
graph_name,
num_nodes,
num_edges,
part_id,
num_parts,
node_map_val,
edge_map_val,
ntypes_map,
etypes_map,
output_dir,
):
""" """
Auxiliary function to create json file for the graph partition metadata Auxiliary function to create json file for the graph partition metadata
...@@ -338,7 +418,7 @@ def create_metadata_json(graph_name, num_nodes, num_edges, part_id, num_parts, n ...@@ -338,7 +418,7 @@ def create_metadata_json(graph_name, num_nodes, num_edges, part_id, num_parts, n
etypes_map : dictionary etypes_map : dictionary
map between edge type(string) and edge_type_id(int) map between edge type(string) and edge_type_id(int)
output_dir : string output_dir : string
directory where the output files are to be stored directory where the output files are to be stored
Returns: Returns:
-------- --------
...@@ -346,22 +426,26 @@ def create_metadata_json(graph_name, num_nodes, num_edges, part_id, num_parts, n ...@@ -346,22 +426,26 @@ def create_metadata_json(graph_name, num_nodes, num_edges, part_id, num_parts, n
map describing the graph information map describing the graph information
""" """
part_metadata = {'graph_name': graph_name, part_metadata = {
'num_nodes': num_nodes, "graph_name": graph_name,
'num_edges': num_edges, "num_nodes": num_nodes,
'part_method': 'metis', "num_edges": num_edges,
'num_parts': num_parts, "part_method": "metis",
'halo_hops': 1, "num_parts": num_parts,
'node_map': node_map_val, "halo_hops": 1,
'edge_map': edge_map_val, "node_map": node_map_val,
'ntypes': ntypes_map, "edge_map": edge_map_val,
'etypes': etypes_map} "ntypes": ntypes_map,
"etypes": etypes_map,
part_dir = 'part' + str(part_id) }
part_dir = "part" + str(part_id)
node_feat_file = os.path.join(part_dir, "node_feat.dgl") node_feat_file = os.path.join(part_dir, "node_feat.dgl")
edge_feat_file = os.path.join(part_dir, "edge_feat.dgl") edge_feat_file = os.path.join(part_dir, "edge_feat.dgl")
part_graph_file = os.path.join(part_dir, "graph.dgl") part_graph_file = os.path.join(part_dir, "graph.dgl")
part_metadata['part-{}'.format(part_id)] = {'node_feats': node_feat_file, part_metadata["part-{}".format(part_id)] = {
'edge_feats': edge_feat_file, "node_feats": node_feat_file,
'part_graph': part_graph_file} "edge_feats": edge_feat_file,
"part_graph": part_graph_file,
}
return part_metadata return part_metadata
tools/distpartitioning/data_proc_pipeline.py
View file @ aa42aaeb
...@@ -8,61 +8,104 @@ import torch.multiprocessing as mp ...@@ -8,61 +8,104 @@ import torch.multiprocessing as mp
from data_shuffle import multi_machine_run, single_machine_run from data_shuffle import multi_machine_run, single_machine_run
def log_params(params):
""" Print all the command line arguments for debugging purposes. def log_params(params):
"""Print all the command line arguments for debugging purposes.
Parameters: Parameters:
----------- -----------
params: argparse object params: argparse object
Argument Parser structure listing all the pre-defined parameters Argument Parser structure listing all the pre-defined parameters
""" """
print('Input Dir: ', params.input_dir) print("Input Dir: ", params.input_dir)
print('Graph Name: ', params.graph_name) print("Graph Name: ", params.graph_name)
print('Schema File: ', params.schema) print("Schema File: ", params.schema)
print('No. partitions: ', params.num_parts) print("No. partitions: ", params.num_parts)
print('Output Dir: ', params.output) print("Output Dir: ", params.output)
print('WorldSize: ', params.world_size) print("WorldSize: ", params.world_size)
print('Metis partitions: ', params.partitions_file) print("Metis partitions: ", params.partitions_file)
if __name__ == "__main__": if __name__ == "__main__":
""" """
Start of execution from this point. Start of execution from this point.
Invoke the appropriate function to begin execution Invoke the appropriate function to begin execution
""" """
#arguments which are already needed by the existing implementation of convert_partition.py # arguments which are already needed by the existing implementation of convert_partition.py
parser = argparse.ArgumentParser(description='Construct graph partitions') parser = argparse.ArgumentParser(description="Construct graph partitions")
parser.add_argument('--input-dir', required=True, type=str, parser.add_argument(
help='The directory path that contains the partition results.') "--input-dir",
parser.add_argument('--graph-name', required=True, type=str, required=True,
help='The graph name') type=str,
parser.add_argument('--schema', required=True, type=str, help="The directory path that contains the partition results.",
help='The schema of the graph') )
parser.add_argument('--num-parts', required=True, type=int, parser.add_argument(
help='The number of partitions') "--graph-name", required=True, type=str, help="The graph name"
parser.add_argument('--output', required=True, type=str, )
help='The output directory of the partitioned results') parser.add_argument(
parser.add_argument('--partitions-dir', help='directory of the partition-ids for each node type', "--schema", required=True, type=str, help="The schema of the graph"
default=None, type=str) )
parser.add_argument('--log-level', type=str, default="info", parser.add_argument(
help='To enable log level for debugging purposes. Available options: \ "--num-parts", required=True, type=int, help="The number of partitions"
)
parser.add_argument(
"--output",
required=True,
type=str,
help="The output directory of the partitioned results",
)
parser.add_argument(
"--partitions-dir",
help="directory of the partition-ids for each node type",
default=None,
type=str,
)
parser.add_argument(
"--log-level",
type=str,
default="info",
help="To enable log level for debugging purposes. Available options: \
(Critical, Error, Warning, Info, Debug, Notset), default value \ (Critical, Error, Warning, Info, Debug, Notset), default value \
is: Info') is: Info",
)
#arguments needed for the distributed implementation # arguments needed for the distributed implementation
parser.add_argument('--world-size', help='no. of processes to spawn', parser.add_argument(
default=1, type=int, required=True) "--world-size",
parser.add_argument('--process-group-timeout', required=True, type=int, help="no. of processes to spawn",
help='timeout[seconds] for operations executed against the process group ' default=1,
'(see torch.distributed.init_process_group)') type=int,
parser.add_argument('--save-orig-nids', action='store_true', required=True,
help='Save original node IDs into files') )
parser.add_argument('--save-orig-eids', action='store_true', parser.add_argument(
help='Save original edge IDs into files') "--process-group-timeout",
parser.add_argument('--graph-formats', default=None, type=str, required=True,
help='Save partitions in specified formats.') type=int,
help="timeout[seconds] for operations executed against the process group "
"(see torch.distributed.init_process_group)",
)
parser.add_argument(
"--save-orig-nids",
action="store_true",
help="Save original node IDs into files",
)
parser.add_argument(
"--save-orig-eids",
action="store_true",
help="Save original edge IDs into files",
)
parser.add_argument(
"--graph-formats",
default=None,
type=str,
help="Save partitions in specified formats.",
)
params = parser.parse_args() params = parser.parse_args()
#invoke the pipeline function # invoke the pipeline function
numeric_level = getattr(logging, params.log_level.upper(), None) numeric_level = getattr(logging, params.log_level.upper(), None)
logging.basicConfig(level=numeric_level, format=f"[{platform.node()} %(levelname)s %(asctime)s PID:%(process)d] %(message)s") logging.basicConfig(
level=numeric_level,
format=f"[{platform.node()} %(levelname)s %(asctime)s PID:%(process)d] %(message)s",
)
multi_machine_run(params) multi_machine_run(params)
tools/distpartitioning/data_shuffle.py
View file @ aa42aaeb
...@@ -6,28 +6,43 @@ import sys ...@@ -6,28 +6,43 @@ import sys
from datetime import timedelta from datetime import timedelta
from timeit import default_timer as timer from timeit import default_timer as timer
import constants
import dgl import dgl
import numpy as np import numpy as np
import torch import torch
import torch.distributed as dist import torch.distributed as dist
import torch.multiprocessing as mp import torch.multiprocessing as mp
import constants
from convert_partition import create_dgl_object, create_metadata_json from convert_partition import create_dgl_object, create_metadata_json
from dataset_utils import get_dataset from dataset_utils import get_dataset
from dist_lookup import DistLookupService from dist_lookup import DistLookupService
from globalids import (assign_shuffle_global_nids_edges, from globalids import (
assign_shuffle_global_nids_nodes, assign_shuffle_global_nids_edges,
lookup_shuffle_global_nids_edges) assign_shuffle_global_nids_nodes,
lookup_shuffle_global_nids_edges,
)
from gloo_wrapper import allgather_sizes, alltoallv_cpu, gather_metadata_json from gloo_wrapper import allgather_sizes, alltoallv_cpu, gather_metadata_json
from utils import (augment_edge_data, get_edge_types, get_etype_featnames, from utils import (
get_gnid_range_map, get_idranges, get_node_types, augment_edge_data,
get_ntype_featnames, memory_snapshot, read_json, get_edge_types,
read_ntype_partition_files, write_dgl_objects, get_etype_featnames,
write_metadata_json, map_partid_rank) get_gnid_range_map,
get_idranges,
def gen_node_data(rank, world_size, num_parts, id_lookup, ntid_ntype_map, schema_map): get_node_types,
''' get_ntype_featnames,
map_partid_rank,
memory_snapshot,
read_json,
read_ntype_partition_files,
write_dgl_objects,
write_metadata_json,
)
def gen_node_data(
rank, world_size, num_parts, id_lookup, ntid_ntype_map, schema_map
):
"""
For this data processing pipeline, reading node files is not needed. All the needed information about For this data processing pipeline, reading node files is not needed. All the needed information about
the nodes can be found in the metadata json file. This function generates the nodes owned by a given the nodes can be found in the metadata json file. This function generates the nodes owned by a given
process, using metis partitions. process, using metis partitions.
...@@ -41,9 +56,9 @@ def gen_node_data(rank, world_size, num_parts, id_lookup, ntid_ntype_map, schema ...@@ -41,9 +56,9 @@ def gen_node_data(rank, world_size, num_parts, id_lookup, ntid_ntype_map, schema
num_parts : int num_parts : int
total no. of partitions total no. of partitions
id_lookup : instance of class DistLookupService id_lookup : instance of class DistLookupService
Distributed lookup service used to map global-nids to respective partition-ids and Distributed lookup service used to map global-nids to respective partition-ids and
shuffle-global-nids shuffle-global-nids
ntid_ntype_map : ntid_ntype_map :
a dictionary where keys are node_type ids(integers) and values are node_type names(strings). a dictionary where keys are node_type ids(integers) and values are node_type names(strings).
schema_map: schema_map:
dictionary formed by reading the input metadata json file for the input dataset. dictionary formed by reading the input metadata json file for the input dataset.
...@@ -92,44 +107,63 @@ def gen_node_data(rank, world_size, num_parts, id_lookup, ntid_ntype_map, schema ...@@ -92,44 +107,63 @@ def gen_node_data(rank, world_size, num_parts, id_lookup, ntid_ntype_map, schema
dictionary where keys are column names and values are numpy arrays, these arrays are generated by dictionary where keys are column names and values are numpy arrays, these arrays are generated by
using information present in the metadata json file using information present in the metadata json file
''' """
local_node_data = {} local_node_data = {}
for local_part_id in range(num_parts//world_size): for local_part_id in range(num_parts // world_size):
local_node_data[constants.GLOBAL_NID+"/"+str(local_part_id)] = [] local_node_data[constants.GLOBAL_NID + "/" + str(local_part_id)] = []
local_node_data[constants.NTYPE_ID+"/"+str(local_part_id)] = [] local_node_data[constants.NTYPE_ID + "/" + str(local_part_id)] = []
local_node_data[constants.GLOBAL_TYPE_NID+"/"+str(local_part_id)] = [] local_node_data[
constants.GLOBAL_TYPE_NID + "/" + str(local_part_id)
] = []
# Note that `get_idranges` always returns two dictionaries. Keys in these # Note that `get_idranges` always returns two dictionaries. Keys in these
# dictionaries are type names for nodes and edges and values are # dictionaries are type names for nodes and edges and values are
# `num_parts` number of tuples indicating the range of type-ids in first # `num_parts` number of tuples indicating the range of type-ids in first
# dictionary and range of global-nids in the second dictionary. # dictionary and range of global-nids in the second dictionary.
type_nid_dict, global_nid_dict = get_idranges(schema_map[constants.STR_NODE_TYPE], type_nid_dict, global_nid_dict = get_idranges(
schema_map[constants.STR_NUM_NODES_PER_CHUNK], schema_map[constants.STR_NODE_TYPE],
num_chunks=num_parts) schema_map[constants.STR_NUM_NODES_PER_CHUNK],
num_chunks=num_parts,
)
for ntype_id, ntype_name in ntid_ntype_map.items(): for ntype_id, ntype_name in ntid_ntype_map.items():
type_start, type_end = type_nid_dict[ntype_name][0][0], type_nid_dict[ntype_name][-1][1] type_start, type_end = (
gnid_start, gnid_end = global_nid_dict[ntype_name][0, 0], global_nid_dict[ntype_name][0, 1] type_nid_dict[ntype_name][0][0],
type_nid_dict[ntype_name][-1][1],
node_partid_slice = id_lookup.get_partition_ids(np.arange(gnid_start, gnid_end, dtype=np.int64)) #exclusive )
gnid_start, gnid_end = (
for local_part_id in range(num_parts//world_size): global_nid_dict[ntype_name][0, 0],
cond = node_partid_slice == (rank + local_part_id*world_size) global_nid_dict[ntype_name][0, 1],
)
node_partid_slice = id_lookup.get_partition_ids(
np.arange(gnid_start, gnid_end, dtype=np.int64)
) # exclusive
for local_part_id in range(num_parts // world_size):
cond = node_partid_slice == (rank + local_part_id * world_size)
own_gnids = np.arange(gnid_start, gnid_end, dtype=np.int64) own_gnids = np.arange(gnid_start, gnid_end, dtype=np.int64)
own_gnids = own_gnids[cond] own_gnids = own_gnids[cond]
own_tnids = np.arange(type_start, type_end, dtype=np.int64) own_tnids = np.arange(type_start, type_end, dtype=np.int64)
own_tnids = own_tnids[cond] own_tnids = own_tnids[cond]
local_node_data[constants.NTYPE_ID+"/"+str(local_part_id)].append(np.ones(own_gnids.shape, dtype=np.int64)*ntype_id) local_node_data[
local_node_data[constants.GLOBAL_NID+"/"+str(local_part_id)].append(own_gnids) constants.NTYPE_ID + "/" + str(local_part_id)
local_node_data[constants.GLOBAL_TYPE_NID+"/"+str(local_part_id)].append(own_tnids) ].append(np.ones(own_gnids.shape, dtype=np.int64) * ntype_id)
local_node_data[
constants.GLOBAL_NID + "/" + str(local_part_id)
].append(own_gnids)
local_node_data[
constants.GLOBAL_TYPE_NID + "/" + str(local_part_id)
].append(own_tnids)
for k in local_node_data.keys(): for k in local_node_data.keys():
local_node_data[k] = np.concatenate(local_node_data[k]) local_node_data[k] = np.concatenate(local_node_data[k])
return local_node_data return local_node_data
def exchange_edge_data(rank, world_size, num_parts, edge_data): def exchange_edge_data(rank, world_size, num_parts, edge_data):
""" """
Exchange edge_data among processes in the world. Exchange edge_data among processes in the world.
...@@ -153,37 +187,59 @@ def exchange_edge_data(rank, world_size, num_parts, edge_data): ...@@ -153,37 +187,59 @@ def exchange_edge_data(rank, world_size, num_parts, edge_data):
in the world. in the world.
""" """
# Prepare data for each rank in the cluster. # Prepare data for each rank in the cluster.
start = timer() start = timer()
for local_part_id in range(num_parts//world_size): for local_part_id in range(num_parts // world_size):
input_list = [] input_list = []
for idx in range(world_size): for idx in range(world_size):
send_idx = (edge_data[constants.OWNER_PROCESS] == (idx + local_part_id*world_size)) send_idx = edge_data[constants.OWNER_PROCESS] == (
send_idx = send_idx.reshape(edge_data[constants.GLOBAL_SRC_ID].shape[0]) idx + local_part_id * world_size
filt_data = np.column_stack((edge_data[constants.GLOBAL_SRC_ID][send_idx == 1], \ )
edge_data[constants.GLOBAL_DST_ID][send_idx == 1], \ send_idx = send_idx.reshape(
edge_data[constants.GLOBAL_TYPE_EID][send_idx == 1], \ edge_data[constants.GLOBAL_SRC_ID].shape[0]
edge_data[constants.ETYPE_ID][send_idx == 1], \ )
edge_data[constants.GLOBAL_EID][send_idx == 1])) filt_data = np.column_stack(
if(filt_data.shape[0] <= 0): (
input_list.append(torch.empty((0,5), dtype=torch.int64)) edge_data[constants.GLOBAL_SRC_ID][send_idx == 1],
edge_data[constants.GLOBAL_DST_ID][send_idx == 1],
edge_data[constants.GLOBAL_TYPE_EID][send_idx == 1],
edge_data[constants.ETYPE_ID][send_idx == 1],
edge_data[constants.GLOBAL_EID][send_idx == 1],
)
)
if filt_data.shape[0] <= 0:
input_list.append(torch.empty((0, 5), dtype=torch.int64))
else: else:
input_list.append(torch.from_numpy(filt_data)) input_list.append(torch.from_numpy(filt_data))
dist.barrier () dist.barrier()
output_list = alltoallv_cpu(rank, world_size, input_list, retain_nones=False) output_list = alltoallv_cpu(
rank, world_size, input_list, retain_nones=False
)
#Replace the values of the edge_data, with the received data from all the other processes. # Replace the values of the edge_data, with the received data from all the other processes.
rcvd_edge_data = torch.cat(output_list).numpy() rcvd_edge_data = torch.cat(output_list).numpy()
edge_data[constants.GLOBAL_SRC_ID+"/"+str(local_part_id)] = rcvd_edge_data[:,0] edge_data[
edge_data[constants.GLOBAL_DST_ID+"/"+str(local_part_id)] = rcvd_edge_data[:,1] constants.GLOBAL_SRC_ID + "/" + str(local_part_id)
edge_data[constants.GLOBAL_TYPE_EID+"/"+str(local_part_id)] = rcvd_edge_data[:,2] ] = rcvd_edge_data[:, 0]
edge_data[constants.ETYPE_ID+"/"+str(local_part_id)] = rcvd_edge_data[:,3] edge_data[
edge_data[constants.GLOBAL_EID+"/"+str(local_part_id)] = rcvd_edge_data[:,4] constants.GLOBAL_DST_ID + "/" + str(local_part_id)
] = rcvd_edge_data[:, 1]
edge_data[
constants.GLOBAL_TYPE_EID + "/" + str(local_part_id)
] = rcvd_edge_data[:, 2]
edge_data[
constants.ETYPE_ID + "/" + str(local_part_id)
] = rcvd_edge_data[:, 3]
edge_data[
constants.GLOBAL_EID + "/" + str(local_part_id)
] = rcvd_edge_data[:, 4]
end = timer() end = timer()
logging.info(f'[Rank: {rank}] Time to send/rcv edge data: {timedelta(seconds=end-start)}') logging.info(
f"[Rank: {rank}] Time to send/rcv edge data: {timedelta(seconds=end-start)}"
)
# Clean up. # Clean up.
edge_data.pop(constants.OWNER_PROCESS) edge_data.pop(constants.OWNER_PROCESS)
...@@ -195,30 +251,45 @@ def exchange_edge_data(rank, world_size, num_parts, edge_data): ...@@ -195,30 +251,45 @@ def exchange_edge_data(rank, world_size, num_parts, edge_data):
return edge_data return edge_data
def exchange_feature(rank, data, id_lookup, feat_type, feat_key, featdata_key, gid_start,
gid_end, type_id_start, type_id_end, local_part_id, world_size, num_parts, def exchange_feature(
cur_features, cur_global_ids): rank,
"""This function is used to send/receive one feature for either nodes or data,
id_lookup,
feat_type,
feat_key,
featdata_key,
gid_start,
gid_end,
type_id_start,
type_id_end,
local_part_id,
world_size,
num_parts,
cur_features,
cur_global_ids,
):
"""This function is used to send/receive one feature for either nodes or
edges of the input graph dataset. edges of the input graph dataset.
Parameters: Parameters:
----------- -----------
rank : int rank : int
integer, unique id assigned to the current process integer, unique id assigned to the current process
data: dicitonary data: dicitonary
dictionry in which node or edge features are stored and this information dictionry in which node or edge features are stored and this information
is read from the appropriate node features file which belongs to the is read from the appropriate node features file which belongs to the
current process current process
id_lookup : instance of DistLookupService id_lookup : instance of DistLookupService
instance of an implementation of dist. lookup service to retrieve values instance of an implementation of dist. lookup service to retrieve values
for keys for keys
feat_type : string feat_type : string
this is used to distinguish which features are being exchanged. Please this is used to distinguish which features are being exchanged. Please
note that for nodes ownership is clearly defined and for edges it is note that for nodes ownership is clearly defined and for edges it is
always assumed that destination end point of the edge defines the always assumed that destination end point of the edge defines the
ownership of that particular edge ownership of that particular edge
feat_key : string feat_key : string
this string is used as a key in the dictionary to store features, as this string is used as a key in the dictionary to store features, as
tensors, in local dictionaries tensors, in local dictionaries
featdata_key : numpy array featdata_key : numpy array
features associated with this feature key being processed features associated with this feature key being processed
...@@ -238,22 +309,22 @@ def exchange_feature(rank, data, id_lookup, feat_type, feat_key, featdata_key, g ...@@ -238,22 +309,22 @@ def exchange_feature(rank, data, id_lookup, feat_type, feat_key, featdata_key, g
num_parts : int num_parts : int
total number of partitions total number of partitions
cur_features : dictionary cur_features : dictionary
dictionary to store the feature data which belongs to the current dictionary to store the feature data which belongs to the current
process process
cur_global_ids : dictionary cur_global_ids : dictionary
dictionary to store global ids, of either nodes or edges, for which dictionary to store global ids, of either nodes or edges, for which
the features stored in the cur_features dictionary the features stored in the cur_features dictionary
Returns: Returns:
------- -------
dictionary : dictionary :
a dictionary is returned where keys are type names and a dictionary is returned where keys are type names and
feature data are the values feature data are the values
list : list :
a dictionary of global_ids either nodes or edges whose features are a dictionary of global_ids either nodes or edges whose features are
received during the data shuffle process received during the data shuffle process
""" """
#type_ids for this feature subset on the current rank # type_ids for this feature subset on the current rank
gids_feat = np.arange(gid_start, gid_end) gids_feat = np.arange(gid_start, gid_end)
tids_feat = np.arange(type_id_start, type_id_end) tids_feat = np.arange(type_id_start, type_id_end)
local_idx = np.arange(0, type_id_end - type_id_start) local_idx = np.arange(0, type_id_end - type_id_start)
...@@ -263,54 +334,64 @@ def exchange_feature(rank, data, id_lookup, feat_type, feat_key, featdata_key, g ...@@ -263,54 +334,64 @@ def exchange_feature(rank, data, id_lookup, feat_type, feat_key, featdata_key, g
tokens = feat_key.split("/") tokens = feat_key.split("/")
assert len(tokens) == 3 assert len(tokens) == 3
local_feat_key = "/".join(tokens[:-1]) +"/"+ str(local_part_id) local_feat_key = "/".join(tokens[:-1]) + "/" + str(local_part_id)
for idx in range(world_size): for idx in range(world_size):
# Get the partition ids for the range of global nids. # Get the partition ids for the range of global nids.
if feat_type == constants.STR_NODE_FEATURES: if feat_type == constants.STR_NODE_FEATURES:
# Retrieve the partition ids for the node features. # Retrieve the partition ids for the node features.
# Each partition id will be in the range [0, num_parts). # Each partition id will be in the range [0, num_parts).
partid_slice = id_lookup.get_partition_ids(np.arange(gid_start, gid_end, dtype=np.int64)) partid_slice = id_lookup.get_partition_ids(
np.arange(gid_start, gid_end, dtype=np.int64)
)
else: else:
#Edge data case. # Edge data case.
#Ownership is determined by the destination node. # Ownership is determined by the destination node.
assert data is not None assert data is not None
global_eids = np.arange(gid_start, gid_end, dtype=np.int64) global_eids = np.arange(gid_start, gid_end, dtype=np.int64)
#Now use `data` to extract destination nodes' global id # Now use `data` to extract destination nodes' global id
#and use that to get the ownership # and use that to get the ownership
common, idx1, idx2 = np.intersect1d(data[constants.GLOBAL_EID], global_eids, return_indices=True) common, idx1, idx2 = np.intersect1d(
data[constants.GLOBAL_EID], global_eids, return_indices=True
)
assert common.shape[0] == idx2.shape[0] assert common.shape[0] == idx2.shape[0]
global_dst_nids = data[constants.GLOBAL_DST_ID][idx1] global_dst_nids = data[constants.GLOBAL_DST_ID][idx1]
assert np.all(global_eids == data[constants.GLOBAL_EID][idx1]) assert np.all(global_eids == data[constants.GLOBAL_EID][idx1])
partid_slice = id_lookup.get_partition_ids(global_dst_nids) partid_slice = id_lookup.get_partition_ids(global_dst_nids)
cond = (partid_slice == (idx + local_part_id*world_size)) cond = partid_slice == (idx + local_part_id * world_size)
gids_per_partid = gids_feat[cond] gids_per_partid = gids_feat[cond]
tids_per_partid = tids_feat[cond] tids_per_partid = tids_feat[cond]
local_idx_partid = local_idx[cond] local_idx_partid = local_idx[cond]
if (gids_per_partid.shape[0] == 0): if gids_per_partid.shape[0] == 0:
feats_per_rank.append(torch.empty((0,1), dtype=torch.float)) feats_per_rank.append(torch.empty((0, 1), dtype=torch.float))
global_id_per_rank.append(torch.empty((0,1), dtype=torch.int64)) global_id_per_rank.append(torch.empty((0, 1), dtype=torch.int64))
else: else:
feats_per_rank.append(featdata_key[local_idx_partid]) feats_per_rank.append(featdata_key[local_idx_partid])
global_id_per_rank.append(torch.from_numpy(gids_per_partid).type(torch.int64)) global_id_per_rank.append(
torch.from_numpy(gids_per_partid).type(torch.int64)
#features (and global nids) per rank to be sent out are ready )
#for transmission, perform alltoallv here.
output_feat_list = alltoallv_cpu(rank, world_size, feats_per_rank, retain_nones=False) # features (and global nids) per rank to be sent out are ready
output_id_list = alltoallv_cpu(rank, world_size, global_id_per_rank, retain_nones=False) # for transmission, perform alltoallv here.
output_feat_list = alltoallv_cpu(
rank, world_size, feats_per_rank, retain_nones=False
)
output_id_list = alltoallv_cpu(
rank, world_size, global_id_per_rank, retain_nones=False
)
assert len(output_feat_list) == len(output_id_list), ( assert len(output_feat_list) == len(output_id_list), (
"Length of feature list and id list are expected to be equal while " "Length of feature list and id list are expected to be equal while "
f"got {len(output_feat_list)} and {len(output_id_list)}." f"got {len(output_feat_list)} and {len(output_id_list)}."
) )
#stitch node_features together to form one large feature tensor # stitch node_features together to form one large feature tensor
if len(output_feat_list) > 0: if len(output_feat_list) > 0:
output_feat_list = torch.cat(output_feat_list) output_feat_list = torch.cat(output_feat_list)
output_id_list = torch.cat(output_id_list) output_id_list = torch.cat(output_id_list)
if local_feat_key in cur_features: if local_feat_key in cur_features:
temp = cur_features[local_feat_key] temp = cur_features[local_feat_key]
cur_features[local_feat_key] = torch.cat([temp, output_feat_list]) cur_features[local_feat_key] = torch.cat([temp, output_feat_list])
temp = cur_global_ids[local_feat_key] temp = cur_global_ids[local_feat_key]
...@@ -322,7 +403,17 @@ def exchange_feature(rank, data, id_lookup, feat_type, feat_key, featdata_key, g ...@@ -322,7 +403,17 @@ def exchange_feature(rank, data, id_lookup, feat_type, feat_key, featdata_key, g
return cur_features, cur_global_ids return cur_features, cur_global_ids
def exchange_features(rank, world_size, num_parts, feature_tids, type_id_map, id_lookup, feature_data, feat_type, data): def exchange_features(
rank,
world_size,
num_parts,
feature_tids,
type_id_map,
id_lookup,
feature_data,
feat_type,
data,
):
""" """
This function is used to shuffle node features so that each process will receive This function is used to shuffle node features so that each process will receive
all the node features whose corresponding nodes are owned by the same process. all the node features whose corresponding nodes are owned by the same process.
...@@ -348,35 +439,35 @@ def exchange_features(rank, world_size, num_parts, feature_tids, type_id_map, id ...@@ -348,35 +439,35 @@ def exchange_features(rank, world_size, num_parts, feature_tids, type_id_map, id
world_size : int world_size : int
total no. of participating processes. total no. of participating processes.
feature_tids : dictionary feature_tids : dictionary
dictionary with keys as node-type names with suffixes as feature names dictionary with keys as node-type names with suffixes as feature names
and value is a dictionary. This dictionary contains information about and value is a dictionary. This dictionary contains information about
node-features associated with a given node-type and value is a list. node-features associated with a given node-type and value is a list.
This list contains a of indexes, like [starting-idx, ending-idx) which This list contains a of indexes, like [starting-idx, ending-idx) which
can be used to index into the node feature tensors read from can be used to index into the node feature tensors read from
corresponding input files. corresponding input files.
type_id_map : dictionary type_id_map : dictionary
mapping between type names and global_ids, of either nodes or edges, mapping between type names and global_ids, of either nodes or edges,
which belong to the keys in this dictionary which belong to the keys in this dictionary
id_lookup : instance of class DistLookupService id_lookup : instance of class DistLookupService
Distributed lookup service used to map global-nids to respective Distributed lookup service used to map global-nids to respective
partition-ids and shuffle-global-nids partition-ids and shuffle-global-nids
feat_type : string feat_type : string
this is used to distinguish which features are being exchanged. Please this is used to distinguish which features are being exchanged. Please
note that for nodes ownership is clearly defined and for edges it is note that for nodes ownership is clearly defined and for edges it is
always assumed that destination end point of the edge defines the always assumed that destination end point of the edge defines the
ownership of that particular edge ownership of that particular edge
data: dicitonary data: dicitonary
dictionry in which node or edge features are stored and this information dictionry in which node or edge features are stored and this information
is read from the appropriate node features file which belongs to the is read from the appropriate node features file which belongs to the
current process current process
Returns: Returns:
-------- --------
dictionary : dictionary :
a dictionary is returned where keys are type names and a dictionary is returned where keys are type names and
feature data are the values feature data are the values
list : list :
a dictionary of global_ids either nodes or edges whose features are a dictionary of global_ids either nodes or edges whose features are
received during the data shuffle process received during the data shuffle process
""" """
start = timer() start = timer()
...@@ -389,21 +480,21 @@ def exchange_features(rank, world_size, num_parts, feature_tids, type_id_map, id ...@@ -389,21 +480,21 @@ def exchange_features(rank, world_size, num_parts, feature_tids, type_id_map, id
# To iterate over the feature data, of a given (node or edge )type # To iterate over the feature data, of a given (node or edge )type
# type_info is a list of 3 elements (as shown below): # type_info is a list of 3 elements (as shown below):
# [feature-name, starting-idx, ending-idx] # [feature-name, starting-idx, ending-idx]
# feature-name is the name given to the feature-data, # feature-name is the name given to the feature-data,
# read from the input metadata file # read from the input metadata file
# [starting-idx, ending-idx) specifies the range of indexes # [starting-idx, ending-idx) specifies the range of indexes
# associated with the features data # associated with the features data
# Determine the owner process for these features. # Determine the owner process for these features.
# Note that the keys in the node features (and similarly edge features) # Note that the keys in the node features (and similarly edge features)
# dictionary is of the following format: # dictionary is of the following format:
# `node_type/feature_name/local_part_id`: # `node_type/feature_name/local_part_id`:
# where node_type and feature_name are self-explanatory and # where node_type and feature_name are self-explanatory and
# local_part_id denotes the partition-id, in the local process, # local_part_id denotes the partition-id, in the local process,
# which will be used a suffix to store all the information of a # which will be used a suffix to store all the information of a
# given partition which is processed by the current process. Its # given partition which is processed by the current process. Its
# values start from 0 onwards, for instance 0, 1, 2 ... etc. # values start from 0 onwards, for instance 0, 1, 2 ... etc.
# local_part_id can be easily mapped to global partition id very # local_part_id can be easily mapped to global partition id very
# easily, using cyclic ordering. All local_part_ids = 0 from all # easily, using cyclic ordering. All local_part_ids = 0 from all
# processes will form global partition-ids between 0 and world_size-1. # processes will form global partition-ids between 0 and world_size-1.
# Similarly all local_part_ids = 1 from all processes will form # Similarly all local_part_ids = 1 from all processes will form
# global partition ids in the range [world_size, 2*world_size-1] and # global partition ids in the range [world_size, 2*world_size-1] and
...@@ -412,7 +503,7 @@ def exchange_features(rank, world_size, num_parts, feature_tids, type_id_map, id ...@@ -412,7 +503,7 @@ def exchange_features(rank, world_size, num_parts, feature_tids, type_id_map, id
assert len(tokens) == 3 assert len(tokens) == 3
type_name = tokens[0] type_name = tokens[0]
feat_name = tokens[1] feat_name = tokens[1]
logging.info(f'[Rank: {rank}] processing feature: {feat_key}') logging.info(f"[Rank: {rank}] processing feature: {feat_key}")
for feat_info in type_info: for feat_info in type_info:
# Compute the global_id range for this feature data # Compute the global_id range for this feature data
...@@ -425,24 +516,51 @@ def exchange_features(rank, world_size, num_parts, feature_tids, type_id_map, id ...@@ -425,24 +516,51 @@ def exchange_features(rank, world_size, num_parts, feature_tids, type_id_map, id
# Check if features exist for this type_name + feat_name. # Check if features exist for this type_name + feat_name.
# This check should always pass, because feature_tids are built # This check should always pass, because feature_tids are built
# by reading the input metadata json file for existing features. # by reading the input metadata json file for existing features.
assert(feat_key in feature_data) assert feat_key in feature_data
for local_part_id in range(num_parts//world_size): for local_part_id in range(num_parts // world_size):
featdata_key = feature_data[feat_key] featdata_key = feature_data[feat_key]
own_features, own_global_ids = exchange_feature(rank, data, id_lookup, own_features, own_global_ids = exchange_feature(
feat_type, feat_key, featdata_key, gid_start, gid_end, type_id_start, rank,
type_id_end, local_part_id, world_size, num_parts, own_features, data,
own_global_ids) id_lookup,
feat_type,
feat_key,
featdata_key,
gid_start,
gid_end,
type_id_start,
type_id_end,
local_part_id,
world_size,
num_parts,
own_features,
own_global_ids,
)
end = timer() end = timer()
logging.info(f'[Rank: {rank}] Total time for feature exchange: {timedelta(seconds = end - start)}') logging.info(
f"[Rank: {rank}] Total time for feature exchange: {timedelta(seconds = end - start)}"
)
return own_features, own_global_ids return own_features, own_global_ids
def exchange_graph_data(rank, world_size, num_parts, node_features, edge_features,
node_feat_tids, edge_feat_tids, def exchange_graph_data(
edge_data, id_lookup, ntypes_ntypeid_map, rank,
ntypes_gnid_range_map, etypes_geid_range_map, world_size,
ntid_ntype_map, schema_map): num_parts,
node_features,
edge_features,
node_feat_tids,
edge_feat_tids,
edge_data,
id_lookup,
ntypes_ntypeid_map,
ntypes_gnid_range_map,
etypes_geid_range_map,
ntid_ntype_map,
schema_map,
):
""" """
Wrapper function which is used to shuffle graph data on all the processes. Wrapper function which is used to shuffle graph data on all the processes.
...@@ -467,13 +585,13 @@ def exchange_graph_data(rank, world_size, num_parts, node_features, edge_feature ...@@ -467,13 +585,13 @@ def exchange_graph_data(rank, world_size, num_parts, node_features, edge_feature
hence each key may have several triplets. hence each key may have several triplets.
edge_feat_tids : dictionary edge_feat_tids : dictionary
a dictionary in which keys are edge-type names and values are triplets of the format a dictionary in which keys are edge-type names and values are triplets of the format
<feat-name, start-per-type-idx, end-per-type-idx>. This triplet is used to identify <feat-name, start-per-type-idx, end-per-type-idx>. This triplet is used to identify
the chunk of feature data for which current process is responsible for the chunk of feature data for which current process is responsible for
edge_data : dictionary edge_data : dictionary
dictionary which is used to store edge information as read from appropriate files assigned dictionary which is used to store edge information as read from appropriate files assigned
to each process. to each process.
id_lookup : instance of class DistLookupService id_lookup : instance of class DistLookupService
Distributed lookup service used to map global-nids to respective partition-ids and Distributed lookup service used to map global-nids to respective partition-ids and
shuffle-global-nids shuffle-global-nids
ntypes_ntypeid_map : dictionary ntypes_ntypeid_map : dictionary
mappings between node type names and node type ids mappings between node type names and node type ids
...@@ -501,7 +619,7 @@ def exchange_graph_data(rank, world_size, num_parts, node_features, edge_feature ...@@ -501,7 +619,7 @@ def exchange_graph_data(rank, world_size, num_parts, node_features, edge_feature
dictionary : dictionary :
the input argument, edge_data dictionary, is updated with the edge data received from other processes the input argument, edge_data dictionary, is updated with the edge data received from other processes
in the world. The edge data is received by each rank in the process of data shuffling. in the world. The edge data is received by each rank in the process of data shuffling.
dictionary : dictionary :
edge features dictionary which has edge features. These destination end points of these edges edge features dictionary which has edge features. These destination end points of these edges
are owned by the current process are owned by the current process
dictionary : dictionary :
...@@ -509,23 +627,49 @@ def exchange_graph_data(rank, world_size, num_parts, node_features, edge_feature ...@@ -509,23 +627,49 @@ def exchange_graph_data(rank, world_size, num_parts, node_features, edge_feature
was performed in the `exchange_features` function call was performed in the `exchange_features` function call
""" """
memory_snapshot("ShuffleNodeFeaturesBegin: ", rank) memory_snapshot("ShuffleNodeFeaturesBegin: ", rank)
rcvd_node_features, rcvd_global_nids = exchange_features(rank, world_size, num_parts, node_feat_tids, rcvd_node_features, rcvd_global_nids = exchange_features(
ntypes_gnid_range_map, id_lookup, node_features, rank,
constants.STR_NODE_FEATURES, None) world_size,
num_parts,
node_feat_tids,
ntypes_gnid_range_map,
id_lookup,
node_features,
constants.STR_NODE_FEATURES,
None,
)
memory_snapshot("ShuffleNodeFeaturesComplete: ", rank) memory_snapshot("ShuffleNodeFeaturesComplete: ", rank)
logging.info(f'[Rank: {rank}] Done with node features exchange.') logging.info(f"[Rank: {rank}] Done with node features exchange.")
rcvd_edge_features, rcvd_global_eids = exchange_features(rank, world_size, num_parts, edge_feat_tids, rcvd_edge_features, rcvd_global_eids = exchange_features(
etypes_geid_range_map, id_lookup, edge_features, rank,
constants.STR_EDGE_FEATURES, edge_data) world_size,
logging.info(f'[Rank: {rank}] Done with edge features exchange.') num_parts,
edge_feat_tids,
etypes_geid_range_map,
id_lookup,
edge_features,
constants.STR_EDGE_FEATURES,
edge_data,
)
logging.info(f"[Rank: {rank}] Done with edge features exchange.")
node_data = gen_node_data(rank, world_size, num_parts, id_lookup, ntid_ntype_map, schema_map) node_data = gen_node_data(
rank, world_size, num_parts, id_lookup, ntid_ntype_map, schema_map
)
memory_snapshot("NodeDataGenerationComplete: ", rank) memory_snapshot("NodeDataGenerationComplete: ", rank)
edge_data = exchange_edge_data(rank, world_size, num_parts, edge_data) edge_data = exchange_edge_data(rank, world_size, num_parts, edge_data)
memory_snapshot("ShuffleEdgeDataComplete: ", rank) memory_snapshot("ShuffleEdgeDataComplete: ", rank)
return node_data, rcvd_node_features, rcvd_global_nids, edge_data, rcvd_edge_features, rcvd_global_eids return (
node_data,
rcvd_node_features,
rcvd_global_nids,
edge_data,
rcvd_edge_features,
rcvd_global_eids,
)
def read_dataset(rank, world_size, id_lookup, params, schema_map): def read_dataset(rank, world_size, id_lookup, params, schema_map):
""" """
...@@ -542,7 +686,7 @@ def read_dataset(rank, world_size, id_lookup, params, schema_map): ...@@ -542,7 +686,7 @@ def read_dataset(rank, world_size, id_lookup, params, schema_map):
world_size : int world_size : int
total no. of processes instantiated total no. of processes instantiated
id_lookup : instance of class DistLookupService id_lookup : instance of class DistLookupService
Distributed lookup service used to map global-nids to respective partition-ids and Distributed lookup service used to map global-nids to respective partition-ids and
shuffle-global-nids shuffle-global-nids
params : argparser object params : argparser object
argument parser object to access command line arguments argument parser object to access command line arguments
...@@ -571,20 +715,47 @@ def read_dataset(rank, world_size, id_lookup, params, schema_map): ...@@ -571,20 +715,47 @@ def read_dataset(rank, world_size, id_lookup, params, schema_map):
a dictionary in which keys are edge-type names and values are tuples indicating the range of ids a dictionary in which keys are edge-type names and values are tuples indicating the range of ids
for edges read by the current process. for edges read by the current process.
dictionary dictionary
a dictionary in which keys are edge-type names and values are triplets, a dictionary in which keys are edge-type names and values are triplets,
(edge-feature-name, start_type_id, end_type_id). These type_ids are indices in the edge-features (edge-feature-name, start_type_id, end_type_id). These type_ids are indices in the edge-features
read by the current process. Note that each edge-type may have several edge-features. read by the current process. Note that each edge-type may have several edge-features.
""" """
edge_features = {} edge_features = {}
#node_tids, node_features, edge_datadict, edge_tids # node_tids, node_features, edge_datadict, edge_tids
node_tids, node_features, node_feat_tids, edge_data, edge_tids, edge_features, edge_feat_tids = \ (
get_dataset(params.input_dir, params.graph_name, rank, world_size, params.num_parts, schema_map) node_tids,
logging.info(f'[Rank: {rank}] Done reading dataset {params.input_dir}') node_features,
node_feat_tids,
edge_data,
edge_tids,
edge_features,
edge_feat_tids,
) = get_dataset(
params.input_dir,
params.graph_name,
rank,
world_size,
params.num_parts,
schema_map,
)
logging.info(f"[Rank: {rank}] Done reading dataset {params.input_dir}")
edge_data = augment_edge_data(edge_data, id_lookup, edge_tids, rank, world_size, params.num_parts) edge_data = augment_edge_data(
logging.info(f'[Rank: {rank}] Done augmenting edge_data: {len(edge_data)}, {edge_data[constants.GLOBAL_SRC_ID].shape}') edge_data, id_lookup, edge_tids, rank, world_size, params.num_parts
)
logging.info(
f"[Rank: {rank}] Done augmenting edge_data: {len(edge_data)}, {edge_data[constants.GLOBAL_SRC_ID].shape}"
)
return (
node_tids,
node_features,
node_feat_tids,
edge_data,
edge_features,
edge_tids,
edge_feat_tids,
)
return node_tids, node_features, node_feat_tids, edge_data, edge_features, edge_tids, edge_feat_tids
def gen_dist_partitions(rank, world_size, params): def gen_dist_partitions(rank, world_size, params):
""" """
...@@ -711,52 +882,90 @@ def gen_dist_partitions(rank, world_size, params): ...@@ -711,52 +882,90 @@ def gen_dist_partitions(rank, world_size, params):
this object, key value pairs, provides access to the command line arguments from the runtime environment this object, key value pairs, provides access to the command line arguments from the runtime environment
""" """
global_start = timer() global_start = timer()
logging.info(f'[Rank: {rank}] Starting distributed data processing pipeline...') logging.info(
f"[Rank: {rank}] Starting distributed data processing pipeline..."
)
memory_snapshot("Pipeline Begin: ", rank) memory_snapshot("Pipeline Begin: ", rank)
#init processing # init processing
schema_map = read_json(os.path.join(params.input_dir, params.schema)) schema_map = read_json(os.path.join(params.input_dir, params.schema))
#Initialize distributed lookup service for partition-id and shuffle-global-nids mappings # Initialize distributed lookup service for partition-id and shuffle-global-nids mappings
#for global-nids # for global-nids
_, global_nid_ranges = get_idranges(schema_map[constants.STR_NODE_TYPE], _, global_nid_ranges = get_idranges(
schema_map[constants.STR_NUM_NODES_PER_CHUNK], params.num_parts) schema_map[constants.STR_NODE_TYPE],
schema_map[constants.STR_NUM_NODES_PER_CHUNK],
params.num_parts,
)
id_map = dgl.distributed.id_map.IdMap(global_nid_ranges) id_map = dgl.distributed.id_map.IdMap(global_nid_ranges)
# The resources, which are node-id to partition-id mappings, are split # The resources, which are node-id to partition-id mappings, are split
# into `world_size` number of parts, where each part can be mapped to # into `world_size` number of parts, where each part can be mapped to
# each physical node. # each physical node.
id_lookup = DistLookupService(os.path.join(params.input_dir, params.partitions_dir),\ id_lookup = DistLookupService(
schema_map[constants.STR_NODE_TYPE],\ os.path.join(params.input_dir, params.partitions_dir),
id_map, rank, world_size) schema_map[constants.STR_NODE_TYPE],
id_map,
rank,
world_size,
)
ntypes_ntypeid_map, ntypes, ntypeid_ntypes_map = get_node_types(schema_map) ntypes_ntypeid_map, ntypes, ntypeid_ntypes_map = get_node_types(schema_map)
etypes_etypeid_map, etypes, etypeid_etypes_map = get_edge_types(schema_map) etypes_etypeid_map, etypes, etypeid_etypes_map = get_edge_types(schema_map)
logging.info(f'[Rank: {rank}] Initialized metis partitions and node_types map...') logging.info(
f"[Rank: {rank}] Initialized metis partitions and node_types map..."
)
#read input graph files and augment these datastructures with # read input graph files and augment these datastructures with
#appropriate information (global_nid and owner process) for node and edge data # appropriate information (global_nid and owner process) for node and edge data
node_tids, node_features, node_feat_tids, edge_data, edge_features, edge_tids, edge_feat_tids = \ (
read_dataset(rank, world_size, id_lookup, params, schema_map) node_tids,
logging.info(f'[Rank: {rank}] Done augmenting file input data with auxilary columns') node_features,
node_feat_tids,
edge_data,
edge_features,
edge_tids,
edge_feat_tids,
) = read_dataset(rank, world_size, id_lookup, params, schema_map)
logging.info(
f"[Rank: {rank}] Done augmenting file input data with auxilary columns"
)
memory_snapshot("DatasetReadComplete: ", rank) memory_snapshot("DatasetReadComplete: ", rank)
#send out node and edge data --- and appropriate features. # send out node and edge data --- and appropriate features.
#this function will also stitch the data recvd from other processes # this function will also stitch the data recvd from other processes
#and return the aggregated data # and return the aggregated data
ntypes_gnid_range_map = get_gnid_range_map(node_tids) ntypes_gnid_range_map = get_gnid_range_map(node_tids)
etypes_geid_range_map = get_gnid_range_map(edge_tids) etypes_geid_range_map = get_gnid_range_map(edge_tids)
node_data, rcvd_node_features, rcvd_global_nids, edge_data, rcvd_edge_features, rcvd_global_eids = \ (
exchange_graph_data(rank, world_size, params.num_parts, node_features, edge_features, \ node_data,
node_feat_tids, edge_feat_tids, edge_data, id_lookup, ntypes_ntypeid_map, \ rcvd_node_features,
ntypes_gnid_range_map, etypes_geid_range_map, \ rcvd_global_nids,
ntypeid_ntypes_map, schema_map) edge_data,
rcvd_edge_features,
rcvd_global_eids,
) = exchange_graph_data(
rank,
world_size,
params.num_parts,
node_features,
edge_features,
node_feat_tids,
edge_feat_tids,
edge_data,
id_lookup,
ntypes_ntypeid_map,
ntypes_gnid_range_map,
etypes_geid_range_map,
ntypeid_ntypes_map,
schema_map,
)
gc.collect() gc.collect()
logging.info(f'[Rank: {rank}] Done with data shuffling...') logging.info(f"[Rank: {rank}] Done with data shuffling...")
memory_snapshot("DataShuffleComplete: ", rank) memory_snapshot("DataShuffleComplete: ", rank)
#sort node_data by ntype # sort node_data by ntype
for local_part_id in range(params.num_parts//world_size): for local_part_id in range(params.num_parts // world_size):
idx = node_data[constants.NTYPE_ID+"/"+str(local_part_id)].argsort() idx = node_data[constants.NTYPE_ID + "/" + str(local_part_id)].argsort()
for k, v in node_data.items(): for k, v in node_data.items():
tokens = k.split("/") tokens = k.split("/")
assert len(tokens) == 2 assert len(tokens) == 2
...@@ -764,34 +973,48 @@ def gen_dist_partitions(rank, world_size, params): ...@@ -764,34 +973,48 @@ def gen_dist_partitions(rank, world_size, params):
node_data[k] = v[idx] node_data[k] = v[idx]
idx = None idx = None
gc.collect() gc.collect()
logging.info(f'[Rank: {rank}] Sorted node_data by node_type') logging.info(f"[Rank: {rank}] Sorted node_data by node_type")
#resolve global_ids for nodes # resolve global_ids for nodes
assign_shuffle_global_nids_nodes(rank, world_size, params.num_parts, node_data) assign_shuffle_global_nids_nodes(
logging.info(f'[Rank: {rank}] Done assigning global-ids to nodes...') rank, world_size, params.num_parts, node_data
)
logging.info(f"[Rank: {rank}] Done assigning global-ids to nodes...")
memory_snapshot("ShuffleGlobalID_Nodes_Complete: ", rank) memory_snapshot("ShuffleGlobalID_Nodes_Complete: ", rank)
#shuffle node feature according to the node order on each rank. # shuffle node feature according to the node order on each rank.
for ntype_name in ntypes: for ntype_name in ntypes:
featnames = get_ntype_featnames(ntype_name, schema_map) featnames = get_ntype_featnames(ntype_name, schema_map)
for featname in featnames: for featname in featnames:
#if a feature name exists for a node-type, then it should also have # if a feature name exists for a node-type, then it should also have
#feature data as well. Hence using the assert statement. # feature data as well. Hence using the assert statement.
for local_part_id in range(params.num_parts//world_size): for local_part_id in range(params.num_parts // world_size):
feature_key = ntype_name+'/'+featname+"/"+str(local_part_id) feature_key = (
assert(feature_key in rcvd_global_nids) ntype_name + "/" + featname + "/" + str(local_part_id)
)
assert feature_key in rcvd_global_nids
global_nids = rcvd_global_nids[feature_key] global_nids = rcvd_global_nids[feature_key]
_, idx1, _ = np.intersect1d(node_data[constants.GLOBAL_NID+"/"+str(local_part_id)], global_nids, return_indices=True) _, idx1, _ = np.intersect1d(
shuffle_global_ids = node_data[constants.SHUFFLE_GLOBAL_NID+"/"+str(local_part_id)][idx1] node_data[constants.GLOBAL_NID + "/" + str(local_part_id)],
global_nids,
return_indices=True,
)
shuffle_global_ids = node_data[
constants.SHUFFLE_GLOBAL_NID + "/" + str(local_part_id)
][idx1]
feature_idx = shuffle_global_ids.argsort() feature_idx = shuffle_global_ids.argsort()
rcvd_node_features[feature_key] = rcvd_node_features[feature_key][feature_idx] rcvd_node_features[feature_key] = rcvd_node_features[
feature_key
][feature_idx]
memory_snapshot("ReorderNodeFeaturesComplete: ", rank) memory_snapshot("ReorderNodeFeaturesComplete: ", rank)
#sort edge_data by etype # sort edge_data by etype
for local_part_id in range(params.num_parts//world_size): for local_part_id in range(params.num_parts // world_size):
sorted_idx = edge_data[constants.ETYPE_ID+"/"+str(local_part_id)].argsort() sorted_idx = edge_data[
constants.ETYPE_ID + "/" + str(local_part_id)
].argsort()
for k, v in edge_data.items(): for k, v in edge_data.items():
tokens = k.split("/") tokens = k.split("/")
assert len(tokens) == 2 assert len(tokens) == 2
...@@ -800,91 +1023,156 @@ def gen_dist_partitions(rank, world_size, params): ...@@ -800,91 +1023,156 @@ def gen_dist_partitions(rank, world_size, params):
sorted_idx = None sorted_idx = None
gc.collect() gc.collect()
shuffle_global_eid_offsets = assign_shuffle_global_nids_edges(rank, world_size, params.num_parts, edge_data) shuffle_global_eid_offsets = assign_shuffle_global_nids_edges(
logging.info(f'[Rank: {rank}] Done assigning global_ids to edges ...') rank, world_size, params.num_parts, edge_data
)
logging.info(f"[Rank: {rank}] Done assigning global_ids to edges ...")
memory_snapshot("ShuffleGlobalID_Edges_Complete: ", rank) memory_snapshot("ShuffleGlobalID_Edges_Complete: ", rank)
#Shuffle edge features according to the edge order on each rank. # Shuffle edge features according to the edge order on each rank.
for etype_name in etypes: for etype_name in etypes:
featnames = get_etype_featnames(etype_name, schema_map) featnames = get_etype_featnames(etype_name, schema_map)
for featname in featnames: for featname in featnames:
for local_part_id in range(params.num_parts//world_size): for local_part_id in range(params.num_parts // world_size):
feature_key = etype_name+'/'+featname+"/"+str(local_part_id) feature_key = (
etype_name + "/" + featname + "/" + str(local_part_id)
)
assert feature_key in rcvd_global_eids assert feature_key in rcvd_global_eids
global_eids = rcvd_global_eids[feature_key] global_eids = rcvd_global_eids[feature_key]
_, idx1, _ = np.intersect1d(edge_data[constants.GLOBAL_EID+"/"+str(local_part_id)], global_eids, return_indices=True) _, idx1, _ = np.intersect1d(
shuffle_global_ids = edge_data[constants.SHUFFLE_GLOBAL_EID+"/"+str(local_part_id)][idx1] edge_data[constants.GLOBAL_EID + "/" + str(local_part_id)],
global_eids,
return_indices=True,
)
shuffle_global_ids = edge_data[
constants.SHUFFLE_GLOBAL_EID + "/" + str(local_part_id)
][idx1]
feature_idx = shuffle_global_ids.argsort() feature_idx = shuffle_global_ids.argsort()
rcvd_edge_features[feature_key] = rcvd_edge_features[feature_key][feature_idx] rcvd_edge_features[feature_key] = rcvd_edge_features[
feature_key
][feature_idx]
#determine global-ids for edge end-points # determine global-ids for edge end-points
edge_data = lookup_shuffle_global_nids_edges(rank, world_size, params.num_parts, edge_data, id_lookup, node_data) edge_data = lookup_shuffle_global_nids_edges(
logging.info(f'[Rank: {rank}] Done resolving orig_node_id for local node_ids...') rank, world_size, params.num_parts, edge_data, id_lookup, node_data
)
logging.info(
f"[Rank: {rank}] Done resolving orig_node_id for local node_ids..."
)
memory_snapshot("ShuffleGlobalID_Lookup_Complete: ", rank) memory_snapshot("ShuffleGlobalID_Lookup_Complete: ", rank)
def prepare_local_data(src_data, local_part_id): def prepare_local_data(src_data, local_part_id):
local_data = {} local_data = {}
for k, v in src_data.items(): for k, v in src_data.items():
tokens = k.split("/") tokens = k.split("/")
if tokens[len(tokens)-1] == str(local_part_id): if tokens[len(tokens) - 1] == str(local_part_id):
local_data["/".join(tokens[:-1])] = v local_data["/".join(tokens[:-1])] = v
return local_data return local_data
#create dgl objects here # create dgl objects here
output_meta_json = {} output_meta_json = {}
start = timer() start = timer()
graph_formats = None graph_formats = None
if params.graph_formats: if params.graph_formats:
graph_formats = params.graph_formats.split(',') graph_formats = params.graph_formats.split(",")
for local_part_id in range(params.num_parts//world_size): for local_part_id in range(params.num_parts // world_size):
num_edges = shuffle_global_eid_offsets[local_part_id] num_edges = shuffle_global_eid_offsets[local_part_id]
node_count = len(node_data[constants.NTYPE_ID+"/"+str(local_part_id)]) node_count = len(
edge_count = len(edge_data[constants.ETYPE_ID+"/"+str(local_part_id)]) node_data[constants.NTYPE_ID + "/" + str(local_part_id)]
)
edge_count = len(
edge_data[constants.ETYPE_ID + "/" + str(local_part_id)]
)
local_node_data = prepare_local_data(node_data, local_part_id) local_node_data = prepare_local_data(node_data, local_part_id)
local_edge_data = prepare_local_data(edge_data, local_part_id) local_edge_data = prepare_local_data(edge_data, local_part_id)
graph_obj, ntypes_map_val, etypes_map_val, ntypes_map, etypes_map, \ (
orig_nids, orig_eids = create_dgl_object(schema_map, rank+local_part_id*world_size, graph_obj,
local_node_data, local_edge_data, ntypes_map_val,
num_edges, params.save_orig_nids, params.save_orig_eids) etypes_map_val,
ntypes_map,
etypes_map,
orig_nids,
orig_eids,
) = create_dgl_object(
schema_map,
rank + local_part_id * world_size,
local_node_data,
local_edge_data,
num_edges,
params.save_orig_nids,
params.save_orig_eids,
)
sort_etypes = len(etypes_map) > 1 sort_etypes = len(etypes_map) > 1
local_node_features = prepare_local_data(rcvd_node_features, local_part_id) local_node_features = prepare_local_data(
local_edge_features = prepare_local_data(rcvd_edge_features, local_part_id) rcvd_node_features, local_part_id
write_dgl_objects(graph_obj, )
local_node_features, local_edge_features, local_edge_features = prepare_local_data(
params.output, rcvd_edge_features, local_part_id
rank + (local_part_id*world_size), )
orig_nids, orig_eids, graph_formats, sort_etypes) write_dgl_objects(
graph_obj,
local_node_features,
local_edge_features,
params.output,
rank + (local_part_id * world_size),
orig_nids,
orig_eids,
graph_formats,
sort_etypes,
)
memory_snapshot("DiskWriteDGLObjectsComplete: ", rank) memory_snapshot("DiskWriteDGLObjectsComplete: ", rank)
#get the meta-data # get the meta-data
json_metadata = create_metadata_json(params.graph_name, node_count, edge_count, \ json_metadata = create_metadata_json(
local_part_id * world_size + rank, params.num_parts, ntypes_map_val, \ params.graph_name,
etypes_map_val, ntypes_map, etypes_map, params.output) node_count,
output_meta_json["local-part-id-"+str(local_part_id*world_size + rank)] = json_metadata edge_count,
local_part_id * world_size + rank,
params.num_parts,
ntypes_map_val,
etypes_map_val,
ntypes_map,
etypes_map,
params.output,
)
output_meta_json[
"local-part-id-" + str(local_part_id * world_size + rank)
] = json_metadata
memory_snapshot("MetadataCreateComplete: ", rank) memory_snapshot("MetadataCreateComplete: ", rank)
if (rank == 0): if rank == 0:
#get meta-data from all partitions and merge them on rank-0 # get meta-data from all partitions and merge them on rank-0
metadata_list = gather_metadata_json(output_meta_json, rank, world_size) metadata_list = gather_metadata_json(output_meta_json, rank, world_size)
metadata_list[0] = output_meta_json metadata_list[0] = output_meta_json
write_metadata_json(metadata_list, params.output, params.graph_name, world_size, params.num_parts) write_metadata_json(
metadata_list,
params.output,
params.graph_name,
world_size,
params.num_parts,
)
else: else:
#send meta-data to Rank-0 process # send meta-data to Rank-0 process
gather_metadata_json(output_meta_json, rank, world_size) gather_metadata_json(output_meta_json, rank, world_size)
end = timer() end = timer()
logging.info(f'[Rank: {rank}] Time to create dgl objects: {timedelta(seconds = end - start)}') logging.info(
f"[Rank: {rank}] Time to create dgl objects: {timedelta(seconds = end - start)}"
)
memory_snapshot("MetadataWriteComplete: ", rank) memory_snapshot("MetadataWriteComplete: ", rank)
global_end = timer() global_end = timer()
logging.info(f'[Rank: {rank}] Total execution time of the program: {timedelta(seconds = global_end - global_start)}') logging.info(
f"[Rank: {rank}] Total execution time of the program: {timedelta(seconds = global_end - global_start)}"
)
memory_snapshot("PipelineComplete: ", rank) memory_snapshot("PipelineComplete: ", rank)
def single_machine_run(params): def single_machine_run(params):
""" Main function for distributed implementation on a single machine """Main function for distributed implementation on a single machine
Parameters: Parameters:
----------- -----------
...@@ -896,16 +1184,20 @@ def single_machine_run(params): ...@@ -896,16 +1184,20 @@ def single_machine_run(params):
processes = [] processes = []
mp.set_start_method("spawn") mp.set_start_method("spawn")
#Invoke `target` function from each of the spawned process for distributed # Invoke `target` function from each of the spawned process for distributed
#implementation # implementation
for rank in range(params.world_size): for rank in range(params.world_size):
p = mp.Process(target=run, args=(rank, params.world_size, gen_dist_partitions, params)) p = mp.Process(
target=run,
args=(rank, params.world_size, gen_dist_partitions, params),
)
p.start() p.start()
processes.append(p) processes.append(p)
for p in processes: for p in processes:
p.join() p.join()
def run(rank, world_size, func_exec, params, backend="gloo"): def run(rank, world_size, func_exec, params, backend="gloo"):
""" """
Init. function which is run by each process in the Gloo ProcessGroup Init. function which is run by each process in the Gloo ProcessGroup
...@@ -923,15 +1215,21 @@ def run(rank, world_size, func_exec, params, backend="gloo"): ...@@ -923,15 +1215,21 @@ def run(rank, world_size, func_exec, params, backend="gloo"):
backend : string backend : string
string specifying the type of backend to use for communication string specifying the type of backend to use for communication
""" """
os.environ["MASTER_ADDR"] = '127.0.0.1' os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = '29500' os.environ["MASTER_PORT"] = "29500"
#create Gloo Process Group # create Gloo Process Group
dist.init_process_group(backend, rank=rank, world_size=world_size, timeout=timedelta(seconds=5*60)) dist.init_process_group(
backend,
rank=rank,
world_size=world_size,
timeout=timedelta(seconds=5 * 60),
)
#Invoke the main function to kick-off each process # Invoke the main function to kick-off each process
func_exec(rank, world_size, params) func_exec(rank, world_size, params)
def multi_machine_run(params): def multi_machine_run(params):
""" """
Function to be invoked when executing data loading pipeline on multiple machines Function to be invoked when executing data loading pipeline on multiple machines
...@@ -943,14 +1241,17 @@ def multi_machine_run(params): ...@@ -943,14 +1241,17 @@ def multi_machine_run(params):
""" """
rank = int(os.environ["RANK"]) rank = int(os.environ["RANK"])
#init the gloo process group here. # init the gloo process group here.
dist.init_process_group( dist.init_process_group(
backend="gloo", backend="gloo",
rank=rank, rank=rank,
world_size=params.world_size, world_size=params.world_size,
timeout=timedelta(seconds=params.process_group_timeout)) timeout=timedelta(seconds=params.process_group_timeout),
logging.info(f'[Rank: {rank}] Done with process group initialization...') )
logging.info(f"[Rank: {rank}] Done with process group initialization...")
#invoke the main function here. # invoke the main function here.
gen_dist_partitions(rank, params.world_size, params) gen_dist_partitions(rank, params.world_size, params)
logging.info(f'[Rank: {rank}] Done with Distributed data processing pipeline processing.') logging.info(
f"[Rank: {rank}] Done with Distributed data processing pipeline processing."
)
tools/distpartitioning/dataset_utils.py
View file @ aa42aaeb
import gc
import logging import logging
import os import os
import gc
import array_readwriter
import constants
import numpy as np import numpy as np
import pyarrow import pyarrow
import pyarrow.parquet as pq import pyarrow.parquet as pq
import torch import torch
import torch.distributed as dist import torch.distributed as dist
import array_readwriter
import constants
from utils import get_idranges, map_partid_rank, generate_read_list
from gloo_wrapper import alltoallv_cpu from gloo_wrapper import alltoallv_cpu
from utils import generate_read_list, get_idranges, map_partid_rank
DATA_TYPE_ID = { DATA_TYPE_ID = {
data_type: id for id, data_type in enumerate([ data_type: id
torch.float32, for id, data_type in enumerate(
torch.float64, [
torch.float16, torch.float32,
torch.uint8, torch.float64,
torch.int8, torch.float16,
torch.int16, torch.uint8,
torch.int32, torch.int8,
torch.int64, torch.int16,
torch.bool, torch.int32,
]) torch.int64,
torch.bool,
]
)
} }
REV_DATA_TYPE_ID = { REV_DATA_TYPE_ID = {id: data_type for data_type, id in DATA_TYPE_ID.items()}
id: data_type for data_type, id in DATA_TYPE_ID.items()
}
def _shuffle_data(data, rank, world_size, tids, num_parts): def _shuffle_data(data, rank, world_size, tids, num_parts):
'''Each process scatters loaded data to all processes in a group and """Each process scatters loaded data to all processes in a group and
return gathered data. return gathered data.
Parameters Parameters
...@@ -57,14 +58,15 @@ def _shuffle_data(data, rank, world_size, tids, num_parts): ...@@ -57,14 +58,15 @@ def _shuffle_data(data, rank, world_size, tids, num_parts):
shuffled_data: tensor shuffled_data: tensor
Shuffled node or edge data. Shuffled node or edge data.
''' """
# Broadcast basic information of loaded data: # Broadcast basic information of loaded data:
# 1. number of data lines # 1. number of data lines
# 2. data dimension # 2. data dimension
# 3. data type # 3. data type
assert len(data.shape) in [1, 2], ( assert len(data.shape) in [
f"Data is expected to be 1-D or 2-D but got {data.shape}." 1,
) 2,
], f"Data is expected to be 1-D or 2-D but got {data.shape}."
data_shape = list(data.shape) data_shape = list(data.shape)
if len(data_shape) == 1: if len(data_shape) == 1:
data_shape.append(1) data_shape.append(1)
...@@ -127,7 +129,7 @@ def _shuffle_data(data, rank, world_size, tids, num_parts): ...@@ -127,7 +129,7 @@ def _shuffle_data(data, rank, world_size, tids, num_parts):
def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map): def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map):
""" """
Function to read the multiple file formatted dataset. Function to read the multiple file formatted dataset.
Parameters: Parameters:
----------- -----------
...@@ -153,18 +155,18 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map): ...@@ -153,18 +155,18 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map):
data for each node type is split into "p" files and each one of these "p" files are data for each node type is split into "p" files and each one of these "p" files are
read a process in the distributed graph partitioning pipeline read a process in the distributed graph partitioning pipeline
dictionary dictionary
Data read from numpy files for all the node features in this dataset. Dictionary built Data read from numpy files for all the node features in this dataset. Dictionary built
using this data has keys as node feature names and values as tensor data representing using this data has keys as node feature names and values as tensor data representing
node features node features
dictionary dictionary
in which keys are node-type and values are a triplet. This triplet has node-feature name, in which keys are node-type and values are a triplet. This triplet has node-feature name,
and range of tids for the node feature data read from files by the current process. Each and range of tids for the node feature data read from files by the current process. Each
node-type may have mutiple feature(s) and associated tensor data. node-type may have mutiple feature(s) and associated tensor data.
dictionary dictionary
Data read from edges.txt file and used to build a dictionary with keys as column names Data read from edges.txt file and used to build a dictionary with keys as column names
and values as columns in the csv file. and values as columns in the csv file.
dictionary dictionary
in which keys are edge-type names and values are triplets. This triplet has edge-feature name, in which keys are edge-type names and values are triplets. This triplet has edge-feature name,
and range of tids for theedge feature data read from the files by the current process. Each and range of tids for theedge feature data read from the files by the current process. Each
edge-type may have several edge features and associated tensor data. edge-type may have several edge features and associated tensor data.
dictionary dictionary
...@@ -173,21 +175,21 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map): ...@@ -173,21 +175,21 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map):
dictionary dictionary
This dictionary is used for identifying the global-id range for the associated edge features This dictionary is used for identifying the global-id range for the associated edge features
present in the previous return value. The keys are edge-type names and values are triplets. present in the previous return value. The keys are edge-type names and values are triplets.
Each triplet consists of edge-feature name and starting and ending points of the range of Each triplet consists of edge-feature name and starting and ending points of the range of
tids representing the corresponding edge feautres. tids representing the corresponding edge feautres.
""" """
#node features dictionary # node features dictionary
#TODO: With the new file format, It is guaranteed that the input dataset will have # TODO: With the new file format, It is guaranteed that the input dataset will have
#no. of nodes with features (node-features) files and nodes metadata will always be the same. # no. of nodes with features (node-features) files and nodes metadata will always be the same.
#This means the dimension indicating the no. of nodes in any node-feature files and the no. of # This means the dimension indicating the no. of nodes in any node-feature files and the no. of
#nodes in the corresponding nodes metadata file will always be the same. With this guarantee, # nodes in the corresponding nodes metadata file will always be the same. With this guarantee,
#we can eliminate the `node_feature_tids` dictionary since the same information is also populated # we can eliminate the `node_feature_tids` dictionary since the same information is also populated
#in the `node_tids` dictionary. This will be remnoved in the next iteration of code changes. # in the `node_tids` dictionary. This will be remnoved in the next iteration of code changes.
node_features = {} node_features = {}
node_feature_tids = {} node_feature_tids = {}
''' """
The structure of the node_data is as follows, which is present in the input metadata json file. The structure of the node_data is as follows, which is present in the input metadata json file.
"node_data" : { "node_data" : {
"ntype0-name" : { "ntype0-name" : {
...@@ -244,23 +246,27 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map): ...@@ -244,23 +246,27 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map):
which are owned by that particular rank. And using the "num_nodes_per_chunk" information each which are owned by that particular rank. And using the "num_nodes_per_chunk" information each
process can easily compute any nodes per-type node_id and global node_id. process can easily compute any nodes per-type node_id and global node_id.
The node-ids are treated as int64's in order to support billions of nodes in the input graph. The node-ids are treated as int64's in order to support billions of nodes in the input graph.
''' """
#read my nodes for each node type # read my nodes for each node type
node_tids, ntype_gnid_offset = get_idranges(schema_map[constants.STR_NODE_TYPE], node_tids, ntype_gnid_offset = get_idranges(
schema_map[constants.STR_NUM_NODES_PER_CHUNK], schema_map[constants.STR_NODE_TYPE],
num_chunks=num_parts) schema_map[constants.STR_NUM_NODES_PER_CHUNK],
num_chunks=num_parts,
)
#iterate over the "node_data" dictionary in the schema_map # iterate over the "node_data" dictionary in the schema_map
#read the node features if exists # read the node features if exists
#also keep track of the type_nids for which the node_features are read. # also keep track of the type_nids for which the node_features are read.
dataset_features = schema_map[constants.STR_NODE_DATA] dataset_features = schema_map[constants.STR_NODE_DATA]
if((dataset_features is not None) and (len(dataset_features) > 0)): if (dataset_features is not None) and (len(dataset_features) > 0):
for ntype_name, ntype_feature_data in dataset_features.items(): for ntype_name, ntype_feature_data in dataset_features.items():
for feat_name, feat_data in ntype_feature_data.items(): for feat_name, feat_data in ntype_feature_data.items():
assert (feat_data[constants.STR_FORMAT][constants.STR_NAME] assert feat_data[constants.STR_FORMAT][constants.STR_NAME] in [
in [constants.STR_NUMPY, constants.STR_PARQUET]) constants.STR_NUMPY,
# It is guaranteed that num_chunks is always greater constants.STR_PARQUET,
]
# It is guaranteed that num_chunks is always greater
# than num_partitions. # than num_partitions.
node_data = [] node_data = []
num_files = len(feat_data[constants.STR_DATA]) num_files = len(feat_data[constants.STR_DATA])
...@@ -291,7 +297,8 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map): ...@@ -291,7 +297,8 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map):
rank, rank,
world_size, world_size,
node_tids[ntype_name], node_tids[ntype_name],
num_parts) num_parts,
)
# collect data on current rank. # collect data on current rank.
offset = 0 offset = 0
...@@ -299,7 +306,7 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map): ...@@ -299,7 +306,7 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map):
if map_partid_rank(local_part_id, world_size) == rank: if map_partid_rank(local_part_id, world_size) == rank:
nfeat = [] nfeat = []
nfeat_tids = [] nfeat_tids = []
start, end = node_tids[ntype_name][local_part_id] start, end = node_tids[ntype_name][local_part_id]
nfeat = node_data[offset : offset + end - start] nfeat = node_data[offset : offset + end - start]
data_key = f"{ntype_name}/{feat_name}/{local_part_id//world_size}" data_key = f"{ntype_name}/{feat_name}/{local_part_id//world_size}"
node_features[data_key] = nfeat node_features[data_key] = nfeat
...@@ -307,19 +314,23 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map): ...@@ -307,19 +314,23 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map):
node_feature_tids[data_key] = nfeat_tids node_feature_tids[data_key] = nfeat_tids
offset += end - start offset += end - start
#done building node_features locally. # done building node_features locally.
if len(node_features) <= 0: if len(node_features) <= 0:
logging.info(f'[Rank: {rank}] This dataset does not have any node features') logging.info(
f"[Rank: {rank}] This dataset does not have any node features"
)
else: else:
assert len(node_features) == len(node_feature_tids) assert len(node_features) == len(node_feature_tids)
# Note that the keys in the node_features dictionary are as follows: # Note that the keys in the node_features dictionary are as follows:
# `ntype_name/feat_name/local_part_id`. # `ntype_name/feat_name/local_part_id`.
# where ntype_name and feat_name are self-explanatory, and # where ntype_name and feat_name are self-explanatory, and
# local_part_id indicates the partition-id, in the context of current # local_part_id indicates the partition-id, in the context of current
# process which take the values 0, 1, 2, .... # process which take the values 0, 1, 2, ....
for feat_name, feat_info in node_features.items(): for feat_name, feat_info in node_features.items():
logging.info(f'[Rank: {rank}] node feature name: {feat_name}, feature data shape: {feat_info.size()}') logging.info(
f"[Rank: {rank}] node feature name: {feat_name}, feature data shape: {feat_info.size()}"
)
tokens = feat_name.split("/") tokens = feat_name.split("/")
assert len(tokens) == 3 assert len(tokens) == 3
...@@ -330,10 +341,11 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map): ...@@ -330,10 +341,11 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map):
# Iterate over the range of type ids for the current node feature # Iterate over the range of type ids for the current node feature
# and count the number of features for this feature name. # and count the number of features for this feature name.
count = tids[0][1] - tids[0][0] count = tids[0][1] - tids[0][0]
assert count == feat_info.size()[0], f"{feat_name}, {count} vs {feat_info.size()[0]}." assert (
count == feat_info.size()[0]
), f"{feat_name}, {count} vs {feat_info.size()[0]}."
"""
'''
Reading edge features now. Reading edge features now.
The structure of the edge_data is as follows, which is present in the input metadata json file. The structure of the edge_data is as follows, which is present in the input metadata json file.
"edge_data" : { "edge_data" : {
...@@ -369,13 +381,16 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map): ...@@ -369,13 +381,16 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map):
Data read from each of the node features file is a multi-dimensional tensor data and is read Data read from each of the node features file is a multi-dimensional tensor data and is read
in numpy format, which is also the storage format of node features on the permanent storage. in numpy format, which is also the storage format of node features on the permanent storage.
''' """
edge_features = {} edge_features = {}
edge_feature_tids = {} edge_feature_tids = {}
# Read edges for each edge type that are processed by the currnet process. # Read edges for each edge type that are processed by the currnet process.
edge_tids, _ = get_idranges(schema_map[constants.STR_EDGE_TYPE], edge_tids, _ = get_idranges(
schema_map[constants.STR_NUM_EDGES_PER_CHUNK], num_parts) schema_map[constants.STR_EDGE_TYPE],
schema_map[constants.STR_NUM_EDGES_PER_CHUNK],
num_parts,
)
# Iterate over the "edge_data" dictionary in the schema_map. # Iterate over the "edge_data" dictionary in the schema_map.
# Read the edge features if exists. # Read the edge features if exists.
...@@ -384,8 +399,10 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map): ...@@ -384,8 +399,10 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map):
if dataset_features and (len(dataset_features) > 0): if dataset_features and (len(dataset_features) > 0):
for etype_name, etype_feature_data in dataset_features.items(): for etype_name, etype_feature_data in dataset_features.items():
for feat_name, feat_data in etype_feature_data.items(): for feat_name, feat_data in etype_feature_data.items():
assert (feat_data[constants.STR_FORMAT][constants.STR_NAME] assert feat_data[constants.STR_FORMAT][constants.STR_NAME] in [
in [constants.STR_NUMPY, constants.STR_PARQUET]) constants.STR_NUMPY,
constants.STR_PARQUET,
]
edge_data = [] edge_data = []
num_files = len(feat_data[constants.STR_DATA]) num_files = len(feat_data[constants.STR_DATA])
...@@ -416,7 +433,8 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map): ...@@ -416,7 +433,8 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map):
rank, rank,
world_size, world_size,
edge_tids[etype_name], edge_tids[etype_name],
num_parts) num_parts,
)
# collect data on current rank. # collect data on current rank.
offset = 0 offset = 0
...@@ -432,19 +450,23 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map): ...@@ -432,19 +450,23 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map):
edge_feature_tids[data_key] = efeat_tids edge_feature_tids[data_key] = efeat_tids
offset += end - start offset += end - start
# Done with building node_features locally. # Done with building node_features locally.
if len(edge_features) <= 0: if len(edge_features) <= 0:
logging.info(f'[Rank: {rank}] This dataset does not have any edge features') logging.info(
f"[Rank: {rank}] This dataset does not have any edge features"
)
else: else:
assert len(edge_features) == len(edge_feature_tids) assert len(edge_features) == len(edge_feature_tids)
for k, v in edge_features.items(): for k, v in edge_features.items():
logging.info(f'[Rank: {rank}] edge feature name: {k}, feature data shape: {v.shape}') logging.info(
f"[Rank: {rank}] edge feature name: {k}, feature data shape: {v.shape}"
)
tids = edge_feature_tids[k] tids = edge_feature_tids[k]
count = tids[0][1] - tids[0][0] count = tids[0][1] - tids[0][0]
assert count == v.size()[0] assert count == v.size()[0]
''' """
Code below is used to read edges from the input dataset with the help of the metadata json file Code below is used to read edges from the input dataset with the help of the metadata json file
for the input graph dataset. for the input graph dataset.
In the metadata json file, we expect the following key-value pairs to help read the edges of the In the metadata json file, we expect the following key-value pairs to help read the edges of the
...@@ -484,27 +506,33 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map): ...@@ -484,27 +506,33 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map):
Each edge file contains two columns representing the source per-type node_ids and destination per-type node_ids Each edge file contains two columns representing the source per-type node_ids and destination per-type node_ids
of any given edge. Since these are node-ids as well they are read in as int64's. of any given edge. Since these are node-ids as well they are read in as int64's.
''' """
#read my edges for each edge type # read my edges for each edge type
etype_names = schema_map[constants.STR_EDGE_TYPE] etype_names = schema_map[constants.STR_EDGE_TYPE]
etype_name_idmap = {e : idx for idx, e in enumerate(etype_names)} etype_name_idmap = {e: idx for idx, e in enumerate(etype_names)}
edge_tids, _ = get_idranges(schema_map[constants.STR_EDGE_TYPE], edge_tids, _ = get_idranges(
schema_map[constants.STR_NUM_EDGES_PER_CHUNK], schema_map[constants.STR_EDGE_TYPE],
num_chunks=num_parts) schema_map[constants.STR_NUM_EDGES_PER_CHUNK],
num_chunks=num_parts,
)
edge_datadict = {} edge_datadict = {}
edge_data = schema_map[constants.STR_EDGES] edge_data = schema_map[constants.STR_EDGES]
#read the edges files and store this data in memory. # read the edges files and store this data in memory.
for col in [constants.GLOBAL_SRC_ID, constants.GLOBAL_DST_ID, \ for col in [
constants.GLOBAL_TYPE_EID, constants.ETYPE_ID]: constants.GLOBAL_SRC_ID,
constants.GLOBAL_DST_ID,
constants.GLOBAL_TYPE_EID,
constants.ETYPE_ID,
]:
edge_datadict[col] = [] edge_datadict[col] = []
for etype_name, etype_info in edge_data.items(): for etype_name, etype_info in edge_data.items():
edge_info = etype_info[constants.STR_DATA] edge_info = etype_info[constants.STR_DATA]
#edgetype strings are in canonical format, src_node_type:edge_type:dst_node_type # edgetype strings are in canonical format, src_node_type:edge_type:dst_node_type
tokens = etype_name.split(":") tokens = etype_name.split(":")
assert len(tokens) == 3 assert len(tokens) == 3
...@@ -525,50 +553,102 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map): ...@@ -525,50 +553,102 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map):
edge_file = edge_info[idx] edge_file = edge_info[idx]
if not os.path.isabs(edge_file): if not os.path.isabs(edge_file):
edge_file = os.path.join(input_dir, edge_file) edge_file = os.path.join(input_dir, edge_file)
logging.info(f'Loading edges of etype[{etype_name}] from {edge_file}') logging.info(
f"Loading edges of etype[{etype_name}] from {edge_file}"
if etype_info[constants.STR_FORMAT][constants.STR_NAME] == constants.STR_CSV: )
read_options=pyarrow.csv.ReadOptions(use_threads=True, block_size=4096, autogenerate_column_names=True)
parse_options=pyarrow.csv.ParseOptions(delimiter=' ') if (
with pyarrow.csv.open_csv(edge_file, read_options=read_options, parse_options=parse_options) as reader: etype_info[constants.STR_FORMAT][constants.STR_NAME]
== constants.STR_CSV
):
read_options = pyarrow.csv.ReadOptions(
use_threads=True,
block_size=4096,
autogenerate_column_names=True,
)
parse_options = pyarrow.csv.ParseOptions(delimiter=" ")
with pyarrow.csv.open_csv(
edge_file,
read_options=read_options,
parse_options=parse_options,
) as reader:
for next_chunk in reader: for next_chunk in reader:
if next_chunk is None: if next_chunk is None:
break break
next_table = pyarrow.Table.from_batches([next_chunk]) next_table = pyarrow.Table.from_batches([next_chunk])
src_ids.append(next_table['f0'].to_numpy()) src_ids.append(next_table["f0"].to_numpy())
dst_ids.append(next_table['f1'].to_numpy()) dst_ids.append(next_table["f1"].to_numpy())
elif etype_info[constants.STR_FORMAT][constants.STR_NAME] == constants.STR_PARQUET: elif (
etype_info[constants.STR_FORMAT][constants.STR_NAME]
== constants.STR_PARQUET
):
data_df = pq.read_table(edge_file) data_df = pq.read_table(edge_file)
data_df = data_df.rename_columns(["f0", "f1"]) data_df = data_df.rename_columns(["f0", "f1"])
src_ids.append(data_df['f0'].to_numpy()) src_ids.append(data_df["f0"].to_numpy())
dst_ids.append(data_df['f1'].to_numpy()) dst_ids.append(data_df["f1"].to_numpy())
else: else:
raise ValueError(f'Unknown edge format {etype_info[constants.STR_FORMAT][constants.STR_NAME]} for edge type {etype_name}') raise ValueError(
f"Unknown edge format {etype_info[constants.STR_FORMAT][constants.STR_NAME]} for edge type {etype_name}"
)
src_ids = np.concatenate(src_ids) src_ids = np.concatenate(src_ids)
dst_ids = np.concatenate(dst_ids) dst_ids = np.concatenate(dst_ids)
#currently these are just type_edge_ids... which will be converted to global ids # currently these are just type_edge_ids... which will be converted to global ids
edge_datadict[constants.GLOBAL_SRC_ID].append(src_ids + ntype_gnid_offset[src_ntype_name][0, 0]) edge_datadict[constants.GLOBAL_SRC_ID].append(
edge_datadict[constants.GLOBAL_DST_ID].append(dst_ids + ntype_gnid_offset[dst_ntype_name][0, 0]) src_ids + ntype_gnid_offset[src_ntype_name][0, 0]
edge_datadict[constants.ETYPE_ID].append(etype_name_idmap[etype_name] * \ )
np.ones(shape=(src_ids.shape), dtype=np.int64)) edge_datadict[constants.GLOBAL_DST_ID].append(
dst_ids + ntype_gnid_offset[dst_ntype_name][0, 0]
)
edge_datadict[constants.ETYPE_ID].append(
etype_name_idmap[etype_name]
* np.ones(shape=(src_ids.shape), dtype=np.int64)
)
for local_part_id in range(num_parts): for local_part_id in range(num_parts):
if (map_partid_rank(local_part_id, world_size) == rank): if map_partid_rank(local_part_id, world_size) == rank:
edge_datadict[constants.GLOBAL_TYPE_EID].append(np.arange(edge_tids[etype_name][local_part_id][0],\ edge_datadict[constants.GLOBAL_TYPE_EID].append(
edge_tids[etype_name][local_part_id][1] ,dtype=np.int64)) np.arange(
edge_tids[etype_name][local_part_id][0],
#stitch together to create the final data on the local machine edge_tids[etype_name][local_part_id][1],
for col in [constants.GLOBAL_SRC_ID, constants.GLOBAL_DST_ID, constants.GLOBAL_TYPE_EID, constants.ETYPE_ID]: dtype=np.int64,
)
)
# stitch together to create the final data on the local machine
for col in [
constants.GLOBAL_SRC_ID,
constants.GLOBAL_DST_ID,
constants.GLOBAL_TYPE_EID,
constants.ETYPE_ID,
]:
edge_datadict[col] = np.concatenate(edge_datadict[col]) edge_datadict[col] = np.concatenate(edge_datadict[col])
assert edge_datadict[constants.GLOBAL_SRC_ID].shape == edge_datadict[constants.GLOBAL_DST_ID].shape assert (
assert edge_datadict[constants.GLOBAL_DST_ID].shape == edge_datadict[constants.GLOBAL_TYPE_EID].shape edge_datadict[constants.GLOBAL_SRC_ID].shape
assert edge_datadict[constants.GLOBAL_TYPE_EID].shape == edge_datadict[constants.ETYPE_ID].shape == edge_datadict[constants.GLOBAL_DST_ID].shape
logging.info(f'[Rank: {rank}] Done reading edge_file: {len(edge_datadict)}, {edge_datadict[constants.GLOBAL_SRC_ID].shape}') )
logging.info(f'Rank: {rank} edge_feat_tids: {edge_feature_tids}') assert (
edge_datadict[constants.GLOBAL_DST_ID].shape
return node_tids, node_features, node_feature_tids, edge_datadict, edge_tids, edge_features, edge_feature_tids == edge_datadict[constants.GLOBAL_TYPE_EID].shape
)
assert (
edge_datadict[constants.GLOBAL_TYPE_EID].shape
== edge_datadict[constants.ETYPE_ID].shape
)
logging.info(
f"[Rank: {rank}] Done reading edge_file: {len(edge_datadict)}, {edge_datadict[constants.GLOBAL_SRC_ID].shape}"
)
logging.info(f"Rank: {rank} edge_feat_tids: {edge_feature_tids}")
return (
node_tids,
node_features,
node_feature_tids,
edge_datadict,
edge_tids,
edge_features,
edge_feature_tids,
)
tools/distpartitioning/dist_lookup.py
View file @ aa42aaeb
import copy
import logging import logging
import os import os
import numpy as np import numpy as np
import pyarrow import pyarrow
import torch import torch
import copy from gloo_wrapper import alltoallv_cpu
from pyarrow import csv from pyarrow import csv
from gloo_wrapper import alltoallv_cpu
from utils import map_partid_rank from utils import map_partid_rank
class DistLookupService: class DistLookupService:
''' """
This is an implementation of a Distributed Lookup Service to provide the following This is an implementation of a Distributed Lookup Service to provide the following
services to its users. Map 1) global node-ids to partition-ids, and 2) global node-ids services to its users. Map 1) global node-ids to partition-ids, and 2) global node-ids
to shuffle global node-ids (contiguous, within each node for a give node_type and across to shuffle global node-ids (contiguous, within each node for a give node_type and across
all the partitions) all the partitions)
This services initializes itself with the node-id to partition-id mappings, which are inputs This services initializes itself with the node-id to partition-id mappings, which are inputs
...@@ -44,7 +45,7 @@ class DistLookupService: ...@@ -44,7 +45,7 @@ class DistLookupService:
integer indicating the rank of a given process integer indicating the rank of a given process
world_size : integer world_size : integer
integer indicating the total no. of processes integer indicating the total no. of processes
''' """
def __init__(self, input_dir, ntype_names, id_map, rank, world_size): def __init__(self, input_dir, ntype_names, id_map, rank, world_size):
assert os.path.isdir(input_dir) assert os.path.isdir(input_dir)
...@@ -60,19 +61,28 @@ class DistLookupService: ...@@ -60,19 +61,28 @@ class DistLookupService:
# Iterate over the node types and extract the partition id mappings. # Iterate over the node types and extract the partition id mappings.
for ntype in ntype_names: for ntype in ntype_names:
filename = f'{ntype}.txt' filename = f"{ntype}.txt"
logging.info(f'[Rank: {rank}] Reading file: {os.path.join(input_dir, filename)}') logging.info(
f"[Rank: {rank}] Reading file: {os.path.join(input_dir, filename)}"
read_options=pyarrow.csv.ReadOptions(use_threads=True, block_size=4096, autogenerate_column_names=True) )
parse_options=pyarrow.csv.ParseOptions(delimiter=' ')
read_options = pyarrow.csv.ReadOptions(
use_threads=True,
block_size=4096,
autogenerate_column_names=True,
)
parse_options = pyarrow.csv.ParseOptions(delimiter=" ")
ntype_partids = [] ntype_partids = []
with pyarrow.csv.open_csv(os.path.join(input_dir, '{}.txt'.format(ntype)), with pyarrow.csv.open_csv(
read_options=read_options, parse_options=parse_options) as reader: os.path.join(input_dir, "{}.txt".format(ntype)),
read_options=read_options,
parse_options=parse_options,
) as reader:
for next_chunk in reader: for next_chunk in reader:
if next_chunk is None: if next_chunk is None:
break break
next_table = pyarrow.Table.from_batches([next_chunk]) next_table = pyarrow.Table.from_batches([next_chunk])
ntype_partids.append(next_table['f0'].to_numpy()) ntype_partids.append(next_table["f0"].to_numpy())
ntype_partids = np.concatenate(ntype_partids) ntype_partids = np.concatenate(ntype_partids)
count = len(ntype_partids) count = len(ntype_partids)
...@@ -80,9 +90,12 @@ class DistLookupService: ...@@ -80,9 +90,12 @@ class DistLookupService:
# Each rank assumes a contiguous set of partition-ids which are equally split # Each rank assumes a contiguous set of partition-ids which are equally split
# across all the processes. # across all the processes.
split_size = np.ceil(count/np.int64(world_size)).astype(np.int64) split_size = np.ceil(count / np.int64(world_size)).astype(np.int64)
start, end = np.int64(rank)*split_size, np.int64(rank+1)*split_size start, end = (
if rank == (world_size-1): np.int64(rank) * split_size,
np.int64(rank + 1) * split_size,
)
if rank == (world_size - 1):
end = count end = count
type_nid_begin.append(start) type_nid_begin.append(start)
type_nid_end.append(end) type_nid_end.append(end)
...@@ -102,14 +115,13 @@ class DistLookupService: ...@@ -102,14 +115,13 @@ class DistLookupService:
self.rank = rank self.rank = rank
self.world_size = world_size self.world_size = world_size
def get_partition_ids(self, global_nids): def get_partition_ids(self, global_nids):
''' """
This function is used to get the partition-ids for a given set of global node ids This function is used to get the partition-ids for a given set of global node ids
global_nids <-> partition-ids mappings are deterministically distributed across global_nids <-> partition-ids mappings are deterministically distributed across
all the participating processes, within the service. A contiguous global-nids all the participating processes, within the service. A contiguous global-nids
(ntype-ids, per-type-nids) are stored within each process and this is determined (ntype-ids, per-type-nids) are stored within each process and this is determined
by the total no. of nodes of a given ntype-id and the rank of the process. by the total no. of nodes of a given ntype-id and the rank of the process.
Process, where the global_nid <-> partition-id mapping is stored can be easily computed Process, where the global_nid <-> partition-id mapping is stored can be easily computed
...@@ -118,7 +130,7 @@ class DistLookupService: ...@@ -118,7 +130,7 @@ class DistLookupService:
partition-ids using locally stored lookup tables. It builds responses to all the other partition-ids using locally stored lookup tables. It builds responses to all the other
processes and performs alltoallv. processes and performs alltoallv.
Once the response, partition-ids, is received, they are re-ordered corresponding to the Once the response, partition-ids, is received, they are re-ordered corresponding to the
incoming global-nids order and returns to the caller. incoming global-nids order and returns to the caller.
Parameters: Parameters:
...@@ -126,33 +138,35 @@ class DistLookupService: ...@@ -126,33 +138,35 @@ class DistLookupService:
self : instance of this class self : instance of this class
instance of this class, which is passed by the runtime implicitly instance of this class, which is passed by the runtime implicitly
global_nids : numpy array global_nids : numpy array
an array of global node-ids for which partition-ids are to be retrieved by an array of global node-ids for which partition-ids are to be retrieved by
the distributed lookup service. the distributed lookup service.
Returns: Returns:
-------- --------
list of integers : list of integers :
list of integers, which are the partition-ids of the global-node-ids (which is the list of integers, which are the partition-ids of the global-node-ids (which is the
function argument) function argument)
''' """
# Find the process where global_nid --> partition-id(owner) is stored. # Find the process where global_nid --> partition-id(owner) is stored.
ntype_ids, type_nids = self.id_map(global_nids) ntype_ids, type_nids = self.id_map(global_nids)
ntype_ids, type_nids = ntype_ids.numpy(), type_nids.numpy() ntype_ids, type_nids = ntype_ids.numpy(), type_nids.numpy()
assert len(ntype_ids) == len(global_nids) assert len(ntype_ids) == len(global_nids)
# For each node-type, the per-type-node-id <-> partition-id mappings are # For each node-type, the per-type-node-id <-> partition-id mappings are
# stored as contiguous chunks by this lookup service. # stored as contiguous chunks by this lookup service.
# The no. of these mappings stored by each process, in the lookup service, are # The no. of these mappings stored by each process, in the lookup service, are
# equally split among all the processes in the lookup service, deterministically. # equally split among all the processes in the lookup service, deterministically.
typeid_counts = self.ntype_count[ntype_ids] typeid_counts = self.ntype_count[ntype_ids]
chunk_sizes = np.ceil(typeid_counts/self.world_size).astype(np.int64) chunk_sizes = np.ceil(typeid_counts / self.world_size).astype(np.int64)
service_owners = np.floor_divide(type_nids, chunk_sizes).astype(np.int64) service_owners = np.floor_divide(type_nids, chunk_sizes).astype(
np.int64
)
# Now `service_owners` is a list of ranks (process-ids) which own the corresponding # Now `service_owners` is a list of ranks (process-ids) which own the corresponding
# global-nid <-> partition-id mapping. # global-nid <-> partition-id mapping.
# Split the input global_nids into a list of lists where each list will be # Split the input global_nids into a list of lists where each list will be
# sent to the respective rank/process # sent to the respective rank/process
# We also need to store the indices, in the indices_list, so that we can re-order # We also need to store the indices, in the indices_list, so that we can re-order
# the final result (partition-ids) in the same order as the global-nids (function argument) # the final result (partition-ids) in the same order as the global-nids (function argument)
...@@ -164,12 +178,14 @@ class DistLookupService: ...@@ -164,12 +178,14 @@ class DistLookupService:
send_list.append(torch.from_numpy(ll)) send_list.append(torch.from_numpy(ll))
indices_list.append(idxes[0]) indices_list.append(idxes[0])
assert len(np.concatenate(indices_list)) == len(global_nids) assert len(np.concatenate(indices_list)) == len(global_nids)
assert np.all(np.sort(np.concatenate(indices_list)) == np.arange(len(global_nids))) assert np.all(
np.sort(np.concatenate(indices_list)) == np.arange(len(global_nids))
)
# Send the request to everyone else. # Send the request to everyone else.
# As a result of this operation, the current process also receives a list of lists # As a result of this operation, the current process also receives a list of lists
# from all the other processes. # from all the other processes.
# These lists are global-node-ids whose global-node-ids <-> partition-id mappings # These lists are global-node-ids whose global-node-ids <-> partition-id mappings
# are owned/stored by the current process # are owned/stored by the current process
owner_req_list = alltoallv_cpu(self.rank, self.world_size, send_list) owner_req_list = alltoallv_cpu(self.rank, self.world_size, send_list)
...@@ -201,12 +217,15 @@ class DistLookupService: ...@@ -201,12 +217,15 @@ class DistLookupService:
local_type_nids = global_type_nids - self.type_nid_begin[tid] local_type_nids = global_type_nids - self.type_nid_begin[tid]
assert np.all(local_type_nids >= 0) assert np.all(local_type_nids >= 0)
assert np.all(local_type_nids <= (self.type_nid_end[tid] + 1 - self.type_nid_begin[tid])) assert np.all(
local_type_nids
<= (self.type_nid_end[tid] + 1 - self.type_nid_begin[tid])
)
cur_owners = self.partid_list[tid][local_type_nids] cur_owners = self.partid_list[tid][local_type_nids]
type_id_lookups.append(cur_owners) type_id_lookups.append(cur_owners)
# Reorder the partition-ids, so that it agrees with the input order -- # Reorder the partition-ids, so that it agrees with the input order --
# which is the order in which the incoming message is received. # which is the order in which the incoming message is received.
if len(type_id_lookups) <= 0: if len(type_id_lookups) <= 0:
out_list.append(torch.empty((0,), dtype=torch.int64)) out_list.append(torch.empty((0,), dtype=torch.int64))
...@@ -219,14 +238,16 @@ class DistLookupService: ...@@ -219,14 +238,16 @@ class DistLookupService:
# Send the partition-ids to their respective requesting processes. # Send the partition-ids to their respective requesting processes.
owner_resp_list = alltoallv_cpu(self.rank, self.world_size, out_list) owner_resp_list = alltoallv_cpu(self.rank, self.world_size, out_list)
# Owner_resp_list, is a list of lists of numpy arrays where each list # Owner_resp_list, is a list of lists of numpy arrays where each list
# is a list of partition-ids which the current process requested # is a list of partition-ids which the current process requested
# Now we need to re-order so that the parition-ids correspond to the # Now we need to re-order so that the parition-ids correspond to the
# global_nids which are passed into this function. # global_nids which are passed into this function.
# Order according to the requesting order. # Order according to the requesting order.
# Owner_resp_list is the list of owner-ids for global_nids (function argument). # Owner_resp_list is the list of owner-ids for global_nids (function argument).
owner_ids = torch.cat([x for x in owner_resp_list if x is not None]).numpy() owner_ids = torch.cat(
[x for x in owner_resp_list if x is not None]
).numpy()
assert len(owner_ids) == len(global_nids) assert len(owner_ids) == len(global_nids)
global_nids_order = np.concatenate(indices_list) global_nids_order = np.concatenate(indices_list)
...@@ -238,16 +259,18 @@ class DistLookupService: ...@@ -238,16 +259,18 @@ class DistLookupService:
# Now the owner_ids (partition-ids) which corresponding to the global_nids. # Now the owner_ids (partition-ids) which corresponding to the global_nids.
return owner_ids return owner_ids
def get_shuffle_nids(self, global_nids, my_global_nids, my_shuffle_global_nids, world_size): def get_shuffle_nids(
''' self, global_nids, my_global_nids, my_shuffle_global_nids, world_size
):
"""
This function is used to retrieve shuffle_global_nids for a given set of incoming This function is used to retrieve shuffle_global_nids for a given set of incoming
global_nids. Note that global_nids are of random order and will contain duplicates global_nids. Note that global_nids are of random order and will contain duplicates
This function first retrieves the partition-ids of the incoming global_nids. This function first retrieves the partition-ids of the incoming global_nids.
These partition-ids which are also the ranks of processes which own the respective These partition-ids which are also the ranks of processes which own the respective
global-nids as well as shuffle-global-nids. alltoallv is performed to send the global-nids as well as shuffle-global-nids. alltoallv is performed to send the
global-nids to respective ranks/partition-ids where the mapping global-nids to respective ranks/partition-ids where the mapping
global-nids <-> shuffle-global-nid is located. global-nids <-> shuffle-global-nid is located.
On the receiving side, once the global-nids are received associated shuffle-global-nids On the receiving side, once the global-nids are received associated shuffle-global-nids
are retrieved and an alltoallv is performed to send the responses to all the other are retrieved and an alltoallv is performed to send the responses to all the other
...@@ -261,7 +284,7 @@ class DistLookupService: ...@@ -261,7 +284,7 @@ class DistLookupService:
self : instance of this class self : instance of this class
instance of this class, which is passed by the runtime implicitly instance of this class, which is passed by the runtime implicitly
global_nids : numpy array global_nids : numpy array
an array of global node-ids for which partition-ids are to be retrieved by an array of global node-ids for which partition-ids are to be retrieved by
the distributed lookup service. the distributed lookup service.
my_global_nids: numpy ndarray my_global_nids: numpy ndarray
array of global_nids which are owned by the current partition/rank/process array of global_nids which are owned by the current partition/rank/process
...@@ -276,17 +299,17 @@ class DistLookupService: ...@@ -276,17 +299,17 @@ class DistLookupService:
list of integers: list of integers:
list of shuffle_global_nids which correspond to the incoming node-ids in the list of shuffle_global_nids which correspond to the incoming node-ids in the
global_nids. global_nids.
''' """
# Get the owner_ids (partition-ids or rank). # Get the owner_ids (partition-ids or rank).
owner_ids = self.get_partition_ids(global_nids) owner_ids = self.get_partition_ids(global_nids)
# These owner_ids, which are also partition ids of the nodes in the # These owner_ids, which are also partition ids of the nodes in the
# input graph, are in the range 0 - (num_partitions - 1). # input graph, are in the range 0 - (num_partitions - 1).
# These ids are generated using some kind of graph partitioning method. # These ids are generated using some kind of graph partitioning method.
# Distribuged lookup service, as used by the graph partitioning # Distribuged lookup service, as used by the graph partitioning
# pipeline, is used to store ntype-ids (also type_nids) and their # pipeline, is used to store ntype-ids (also type_nids) and their
# mapping to the associated partition-id. # mapping to the associated partition-id.
# These ids are split into `num_process` chunks and processes in the # These ids are split into `num_process` chunks and processes in the
# dist. lookup service are assigned the owernship of these chunks. # dist. lookup service are assigned the owernship of these chunks.
# The pipeline also enforeces the following constraint among the # The pipeline also enforeces the following constraint among the
...@@ -318,8 +341,15 @@ class DistLookupService: ...@@ -318,8 +341,15 @@ class DistLookupService:
shuffle_nids_list.append(torch.empty((0,), dtype=torch.int64)) shuffle_nids_list.append(torch.empty((0,), dtype=torch.int64))
continue continue
uniq_ids, inverse_idx = np.unique(cur_global_nids[idx], return_inverse=True) uniq_ids, inverse_idx = np.unique(
common, idx1, idx2 = np.intersect1d(uniq_ids, my_global_nids, assume_unique=True, return_indices=True) cur_global_nids[idx], return_inverse=True
)
common, idx1, idx2 = np.intersect1d(
uniq_ids,
my_global_nids,
assume_unique=True,
return_indices=True,
)
assert len(common) == len(uniq_ids) assert len(common) == len(uniq_ids)
req_shuffle_global_nids = my_shuffle_global_nids[idx2][inverse_idx] req_shuffle_global_nids = my_shuffle_global_nids[idx2][inverse_idx]
...@@ -327,7 +357,9 @@ class DistLookupService: ...@@ -327,7 +357,9 @@ class DistLookupService:
shuffle_nids_list.append(torch.from_numpy(req_shuffle_global_nids)) shuffle_nids_list.append(torch.from_numpy(req_shuffle_global_nids))
# Send the shuffle-global-nids to their respective ranks. # Send the shuffle-global-nids to their respective ranks.
mapped_global_nids = alltoallv_cpu(self.rank, self.world_size, shuffle_nids_list) mapped_global_nids = alltoallv_cpu(
self.rank, self.world_size, shuffle_nids_list
)
for idx in range(len(mapped_global_nids)): for idx in range(len(mapped_global_nids)):
if mapped_global_nids[idx] == None: if mapped_global_nids[idx] == None:
mapped_global_nids[idx] = torch.empty((0,), dtype=torch.int64) mapped_global_nids[idx] = torch.empty((0,), dtype=torch.int64)
...@@ -338,7 +370,7 @@ class DistLookupService: ...@@ -338,7 +370,7 @@ class DistLookupService:
assert len(shuffle_global_nids) == len(global_nids) assert len(shuffle_global_nids) == len(global_nids)
sorted_idx = np.argsort(global_nids_order) sorted_idx = np.argsort(global_nids_order)
shuffle_global_nids = shuffle_global_nids[ sorted_idx ] shuffle_global_nids = shuffle_global_nids[sorted_idx]
global_nids_ordered = global_nids_order[sorted_idx] global_nids_ordered = global_nids_order[sorted_idx]
assert np.all(global_nids_ordered == np.arange(len(global_nids))) assert np.all(global_nids_ordered == np.arange(len(global_nids)))
......
tools/distpartitioning/globalids.py
View file @ aa42aaeb
import itertools import itertools
import operator import operator
import constants
import numpy as np import numpy as np
import torch import torch
import constants
from dist_lookup import DistLookupService from dist_lookup import DistLookupService
from gloo_wrapper import allgather_sizes, alltoallv_cpu from gloo_wrapper import allgather_sizes, alltoallv_cpu
from utils import memory_snapshot from utils import memory_snapshot
def get_shuffle_global_nids(rank, world_size, global_nids_ranks, node_data): def get_shuffle_global_nids(rank, world_size, global_nids_ranks, node_data):
""" """
For nodes which are not owned by the current rank, whose global_nid <-> shuffle_global-nid mapping For nodes which are not owned by the current rank, whose global_nid <-> shuffle_global-nid mapping
is not present at the current rank, this function retrieves their shuffle_global_ids from the owner rank is not present at the current rank, this function retrieves their shuffle_global_ids from the owner rank
Parameters: Parameters:
----------- -----------
rank : integer rank : integer
rank of the process rank of the process
...@@ -23,7 +23,7 @@ def get_shuffle_global_nids(rank, world_size, global_nids_ranks, node_data): ...@@ -23,7 +23,7 @@ def get_shuffle_global_nids(rank, world_size, global_nids_ranks, node_data):
total no. of ranks configured total no. of ranks configured
global_nids_ranks : list global_nids_ranks : list
list of numpy arrays (of global_nids), index of the list is the rank of the process list of numpy arrays (of global_nids), index of the list is the rank of the process
where global_nid <-> shuffle_global_nid mapping is located. where global_nid <-> shuffle_global_nid mapping is located.
node_data : dictionary node_data : dictionary
node_data is a dictionary with keys as column names and values as numpy arrays node_data is a dictionary with keys as column names and values as numpy arrays
...@@ -31,36 +31,51 @@ def get_shuffle_global_nids(rank, world_size, global_nids_ranks, node_data): ...@@ -31,36 +31,51 @@ def get_shuffle_global_nids(rank, world_size, global_nids_ranks, node_data):
-------- --------
numpy ndarray numpy ndarray
where the column-0 are global_nids and column-1 are shuffle_global_nids which are retrieved where the column-0 are global_nids and column-1 are shuffle_global_nids which are retrieved
from other processes. from other processes.
""" """
#build a list of sizes (lengths of lists) # build a list of sizes (lengths of lists)
global_nids_ranks = [torch.from_numpy(x) for x in global_nids_ranks] global_nids_ranks = [torch.from_numpy(x) for x in global_nids_ranks]
recv_nodes = alltoallv_cpu(rank, world_size, global_nids_ranks) recv_nodes = alltoallv_cpu(rank, world_size, global_nids_ranks)
# Use node_data to lookup global id to send over. # Use node_data to lookup global id to send over.
send_nodes = [] send_nodes = []
for proc_i_nodes in recv_nodes: for proc_i_nodes in recv_nodes:
#list of node-ids to lookup # list of node-ids to lookup
if proc_i_nodes is not None: if proc_i_nodes is not None:
global_nids = proc_i_nodes.numpy() global_nids = proc_i_nodes.numpy()
if(len(global_nids) != 0): if len(global_nids) != 0:
common, ind1, ind2 = np.intersect1d(node_data[constants.GLOBAL_NID], global_nids, return_indices=True) common, ind1, ind2 = np.intersect1d(
shuffle_global_nids = node_data[constants.SHUFFLE_GLOBAL_NID][ind1] node_data[constants.GLOBAL_NID],
send_nodes.append(torch.from_numpy(shuffle_global_nids).type(dtype=torch.int64)) global_nids,
return_indices=True,
)
shuffle_global_nids = node_data[constants.SHUFFLE_GLOBAL_NID][
ind1
]
send_nodes.append(
torch.from_numpy(shuffle_global_nids).type(
dtype=torch.int64
)
)
else: else:
send_nodes.append(torch.empty((0), dtype=torch.int64)) send_nodes.append(torch.empty((0), dtype=torch.int64))
else: else:
send_nodes.append(torch.empty((0), dtype=torch.int64)) send_nodes.append(torch.empty((0), dtype=torch.int64))
#send receive global-ids # send receive global-ids
recv_shuffle_global_nids = alltoallv_cpu(rank, world_size, send_nodes) recv_shuffle_global_nids = alltoallv_cpu(rank, world_size, send_nodes)
shuffle_global_nids = np.concatenate([x.numpy() if x is not None else [] for x in recv_shuffle_global_nids]) shuffle_global_nids = np.concatenate(
[x.numpy() if x is not None else [] for x in recv_shuffle_global_nids]
)
global_nids = np.concatenate([x for x in global_nids_ranks]) global_nids = np.concatenate([x for x in global_nids_ranks])
ret_val = np.column_stack([global_nids, shuffle_global_nids]) ret_val = np.column_stack([global_nids, shuffle_global_nids])
return ret_val return ret_val
def lookup_shuffle_global_nids_edges(rank, world_size, num_parts, edge_data, id_lookup, node_data):
''' def lookup_shuffle_global_nids_edges(
rank, world_size, num_parts, edge_data, id_lookup, node_data
):
"""
This function is a helper function used to lookup shuffle-global-nids for a given set of This function is a helper function used to lookup shuffle-global-nids for a given set of
global-nids using a distributed lookup service. global-nids using a distributed lookup service.
...@@ -87,56 +102,87 @@ def lookup_shuffle_global_nids_edges(rank, world_size, num_parts, edge_data, id_ ...@@ -87,56 +102,87 @@ def lookup_shuffle_global_nids_edges(rank, world_size, num_parts, edge_data, id_
dictionary : dictionary :
dictionary where keys are column names and values are numpy arrays representing all the dictionary where keys are column names and values are numpy arrays representing all the
edges present in the current graph partition edges present in the current graph partition
''' """
# Make sure that the outgoing message size does not exceed 2GB in size. # Make sure that the outgoing message size does not exceed 2GB in size.
# Even though gloo can handle upto 10GB size of data in the outgoing messages, # Even though gloo can handle upto 10GB size of data in the outgoing messages,
# it needs additional memory to store temporary information into the buffers which will increase # it needs additional memory to store temporary information into the buffers which will increase
# the memory needs of the process. # the memory needs of the process.
MILLION = 1000 * 1000 MILLION = 1000 * 1000
BATCH_SIZE = 250 * MILLION BATCH_SIZE = 250 * MILLION
memory_snapshot("GlobalToShuffleIDMapBegin: ", rank) memory_snapshot("GlobalToShuffleIDMapBegin: ", rank)
local_nids = [] local_nids = []
local_shuffle_nids = [] local_shuffle_nids = []
for local_part_id in range(num_parts//world_size): for local_part_id in range(num_parts // world_size):
local_nids.append(node_data[constants.GLOBAL_NID+"/"+str(local_part_id)]) local_nids.append(
local_shuffle_nids.append(node_data[constants.SHUFFLE_GLOBAL_NID+"/"+str(local_part_id)]) node_data[constants.GLOBAL_NID + "/" + str(local_part_id)]
)
local_shuffle_nids.append(
node_data[constants.SHUFFLE_GLOBAL_NID + "/" + str(local_part_id)]
)
local_nids = np.concatenate(local_nids) local_nids = np.concatenate(local_nids)
local_shuffle_nids = np.concatenate(local_shuffle_nids) local_shuffle_nids = np.concatenate(local_shuffle_nids)
for local_part_id in range(num_parts//world_size): for local_part_id in range(num_parts // world_size):
node_list = edge_data[constants.GLOBAL_SRC_ID+"/"+str(local_part_id)] node_list = edge_data[
constants.GLOBAL_SRC_ID + "/" + str(local_part_id)
]
# Determine the no. of times each process has to send alltoall messages. # Determine the no. of times each process has to send alltoall messages.
all_sizes = allgather_sizes([node_list.shape[0]], world_size, num_parts, return_sizes=True) all_sizes = allgather_sizes(
[node_list.shape[0]], world_size, num_parts, return_sizes=True
)
max_count = np.amax(all_sizes) max_count = np.amax(all_sizes)
num_splits = max_count // BATCH_SIZE + 1 num_splits = max_count // BATCH_SIZE + 1
# Split the message into batches and send. # Split the message into batches and send.
splits = np.array_split(node_list, num_splits) splits = np.array_split(node_list, num_splits)
shuffle_mappings = [] shuffle_mappings = []
for item in splits: for item in splits:
shuffle_ids = id_lookup.get_shuffle_nids(item, local_nids, local_shuffle_nids, world_size) shuffle_ids = id_lookup.get_shuffle_nids(
item, local_nids, local_shuffle_nids, world_size
)
shuffle_mappings.append(shuffle_ids) shuffle_mappings.append(shuffle_ids)
shuffle_ids = np.concatenate(shuffle_mappings) shuffle_ids = np.concatenate(shuffle_mappings)
assert shuffle_ids.shape[0] == node_list.shape[0] assert shuffle_ids.shape[0] == node_list.shape[0]
edge_data[constants.SHUFFLE_GLOBAL_SRC_ID+"/"+str(local_part_id)] = shuffle_ids edge_data[
constants.SHUFFLE_GLOBAL_SRC_ID + "/" + str(local_part_id)
] = shuffle_ids
# Destination end points of edges are owned by the current node and therefore # Destination end points of edges are owned by the current node and therefore
# should have corresponding SHUFFLE_GLOBAL_NODE_IDs. # should have corresponding SHUFFLE_GLOBAL_NODE_IDs.
# Here retrieve SHUFFLE_GLOBAL_NODE_IDs for the destination end points of local edges. # Here retrieve SHUFFLE_GLOBAL_NODE_IDs for the destination end points of local edges.
uniq_ids, inverse_idx = np.unique(edge_data[constants.GLOBAL_DST_ID+"/"+str(local_part_id)], return_inverse=True) uniq_ids, inverse_idx = np.unique(
common, idx1, idx2 = np.intersect1d(uniq_ids, node_data[constants.GLOBAL_NID+"/"+str(local_part_id)], assume_unique=True, return_indices=True) edge_data[constants.GLOBAL_DST_ID + "/" + str(local_part_id)],
return_inverse=True,
)
common, idx1, idx2 = np.intersect1d(
uniq_ids,
node_data[constants.GLOBAL_NID + "/" + str(local_part_id)],
assume_unique=True,
return_indices=True,
)
assert len(common) == len(uniq_ids) assert len(common) == len(uniq_ids)
edge_data[constants.SHUFFLE_GLOBAL_DST_ID+"/"+str(local_part_id)] = node_data[constants.SHUFFLE_GLOBAL_NID+"/"+str(local_part_id)][idx2][inverse_idx] edge_data[
assert len(edge_data[constants.SHUFFLE_GLOBAL_DST_ID+"/"+str(local_part_id)]) == len(edge_data[constants.GLOBAL_DST_ID+"/"+str(local_part_id)]) constants.SHUFFLE_GLOBAL_DST_ID + "/" + str(local_part_id)
] = node_data[constants.SHUFFLE_GLOBAL_NID + "/" + str(local_part_id)][
idx2
][
inverse_idx
]
assert len(
edge_data[
constants.SHUFFLE_GLOBAL_DST_ID + "/" + str(local_part_id)
]
) == len(edge_data[constants.GLOBAL_DST_ID + "/" + str(local_part_id)])
memory_snapshot("GlobalToShuffleIDMap_AfterLookupServiceCalls: ", rank) memory_snapshot("GlobalToShuffleIDMap_AfterLookupServiceCalls: ", rank)
return edge_data return edge_data
def assign_shuffle_global_nids_nodes(rank, world_size, num_parts, node_data): def assign_shuffle_global_nids_nodes(rank, world_size, num_parts, node_data):
""" """
Utility function to assign shuffle global ids to nodes at a given rank Utility function to assign shuffle global ids to nodes at a given rank
...@@ -145,10 +191,10 @@ def assign_shuffle_global_nids_nodes(rank, world_size, num_parts, node_data): ...@@ -145,10 +191,10 @@ def assign_shuffle_global_nids_nodes(rank, world_size, num_parts, node_data):
where shuffle_global_nid : global id of the node after data shuffle where shuffle_global_nid : global id of the node after data shuffle
ntype : node-type as read from xxx_nodes.txt ntype : node-type as read from xxx_nodes.txt
global_type_nid : node-type-id as read from xxx_nodes.txt global_type_nid : node-type-id as read from xxx_nodes.txt
global_nid : node-id as read from xxx_nodes.txt, implicitly global_nid : node-id as read from xxx_nodes.txt, implicitly
this is the line no. in the file this is the line no. in the file
part_local_type_nid : type_nid assigned by the current rank within its scope part_local_type_nid : type_nid assigned by the current rank within its scope
Parameters: Parameters:
----------- -----------
rank : integer rank : integer
...@@ -162,17 +208,27 @@ def assign_shuffle_global_nids_nodes(rank, world_size, num_parts, node_data): ...@@ -162,17 +208,27 @@ def assign_shuffle_global_nids_nodes(rank, world_size, num_parts, node_data):
""" """
# Compute prefix sum to determine node-id offsets # Compute prefix sum to determine node-id offsets
local_row_counts = [] local_row_counts = []
for local_part_id in range(num_parts//world_size): for local_part_id in range(num_parts // world_size):
local_row_counts.append(node_data[constants.GLOBAL_NID+"/"+str(local_part_id)].shape[0]) local_row_counts.append(
node_data[constants.GLOBAL_NID + "/" + str(local_part_id)].shape[0]
)
# Perform allgather to compute the local offsets. # Perform allgather to compute the local offsets.
prefix_sum_nodes = allgather_sizes(local_row_counts, world_size, num_parts) prefix_sum_nodes = allgather_sizes(local_row_counts, world_size, num_parts)
for local_part_id in range(num_parts//world_size): for local_part_id in range(num_parts // world_size):
shuffle_global_nid_start = prefix_sum_nodes[rank + (local_part_id*world_size)] shuffle_global_nid_start = prefix_sum_nodes[
shuffle_global_nid_end = prefix_sum_nodes[rank + 1 + (local_part_id*world_size)] rank + (local_part_id * world_size)
shuffle_global_nids = np.arange(shuffle_global_nid_start, shuffle_global_nid_end, dtype=np.int64) ]
node_data[constants.SHUFFLE_GLOBAL_NID+"/"+str(local_part_id)] = shuffle_global_nids shuffle_global_nid_end = prefix_sum_nodes[
rank + 1 + (local_part_id * world_size)
]
shuffle_global_nids = np.arange(
shuffle_global_nid_start, shuffle_global_nid_end, dtype=np.int64
)
node_data[
constants.SHUFFLE_GLOBAL_NID + "/" + str(local_part_id)
] = shuffle_global_nids
def assign_shuffle_global_nids_edges(rank, world_size, num_parts, edge_data): def assign_shuffle_global_nids_edges(rank, world_size, num_parts, edge_data):
...@@ -198,19 +254,31 @@ def assign_shuffle_global_nids_edges(rank, world_size, num_parts, edge_data): ...@@ -198,19 +254,31 @@ def assign_shuffle_global_nids_edges(rank, world_size, num_parts, edge_data):
shuffle_global_eid_start, which indicates the starting value from which shuffle_global-ids are assigned to edges shuffle_global_eid_start, which indicates the starting value from which shuffle_global-ids are assigned to edges
on this rank on this rank
""" """
#get prefix sum of edge counts per rank to locate the starting point # get prefix sum of edge counts per rank to locate the starting point
#from which global-ids to edges are assigned in the current rank # from which global-ids to edges are assigned in the current rank
local_row_counts = [] local_row_counts = []
for local_part_id in range(num_parts//world_size): for local_part_id in range(num_parts // world_size):
local_row_counts.append(edge_data[constants.GLOBAL_SRC_ID+"/"+str(local_part_id)].shape[0]) local_row_counts.append(
edge_data[constants.GLOBAL_SRC_ID + "/" + str(local_part_id)].shape[
0
]
)
shuffle_global_eid_offset = [] shuffle_global_eid_offset = []
prefix_sum_edges = allgather_sizes(local_row_counts, world_size, num_parts) prefix_sum_edges = allgather_sizes(local_row_counts, world_size, num_parts)
for local_part_id in range(num_parts//world_size): for local_part_id in range(num_parts // world_size):
shuffle_global_eid_start = prefix_sum_edges[rank + (local_part_id*world_size)] shuffle_global_eid_start = prefix_sum_edges[
shuffle_global_eid_end = prefix_sum_edges[rank + 1 + (local_part_id*world_size)] rank + (local_part_id * world_size)
shuffle_global_eids = np.arange(shuffle_global_eid_start, shuffle_global_eid_end, dtype=np.int64) ]
edge_data[constants.SHUFFLE_GLOBAL_EID+"/"+str(local_part_id)] = shuffle_global_eids shuffle_global_eid_end = prefix_sum_edges[
rank + 1 + (local_part_id * world_size)
]
shuffle_global_eids = np.arange(
shuffle_global_eid_start, shuffle_global_eid_end, dtype=np.int64
)
edge_data[
constants.SHUFFLE_GLOBAL_EID + "/" + str(local_part_id)
] = shuffle_global_eids
shuffle_global_eid_offset.append(shuffle_global_eid_start) shuffle_global_eid_offset.append(shuffle_global_eid_start)
return shuffle_global_eid_offset return shuffle_global_eid_offset
tools/distpartitioning/gloo_wrapper.py
View file @ aa42aaeb
...@@ -2,15 +2,16 @@ import numpy as np ...@@ -2,15 +2,16 @@ import numpy as np
import torch import torch
import torch.distributed as dist import torch.distributed as dist
def allgather_sizes(send_data, world_size, num_parts, return_sizes=False): def allgather_sizes(send_data, world_size, num_parts, return_sizes=False):
""" """
Perform all gather on list lengths, used to compute prefix sums Perform all gather on list lengths, used to compute prefix sums
to determine the offsets on each ranks. This is used to allocate to determine the offsets on each ranks. This is used to allocate
global ids for edges/nodes on each ranks. global ids for edges/nodes on each ranks.
Parameters Parameters
---------- ----------
send_data : numpy array send_data : numpy array
Data on which allgather is performed. Data on which allgather is performed.
world_size : integer world_size : integer
No. of processes configured for execution No. of processes configured for execution
...@@ -20,7 +21,7 @@ def allgather_sizes(send_data, world_size, num_parts, return_sizes=False): ...@@ -20,7 +21,7 @@ def allgather_sizes(send_data, world_size, num_parts, return_sizes=False):
Boolean flag to indicate whether to return raw sizes from each process Boolean flag to indicate whether to return raw sizes from each process
or perform prefix sum on the raw sizes. or perform prefix sum on the raw sizes.
Returns : Returns :
--------- ---------
numpy array numpy array
array with the prefix sum array with the prefix sum
...@@ -29,33 +30,35 @@ def allgather_sizes(send_data, world_size, num_parts, return_sizes=False): ...@@ -29,33 +30,35 @@ def allgather_sizes(send_data, world_size, num_parts, return_sizes=False):
# Assert on the world_size, num_parts # Assert on the world_size, num_parts
assert (num_parts % world_size) == 0 assert (num_parts % world_size) == 0
#compute the length of the local data # compute the length of the local data
send_length = len(send_data) send_length = len(send_data)
out_tensor = torch.as_tensor(send_data, dtype=torch.int64) out_tensor = torch.as_tensor(send_data, dtype=torch.int64)
in_tensor = [torch.zeros(send_length, dtype=torch.int64) in_tensor = [
for _ in range(world_size)] torch.zeros(send_length, dtype=torch.int64) for _ in range(world_size)
]
#all_gather message # all_gather message
dist.all_gather(in_tensor, out_tensor) dist.all_gather(in_tensor, out_tensor)
# Return on the raw sizes from each process # Return on the raw sizes from each process
if return_sizes: if return_sizes:
return torch.cat(in_tensor).numpy() return torch.cat(in_tensor).numpy()
#gather sizes in on array to return to the invoking function # gather sizes in on array to return to the invoking function
rank_sizes = np.zeros(num_parts + 1, dtype=np.int64) rank_sizes = np.zeros(num_parts + 1, dtype=np.int64)
part_counts = torch.cat(in_tensor).numpy() part_counts = torch.cat(in_tensor).numpy()
count = rank_sizes[0] count = rank_sizes[0]
idx = 1 idx = 1
for local_part_id in range(num_parts//world_size): for local_part_id in range(num_parts // world_size):
for r in range(world_size): for r in range(world_size):
count += part_counts[r*(num_parts//world_size) + local_part_id] count += part_counts[r * (num_parts // world_size) + local_part_id]
rank_sizes[idx] = count rank_sizes[idx] = count
idx += 1 idx += 1
return rank_sizes return rank_sizes
def __alltoall_cpu(rank, world_size, output_tensor_list, input_tensor_list): def __alltoall_cpu(rank, world_size, output_tensor_list, input_tensor_list):
""" """
Each process scatters list of input tensors to all processes in a cluster Each process scatters list of input tensors to all processes in a cluster
...@@ -72,36 +75,41 @@ def __alltoall_cpu(rank, world_size, output_tensor_list, input_tensor_list): ...@@ -72,36 +75,41 @@ def __alltoall_cpu(rank, world_size, output_tensor_list, input_tensor_list):
input_tensor_list : List of tensor input_tensor_list : List of tensor
The tensors to exchange The tensors to exchange
""" """
input_tensor_list = [tensor.to(torch.device('cpu')) for tensor in input_tensor_list] input_tensor_list = [
tensor.to(torch.device("cpu")) for tensor in input_tensor_list
]
# TODO(#5002): As Boolean data is not supported in # TODO(#5002): As Boolean data is not supported in
# ``torch.distributed.scatter()``, we convert boolean into uint8 before # ``torch.distributed.scatter()``, we convert boolean into uint8 before
# scatter and convert it back afterwards. # scatter and convert it back afterwards.
dtypes = [ t.dtype for t in input_tensor_list] dtypes = [t.dtype for t in input_tensor_list]
for i, dtype in enumerate(dtypes): for i, dtype in enumerate(dtypes):
if dtype == torch.bool: if dtype == torch.bool:
input_tensor_list[i] = input_tensor_list[i].to(torch.int8) input_tensor_list[i] = input_tensor_list[i].to(torch.int8)
output_tensor_list[i] = output_tensor_list[i].to(torch.int8) output_tensor_list[i] = output_tensor_list[i].to(torch.int8)
for i in range(world_size): for i in range(world_size):
dist.scatter(output_tensor_list[i], input_tensor_list if i == rank else [], src=i) dist.scatter(
output_tensor_list[i], input_tensor_list if i == rank else [], src=i
)
# Convert back to original dtype # Convert back to original dtype
for i, dtype in enumerate(dtypes): for i, dtype in enumerate(dtypes):
if dtype == torch.bool: if dtype == torch.bool:
input_tensor_list[i] = input_tensor_list[i].to(dtype) input_tensor_list[i] = input_tensor_list[i].to(dtype)
output_tensor_list[i] = output_tensor_list[i].to(dtype) output_tensor_list[i] = output_tensor_list[i].to(dtype)
def alltoallv_cpu(rank, world_size, input_tensor_list, retain_nones=True): def alltoallv_cpu(rank, world_size, input_tensor_list, retain_nones=True):
""" """
Wrapper function to providing the alltoallv functionality by using underlying alltoall Wrapper function to providing the alltoallv functionality by using underlying alltoall
messaging primitive. This function, in its current implementation, supports exchanging messaging primitive. This function, in its current implementation, supports exchanging
messages of arbitrary dimensions and is not tied to the user of this function. messages of arbitrary dimensions and is not tied to the user of this function.
This function pads all input tensors, except one, so that all the messages are of the same This function pads all input tensors, except one, so that all the messages are of the same
size. Once the messages are padded, It first sends a vector whose first two elements are size. Once the messages are padded, It first sends a vector whose first two elements are
1) actual message size along first dimension, and 2) Message size along first dimension 1) actual message size along first dimension, and 2) Message size along first dimension
which is used for communication. The rest of the dimensions are assumed to be same across which is used for communication. The rest of the dimensions are assumed to be same across
all the input tensors. After receiving the message sizes, the receiving end will create buffers all the input tensors. After receiving the message sizes, the receiving end will create buffers
of appropriate sizes. And then slices the received messages to remove the added padding, if any, of appropriate sizes. And then slices the received messages to remove the added padding, if any,
and returns to the caller. and returns to the caller.
Parameters: Parameters:
----------- -----------
...@@ -116,81 +124,99 @@ def alltoallv_cpu(rank, world_size, input_tensor_list, retain_nones=True): ...@@ -116,81 +124,99 @@ def alltoallv_cpu(rank, world_size, input_tensor_list, retain_nones=True):
Returns: Returns:
-------- --------
list : list :
list of tensors received from other processes during alltoall message list of tensors received from other processes during alltoall message
""" """
#ensure len of input_tensor_list is same as the world_size. # ensure len of input_tensor_list is same as the world_size.
assert input_tensor_list != None assert input_tensor_list != None
assert len(input_tensor_list) == world_size assert len(input_tensor_list) == world_size
#ensure that all the tensors in the input_tensor_list are of same size. # ensure that all the tensors in the input_tensor_list are of same size.
sizes = [list(x.size()) for x in input_tensor_list] sizes = [list(x.size()) for x in input_tensor_list]
for idx in range(1,len(sizes)): for idx in range(1, len(sizes)):
assert len(sizes[idx-1]) == len(sizes[idx]) #no. of dimensions should be same assert len(sizes[idx - 1]) == len(
assert input_tensor_list[idx-1].dtype == input_tensor_list[idx].dtype # dtype should be same sizes[idx]
assert sizes[idx-1][1:] == sizes[idx][1:] #except first dimension remaining dimensions should all be the same ) # no. of dimensions should be same
assert (
#decide how much to pad. input_tensor_list[idx - 1].dtype == input_tensor_list[idx].dtype
#always use the first-dimension for padding. ) # dtype should be same
ll = [ x[0] for x in sizes ] assert (
sizes[idx - 1][1:] == sizes[idx][1:]
#dims of the padding needed, if any ) # except first dimension remaining dimensions should all be the same
#these dims are used for padding purposes.
diff_dims = [ [np.amax(ll) - l[0]] + l[1:] for l in sizes ] # decide how much to pad.
# always use the first-dimension for padding.
#pad the actual message ll = [x[0] for x in sizes]
input_tensor_list = [torch.cat((x, torch.zeros(diff_dims[idx]).type(x.dtype))) for idx, x in enumerate(input_tensor_list)]
# dims of the padding needed, if any
#send useful message sizes to all # these dims are used for padding purposes.
diff_dims = [[np.amax(ll) - l[0]] + l[1:] for l in sizes]
# pad the actual message
input_tensor_list = [
torch.cat((x, torch.zeros(diff_dims[idx]).type(x.dtype)))
for idx, x in enumerate(input_tensor_list)
]
# send useful message sizes to all
send_counts = [] send_counts = []
recv_counts = [] recv_counts = []
for idx in range(world_size): for idx in range(world_size):
#send a vector, of atleast 3 elements, [a, b, ....] where # send a vector, of atleast 3 elements, [a, b, ....] where
#a = useful message dim, b = actual message outgoing message size along the first dimension # a = useful message dim, b = actual message outgoing message size along the first dimension
#and remaining elements are the remaining dimensions of the tensor # and remaining elements are the remaining dimensions of the tensor
send_counts.append(torch.from_numpy(np.array([sizes[idx][0]] + [np.amax(ll)] + sizes[idx][1:] )).type(torch.int64)) send_counts.append(
recv_counts.append(torch.zeros((1 + len(sizes[idx])), dtype=torch.int64)) torch.from_numpy(
np.array([sizes[idx][0]] + [np.amax(ll)] + sizes[idx][1:])
).type(torch.int64)
)
recv_counts.append(
torch.zeros((1 + len(sizes[idx])), dtype=torch.int64)
)
__alltoall_cpu(rank, world_size, recv_counts, send_counts) __alltoall_cpu(rank, world_size, recv_counts, send_counts)
#allocate buffers for receiving message # allocate buffers for receiving message
output_tensor_list = [] output_tensor_list = []
recv_counts = [ tsize.numpy() for tsize in recv_counts] recv_counts = [tsize.numpy() for tsize in recv_counts]
for idx, tsize in enumerate(recv_counts): for idx, tsize in enumerate(recv_counts):
output_tensor_list.append(torch.zeros(tuple(tsize[1:])).type(input_tensor_list[idx].dtype)) output_tensor_list.append(
torch.zeros(tuple(tsize[1:])).type(input_tensor_list[idx].dtype)
)
#send actual message itself. # send actual message itself.
__alltoall_cpu(rank, world_size, output_tensor_list, input_tensor_list) __alltoall_cpu(rank, world_size, output_tensor_list, input_tensor_list)
#extract un-padded message from the output_tensor_list and return it # extract un-padded message from the output_tensor_list and return it
return_vals = [] return_vals = []
for s, t in zip(recv_counts, output_tensor_list): for s, t in zip(recv_counts, output_tensor_list):
if s[0] == 0: if s[0] == 0:
if retain_nones: if retain_nones:
return_vals.append(None) return_vals.append(None)
else: else:
return_vals.append(t[0:s[0]]) return_vals.append(t[0 : s[0]])
return return_vals return return_vals
def gather_metadata_json(metadata, rank, world_size):
""" def gather_metadata_json(metadata, rank, world_size):
"""
Gather an object (json schema on `rank`) Gather an object (json schema on `rank`)
Parameters: Parameters:
----------- -----------
metadata : json dictionary object metadata : json dictionary object
json schema formed on each rank with graph level data. json schema formed on each rank with graph level data.
This will be used as input to the distributed training in the later steps. This will be used as input to the distributed training in the later steps.
Returns: Returns:
-------- --------
list : list of json dictionary objects list : list of json dictionary objects
The result of the gather operation, which is the list of json dicitonary The result of the gather operation, which is the list of json dicitonary
objects from each rank in the world objects from each rank in the world
""" """
#Populate input obj and output obj list on rank-0 and non-rank-0 machines # Populate input obj and output obj list on rank-0 and non-rank-0 machines
input_obj = None if rank == 0 else metadata input_obj = None if rank == 0 else metadata
output_objs = [None for _ in range(world_size)] if rank == 0 else None output_objs = [None for _ in range(world_size)] if rank == 0 else None
#invoke the gloo method to perform gather on rank-0 # invoke the gloo method to perform gather on rank-0
dist.gather_object(input_obj, output_objs, dst=0) dist.gather_object(input_obj, output_objs, dst=0)
return output_objs return output_objs
tools/distpartitioning/parmetis_postprocess.py
View file @ aa42aaeb
...@@ -5,13 +5,13 @@ import platform ...@@ -5,13 +5,13 @@ import platform
import sys import sys
from pathlib import Path from pathlib import Path
import constants
import numpy as np import numpy as np
import pyarrow import pyarrow
import pyarrow.csv as csv import pyarrow.csv as csv
from partition_algo.base import dump_partition_meta, PartitionMeta
import constants
from utils import get_idranges, get_node_types, read_json from utils import get_idranges, get_node_types, read_json
from partition_algo.base import PartitionMeta, dump_partition_meta
def post_process(params): def post_process(params):
......
tools/distpartitioning/parmetis_preprocess.py
View file @ aa42aaeb
...@@ -4,14 +4,14 @@ import os ...@@ -4,14 +4,14 @@ import os
import sys import sys
from pathlib import Path from pathlib import Path
import constants
import numpy as np import numpy as np
import pyarrow import pyarrow
import pyarrow.csv as csv import pyarrow.csv as csv
import pyarrow.parquet as pq import pyarrow.parquet as pq
import torch import torch
import torch.distributed as dist import torch.distributed as dist
import constants
from utils import get_idranges, get_node_types, read_json from utils import get_idranges, get_node_types, read_json
import array_readwriter import array_readwriter
...@@ -33,12 +33,13 @@ def get_proc_info(): ...@@ -33,12 +33,13 @@ def get_proc_info():
# mpich # mpich
if "PMI_RANK" in env_variables: if "PMI_RANK" in env_variables:
return int(env_variables["PMI_RANK"]) return int(env_variables["PMI_RANK"])
#openmpi # openmpi
elif "OMPI_COMM_WORLD_RANK" in env_variables: elif "OMPI_COMM_WORLD_RANK" in env_variables:
return int(env_variables["OMPI_COMM_WORLD_RANK"]) return int(env_variables["OMPI_COMM_WORLD_RANK"])
else: else:
return 0 return 0
def gen_edge_files(schema_map, output): def gen_edge_files(schema_map, output):
"""Function to create edges files to be consumed by ParMETIS """Function to create edges files to be consumed by ParMETIS
for partitioning purposes. for partitioning purposes.
...@@ -106,12 +107,16 @@ def gen_edge_files(schema_map, output): ...@@ -106,12 +107,16 @@ def gen_edge_files(schema_map, output):
options = csv.WriteOptions(include_header=False, delimiter=" ") options = csv.WriteOptions(include_header=False, delimiter=" ")
options.delimiter = " " options.delimiter = " "
csv.write_csv( csv.write_csv(
pyarrow.Table.from_arrays(cols, names=col_names), out_file, options pyarrow.Table.from_arrays(cols, names=col_names),
out_file,
options,
) )
return out_file return out_file
if edges_format == constants.STR_CSV: if edges_format == constants.STR_CSV:
delimiter = etype_info[constants.STR_FORMAT][constants.STR_FORMAT_DELIMITER] delimiter = etype_info[constants.STR_FORMAT][
constants.STR_FORMAT_DELIMITER
]
data_df = csv.read_csv( data_df = csv.read_csv(
edge_data_files[rank], edge_data_files[rank],
read_options=pyarrow.csv.ReadOptions( read_options=pyarrow.csv.ReadOptions(
...@@ -309,16 +314,22 @@ def gen_parmetis_input_args(params, schema_map): ...@@ -309,16 +314,22 @@ def gen_parmetis_input_args(params, schema_map):
) )
# Check if <graph-name>_stats.txt exists, if not create one using metadata. # Check if <graph-name>_stats.txt exists, if not create one using metadata.
# Here stats file will be created in the current directory. # Here stats file will be created in the current directory.
# No. of constraints, third column in the stats file is computed as follows: # No. of constraints, third column in the stats file is computed as follows:
# num_constraints = no. of node types + train_mask + test_mask + val_mask # num_constraints = no. of node types + train_mask + test_mask + val_mask
# Here, (train/test/val) masks will be set to 1 if these masks exist for # Here, (train/test/val) masks will be set to 1 if these masks exist for
# all the node types in the graph, otherwise these flags will be set to 0 # all the node types in the graph, otherwise these flags will be set to 0
assert constants.STR_GRAPH_NAME in schema_map, "Graph name is not present in the json file" assert (
constants.STR_GRAPH_NAME in schema_map
), "Graph name is not present in the json file"
graph_name = schema_map[constants.STR_GRAPH_NAME] graph_name = schema_map[constants.STR_GRAPH_NAME]
if not os.path.isfile(f'{graph_name}_stats.txt'): if not os.path.isfile(f"{graph_name}_stats.txt"):
num_nodes = np.sum(np.concatenate(schema_map[constants.STR_NUM_NODES_PER_CHUNK])) num_nodes = np.sum(
num_edges = np.sum(np.concatenate(schema_map[constants.STR_NUM_EDGES_PER_CHUNK])) np.concatenate(schema_map[constants.STR_NUM_NODES_PER_CHUNK])
)
num_edges = np.sum(
np.concatenate(schema_map[constants.STR_NUM_EDGES_PER_CHUNK])
)
num_ntypes = len(schema_map[constants.STR_NODE_TYPE]) num_ntypes = len(schema_map[constants.STR_NODE_TYPE])
train_mask = test_mask = val_mask = 0 train_mask = test_mask = val_mask = 0
...@@ -335,8 +346,8 @@ def gen_parmetis_input_args(params, schema_map): ...@@ -335,8 +346,8 @@ def gen_parmetis_input_args(params, schema_map):
val_mask = val_mask // num_ntypes val_mask = val_mask // num_ntypes
num_constraints = num_ntypes + train_mask + test_mask + val_mask num_constraints = num_ntypes + train_mask + test_mask + val_mask
with open(f'{graph_name}_stats.txt', 'w') as sf: with open(f"{graph_name}_stats.txt", "w") as sf:
sf.write(f'{num_nodes} {num_edges} {num_constraints}') sf.write(f"{num_nodes} {num_edges} {num_constraints}")
node_files = [] node_files = []
outdir = Path(params.output_dir) outdir = Path(params.output_dir)
......
tools/distpartitioning/utils.py
View file @ aa42aaeb
...@@ -3,56 +3,61 @@ import logging ...@@ -3,56 +3,61 @@ import logging
import os import os
import constants import constants
import dgl
import numpy as np import numpy as np
import psutil import psutil
import pyarrow import pyarrow
from pyarrow import csv
import dgl
from dgl.distributed.partition import _dump_part_config from dgl.distributed.partition import _dump_part_config
from pyarrow import csv
def read_ntype_partition_files(schema_map, input_dir): def read_ntype_partition_files(schema_map, input_dir):
""" """
Utility method to read the partition id mapping for each node. Utility method to read the partition id mapping for each node.
For each node type, there will be an file, in the input directory argument For each node type, there will be an file, in the input directory argument
containing the partition id mapping for a given nodeid. containing the partition id mapping for a given nodeid.
Parameters: Parameters:
----------- -----------
schema_map : dictionary schema_map : dictionary
dictionary created by reading the input metadata json file dictionary created by reading the input metadata json file
input_dir : string input_dir : string
directory in which the node-id to partition-id mappings files are directory in which the node-id to partition-id mappings files are
located for each of the node types in the input graph located for each of the node types in the input graph
Returns: Returns:
-------- --------
numpy array : numpy array :
array of integers representing mapped partition-ids for a given node-id. array of integers representing mapped partition-ids for a given node-id.
The line number, in these files, are used as the type_node_id in each of The line number, in these files, are used as the type_node_id in each of
the files. The index into this array will be the homogenized node-id and the files. The index into this array will be the homogenized node-id and
value will be the partition-id for that node-id (index). Please note that value will be the partition-id for that node-id (index). Please note that
the partition-ids of each node-type are stacked together vertically and the partition-ids of each node-type are stacked together vertically and
in this way heterogenous node-ids are converted to homogenous node-ids. in this way heterogenous node-ids are converted to homogenous node-ids.
""" """
assert os.path.isdir(input_dir) assert os.path.isdir(input_dir)
#iterate over the node types and extract the partition id mappings # iterate over the node types and extract the partition id mappings
part_ids = [] part_ids = []
ntype_names = schema_map[constants.STR_NODE_TYPE] ntype_names = schema_map[constants.STR_NODE_TYPE]
for ntype in ntype_names: for ntype in ntype_names:
df = csv.read_csv(os.path.join(input_dir, '{}.txt'.format(ntype)), \ df = csv.read_csv(
read_options=pyarrow.csv.ReadOptions(autogenerate_column_names=True), \ os.path.join(input_dir, "{}.txt".format(ntype)),
parse_options=pyarrow.csv.ParseOptions(delimiter=' ')) read_options=pyarrow.csv.ReadOptions(
ntype_partids = df['f0'].to_numpy() autogenerate_column_names=True
),
parse_options=pyarrow.csv.ParseOptions(delimiter=" "),
)
ntype_partids = df["f0"].to_numpy()
part_ids.append(ntype_partids) part_ids.append(ntype_partids)
return np.concatenate(part_ids) return np.concatenate(part_ids)
def read_json(json_file): def read_json(json_file):
""" """
Utility method to read a json file schema Utility method to read a json file schema
Parameters: Parameters:
----------- -----------
json_file : string json_file : string
...@@ -67,6 +72,7 @@ def read_json(json_file): ...@@ -67,6 +72,7 @@ def read_json(json_file):
return val return val
def get_etype_featnames(etype_name, schema_map): def get_etype_featnames(etype_name, schema_map):
"""Retrieves edge feature names for a given edge_type """Retrieves edge feature names for a given edge_type
...@@ -81,14 +87,15 @@ def get_etype_featnames(etype_name, schema_map): ...@@ -81,14 +87,15 @@ def get_etype_featnames(etype_name, schema_map):
Returns: Returns:
-------- --------
list : list :
a list of feature names for a given edge_type a list of feature names for a given edge_type
""" """
edge_data = schema_map[constants.STR_EDGE_DATA] edge_data = schema_map[constants.STR_EDGE_DATA]
feats = edge_data.get(etype_name, {}) feats = edge_data.get(etype_name, {})
return [feat for feat in feats] return [feat for feat in feats]
def get_ntype_featnames(ntype_name, schema_map):
def get_ntype_featnames(ntype_name, schema_map):
""" """
Retrieves node feature names for a given node_type Retrieves node feature names for a given node_type
...@@ -103,13 +110,14 @@ def get_ntype_featnames(ntype_name, schema_map): ...@@ -103,13 +110,14 @@ def get_ntype_featnames(ntype_name, schema_map):
Returns: Returns:
-------- --------
list : list :
a list of feature names for a given node_type a list of feature names for a given node_type
""" """
node_data = schema_map[constants.STR_NODE_DATA] node_data = schema_map[constants.STR_NODE_DATA]
feats = node_data.get(ntype_name, {}) feats = node_data.get(ntype_name, {})
return [feat for feat in feats] return [feat for feat in feats]
def get_edge_types(schema_map): def get_edge_types(schema_map):
"""Utility method to extract edge_typename -> edge_type mappings """Utility method to extract edge_typename -> edge_type mappings
as defined by the input schema as defined by the input schema
...@@ -130,12 +138,13 @@ def get_edge_types(schema_map): ...@@ -130,12 +138,13 @@ def get_edge_types(schema_map):
with keys as etype ids (integers) and values as edge type names with keys as etype ids (integers) and values as edge type names
""" """
etypes = schema_map[constants.STR_EDGE_TYPE] etypes = schema_map[constants.STR_EDGE_TYPE]
etype_etypeid_map = {e : i for i, e in enumerate(etypes)} etype_etypeid_map = {e: i for i, e in enumerate(etypes)}
etypeid_etype_map = {i : e for i, e in enumerate(etypes)} etypeid_etype_map = {i: e for i, e in enumerate(etypes)}
return etype_etypeid_map, etypes, etypeid_etype_map return etype_etypeid_map, etypes, etypeid_etype_map
def get_node_types(schema_map): def get_node_types(schema_map):
""" """
Utility method to extract node_typename -> node_type mappings Utility method to extract node_typename -> node_type mappings
as defined by the input schema as defined by the input schema
...@@ -155,11 +164,12 @@ def get_node_types(schema_map): ...@@ -155,11 +164,12 @@ def get_node_types(schema_map):
with keys as ntype ids (integers) and values as node type names with keys as ntype ids (integers) and values as node type names
""" """
ntypes = schema_map[constants.STR_NODE_TYPE] ntypes = schema_map[constants.STR_NODE_TYPE]
ntype_ntypeid_map = {e : i for i, e in enumerate(ntypes)} ntype_ntypeid_map = {e: i for i, e in enumerate(ntypes)}
ntypeid_ntype_map = {i : e for i, e in enumerate(ntypes)} ntypeid_ntype_map = {i: e for i, e in enumerate(ntypes)}
return ntype_ntypeid_map, ntypes, ntypeid_ntype_map return ntype_ntypeid_map, ntypes, ntypeid_ntype_map
def get_gnid_range_map(node_tids):
def get_gnid_range_map(node_tids):
""" """
Retrieves auxiliary dictionaries from the metadata json object Retrieves auxiliary dictionaries from the metadata json object
...@@ -167,28 +177,31 @@ def get_gnid_range_map(node_tids): ...@@ -167,28 +177,31 @@ def get_gnid_range_map(node_tids):
----------- -----------
node_tids: dictionary node_tids: dictionary
This dictionary contains the information about nodes for each node_type. This dictionary contains the information about nodes for each node_type.
Typically this information contains p-entries, where each entry has a file-name, Typically this information contains p-entries, where each entry has a file-name,
starting and ending type_node_ids for the nodes in this file. Keys in this dictionary starting and ending type_node_ids for the nodes in this file. Keys in this dictionary
are the node_type and value is a list of lists. Each individual entry in this list has are the node_type and value is a list of lists. Each individual entry in this list has
three items: file-name, starting type_nid and ending type_nid three items: file-name, starting type_nid and ending type_nid
Returns: Returns:
-------- --------
dictionary : dictionary :
a dictionary where keys are node_type names and values are global_nid range, which is a tuple. a dictionary where keys are node_type names and values are global_nid range, which is a tuple.
""" """
ntypes_gid_range = {} ntypes_gid_range = {}
offset = 0 offset = 0
for k, v in node_tids.items(): for k, v in node_tids.items():
ntypes_gid_range[k] = [offset + int(v[0][0]), offset + int(v[-1][1])] ntypes_gid_range[k] = [offset + int(v[0][0]), offset + int(v[-1][1])]
offset += int(v[-1][1]) offset += int(v[-1][1])
return ntypes_gid_range return ntypes_gid_range
def write_metadata_json(input_list, output_dir, graph_name, world_size, num_parts):
def write_metadata_json(
input_list, output_dir, graph_name, world_size, num_parts
):
""" """
Merge json schema's from each of the rank's on rank-0. Merge json schema's from each of the rank's on rank-0.
This utility function, to be used on rank-0, to create aggregated json file. This utility function, to be used on rank-0, to create aggregated json file.
Parameters: Parameters:
...@@ -201,50 +214,73 @@ def write_metadata_json(input_list, output_dir, graph_name, world_size, num_part ...@@ -201,50 +214,73 @@ def write_metadata_json(input_list, output_dir, graph_name, world_size, num_part
a string specifying the graph name a string specifying the graph name
""" """
# Preprocess the input_list, a list of dictionaries # Preprocess the input_list, a list of dictionaries
# each dictionary will contain num_parts/world_size metadata json # each dictionary will contain num_parts/world_size metadata json
# which correspond to local partitions on the respective ranks. # which correspond to local partitions on the respective ranks.
metadata_list = [] metadata_list = []
for local_part_id in range(num_parts//world_size): for local_part_id in range(num_parts // world_size):
for idx in range(world_size): for idx in range(world_size):
metadata_list.append(input_list[idx]["local-part-id-"+str(local_part_id*world_size + idx)]) metadata_list.append(
input_list[idx][
"local-part-id-" + str(local_part_id * world_size + idx)
]
)
#Initialize global metadata # Initialize global metadata
graph_metadata = {} graph_metadata = {}
#Merge global_edge_ids from each json object in the input list # Merge global_edge_ids from each json object in the input list
edge_map = {} edge_map = {}
x = metadata_list[0]["edge_map"] x = metadata_list[0]["edge_map"]
for k in x: for k in x:
edge_map[k] = [] edge_map[k] = []
for idx in range(len(metadata_list)): for idx in range(len(metadata_list)):
edge_map[k].append([int(metadata_list[idx]["edge_map"][k][0][0]),int(metadata_list[idx]["edge_map"][k][0][1])]) edge_map[k].append(
[
int(metadata_list[idx]["edge_map"][k][0][0]),
int(metadata_list[idx]["edge_map"][k][0][1]),
]
)
graph_metadata["edge_map"] = edge_map graph_metadata["edge_map"] = edge_map
graph_metadata["etypes"] = metadata_list[0]["etypes"] graph_metadata["etypes"] = metadata_list[0]["etypes"]
graph_metadata["graph_name"] = metadata_list[0]["graph_name"] graph_metadata["graph_name"] = metadata_list[0]["graph_name"]
graph_metadata["halo_hops"] = metadata_list[0]["halo_hops"] graph_metadata["halo_hops"] = metadata_list[0]["halo_hops"]
#Merge global_nodeids from each of json object in the input list # Merge global_nodeids from each of json object in the input list
node_map = {} node_map = {}
x = metadata_list[0]["node_map"] x = metadata_list[0]["node_map"]
for k in x: for k in x:
node_map[k] = [] node_map[k] = []
for idx in range(len(metadata_list)): for idx in range(len(metadata_list)):
node_map[k].append([int(metadata_list[idx]["node_map"][k][0][0]), int(metadata_list[idx]["node_map"][k][0][1])]) node_map[k].append(
[
int(metadata_list[idx]["node_map"][k][0][0]),
int(metadata_list[idx]["node_map"][k][0][1]),
]
)
graph_metadata["node_map"] = node_map graph_metadata["node_map"] = node_map
graph_metadata["ntypes"] = metadata_list[0]["ntypes"] graph_metadata["ntypes"] = metadata_list[0]["ntypes"]
graph_metadata["num_edges"] = int(sum([metadata_list[i]["num_edges"] for i in range(len(metadata_list))])) graph_metadata["num_edges"] = int(
graph_metadata["num_nodes"] = int(sum([metadata_list[i]["num_nodes"] for i in range(len(metadata_list))])) sum([metadata_list[i]["num_edges"] for i in range(len(metadata_list))])
)
graph_metadata["num_nodes"] = int(
sum([metadata_list[i]["num_nodes"] for i in range(len(metadata_list))])
)
graph_metadata["num_parts"] = metadata_list[0]["num_parts"] graph_metadata["num_parts"] = metadata_list[0]["num_parts"]
graph_metadata["part_method"] = metadata_list[0]["part_method"] graph_metadata["part_method"] = metadata_list[0]["part_method"]
for i in range(len(metadata_list)): for i in range(len(metadata_list)):
graph_metadata["part-{}".format(i)] = metadata_list[i]["part-{}".format(i)] graph_metadata["part-{}".format(i)] = metadata_list[i][
"part-{}".format(i)
]
_dump_part_config(f"{output_dir}/metadata.json", graph_metadata)
_dump_part_config(f'{output_dir}/metadata.json', graph_metadata)
def augment_edge_data(edge_data, lookup_service, edge_tids, rank, world_size, num_parts): def augment_edge_data(
edge_data, lookup_service, edge_tids, rank, world_size, num_parts
):
""" """
Add partition-id (rank which owns an edge) column to the edge_data. Add partition-id (rank which owns an edge) column to the edge_data.
...@@ -267,22 +303,22 @@ def augment_edge_data(edge_data, lookup_service, edge_tids, rank, world_size, nu ...@@ -267,22 +303,22 @@ def augment_edge_data(edge_data, lookup_service, edge_tids, rank, world_size, nu
Returns: Returns:
-------- --------
dictionary : dictionary :
dictionary with keys as column names and values as numpy arrays and this information is dictionary with keys as column names and values as numpy arrays and this information is
loaded from input dataset files. In addition to this we include additional columns which loaded from input dataset files. In addition to this we include additional columns which
aid this pipelines computation, like constants.OWNER_PROCESS aid this pipelines computation, like constants.OWNER_PROCESS
""" """
#add global_nids to the node_data # add global_nids to the node_data
etype_offset = {} etype_offset = {}
offset = 0 offset = 0
for etype_name, tid_range in edge_tids.items(): for etype_name, tid_range in edge_tids.items():
assert int(tid_range[0][0]) == 0 assert int(tid_range[0][0]) == 0
assert len(tid_range) == num_parts assert len(tid_range) == num_parts
etype_offset[etype_name] = offset + int(tid_range[0][0]) etype_offset[etype_name] = offset + int(tid_range[0][0])
offset += int(tid_range[-1][1]) offset += int(tid_range[-1][1])
global_eids = [] global_eids = []
for etype_name, tid_range in edge_tids.items(): for etype_name, tid_range in edge_tids.items():
for idx in range(num_parts): for idx in range(num_parts):
if map_partid_rank(idx, world_size) == rank: if map_partid_rank(idx, world_size) == rank:
global_eid_start = etype_offset[etype_name] global_eid_start = etype_offset[etype_name]
...@@ -293,12 +329,15 @@ def augment_edge_data(edge_data, lookup_service, edge_tids, rank, world_size, nu ...@@ -293,12 +329,15 @@ def augment_edge_data(edge_data, lookup_service, edge_tids, rank, world_size, nu
assert global_eids.shape[0] == edge_data[constants.ETYPE_ID].shape[0] assert global_eids.shape[0] == edge_data[constants.ETYPE_ID].shape[0]
edge_data[constants.GLOBAL_EID] = global_eids edge_data[constants.GLOBAL_EID] = global_eids
#assign the owner process/rank for each edge # assign the owner process/rank for each edge
edge_data[constants.OWNER_PROCESS] = lookup_service.get_partition_ids(edge_data[constants.GLOBAL_DST_ID]) edge_data[constants.OWNER_PROCESS] = lookup_service.get_partition_ids(
edge_data[constants.GLOBAL_DST_ID]
)
return edge_data return edge_data
def read_edges_file(edge_file, edge_data_dict): def read_edges_file(edge_file, edge_data_dict):
""" """
Utility function to read xxx_edges.txt file Utility function to read xxx_edges.txt file
Parameters: Parameters:
...@@ -310,25 +349,29 @@ def read_edges_file(edge_file, edge_data_dict): ...@@ -310,25 +349,29 @@ def read_edges_file(edge_file, edge_data_dict):
-------- --------
dictionary dictionary
edge data as read from xxx_edges.txt file and columns are stored edge data as read from xxx_edges.txt file and columns are stored
in a dictionary with key-value pairs as column-names and column-data. in a dictionary with key-value pairs as column-names and column-data.
""" """
if edge_file == "" or edge_file == None: if edge_file == "" or edge_file == None:
return None return None
#Read the file from here. # Read the file from here.
#<global_src_id> <global_dst_id> <type_eid> <etype> <attributes> # <global_src_id> <global_dst_id> <type_eid> <etype> <attributes>
# global_src_id -- global idx for the source node ... line # in the graph_nodes.txt # global_src_id -- global idx for the source node ... line # in the graph_nodes.txt
# global_dst_id -- global idx for the destination id node ... line # in the graph_nodes.txt # global_dst_id -- global idx for the destination id node ... line # in the graph_nodes.txt
edge_data_df = csv.read_csv(edge_file, read_options=pyarrow.csv.ReadOptions(autogenerate_column_names=True), edge_data_df = csv.read_csv(
parse_options=pyarrow.csv.ParseOptions(delimiter=' ')) edge_file,
read_options=pyarrow.csv.ReadOptions(autogenerate_column_names=True),
parse_options=pyarrow.csv.ParseOptions(delimiter=" "),
)
edge_data_dict = {} edge_data_dict = {}
edge_data_dict[constants.GLOBAL_SRC_ID] = edge_data_df['f0'].to_numpy() edge_data_dict[constants.GLOBAL_SRC_ID] = edge_data_df["f0"].to_numpy()
edge_data_dict[constants.GLOBAL_DST_ID] = edge_data_df['f1'].to_numpy() edge_data_dict[constants.GLOBAL_DST_ID] = edge_data_df["f1"].to_numpy()
edge_data_dict[constants.GLOBAL_TYPE_EID] = edge_data_df['f2'].to_numpy() edge_data_dict[constants.GLOBAL_TYPE_EID] = edge_data_df["f2"].to_numpy()
edge_data_dict[constants.ETYPE_ID] = edge_data_df['f3'].to_numpy() edge_data_dict[constants.ETYPE_ID] = edge_data_df["f3"].to_numpy()
return edge_data_dict return edge_data_dict
def read_node_features_file(nodes_features_file): def read_node_features_file(nodes_features_file):
""" """
Utility function to load tensors from a file Utility function to load tensors from a file
...@@ -347,8 +390,9 @@ def read_node_features_file(nodes_features_file): ...@@ -347,8 +390,9 @@ def read_node_features_file(nodes_features_file):
node_features = dgl.data.utils.load_tensors(nodes_features_file, False) node_features = dgl.data.utils.load_tensors(nodes_features_file, False)
return node_features return node_features
def read_edge_features_file(edge_features_file): def read_edge_features_file(edge_features_file):
""" """
Utility function to load tensors from a file Utility function to load tensors from a file
Parameters: Parameters:
...@@ -364,6 +408,7 @@ def read_edge_features_file(edge_features_file): ...@@ -364,6 +408,7 @@ def read_edge_features_file(edge_features_file):
edge_features = dgl.data.utils.load_tensors(edge_features_file, True) edge_features = dgl.data.utils.load_tensors(edge_features_file, True)
return edge_features return edge_features
def write_node_features(node_features, node_file): def write_node_features(node_features, node_file):
""" """
Utility function to serialize node_features in node_file file Utility function to serialize node_features in node_file file
...@@ -372,12 +417,13 @@ def write_node_features(node_features, node_file): ...@@ -372,12 +417,13 @@ def write_node_features(node_features, node_file):
----------- -----------
node_features : dictionary node_features : dictionary
dictionary storing ntype <-> list of features dictionary storing ntype <-> list of features
node_file : string node_file : string
File in which the node information is serialized File in which the node information is serialized
""" """
dgl.data.utils.save_tensors(node_file, node_features) dgl.data.utils.save_tensors(node_file, node_features)
def write_edge_features(edge_features, edge_file):
def write_edge_features(edge_features, edge_file):
""" """
Utility function to serialize edge_features in edge_file file Utility function to serialize edge_features in edge_file file
...@@ -385,11 +431,12 @@ def write_edge_features(edge_features, edge_file): ...@@ -385,11 +431,12 @@ def write_edge_features(edge_features, edge_file):
----------- -----------
edge_features : dictionary edge_features : dictionary
dictionary storing etype <-> list of features dictionary storing etype <-> list of features
edge_file : string edge_file : string
File in which the edge information is serialized File in which the edge information is serialized
""" """
dgl.data.utils.save_tensors(edge_file, edge_features) dgl.data.utils.save_tensors(edge_file, edge_features)
def write_graph_dgl(graph_file, graph_obj, formats, sort_etypes): def write_graph_dgl(graph_file, graph_obj, formats, sort_etypes):
""" """
Utility function to serialize graph dgl objects Utility function to serialize graph dgl objects
...@@ -405,11 +452,22 @@ def write_graph_dgl(graph_file, graph_obj, formats, sort_etypes): ...@@ -405,11 +452,22 @@ def write_graph_dgl(graph_file, graph_obj, formats, sort_etypes):
sort_etypes : bool sort_etypes : bool
Whether to sort etypes in csc/csr. Whether to sort etypes in csc/csr.
""" """
dgl.distributed.partition._save_graphs(graph_file, [graph_obj], dgl.distributed.partition._save_graphs(
formats, sort_etypes) graph_file, [graph_obj], formats, sort_etypes
)
def write_dgl_objects(graph_obj, node_features, edge_features, def write_dgl_objects(
output_dir, part_id, orig_nids, orig_eids, formats, sort_etypes): graph_obj,
node_features,
edge_features,
output_dir,
part_id,
orig_nids,
orig_eids,
formats,
sort_etypes,
):
""" """
Wrapper function to write graph, node/edge feature, original node/edge IDs. Wrapper function to write graph, node/edge feature, original node/edge IDs.
...@@ -434,27 +492,33 @@ def write_dgl_objects(graph_obj, node_features, edge_features, ...@@ -434,27 +492,33 @@ def write_dgl_objects(graph_obj, node_features, edge_features,
sort_etypes : bool sort_etypes : bool
Whether to sort etypes in csc/csr. Whether to sort etypes in csc/csr.
""" """
part_dir = output_dir + '/part' + str(part_id) part_dir = output_dir + "/part" + str(part_id)
os.makedirs(part_dir, exist_ok=True) os.makedirs(part_dir, exist_ok=True)
write_graph_dgl(os.path.join(part_dir ,'graph.dgl'), graph_obj, write_graph_dgl(
formats, sort_etypes) os.path.join(part_dir, "graph.dgl"), graph_obj, formats, sort_etypes
)
if node_features != None: if node_features != None:
write_node_features(node_features, os.path.join(part_dir, "node_feat.dgl")) write_node_features(
node_features, os.path.join(part_dir, "node_feat.dgl")
)
if (edge_features != None): if edge_features != None:
write_edge_features(edge_features, os.path.join(part_dir, "edge_feat.dgl")) write_edge_features(
edge_features, os.path.join(part_dir, "edge_feat.dgl")
)
if orig_nids is not None: if orig_nids is not None:
orig_nids_file = os.path.join(part_dir, 'orig_nids.dgl') orig_nids_file = os.path.join(part_dir, "orig_nids.dgl")
dgl.data.utils.save_tensors(orig_nids_file, orig_nids) dgl.data.utils.save_tensors(orig_nids_file, orig_nids)
if orig_eids is not None: if orig_eids is not None:
orig_eids_file = os.path.join(part_dir, 'orig_eids.dgl') orig_eids_file = os.path.join(part_dir, "orig_eids.dgl")
dgl.data.utils.save_tensors(orig_eids_file, orig_eids) dgl.data.utils.save_tensors(orig_eids_file, orig_eids)
def get_idranges(names, counts, num_chunks=None):
def get_idranges(names, counts, num_chunks=None):
""" """
Utility function to compute typd_id/global_id ranges for both nodes and edges. Utility function to compute typd_id/global_id ranges for both nodes and edges.
Parameters: Parameters:
----------- -----------
...@@ -472,8 +536,8 @@ def get_idranges(names, counts, num_chunks=None): ...@@ -472,8 +536,8 @@ def get_idranges(names, counts, num_chunks=None):
-------- --------
dictionary dictionary
dictionary where the keys are node-/edge-type names and values are dictionary where the keys are node-/edge-type names and values are
list of tuples where each tuple indicates the range of values for list of tuples where each tuple indicates the range of values for
corresponding type-ids. corresponding type-ids.
dictionary dictionary
dictionary where the keys are node-/edge-type names and value is a tuple. dictionary where the keys are node-/edge-type names and value is a tuple.
This tuple indicates the global-ids for the associated node-/edge-type. This tuple indicates the global-ids for the associated node-/edge-type.
...@@ -483,7 +547,7 @@ def get_idranges(names, counts, num_chunks=None): ...@@ -483,7 +547,7 @@ def get_idranges(names, counts, num_chunks=None):
tid_dict = {} tid_dict = {}
gid_dict = {} gid_dict = {}
orig_num_chunks = 0 orig_num_chunks = 0
for idx, typename in enumerate(names): for idx, typename in enumerate(names):
type_counts = counts[idx] type_counts = counts[idx]
tid_start = np.cumsum([0] + type_counts[:-1]) tid_start = np.cumsum([0] + type_counts[:-1])
tid_end = np.cumsum(type_counts) tid_end = np.cumsum(type_counts)
...@@ -492,7 +556,7 @@ def get_idranges(names, counts, num_chunks=None): ...@@ -492,7 +556,7 @@ def get_idranges(names, counts, num_chunks=None):
gnid_end += tid_ranges[-1][1] gnid_end += tid_ranges[-1][1]
tid_dict[typename] = tid_ranges tid_dict[typename] = tid_ranges
gid_dict[typename] = np.array([gnid_start, gnid_end]).reshape([1,2]) gid_dict[typename] = np.array([gnid_start, gnid_end]).reshape([1, 2])
gnid_start = gnid_end gnid_start = gnid_end
orig_num_chunks = len(tid_start) orig_num_chunks = len(tid_start)
...@@ -500,14 +564,17 @@ def get_idranges(names, counts, num_chunks=None): ...@@ -500,14 +564,17 @@ def get_idranges(names, counts, num_chunks=None):
if num_chunks is None: if num_chunks is None:
return tid_dict, gid_dict return tid_dict, gid_dict
assert num_chunks <= orig_num_chunks, \ assert (
'Specified number of chunks should be less/euqual than original numbers of ID ranges.' num_chunks <= orig_num_chunks
), "Specified number of chunks should be less/euqual than original numbers of ID ranges."
chunk_list = np.array_split(np.arange(orig_num_chunks), num_chunks) chunk_list = np.array_split(np.arange(orig_num_chunks), num_chunks)
for typename in tid_dict: for typename in tid_dict:
orig_tid_ranges = tid_dict[typename] orig_tid_ranges = tid_dict[typename]
tid_ranges = [] tid_ranges = []
for idx in chunk_list: for idx in chunk_list:
tid_ranges.append((orig_tid_ranges[idx[0]][0], orig_tid_ranges[idx[-1]][-1])) tid_ranges.append(
(orig_tid_ranges[idx[0]][0], orig_tid_ranges[idx[-1]][-1])
)
tid_dict[typename] = tid_ranges tid_dict[typename] = tid_ranges
return tid_dict, gid_dict return tid_dict, gid_dict
...@@ -517,8 +584,8 @@ def memory_snapshot(tag, rank): ...@@ -517,8 +584,8 @@ def memory_snapshot(tag, rank):
""" """
Utility function to take a snapshot of the usage of system resources Utility function to take a snapshot of the usage of system resources
at a given point of time. at a given point of time.
Parameters: Parameters:
----------- -----------
tag : string tag : string
string provided by the user for bookmarking purposes string provided by the user for bookmarking purposes
...@@ -529,19 +596,19 @@ def memory_snapshot(tag, rank): ...@@ -529,19 +596,19 @@ def memory_snapshot(tag, rank):
MB = 1024 * 1024 MB = 1024 * 1024
KB = 1024 KB = 1024
peak = dgl.partition.get_peak_mem()*KB peak = dgl.partition.get_peak_mem() * KB
mem = psutil.virtual_memory() mem = psutil.virtual_memory()
avail = mem.available / MB avail = mem.available / MB
used = mem.used / MB used = mem.used / MB
total = mem.total / MB total = mem.total / MB
mem_string = f'{total:.0f} (MB) total, {peak:.0f} (MB) peak, {used:.0f} (MB) used, {avail:.0f} (MB) avail' mem_string = f"{total:.0f} (MB) total, {peak:.0f} (MB) peak, {used:.0f} (MB) used, {avail:.0f} (MB) avail"
logging.debug(f'[Rank: {rank} MEMORY_SNAPSHOT] {mem_string} - {tag}') logging.debug(f"[Rank: {rank} MEMORY_SNAPSHOT] {mem_string} - {tag}")
def map_partid_rank(partid, world_size): def map_partid_rank(partid, world_size):
"""Auxiliary function to map a given partition id to one of the rank in the """Auxiliary function to map a given partition id to one of the rank in the
MPI_WORLD processes. The range of partition ids is assumed to equal or a MPI_WORLD processes. The range of partition ids is assumed to equal or a
multiple of the total size of MPI_WORLD. In this implementation, we use multiple of the total size of MPI_WORLD. In this implementation, we use
a cyclical mapping procedure to convert partition ids to ranks. a cyclical mapping procedure to convert partition ids to ranks.
...@@ -552,7 +619,7 @@ def map_partid_rank(partid, world_size): ...@@ -552,7 +619,7 @@ def map_partid_rank(partid, world_size):
Returns: Returns:
-------- --------
int : int :
rank of the process, which will be responsible for the given partition rank of the process, which will be responsible for the given partition
id. id.
""" """
......
tools/partition_algo/base.py
View file @ aa42aaeb
import json
from typing import Optional from typing import Optional
import pydantic as dt import pydantic as dt
import json
from dgl import DGLError from dgl import DGLError
class PartitionMeta(dt.BaseModel): class PartitionMeta(dt.BaseModel):
""" Metadata that describes the partition assignment results. """Metadata that describes the partition assignment results.
Regardless of the choice of partitioning algorithm, a metadata JSON file Regardless of the choice of partitioning algorithm, a metadata JSON file
will be created in the output directory which includes the meta information will be created in the output directory which includes the meta information
...@@ -22,15 +24,17 @@ class PartitionMeta(dt.BaseModel): ...@@ -22,15 +24,17 @@ class PartitionMeta(dt.BaseModel):
... part_meta = PartitionMeta(**(json.load(f))) ... part_meta = PartitionMeta(**(json.load(f)))
""" """
# version of metadata JSON. # version of metadata JSON.
version: Optional[str] = '1.0.0' version: Optional[str] = "1.0.0"
# number of partitions. # number of partitions.
num_parts: int num_parts: int
# name of partition algorithm. # name of partition algorithm.
algo_name: str algo_name: str
def dump_partition_meta(part_meta, meta_file): def dump_partition_meta(part_meta, meta_file):
""" Dump partition metadata into json file. """Dump partition metadata into json file.
Parameters Parameters
---------- ----------
...@@ -39,11 +43,12 @@ def dump_partition_meta(part_meta, meta_file): ...@@ -39,11 +43,12 @@ def dump_partition_meta(part_meta, meta_file):
meta_file : str meta_file : str
The target file to save data. The target file to save data.
""" """
with open(meta_file, 'w') as f: with open(meta_file, "w") as f:
json.dump(part_meta.dict(), f, sort_keys=True, indent=4) json.dump(part_meta.dict(), f, sort_keys=True, indent=4)
def load_partition_meta(meta_file): def load_partition_meta(meta_file):
""" Load partition metadata and do sanity check. """Load partition metadata and do sanity check.
Parameters Parameters
---------- ----------
...@@ -60,14 +65,18 @@ def load_partition_meta(meta_file): ...@@ -60,14 +65,18 @@ def load_partition_meta(meta_file):
part_meta = PartitionMeta(**(json.load(f))) part_meta = PartitionMeta(**(json.load(f)))
except dt.ValidationError as e: except dt.ValidationError as e:
raise DGLError( raise DGLError(
f"Invalid partition metadata JSON. Error details: {e.json()}") f"Invalid partition metadata JSON. Error details: {e.json()}"
if part_meta.version != '1.0.0': )
if part_meta.version != "1.0.0":
raise DGLError( raise DGLError(
f"Invalid version[{part_meta.version}]. Supported versions: '1.0.0'") f"Invalid version[{part_meta.version}]. Supported versions: '1.0.0'"
)
if part_meta.num_parts <= 0: if part_meta.num_parts <= 0:
raise DGLError( raise DGLError(
f"num_parts[{part_meta.num_parts}] should be greater than 0.") f"num_parts[{part_meta.num_parts}] should be greater than 0."
if part_meta.algo_name not in ['random', 'metis']: )
if part_meta.algo_name not in ["random", "metis"]:
raise DGLError( raise DGLError(
f"algo_name[{part_meta.num_parts}] is not supported.") f"algo_name[{part_meta.num_parts}] is not supported."
)
return part_meta return part_meta
tools/partition_algo/random_partition.py
View file @ aa42aaeb
...@@ -6,10 +6,11 @@ import os ...@@ -6,10 +6,11 @@ import os
import sys import sys
import numpy as np import numpy as np
from base import PartitionMeta, dump_partition_meta from base import dump_partition_meta, PartitionMeta
from distpartitioning import array_readwriter from distpartitioning import array_readwriter
from files import setdir from files import setdir
def _random_partition(metadata, num_parts): def _random_partition(metadata, num_parts):
num_nodes_per_type = [sum(_) for _ in metadata["num_nodes_per_chunk"]] num_nodes_per_type = [sum(_) for _ in metadata["num_nodes_per_chunk"]]
ntypes = metadata["node_type"] ntypes = metadata["node_type"]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment