[Distributed Example] Check args.num_servers with DGL_NUM_SERVER (dml…

…c#2203) * update * update * update
zeta1999 · Sep 16, 2020 · 3926024 · 3926024
1 parent b25bbe6
commit 3926024
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/examples/pytorch/graphsage/experimental/train_dist.py b/examples/pytorch/graphsage/experimental/train_dist.py
@@ -316,7 +316,9 @@ def main(args):
     parser.add_argument('--close_profiler', action='store_true', help='Close pyinstrument profiler')
     args = parser.parse_args()
     assert args.num_workers == int(os.environ.get('DGL_NUM_SAMPLER')), \
-    'The num_workers should be the same value with num_samplers.'
+    'The num_workers should be the same value with DGL_NUM_SAMPLER.'
+    assert args.num_servers == int(os.environ.get('DGL_NUM_SERVER')), \
+    'The num_servers should be the same value with DGL_NUM_SERVER.'
 
     print(args)
     main(args)
diff --git a/examples/pytorch/graphsage/experimental/train_dist_unsupervised.py b/examples/pytorch/graphsage/experimental/train_dist_unsupervised.py
@@ -486,7 +486,9 @@ def main(args):
         help="whether to remove edges during sampling")
     args = parser.parse_args()
     assert args.num_workers == int(os.environ.get('DGL_NUM_SAMPLER')), \
-    'The num_workers should be the same value with num_samplers.'
-
+    'The num_workers should be the same value with DGL_NUM_SAMPLER.'
+    assert args.num_servers == int(os.environ.get('DGL_NUM_SERVER')), \
+    'The num_servers should be the same value with DGL_NUM_SERVER.'
+
     print(args)
     main(args)