Change the absolute path to relative one in part_config.json (dmlc#…

…2690) * Update partition.py * fix lint. Co-authored-by: Da Zheng <[email protected]>
XINyexun · Feb 25, 2021 · 7ab3d99 · 7ab3d99
1 parent 1686106
commit 7ab3d99
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/python/dgl/distributed/partition.py b/python/dgl/distributed/partition.py
@@ -79,16 +79,19 @@ def load_partition(part_config, part_id):
     List[str]
         The edge types
     '''
+    config_path = os.path.dirname(part_config)
+    relative_to_config = lambda path: os.path.join(config_path, path)
+
     with open(part_config) as conf_f:
         part_metadata = json.load(conf_f)
     assert 'part-{}'.format(part_id) in part_metadata, "part-{} does not exist".format(part_id)
     part_files = part_metadata['part-{}'.format(part_id)]
     assert 'node_feats' in part_files, "the partition does not contain node features."
     assert 'edge_feats' in part_files, "the partition does not contain edge feature."
     assert 'part_graph' in part_files, "the partition does not contain graph structure."
-    node_feats = load_tensors(part_files['node_feats'])
-    edge_feats = load_tensors(part_files['edge_feats'])
-    graph = load_graphs(part_files['part_graph'])[0][0]
+    node_feats = load_tensors(relative_to_config(part_files['node_feats']))
+    edge_feats = load_tensors(relative_to_config(part_files['edge_feats']))
+    graph = load_graphs(relative_to_config(part_files['part_graph']))[0][0]
     # In the old format, the feature name doesn't contain node/edge type.
     # For compatibility, let's add node/edge types to the feature names.
     node_feats1 = {}
@@ -687,9 +690,10 @@ def get_homogeneous(g, balance_ntypes):
         node_feat_file = os.path.join(part_dir, "node_feat.dgl")
         edge_feat_file = os.path.join(part_dir, "edge_feat.dgl")
         part_graph_file = os.path.join(part_dir, "graph.dgl")
-        part_metadata['part-{}'.format(part_id)] = {'node_feats': node_feat_file,
-                                                    'edge_feats': edge_feat_file,
-                                                    'part_graph': part_graph_file}
+        part_metadata['part-{}'.format(part_id)] = {
+            'node_feats': os.path.relpath(node_feat_file, out_path),
+            'edge_feats': os.path.relpath(edge_feat_file, out_path),
+            'part_graph': os.path.relpath(part_graph_file, out_path)}
         os.makedirs(part_dir, mode=0o775, exist_ok=True)
         save_tensors(node_feat_file, node_feats)
         save_tensors(edge_feat_file, edge_feats)