[parameter_server] fix deadlock in multiprocessing by using spawn (#1025

)
pytorch · Jul 18, 2022 · 7ed7ac7 · 7ed7ac7
1 parent 91ccd7a
commit 7ed7ac7
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/distributed/rpc/parameter_server/rpc_parameter_server.py b/distributed/rpc/parameter_server/rpc_parameter_server.py
@@ -266,6 +266,11 @@ def run_worker(rank, world_size, num_gpus, train_loader, test_loader):
     os.environ['MASTER_PORT'] = args.master_port
     processes = []
     world_size = args.world_size
+
+    # Note that Linux uses "fork" by default, which may cause deadlock.
+    # Besides, cuda doesn't support "fork" and Windows only supports "spawn"
+    mp.set_start_method("spawn")
+
     if args.rank == 0:
         p = mp.Process(target=run_parameter_server, args=(0, world_size))
         p.start()