wehub-resource-sync
150 lines
5.1 KiB
Python
150 lines
5.1 KiB
Python
"""Copy the partitions to a cluster of machines."""
|
|
import argparse
|
|
import copy
|
|
import json
|
|
import logging
|
|
import os
|
|
import signal
|
|
import stat
|
|
import subprocess
|
|
import sys
|
|
|
|
|
|
def copy_file(file_name, ip, workspace, param=""):
|
|
print("copy {} to {}".format(file_name, ip + ":" + workspace + "/"))
|
|
cmd = "scp " + param + " " + file_name + " " + ip + ":" + workspace + "/"
|
|
subprocess.check_call(cmd, shell=True)
|
|
|
|
|
|
def exec_cmd(ip, cmd):
|
|
cmd = "ssh -o StrictHostKeyChecking=no " + ip + " '" + cmd + "'"
|
|
subprocess.check_call(cmd, shell=True)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Copy data to the servers.")
|
|
parser.add_argument(
|
|
"--workspace",
|
|
type=str,
|
|
required=True,
|
|
help="Path of user directory of distributed tasks. \
|
|
This is used to specify a destination location where \
|
|
data are copied to on remote machines.",
|
|
)
|
|
parser.add_argument(
|
|
"--rel_data_path",
|
|
type=str,
|
|
required=True,
|
|
help="Relative path in workspace to store the partition data.",
|
|
)
|
|
parser.add_argument(
|
|
"--part_config",
|
|
type=str,
|
|
required=True,
|
|
help="The partition config file. The path is on the local machine.",
|
|
)
|
|
parser.add_argument(
|
|
"--script_folder",
|
|
type=str,
|
|
required=True,
|
|
help="The folder contains all the user code scripts.",
|
|
)
|
|
parser.add_argument(
|
|
"--ip_config",
|
|
type=str,
|
|
required=True,
|
|
help="The file of IP configuration for servers. \
|
|
The path is on the local machine.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
hosts = []
|
|
with open(args.ip_config) as f:
|
|
for line in f:
|
|
res = line.strip().split(" ")
|
|
ip = res[0]
|
|
hosts.append(ip)
|
|
|
|
# We need to update the partition config file so that the paths are relative to
|
|
# the workspace in the remote machines.
|
|
with open(args.part_config) as conf_f:
|
|
part_metadata = json.load(conf_f)
|
|
tmp_part_metadata = copy.deepcopy(part_metadata)
|
|
num_parts = part_metadata["num_parts"]
|
|
assert num_parts == len(
|
|
hosts
|
|
), "The number of partitions needs to be the same as the number of hosts."
|
|
graph_name = part_metadata["graph_name"]
|
|
node_map = part_metadata["node_map"]
|
|
edge_map = part_metadata["edge_map"]
|
|
if not isinstance(node_map, dict):
|
|
assert (
|
|
node_map[-4:] == ".npy"
|
|
), "node map should be stored in a NumPy array."
|
|
tmp_part_metadata["node_map"] = "{}/{}/node_map.npy".format(
|
|
args.workspace, args.rel_data_path
|
|
)
|
|
if not isinstance(edge_map, dict):
|
|
assert (
|
|
edge_map[-4:] == ".npy"
|
|
), "edge map should be stored in a NumPy array."
|
|
tmp_part_metadata["edge_map"] = "{}/{}/edge_map.npy".format(
|
|
args.workspace, args.rel_data_path
|
|
)
|
|
|
|
for part_id in range(num_parts):
|
|
part_files = tmp_part_metadata["part-{}".format(part_id)]
|
|
part_files["edge_feats"] = "{}/part{}/edge_feat.dgl".format(
|
|
args.rel_data_path, part_id
|
|
)
|
|
part_files["node_feats"] = "{}/part{}/node_feat.dgl".format(
|
|
args.rel_data_path, part_id
|
|
)
|
|
part_files["part_graph"] = "{}/part{}/graph.dgl".format(
|
|
args.rel_data_path, part_id
|
|
)
|
|
tmp_part_config = "/tmp/{}.json".format(graph_name)
|
|
with open(tmp_part_config, "w") as outfile:
|
|
json.dump(tmp_part_metadata, outfile, sort_keys=True, indent=4)
|
|
|
|
# Copy ip config.
|
|
for part_id, ip in enumerate(hosts):
|
|
remote_path = "{}/{}".format(args.workspace, args.rel_data_path)
|
|
exec_cmd(ip, "mkdir -p {}".format(remote_path))
|
|
|
|
copy_file(args.ip_config, ip, args.workspace)
|
|
copy_file(
|
|
tmp_part_config,
|
|
ip,
|
|
"{}/{}".format(args.workspace, args.rel_data_path),
|
|
)
|
|
node_map = part_metadata["node_map"]
|
|
edge_map = part_metadata["edge_map"]
|
|
if not isinstance(node_map, dict):
|
|
copy_file(node_map, ip, tmp_part_metadata["node_map"])
|
|
if not isinstance(edge_map, dict):
|
|
copy_file(edge_map, ip, tmp_part_metadata["edge_map"])
|
|
remote_path = "{}/{}/part{}".format(
|
|
args.workspace, args.rel_data_path, part_id
|
|
)
|
|
exec_cmd(ip, "mkdir -p {}".format(remote_path))
|
|
|
|
part_files = part_metadata["part-{}".format(part_id)]
|
|
copy_file(part_files["node_feats"], ip, remote_path)
|
|
copy_file(part_files["edge_feats"], ip, remote_path)
|
|
copy_file(part_files["part_graph"], ip, remote_path)
|
|
# copy script folder
|
|
copy_file(args.script_folder, ip, args.workspace, "-r")
|
|
|
|
|
|
def signal_handler(signal, frame):
|
|
logging.info("Stop copying")
|
|
sys.exit(0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
fmt = "%(asctime)s %(levelname)s %(message)s"
|
|
logging.basicConfig(format=fmt, level=logging.INFO)
|
|
signal.signal(signal.SIGINT, signal_handler)
|
|
main()
|