[DoNotMerge] Integrate GraphBolt with DistDGL

Jenkinsfile

@@ -580,7 +580,6 @@ pipeline {
                   steps {
                     unit_distributed_linux('pytorch', 'cpu')
                   }
-                  when { expression { false } }
                 }
               }
               post {

check_mem_footprint.py

@@ -0,0 +1,52 @@
+import dgl
+from dgl.distributed import load_partition
+import psutil
+import os
+import argparse
+import gc
+
+parser = argparse.ArgumentParser(description="check memory footprint")
+parser.add_argument(
+    "--part_config",
+    type=str,
+    help="partition config file",
+)
+parser.add_argument(
+    "--graphbolt",
+    action="store_true",
+    help="use graphbolt",
+)
+parser.add_argument(
+    "--part_id",
+    type=int,
+    help="partition id",
+)
+
+args = parser.parse_args()
+
+use_graphbolt = args.graphbolt
+part_id = args.part_id
+
+prev_rss = psutil.Process(os.getpid()).memory_info().rss
+(
+    client_g,
+    _,
+    _,
+    gpb,
+    graph_name,
+    ntypes,
+    etypes,
+) = load_partition(
+    args.part_config,
+    part_id,
+    load_feats=False,
+    use_graphbolt=use_graphbolt,
+)
+if not use_graphbolt:
+    graph_format=("csc")
+    client_g = client_g.formats(graph_format)
+    client_g.create_formats_()
+new_rss = psutil.Process(os.getpid()).memory_info().rss
+print(f"[PartID_{part_id}] Loaded {graph_name} with use_graphbolt[{use_graphbolt}] in size[{(new_rss - prev_rss)/1024/1024 : .0f} MB]")
+client_g = None
+gc.collect()

dt.py

@@ -0,0 +1,101 @@
+import dgl
+import dgl.graphbolt as gb
+import numpy as np
+import torch as th
+
+# [TODO][P0] Set up distributed environment.
+
+"""
+num_trainers = 8
+num_servers = 4
+num_samplers = 0
+part_config = ./ogbn-products.json
+ip_config = ./ip_config.txt
+"""
+
+args = {}
+
+# Initialize distributed environment
+dgl.distributed.initialize(args.ip_config)
+th.distributed.init_process_group(backend=args.backend)
+# [TODO][P0] Convert dgl partitioned graphs to graphbolt.CSCSamplingGraph.
+#           done@2023-10-23 16:49:00
+#           see details in: https://github.com/Rhett-Ying/dgl/commits/gb_distdgl
+#                           ddce1d42de016be040cd0f8a5e71f2a10148de82
+'''
+In [1]: part_config='/home/ubuntu/workspace/dgl_2/data/ogbn-mag.json'
+In [3]: dgl.distributed.convert_dgl_partition_to_csc_sampling_graph(part_config, store_orig_nids=True)
+In [7]: !ls data/part0 -lh
+total 1.1G
+-rw-rw-r-- 1 ubuntu ubuntu 207M Oct 23 08:44 csc_sampling_graph.tar
+-rw-rw-r-- 1 ubuntu ubuntu 694M Oct 23 02:47 graph.dgl
+
+In [8]: !ls data/part1 -lh
+total 1.1G
+-rw-rw-r-- 1 ubuntu ubuntu 202M Oct 23 08:44 csc_sampling_graph.tar
+-rw-rw-r-- 1 ubuntu ubuntu 678M Oct 23 02:47 graph.dgl
+'''
+
+# [TODO][P0] Load `CSCSamplingGraph` into `DistGraph`.
+#           done@2023-10-24 15:10:00
+#           see details in: https://github.com/Rhett-Ying/dgl/commits/gb_distdgl
+#                           222dd2bd51084cc4f242148b0a7e6e5d91e0ae80
+## NID/EIDs are required.
+g = dgl.distributed.DistGraph(args.graph_name, part_config=args.part_config)
+
+# Generate train/val/test splits
+##############
+# train/val/test splits could be generated offline, then `train/val/test_masks`
+#   could be offloaded.
+# No change is required as `node_split` requires graph parition book and
+#   masks only.
+# This should be part of `OnDiskDataset::TVT`.
+# [TODO][P1]: Add a standalone API to generate train/val/test splits.
+##############
+gpb = g.get_partition_book()
+train_nids = dgl.distributed.node_split(g.ndata["train_masks"], gpb)
+val_nids = dgl.distributed.node_split(g.ndata["val_masks"], gpb)
+test_nids = dgl.distributed.node_split(g.ndata["test_masks"], gpb)
+all_nids = dgl.distributed.node_split(th.arange(g.num_nodes()), gpb)
+
+# [TODO][P2] How to handle feature data such as 'feat', 'mask'?
+# Just use `g.ndata['feat']` for now. As no more memory could be offloaded.
+# GB: feat_data = gb.OnDiskDataset().feature
+# DistDGL: feat_data = g.ndata['feat'] # DistTensor
+
+
+# Train.
+##############
+# GraphBolt version
+# [TODO][P0] Add `gb.distributed_sample_neighbor` API.
+# [TODO][P0] `remote_sample_neighbor()` returns original global node pairs + eids.
+# [TODO][P0] Upldate `dgl.distributed.merge_graphs` API.
+#     https://github.com/dmlc/dgl/blob/7439b7e73bdb85b4285ab01f704ac5a4f77c927e/python/dgl/distributed/graph_services.py#L440.
+##############
+"""
+datapipe = gb.ItemSampler(item_set, batch_size=batch_size, shuffle=shuffle)
+datapipe = datapipe.sample_neighbor(g._graph, fanouts=fanouts)
+datapipe = datapipe.to_dgl()
+device = th.device("cpu")
+datapipe = datapipe.copy_to(device)
+data_loader = gb.MultiProcessDataLoader(datapipe, num_workers=num_workers)
+"""
+sampler = dgl.dataloading.NeighborSampler([25, 10])
+train_dataloader = dgl.distributed.DistDataLoader(
+    g, train_nids, sampler=sampler, batch_size=args.batch_size, shuffle=True
+)
+model = None
+for mini_batch in train_dataloader:
+    in_feats = g.ndata["feat"][mini_batch.input_nodes]
+    labels = g.ndata["label"][mini_batch.output_nodes]
+    _ = model(mini_batch, in_feats)
+
+# Evaluate.
+model.eval()
+sampler = dgl.dataloading.NeighborSampler([-1])
+val_dataloader = dgl.distributed.DistDataLoader(
+    g, val_nids, sampler=sampler, batch_size=args.batch_size, shuffle=False
+)
+test_dataloader = dgl.distributed.DistDataLoader(
+    g, test_nids, sampler=sampler, batch_size=args.batch_size, shuffle=False
+)

examples/distributed/graphsage/README.md

@@ -1,3 +1,83 @@
+## DistDGL with GraphBolt(Homograph + Node Classification)
+
+### How to partition graph
+
+#### Partition from original dataset with `dgl.distributed.partition_graph()`
+
+```
+DGL_HOME=/home/ubuntu/workspace/dgl_2 DGL_LIBRARY_PATH=$DGL_HOME/build PYTHONPATH=tests:$DGL_HOME/python:tests/python/pytorch/graphbolt:$PYTHONPATH python3 examples/distributed/graphsage/partition_graph.py --dataset ogbn-products --num_parts 2 --balance_train --balance_edges --graphbolt
+```
+
+#### Convert existing partitions into GraphBolt formats
+
+```
+DGL_LIBRARY_PATH=$DGL_HOME/build PYTHONPATH=tests:$DGL_HOME/python:tests/python/pytorch/graphbolt:$PYTHONPATH python3 -c "from dgl.distributed import convert_dgl_partition_to_csc_sampling_graph as f;f('data/ogbn-products.json')"
+```
+
+#### Partition sizes compared between GraphBolt and DistDGL
+
+`csc_sampling_graph.tar` is the GraphBolt partitions.
+`graph.dgl` is the original DistDGL partitions, namely, DGLGraph.
+
+###### ogbn-products
+homogeneous, ~2.4M nodes, ~123.7M edges(reverse edges are added), 2 parts.
+
+| DGL(GB) | GraphBolt w/o EIDs(MB) | GraphBolt w/ EIDs(MB) |
+| --- | ------------------ | ----------------- |
+| 1.6/1.7 | 258/272 | 502/530 |
+
+```
+-rw-rw-r-- 1 ubuntu ubuntu 258M Oct 31 01:56 homo_data/part0/csc_sampling_graph.tar
+-rw-rw-r-- 1 ubuntu ubuntu 502M Oct 31 04:45 homo_data/part0/csc_sampling_graph_eids.tar
+-rw-rw-r-- 1 ubuntu ubuntu   24 Oct 31 00:51 homo_data/part0/edge_feat.dgl
+-rw-rw-r-- 1 ubuntu ubuntu 1.6G Oct 31 00:51 homo_data/part0/graph.dgl
+-rw-rw-r-- 1 ubuntu ubuntu 501M Oct 31 00:51 homo_data/part0/node_feat.dgl
+-rw-rw-r-- 1 ubuntu ubuntu 272M Oct 31 01:56 homo_data/part1/csc_sampling_graph.tar
+-rw-rw-r-- 1 ubuntu ubuntu 530M Oct 31 04:45 homo_data/part1/csc_sampling_graph_eids.tar
+-rw-rw-r-- 1 ubuntu ubuntu   24 Oct 31 00:51 homo_data/part1/edge_feat.dgl
+-rw-rw-r-- 1 ubuntu ubuntu 1.7G Oct 31 00:51 homo_data/part1/graph.dgl
+-rw-rw-r-- 1 ubuntu ubuntu 460M Oct 31 00:51 homo_data/part1/node_feat.dgl
+```
+
+### Train with GraphBolt partitions
+just append `--graphbolt`.
+
+```
+python3 /home/ubuntu/workspace/dgl_2/tools/launch.py \
+    --workspace /home/ubuntu/workspace/dgl_2/examples/distributed/graphsage/ \
+    --num_trainers 4 \
+    --num_servers 2 \
+    --num_samplers 0 \
+    --part_config /home/ubuntu/workspace/dgl_2/homo_data/ogbn-products.json \
+    --ip_config /home/ubuntu/workspace/ip_config.txt \
+    "DGL_LIBRARY_PATH=/home/ubuntu/workspace/dgl_2/build PYTHONPATH=tests:/home/ubuntu/workspace/dgl_2/python:tests/python/pytorch/graphbolt:$PYTHONPATH python3 node_classification.py --graph_name ogbn-products --ip_config /home/ubuntu/workspace/ip_config.txt --num_epochs 3 --eval_every 2 --graphbolt"
+```
+
+#### Results
+`g4dn.metal` x 2, `ogbn-products`.
+
+DistDGL with GraphBolt takes less time for sampling(from **1.8283s** to **1.4470s**) and for whole epoch(from **4.9259s** to **4.4898s**) while keeping comparable accuracies in validation and test.
+
+##### DistDGL
+
+```
+Part 0, Epoch Time(s): 4.9648, sample+data_copy: 1.8283, forward: 0.2912, backward: 1.1307, update: 0.0232, #seeds: 24577, #inputs: 4136843
+
+Summary of node classification(GraphSAGE): GraphName ogbn-products | TrainEpochTime(mean) 4.9259 | TestAccuracy 0.6213
+```
+
+##### DistDGL with GraphBolt
+
+```
+Part 0, Epoch Time(s): 4.4826, sample+data_copy: 1.4470, forward: 0.2517, backward: 0.9081, update: 0.0175, #seeds: 24577, #inputs: 41369
+80
+
+Summary of node classification(GraphSAGE): GraphName ogbn-products | TrainEpochTime(mean) 4.4898 | TestAccuracy 0.6174
+```
+
+---------------------------------------
+
+
 ## Distributed training
 
 This is an example of training GraphSage in a distributed fashion. Before training, please install some python libs by pip:

examples/distributed/graphsage/dgl_cmd.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+python3 /home/ubuntu/workspace/dgl_2/tools/launch.py \
+    --workspace /home/ubuntu/workspace/dgl_2/examples/distributed/graphsage/ \
+    --num_trainers 4 \
+    --num_servers 2 \
+    --num_samplers 0 \
+    --part_config /home/ubuntu/workspace/dgl_2/homo_data/ogbn-products.json \
+    --ip_config /home/ubuntu/workspace/ip_config.txt \
+    "DGL_LIBRARY_PATH=/home/ubuntu/workspace/dgl_2/build PYTHONPATH=tests:/home/ubuntu/workspace/dgl_2/python:tests/python/pytorch/graphbolt:$PYTHONPATH python3 node_classification.py --graph_name ogbn-products --ip_config /home/ubuntu/workspace/ip_config.txt --num_epochs 3 --eval_every 2"

examples/distributed/graphsage/gb_cmd.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+python3 /home/ubuntu/workspace/dgl_2/tools/launch.py \
+    --workspace /home/ubuntu/workspace/dgl_2/examples/distributed/graphsage/ \
+    --num_trainers 4 \
+    --num_servers 2 \
+    --num_samplers 0 \
+    --part_config /home/ubuntu/workspace/dgl_2/homo_data/ogbn-products.json \
+    --ip_config /home/ubuntu/workspace/ip_config.txt \
+    "DGL_LIBRARY_PATH=/home/ubuntu/workspace/dgl_2/build PYTHONPATH=tests:/home/ubuntu/workspace/dgl_2/python:tests/python/pytorch/graphbolt:$PYTHONPATH python3 node_classification.py --graph_name ogbn-products --ip_config /home/ubuntu/workspace/ip_config.txt --num_epochs 3 --eval_every 2 --graphbolt"

examples/distributed/graphsage/node_classification.py

@@ -66,7 +66,7 @@ def forward(self, blocks, x):
                 h = self.dropout(h)
         return h
 
-    def inference(self, g, x, batch_size, device):
+    def inference(self, g, x, batch_size, device, use_graphbolt):
         """
         Distributed layer-wise inference with the GraphSAGE model on full
         neighbors.
@@ -116,6 +116,7 @@ def inference(self, g, x, batch_size, device):
                 batch_size=batch_size,
                 shuffle=False,
                 drop_last=False,
+                use_graphbolt=use_graphbolt,
             )
 
             for input_nodes, output_nodes, blocks in tqdm.tqdm(dataloader):
@@ -155,7 +156,7 @@ def compute_acc(pred, labels):
     return (th.argmax(pred, dim=1) == labels).float().sum() / len(pred)
 
 
-def evaluate(model, g, inputs, labels, val_nid, test_nid, batch_size, device):
+def evaluate(model, g, inputs, labels, val_nid, test_nid, batch_size, device, use_graphbolt):
     """
     Evaluate the model on the validation and test set.
 
@@ -187,7 +188,7 @@ def evaluate(model, g, inputs, labels, val_nid, test_nid, batch_size, device):
     """
     model.eval()
     with th.no_grad():
-        pred = model.inference(g, inputs, batch_size, device)
+        pred = model.inference(g, inputs, batch_size, device, use_graphbolt)
     model.train()
     return compute_acc(pred[val_nid], labels[val_nid]), compute_acc(
         pred[test_nid], labels[test_nid]
@@ -219,6 +220,7 @@ def run(args, device, data):
         batch_size=args.batch_size,
         shuffle=True,
         drop_last=False,
+        use_graphbolt=args.graphbolt,
     )
     model = DistSAGE(
         in_feats,
@@ -325,6 +327,7 @@ def run(args, device, data):
                 test_nid,
                 args.batch_size_eval,
                 device,
+                args.graphbolt,
             )
             print(
                 f"Part {g.rank()}, Val Acc {val_acc:.4f}, "
@@ -338,13 +341,16 @@ def main(args):
     """
     Main function.
     """
+    if args.graphbolt:
+        print("DistDGL with GraphBolt...")
     host_name = socket.gethostname()
     print(f"{host_name}: Initializing DistDGL.")
-    dgl.distributed.initialize(args.ip_config)
+    dgl.distributed.initialize(args.ip_config, use_graphbolt=args.graphbolt)
     print(f"{host_name}: Initializing PyTorch process group.")
     th.distributed.init_process_group(backend=args.backend)
     print(f"{host_name}: Initializing DistGraph.")
-    g = dgl.distributed.DistGraph(args.graph_name, part_config=args.part_config)
+    g = dgl.distributed.DistGraph(args.graph_name, part_config=args.part_config,
+                                  use_graphbolt=args.graphbolt)
     print(f"Rank of {host_name}: {g.rank()}")
 
     # Split train/val/test IDs for each trainer.
@@ -415,6 +421,12 @@ def main(args):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Distributed GraphSAGE.")
+    parser.add_argument(
+        "--graphbolt",
+        default=False,
+        action="store_true",
+        help="train with GraphBolt",
+    )
     parser.add_argument("--graph_name", type=str, help="graph name")
     parser.add_argument(
         "--ip_config", type=str, help="The file for IP configuration"

examples/distributed/graphsage/partition_graph.py

@@ -59,6 +59,11 @@ def load_ogb(name, root="dataset"):
     argparser.add_argument(
         "--part_method", type=str, default="metis", help="the partition method"
     )
+    argparser.add_argument(
+        "--graphbolt",
+        action="store_true",
+        help="convert DGL to GraphBolt partitions.",
+    )
     argparser.add_argument(
         "--balance_train",
         action="store_true",
@@ -127,4 +132,5 @@ def load_ogb(name, root="dataset"):
         balance_ntypes=balance_ntypes,
         balance_edges=args.balance_edges,
         num_trainers_per_machine=args.num_trainers_per_machine,
+        use_graphbolt=args.graphbolt,
     )