[benchmark] fix multi process issue (#5608)

benchmarks/benchmarks/multigpu/bench_multigpu_rgcn.py

@@ -27,14 +27,6 @@ from torch.utils.data import DataLoader
 
 from .. import utils
 
-# import sys
-# import os
-# dir_path = Path(os.path.dirname(__file__))
-# sys.path.insert(0, dir_path.parent)
-
-
-# import utils
-
 
 class EntityClassify(nn.Module):
     def __init__(
@@ -475,7 +467,7 @@ def track_time(data, dgl_sparse):
     n_gpus = len(devices)
     n_cpus = mp.cpu_count()
 
-    ctx = mp.get_context("fork")
+    ctx = mp.get_context("spawn")
     queue = ctx.Queue()
     procs = []
     num_train_seeds = train_idx.shape[0]

benchmarks/benchmarks/multigpu/bench_multigpu_sage.py

@@ -53,8 +53,7 @@ def load_subtensor(nfeat, labels, seeds, input_nodes, dev_id):
 
 
 # Entry point
-
-
+@utils.thread_wrapped_func
 def run(result_queue, proc_id, n_gpus, args, devices, data):
     dev_id = devices[proc_id]
     timing_records = []
@@ -174,11 +173,12 @@ def track_time(data):
     # Pack data
     data = n_classes, train_g, val_g, test_g
 
-    result_queue = mp.Queue()
+    ctx = mp.get_context("spawn")
+    result_queue = ctx.Queue()
     procs = []
     for proc_id in range(n_gpus):
-        p = mp.Process(
-            target=utils.thread_wrapped_func(run),
+        p = ctx.Process(
+            target=run,
             args=(result_queue, proc_id, n_gpus, args, devices, data),
         )
         p.start()

benchmarks/run.sh

@@ -19,6 +19,8 @@ echo "DGL_BENCH_DEVICE=$DGL_BENCH_DEVICE"
 pushd $ROOT/benchmarks
 cat asv.conf.json
 asv machine --yes
-asv run --launch-method=spawn -e -v
+# If --launch-method is specified as 'spawn', multigpu tests will crash with
+# "No module named 'benchmarks' is found".
+asv run -e -v
 asv publish
 popd