dist partition

tools/dispatch_data.py

@@ -75,7 +75,10 @@ def submit_jobs(args) -> str:
     argslist += "--log-level {} ".format(args.log_level)
     argslist += "--save-orig-nids " if args.save_orig_nids else ""
     argslist += "--save-orig-eids " if args.save_orig_eids else ""
-    argslist += "--use-graphbolt" if args.use_graphbolt else ""
+    argslist += "--use-graphbolt " if args.use_graphbolt else ""
+    argslist += "--store-inner-edge " if args.store_inner_edge else ""
+    argslist += "--store-inner-node " if args.store_inner_node else ""
+    argslist += "--store-eids " if args.store_eids else ""
     argslist += (
         f"--graph-formats {args.graph_formats} " if args.graph_formats else ""
     )
@@ -164,6 +167,26 @@ def main():
         action="store_true",
         help="Use GraphBolt for distributed train.",
     )
+    parser.add_argument(
+        "--store-inner-node",
+        action="store_true",
+        default=False,
+        help="Store inner nodes.",
+    )
+
+    parser.add_argument(
+        "--store-inner-edge",
+        action="store_true",
+        default=False,
+        help="Store inner edges.",
+    )
+
+    parser.add_argument(
+        "--store-eids",
+        action="store_true",
+        default=False,
+        help="Store edge IDs.",
+    )
     parser.add_argument(
         "--graph-formats",
         type=str,
@@ -175,7 +198,7 @@ def main():
     )
 
     args, _ = parser.parse_known_args()
-
+    assert args.store_inner_edge is True
     fmt = "%(asctime)s %(levelname)s %(message)s"
     logging.basicConfig(
         format=fmt,

tools/distpartitioning/data_shuffle.py

@@ -13,7 +13,7 @@
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
-from convert_partition import create_dgl_object, create_metadata_json
+from convert_partition import create_graph_object, create_metadata_json
 from dataset_utils import get_dataset
 from dist_lookup import DistLookupService
 from globalids import (
@@ -1323,7 +1323,7 @@ def prepare_local_data(src_data, local_part_id):
             etypes_map,
             orig_nids,
             orig_eids,
-        ) = create_dgl_object(
+        ) = create_graph_object(
             schema_map,
             rank + local_part_id * world_size,
             local_node_data,
@@ -1334,9 +1334,12 @@ def prepare_local_data(src_data, local_part_id):
                 schema_map[constants.STR_NUM_NODES_PER_TYPE],
             ),
             edge_typecounts,
-            params.save_orig_nids,
-            params.save_orig_eids,
-            params.use_graphbolt,
+            return_orig_nids=params.save_orig_nids,
+            return_orig_eids=params.save_orig_eids,
+            use_graphbolt=params.use_graphbolt,
+            store_inner_node=params.store_inner_node, 
+            store_inner_edge=params.store_inner_edge,
+            store_eids=params.store_eids,
         )
         sort_etypes = len(etypes_map) > 1
         local_node_features = prepare_local_data(