propagate ProcessGroup timeout to Store

Author: mrshenli

torch/distributed/distributed_c10d.py

@@ -351,7 +351,7 @@ def init_process_group(backend,
         elif world_size != -1:
             url += "?world_size={}".format(world_size)
 
-        store, rank, world_size = next(rendezvous(url))
+        store, rank, world_size = next(rendezvous(url, timeout=timeout))
         if backend == Backend.GLOO:
             _default_pg = ProcessGroupGloo(
                 store,

torch/distributed/rendezvous.py

@@ -8,6 +8,7 @@
 
 
 _rendezvous_handlers = {}
+_default_store_timeout = timedelta(minutes=5)
 
 
 def register_rendezvous_handler(scheme, handler):
@@ -53,7 +54,7 @@ def _rendezvous_error(msg):
     return ValueError("Error initializing torch.distributed using " + msg)
 
 
-def _file_rendezvous_handler(url):
+def _file_rendezvous_handler(url, timeout=_default_store_timeout):
     def _error(msg):
         return _rendezvous_error("file:// rendezvous: " + msg)
 
@@ -69,14 +70,14 @@ def _error(msg):
 
     rank = int(query["rank"])
     world_size = int(query["world_size"])
-    store = FileStore(path, world_size)
+    store = FileStore(path, world_size, timeout)
     yield (store, rank, world_size)
 
     # If this configuration is invalidated, there is nothing we can do about it
     raise RuntimeError("Unable to perform rerendezvous using file:// method")
 
 
-def _tcp_rendezvous_handler(url):
+def _tcp_rendezvous_handler(url, timeout=_default_store_timeout):
     def _error(msg):
         return _rendezvous_error("tcp:// rendezvous: " + msg)
 
@@ -92,14 +93,14 @@ def _error(msg):
     rank = int(query["rank"])
     world_size = int(query["world_size"])
     start_daemon = rank == 0
-    store = TCPStore(result.hostname, result.port, world_size, start_daemon)
+    store = TCPStore(result.hostname, result.port, world_size, start_daemon, timeout)
     yield (store, rank, world_size)
 
     # If this configuration is invalidated, there is nothing we can do about it
     raise RuntimeError("Unable to perform rerendezvous using tcp:// method")
 
 
-def _env_rendezvous_handler(url):
+def _env_rendezvous_handler(url, timeout=_default_store_timeout):
     def _error(msg):
         return _rendezvous_error("env:// rendezvous: " + msg)
 
@@ -140,7 +141,7 @@ def _env_error(var):
 
     # Now start the TCP store daemon on the rank 0
     start_daemon = rank == 0
-    store = TCPStore(master_addr, master_port, world_size, start_daemon)
+    store = TCPStore(master_addr, master_port, world_size, start_daemon, timeout)
     yield (store, rank, world_size)
 
     # If this configuration is invalidated, there is nothing we can do about it

torch/lib/c10d/FileStore.cpp

@@ -206,8 +206,11 @@ off_t refresh(
 
 } // namespace
 
-FileStore::FileStore(const std::string& path, int numWorkers)
-    : Store(),
+FileStore::FileStore(
+    const std::string& path,
+    int numWorkers,
+    std::chrono::milliseconds timeout)
+    : Store(timeout),
       path_(path),
       pos_(0),
       numWorkers_(numWorkers),

torch/lib/c10d/FileStore.hpp

@@ -10,7 +10,10 @@ namespace c10d {
 
 class FileStore : public Store {
  public:
-  explicit FileStore(const std::string& path, int numWorkers);
+  explicit FileStore(
+      const std::string& path,
+      int numWorkers,
+      std::chrono::milliseconds timeout=kDefaultTimeout);
 
   virtual ~FileStore();
 

torch/lib/c10d/PrefixStore.hpp

@@ -6,7 +6,9 @@ namespace c10d {
 
 class PrefixStore : public Store {
  public:
-  explicit PrefixStore(const std::string& prefix, Store& store);
+  explicit PrefixStore(
+      const std::string& prefix,
+      Store& store);
 
   virtual ~PrefixStore(){};
 

torch/lib/c10d/ProcessGroupGloo.cpp

@@ -66,7 +66,7 @@ class GlooStore : public ::gloo::rendezvous::Store {
   }
 
   void wait(const std::vector<std::string>& keys) override {
-    store_->wait(keys, Store::kDefaultTimeout);
+    store_->wait(keys);
   }
 
   void wait(

torch/lib/c10d/Store.hpp

@@ -15,7 +15,8 @@ class Store {
   static constexpr std::chrono::milliseconds kNoTimeout =
       std::chrono::milliseconds::zero();
 
-  Store() : timeout_(kDefaultTimeout) {}
+  Store(std::chrono::milliseconds timeout=kDefaultTimeout)
+      : timeout_(timeout) {}
 
   virtual ~Store();
 

torch/lib/c10d/TCPStore.cpp

@@ -279,8 +279,10 @@ TCPStore::TCPStore(
     const std::string& masterAddr,
     PortType masterPort,
     int numWorkers,
-    bool isServer)
-    : isServer_(isServer),
+    bool isServer,
+    std::chrono::milliseconds timeout)
+    : Store(timeout),
+      isServer_(isServer),
       tcpStoreAddr_(masterAddr),
       tcpStorePort_(masterPort),
       numWorkers_(numWorkers),

torch/lib/c10d/TCPStore.hpp

@@ -49,7 +49,8 @@ class TCPStore : public Store {
       const std::string& masterAddr,
       PortType masterPort,
       int numWorkers,
-      bool isServer = false);
+      bool isServer = false,
+      std::chrono::milliseconds timeout=kDefaultTimeout);
 
   virtual ~TCPStore();