Unverified Commit 6bc82161 authored by Hongzhi (Steve), Chen's avatar Hongzhi (Steve), Chen Committed by GitHub
Browse files

[Misc] auto-format tools. (#5321)


Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-28-63.ap-northeast-1.compute.internal>
parent 63ac788f
...@@ -533,7 +533,6 @@ def exchange_features( ...@@ -533,7 +533,6 @@ def exchange_features(
# To iterate over the node_types and associated node_features # To iterate over the node_types and associated node_features
for feat_key, type_info in feature_tids.items(): for feat_key, type_info in feature_tids.items():
# 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]
......
...@@ -60,7 +60,6 @@ class DistLookupService: ...@@ -60,7 +60,6 @@ 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( logging.info(
f"[Rank: {rank}] Reading file: {os.path.join(input_dir, filename)}" f"[Rank: {rank}] Reading file: {os.path.join(input_dir, filename)}"
......
...@@ -4,6 +4,8 @@ import os ...@@ -4,6 +4,8 @@ import os
import sys import sys
from pathlib import Path from pathlib import Path
import array_readwriter
import constants import constants
import numpy as np import numpy as np
...@@ -13,7 +15,6 @@ import pyarrow.parquet as pq ...@@ -13,7 +15,6 @@ import pyarrow.parquet as pq
import torch import torch
import torch.distributed as dist import torch.distributed as dist
from utils import get_idranges, get_node_types, read_json from utils import get_idranges, get_node_types, read_json
import array_readwriter
def get_proc_info(): def get_proc_info():
...@@ -79,7 +80,6 @@ def gen_edge_files(schema_map, output): ...@@ -79,7 +80,6 @@ def gen_edge_files(schema_map, output):
edge_files = [] edge_files = []
num_parts = len(schema_map[constants.STR_NUM_EDGES_PER_CHUNK][0]) num_parts = len(schema_map[constants.STR_NUM_EDGES_PER_CHUNK][0])
for etype_name, etype_info in edge_data.items(): for etype_name, etype_info in edge_data.items():
edges_format = etype_info[constants.STR_FORMAT][constants.STR_NAME] edges_format = etype_info[constants.STR_FORMAT][constants.STR_NAME]
edge_data_files = etype_info[constants.STR_DATA] edge_data_files = etype_info[constants.STR_DATA]
...@@ -170,14 +170,19 @@ def read_node_features(schema_map, tgt_ntype_name, feat_names, input_dir): ...@@ -170,14 +170,19 @@ def read_node_features(schema_map, tgt_ntype_name, feat_names, input_dir):
if feat_name in feat_names: if feat_name in feat_names:
feat_data_fname = feat_data[constants.STR_DATA][rank] feat_data_fname = feat_data[constants.STR_DATA][rank]
if not os.path.isabs(feat_data_fname): if not os.path.isabs(feat_data_fname):
feat_data_fname = os.path.join(input_dir, feat_data_fname) feat_data_fname = os.path.join(
input_dir, feat_data_fname
)
logging.info(f"Reading: {feat_data_fname}") logging.info(f"Reading: {feat_data_fname}")
file_suffix = Path(feat_data_fname).suffix file_suffix = Path(feat_data_fname).suffix
reader_fmt_meta = { reader_fmt_meta = {"name": file_suffix[1:]}
"name": file_suffix[1:] node_features[
} feat_name
node_features[feat_name] = array_readwriter.get_array_parser( ] = array_readwriter.get_array_parser(
**reader_fmt_meta).read(feat_data_fname) **reader_fmt_meta
).read(
feat_data_fname
)
return node_features return node_features
...@@ -248,8 +253,10 @@ def gen_node_weights_files(schema_map, input_dir, output): ...@@ -248,8 +253,10 @@ def gen_node_weights_files(schema_map, input_dir, output):
# Add train/test/validation masks if present. node-degree will be added when this file # Add train/test/validation masks if present. node-degree will be added when this file
# is read by ParMETIS to mimic the exisiting single process pipeline present in dgl. # is read by ParMETIS to mimic the exisiting single process pipeline present in dgl.
node_feats = read_node_features( node_feats = read_node_features(
schema_map, ntype_name, set(["train_mask", "val_mask", "test_mask"]), schema_map,
input_dir ntype_name,
set(["train_mask", "val_mask", "test_mask"]),
input_dir,
) )
for k, v in node_feats.items(): for k, v in node_feats.items():
assert sz == v.shape assert sz == v.shape
......
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