Repository: pytorch
Issue: Callbacks of Futures shouldn't wait inline on another future
Labels: oncall: distributed,module: ddp

These are all examples of a (blocking) wait on a Future inside the callback of another Future.

https://github.com/pytorch/pytorch/blob/bc2c1d7d59ae16b1a35bf2d383cdd6d2b351c8a9/torch/distributed/algorithms/ddp_comm_hooks/powerSGD_hook.py#L128 (which, in fact, should be a `wait()` rather than a `value()`, but that's a separate bug)

https://github.com/pytorch/pytorch/blob/bc2c1d7d59ae16b1a35bf2d383cdd6d2b351c8a9/torch/distributed/algorithms/ddp_comm_hooks/quantization_hooks.py#L100

https://github.com/pytorch/pytorch/blob/bc2c1d7d59ae16b1a35bf2d383cdd6d2b351c8a9/torch/distributed/algorithms/ddp_comm_hooks/quantization_hooks.py#L195

These are problematic. Apart from the "antipattern" that in async programming one should just chain callbacks and only wait at the very end, there's a concrete risk of deadlock. Suppose a single event loop is driving things, which means the callbacks would be invoked in its thread, but any futures created by the callback will also be served by the same event loop, so we end up with the callback waiting for the future to be finished, and blocking the event loop by doing so, which prevents the future from ever being completed.

So far this has worked because I think those hooks were only used with NCCL which has the unusual property that futures are always complete immediately upon creation and thus callbacks are always called inline. However this won't always hold and in those cases that code would break.

Here's what I think is a working approach.

If the original (buggy) code is:
```py
original_fut = ...

def blocking_cb(original_fut):
    follow_up_fut = ...
    return follow_up_fut.wait()

result_fut = original_fut.then(blocking_cb)
```
Then it could be transformed into:
```py
original_fut = ...

result_fut = torch.futures.Future()

def on_follow_up_complete(follow_up_fut):
    result_fut.set_result(follow_up_fut.value())

def non_blocking_cb(original_fut):
    follow_up_fut = ...
    follow_up_fut.then(on_follow_up_complete)
    return None

original_fut.then(non_blocking_cb)
```

As you can see, it's a bit more convoluted and less readable than the (buggy) one above. If we think that this is a common use case and that people will risk falling into this mistake then we might want to provide some helper to do this.

In principle, though, we cannot prevent users from doing blocking waits (it's their opaque callback and they can do what they want). So the best we can do is to raise awareness of this issue, try our best to avoid it in our own code, and lower the barrier for doing the right thing.

One option to support this is to give a second `then()` method to Futures, let's call it `thenFut()`, which takes a callback that returns another Future, rather than just a "plain value". In this case, the Future returned by `thenFut()` is _not_ marked complete when the callback returns, and instead it's marked complete when the Future returned by the callback completes in turn. (We could even merge `thenFut()` into `then()` and change behavior based on the return type of the callback, although I'm not fully convinced by that)

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @agolynski @SciPioneer @H-Huang @mrzzd