ProcessGroupAgent failed to join with long async forward chain

[mrshenli]

Apart from the deadlock issue mentioned in #26362, there could be two other reasons that causes join() to occasionally throw timeout errors.

1. There is a period of time that the sender has finished sending and the receiver has finished receiving, but the received message is still in the listenLoop() and hasn't been enqueued into the ThreadPool yet. It means that the following sync/join code could exit before all communications finish.

pytorch/torch/csrc/distributed/rpc/process_group_agent.cpp

Lines 162 to 173 in 7e95439

	void ProcessGroupAgent::sync() {
	// Block until all processes wants to sync. This is necessary before acquiring
	// the lock below, because other processes might not enter sync() until it
	// gets some response from this RpcAgent.
	pg_->barrier()->wait();
	// Wait until the all send works are done.
	// NB: There might be additional send works inserted while waiting.
	threadPool_.waitWorkComplete();
	// Use another barrier in case different RpcAgent handles different amounts of
	// workloads.
	pg_->barrier()->wait();
	}

As a result, the SHUTDOWN message could be sent before regular messages, leading to the timeout.

2. Fix termination crash when having unresolved future #24074 and Add Python RRef as args and return value #25499 attempts to address the above problem by waiting for all futures to settle on each individual worker. However, this might not be sufficient. When we have many layers of nested rpc_async/remote calls, received messages could trigger more sends. It means even if one worker didn't see any send task in the thread pool and no unsettled futures at that time, it doesn't mean the messages it received later won't create new sends.

These might be one reason for the flakiness @rohan-varma saw earlier today in #26570.

Proposed Solution

How about using ProcessGroup::allreduce(MAX) to globally check if there is any unfinished future on any worker, and use a while loop to ensure that SHUTDOWN message is only sent after that? Something like:

while(true) {
  auto t = torch::tensor({futures_.size()});
  pg_->allreduce(t, ReduceOp::MAX)->wait();
  if (t.storage().data<int64_t>()[0] == 0) {
    break;
  }
  std::unique_lock<std::mutex> lock(futureMutex_);
  futureCV_.wait(lock, [this] {
    return futures_.empty();
  });
  lock.unlock();
}

However, future p2p RpcAgent implementations does not have the luxury to use collective communication APIs. We might eventually need to build a more complex termination detection algorithm on top of RpcAgent.

@xush6528 @satgera @aazzolini @rohan-varma @pietern Does this make sense?

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini

[pietern]

The form you propose it will be racy as well. The allreduce acts as a barrier, and by the time it goes through, the local snapshot of futures_.size() may already have changed again. You can imagine getting an (inconsistent) cut where the snapshots of futures_.size() are 0 for all workers, but there is still work happening. To get a consistent cut, you must hold the future lock while running the allreduce. This is not practical because it means no work can proceed until all workers check in.

See https://www.cs.usfca.edu/~srollins/courses/cs682-s08/web/notes/timeandstates.html on Global States and Cuts. We could settle for a solution that kinda works, but we'll get bitten by it eventually. Instead, I believe a deeper analysis is warranted and we can come up with a provably correct solution.

I think there must be some solution where we use a pairwise (for peers) counter for messages that are posted and message that have been executed. Let's say worker A sends a message to B. It will increment counter[A_B] on worker A. Only when B has finished acting on that message will it increment its own local counter[A_B]. If acting on that message by B led to more messages being sent to other workers, it will have incremented counter[B_C], counter[B_D], etc. On a local level, these counters represent just how many messages have been sent out and have been received and acted on. But on a global level we can infer if there is still work ongoing or not.

If we take a snapshot of these counters on every worker and distribute copies to every worker (allgather), then every worker can determine whether or not all messages have drained by comparing the pairwise counters, namely if counter[X_Y]_X == counter[X_Y]_Y for every worker X and Y. If this is not the case, there is at least 1 message in flight and has not been acted upon (or according to the snapshot). Then, sync can be implemented by checking this in a loop until the condition holds.

Intuitively, this should work. What do you think?

[mrshenli]

@pietern

The form you propose it will be racy as well. The allreduce acts as a barrier, and by the time it goes through, the local snapshot of futures_.size() may already have changed again. You can imagine getting an (inconsistent) cut where the snapshots of futures_.size() are 0 for all workers, but there is still work happening. To get a consistent cut, you must hold the future lock while running the allreduce. This is not practical because it means no work can proceed until all workers check in.

That's correct. The above proposed solution does not solve the flakiness. The figure below shows one example of such case.

And I like the idea of using a p2p message counters. I think that will cover the cases where nested RPCs that are immediately sent out (either in another RPC or in response callback). And if the application uses some timer to send out RPC, it should be the application's responsibility to make sure that all timers are fired before calling join. Let me add a PR for the message counters.

[pietern]

It's beyond the scope to deal with outside factors that might trigger a message being sent IMO. If that needs to be taken into account, those timers would need to be registered with the RPC agent itself. I think we're safe to ignore that as long as there clear documentation (stating the obvious, tbh).

Note that this algorithm has implications for the retry logic we talked about in #26759 and that it only works if we take the low level message dedup route.

cc @gqchen