poll for timed out futures in process group agent (pytorch#29601)

Summary:
Pull Request resolved: pytorch#29601

Follow up from pytorch#28392. Adds a background thread to `ProcessGroupAgent` that polls for timed out RPCs at a pre-set interval, and marks them as completed with a timeout exception if they have timed out. Also deletes the futures from the corresponding maps `futures_` and `futureTimeouts`. Unit tests are added to ensure that timed out RPCs are appropriately cleaned up.

Also adds a `shutdown` variable to process group agent to control the shutting down of this background thread, which can eventually be extended to use for controlling a clean shutdown of process group agent.
ghstack-source-id: 9417513

Test Plan: Added unit tests

Differential Revision: D18434215

fbshipit-source-id: c48abdb8759fe1447200ec66bb9d4b1c50ec4535

Author: rohan-varma

test/rpc_test.py

@@ -125,6 +125,10 @@ def my_function(a, b, c):
 def my_tensor_function(a, b):
     return a + b
 
+def my_sleep_func(seconds=1):
+    import time
+    time.sleep(seconds)
+
 
 def my_complex_tensor_function(list_input, tensor_class_input, dict_input):
     res = list_input[0]
@@ -1084,13 +1088,8 @@ def test_call_method_on_rref(self):
 
         self.assertEqual(result, sum(vals))
 
-    @dist_init
-    def test_get_default_rpc_timeout(self):
-        timeout = rpc.get_rpc_timeout()
-        self.assertEqual(timeout, rpc.constants.DEFAULT_RPC_TIMEOUT)
-
     @dist_init(setup_rpc=False)
-    def test_set_rpc_timeout(self):
+    def test_get_rpc_timeout(self):
         timeout = timedelta(seconds=1)
 
         # A new `RpcAgentOptions` is constructed
@@ -1110,6 +1109,35 @@ def test_set_rpc_timeout(self):
         self.assertEqual(timeout, set_timeout)
         rpc.join_rpc()
 
+    @dist_init
+    @requires_process_group_agent("PROCESS_GROUP rpc backend specific test, skip")
+    def test_rpc_timeouts(self):
+        dst_rank = (self.rank + 1) % self.world_size
+        rpc._set_rpc_timeout(timedelta(milliseconds=1))
+        # futures should time out and be marked with an exception indicating it as such.
+        futs = [rpc.rpc_async("worker{}".format(dst_rank), my_sleep_func, args=()) for _ in range(10)]
+        for fut in futs:
+            with self.assertRaisesRegex(RuntimeError, "RPC ran for more than"):
+                fut.wait()
+
+        # ensure that if a new timeout is set old futures don't time out but new ones do.
+        rpc._set_rpc_timeout(timedelta(seconds=200))
+        # create a longstanding RPC.
+        fut1 = rpc.rpc_async("worker{}".format(dst_rank), my_sleep_func, args=(1,))
+        # now, set a short timeout.
+        rpc._set_rpc_timeout(timedelta(milliseconds=1))
+        # f2 should time out, f should not.
+        fut2 = rpc.rpc_async("worker{}".format(dst_rank), my_sleep_func, args=(1,))
+        with self.assertRaises(RuntimeError):
+            fut2.wait()
+        fut1.wait()
+
+        # future should run to completion if the timeout is zero.
+        rpc._set_rpc_timeout(timedelta(seconds=0))
+        rpc.rpc_async("worker{}".format(dst_rank), my_sleep_func, args=()).wait()
+
+        # reset to default timeout so shutdown messages can process cleanly.
+        rpc._set_rpc_timeout(rpc.constants.DEFAULT_RPC_TIMEOUT)
 
     def test_requires_process_group_agent_decorator(self):
         @requires_process_group_agent("test_func did not run")

torch/csrc/distributed/rpc/init.cpp

@@ -200,6 +200,16 @@ PyObject* rpc_init(PyObject* /* unused */) {
             `datetime.timedelta` instance indicating the RPC timeout.
       )");
 
+  module.def(
+      "_set_rpc_timeout",
+      [](const std::chrono::milliseconds& rpcTimeout) {
+        RpcAgent::getDefaultRpcAgent()->setRpcTimeout(rpcTimeout);
+      },
+      R"(
+          Set the timeout for all RPCs. If an RPC is not completed within this
+          time, an exception indicating it has timed out will be raised.
+      )");
+
   Py_RETURN_TRUE;
 }
 

torch/csrc/distributed/rpc/process_group_agent.cpp

@@ -84,6 +84,9 @@ std::vector<int64_t> ProcessGroupAgent::MessageCounter::snapshot() {
 
 ////////////////////////  ProcessGroupAgent  /////////////////////////////////
 
+const std::chrono::milliseconds INFINITE_TIMEOUT =
+    std::chrono::milliseconds::max();
+
 void ProcessGroupAgent::collectNames() {
   const std::string& workerName = workerInfo_.name_;
   const auto worldSize = pg_->getSize();
@@ -124,6 +127,7 @@ ProcessGroupAgent::ProcessGroupAgent(
           WorkerInfo(std::move(workerName), pg->getRank()),
           c10::guts::make_unique<RequestCallbackImpl>(),
           rpcTimeout),
+      shutdown_{false},
       pg_(std::move(pg)),
       sendCounts_(pg_->getSize()),
       recvCounts_(pg_->getSize()),
@@ -182,7 +186,12 @@ void ProcessGroupAgent::join() {
   //    feed it a message or kill the thread.
   // 2. A GLOO process cannot send message to itself. (there is an ongoing
   //    effort to fix this problem).
+  shutdown_.store(true);
   sync();
+  // This is needed in case no futures were created, otherwise the future
+  // timeout watchdog would sleep forever.
+
+  futureTimeoutCV_.notify_one();
   std::unique_lock<std::mutex> lock(futureMutex_);
   futureCV_.wait(
       lock, [this] { return futures_.empty() && futureTimeouts_.empty(); });
@@ -193,6 +202,7 @@ void ProcessGroupAgent::join() {
       SendWork(allWorkerInfo_[dst], Message({}, {}, MessageType::SHUTDOWN)));
   threadPool_.waitWorkComplete();
   listenerThread_.join();
+  futureTimeoutThread_.join();
   PythonRpcHandler::getInstance().cleanup();
 }
 
@@ -260,6 +270,8 @@ void ProcessGroupAgent::sync() {
 
 void ProcessGroupAgent::start() {
   listenerThread_ = std::thread(&ProcessGroupAgent::listenLoop, this);
+  futureTimeoutThread_ =
+      std::thread(&ProcessGroupAgent::pollTimedOutRPCs, this);
 }
 
 std::shared_ptr<FutureMessage> ProcessGroupAgent::send(
@@ -281,9 +293,22 @@ std::shared_ptr<FutureMessage> ProcessGroupAgent::send(
             std::chrono::steady_clock::now().time_since_epoch());
     {
       std::lock_guard<std::mutex> lock{futureMutex_};
-      futures_[requestId] = std::make_pair(future, futureStartTime);
+      // Set infinite timeout if specified.
+      auto timeout = rpcTimeout_.load();
+      if (timeout.count() == 0) {
+        timeout = INFINITE_TIMEOUT;
+      }
+      auto futureInfo = FutureInfo(future, futureStartTime, to.id_, timeout);
+      futures_[requestId] = futureInfo;
+      auto rpcEndTime = getRPCEndTime(futureInfo);
       // insert future into timeouts map to keep track of its timeout
-      futureTimeouts_[futureStartTime].push_back(requestId);
+      futureTimeouts_[rpcEndTime].push_back(requestId);
+      // Signal the watchdog to monitor future timeouts if this is the first
+      // future created or if an RPC with a shorter TTL has been created.
+      if (futures_.size() == 1 ||
+          futureTimeouts_.begin()->first == rpcEndTime) {
+        futureTimeoutCV_.notify_one();
+      }
     }
     message.setId(requestId);
   } else {
@@ -389,27 +414,35 @@ void ProcessGroupAgent::enqueueRecv(RecvWork work) {
         } else if (message.isResponse()) {
           auto id = message.id();
           std::shared_ptr<FutureMessage> fm = nullptr;
-          std::chrono::milliseconds futureStartTime;
-          {
-            std::lock_guard<std::mutex> lock{futureMutex_};
-            std::tie(fm, futureStartTime) = futures_[id];
-          }
-          // Not holding lock on markCompleted as this could run callbacks that
-          // call agent_->send
-          fm->markCompleted(std::move(message));
           {
             std::lock_guard<std::mutex> lock{futureMutex_};
+            const auto& futureInfo = futures_.find(id);
+            if (futureInfo == futures_.end()) {
+              // Received a completion for a timed out future, drop the recv.
+              // RecvCounts will not be incremented here, it will be incremented
+              // by the sender who has determined the future has timed out.
+              return;
+            }
+
+            fm = futureInfo->second.future_;
+            auto rpcEndTime = getRPCEndTime(futureInfo->second);
             futures_.erase(id);
             // look up the corresponding future by its time out and request ID,
             // and remove it from the timeouts map
-            auto& futuresAtTime = futureTimeouts_[futureStartTime];
-            futuresAtTime.erase(
-                std::find(futuresAtTime.begin(), futuresAtTime.end(), id));
-            if (futuresAtTime.size() == 0) {
+            auto& futuresAtTime = futureTimeouts_[rpcEndTime];
+            auto it = std::find(futuresAtTime.begin(), futuresAtTime.end(), id);
+            TORCH_INTERNAL_ASSERT(
+                it != futuresAtTime.end(),
+                "Error: could not find future in futureTimeouts map, race condition.");
+            futuresAtTime.erase(it);
+            if (futuresAtTime.empty()) {
               // remove the key from futureTimeouts_
-              futureTimeouts_.erase(futureStartTime);
+              futureTimeouts_.erase(rpcEndTime);
             }
           }
+          // Not holding lock on markCompleted as this could run callbacks that
+          // call agent_->send
+          fm->markCompleted(std::move(message));
           futureCV_.notify_all();
         } else {
           // TODO: pass the error back to the caller instead of crashing here.
@@ -449,6 +482,98 @@ void ProcessGroupAgent::listenLoop() {
   }
 }
 
+void ProcessGroupAgent::pollTimedOutRPCs() {
+  while (!shutdown_.load()) {
+    std::chrono::milliseconds sleepTime;
+    std::unique_lock<std::mutex> lock{futureMutex_};
+    // Estimate amount of time the first future will time out in, and sleep
+    // for that long.
+    // if there are no futures or the first future's RPC timeout is set to 0
+    // (meaning no timeout), then sleep for a set "infinity" time.
+    if (futureTimeouts_.empty() ||
+        futureTimeouts_.begin()->first == INFINITE_TIMEOUT) {
+      sleepTime = INFINITE_TIMEOUT;
+    } else {
+      const auto minFutureExpirationTime = futureTimeouts_.begin()->first;
+      const auto remainingTime = getRPCRemainingTime(minFutureExpirationTime);
+      sleepTime = std::max(remainingTime, std::chrono::milliseconds(0));
+    }
+
+    if (sleepTime == INFINITE_TIMEOUT) {
+      futureTimeoutCV_.wait(lock);
+    } else {
+      futureTimeoutCV_.wait_for(lock, sleepTime);
+    }
+
+    if (shutdown_.load()) {
+      return;
+    }
+
+    const auto timedOutFutures = processTimedOutFutures();
+
+    // Do not hold the lock while marking futures completed, as markCompleted()
+    // could invoke callbacks.
+    lock.unlock();
+    for (const auto& timedOutFuture : timedOutFutures) {
+      std::ostringstream ss;
+      ss << "RPC ran for more than " << timedOutFuture.timeout_.count()
+         << " milliseconds and timed out.";
+      const auto exceptionMsg = createExceptionResponse(
+          Message({}, {}, MessageType::EXCEPTION), ss.str());
+      timedOutFuture.future_->markCompleted(exceptionMsg);
+
+      const int dst = timedOutFuture.dstRank_;
+      recvCounts_.increment(dst);
+      futureCV_.notify_all();
+    }
+  }
+}
+
+const std::vector<ProcessGroupAgent::FutureInfo> ProcessGroupAgent::
+    processTimedOutFutures() {
+  std::vector<FutureInfo> timedOutFutures;
+  for (auto it = futureTimeouts_.begin(); it != futureTimeouts_.end();
+       /* intentional no increment */) {
+    const auto& endTime = it->first;
+    const auto remainingTime = getRPCRemainingTime(endTime);
+
+    if (remainingTime.count() > 0) {
+      // Since the futureTimeouts_ map is ordered by timeout, we don't need
+      // to check the remaining futures.
+      break;
+    } else {
+      const std::vector<int64_t>& futureIDs = it->second;
+      for (const auto& futureID : futureIDs) {
+        auto futureIt = futures_.find(futureID);
+        TORCH_INTERNAL_ASSERT(
+            futureIt != futures_.end(),
+            "Race Condition - Expected future does not exist in map");
+        const auto futInfo = futureIt->second;
+        timedOutFutures.push_back(futInfo);
+        futures_.erase(futureID);
+      }
+      it = futureTimeouts_.erase(it);
+    }
+  }
+  return timedOutFutures;
+}
+
+const std::chrono::milliseconds ProcessGroupAgent::getRPCRemainingTime(
+    const std::chrono::milliseconds& rpcEndTime) const {
+  const auto remainingTime =
+      rpcEndTime -
+      std::chrono::duration_cast<std::chrono::milliseconds>(
+          std::chrono::steady_clock::now().time_since_epoch());
+  return remainingTime;
+}
+
+const std::chrono::milliseconds ProcessGroupAgent::getRPCEndTime(
+    const FutureInfo& futureInfo) const {
+  return futureInfo.timeout_ == INFINITE_TIMEOUT
+      ? INFINITE_TIMEOUT
+      : futureInfo.startTime_ + futureInfo.timeout_;
+}
+
 } // namespace rpc
 } // namespace distributed
 } // namespace torch