Repository: pytorch
Issue: [RFC] Manage CUDA Stream in TensorPipe RPC Agent
Labels: feature,triaged,module: rpc,module: tensorpipe

with @lw @beauby 

## API

There is no new API, but has a new behavior for sending/receiving CUDA tensors. The TensorPipe RPC agent will maintain a stream pool, and grab streams from the pool to 1) send tensors 2) receive tensors 3) run user functions. 

The guarantees we offer are:

1. RPC user functions are launched when all comm ops are enqueued to the current stream for all involved devices, but these is no guarantee any of the comm ops are done. This should be fine if user functions do not switch streams. However, if they do use different stream, they will need to explicitly synchronize. 
2. `NCCL_BLOCKING_WAIT` dictates the behavior of `Future.wait()` when the response contains CUDA tensors. If `NCCL_BLOCKING_WAIT=1` at the RPC initialization time, `Future.wait()` means that all ops are done. Otherwise, `Future.wait()` only means all CUDA ops are enqueued to the stream. 

## Design 

two principles:

* No RPC CUDA tensor transfer should block ops in the current stream. 
* All ops in a stream must be cleared before it is returned to the pool to avoid pollute next user of this stream.

### Request on the Caller

* Determine the set of devices used by the CUDA tensors in the message.
* Get the current streams for all these devices.
* Record an event in each of these current streams (and then we won’t do anything more on these streams)
* Pick a new fresh stream for each device from the stream pool
* On each of these new streams, enqueue a wait for the above events
* Pass these new streams to TensorPipe in the form of `cudaStream_t` (as TensorPipe cannot depend on PyTorch and hence cannot use PyTorch's Stream type) 
* In the on-complete callback of `pipe->write`, either capture the stream in another lambda for `cudaStreamAddCallback` or use a thread to synchronize the stream before destructing it (returning it to the pool).


### Request on the Callee


* Let `TensorPipe::Tensor.metadata` capture the device information, so that callee can know which device it will use from the descriptor before calling `pipe->read`.
* Grab one stream for each distinct device, and pass those streams to TensorPipe `pipe->read`.
* In `pipe->read`'s on-complete callback, set those streams from the pool as current, and then directly launch user functions without synchronization. All comm ops are already enqueued into the current stream. If user functions need to use different streams, they need to explicitly synchronize the streams. 


### Response on the callee

* Once the user function returns, preserve the current streams (they should still be the same ones from pool).
* Do not synchronize on these streams and directly pass them to TensorPipe to send response message.
* After the write callback from TensorPipe `pipe->write` fires, we synchronize on those streams and only return them to their pools once all the enqueued work on them has completed.

### Response on the caller
    

* When the response is coming in, pick a stream for each relevant device from the pool.
* Pass these streams to TensorPipe to use for receiving the response tensors
* After the read callback fires, synchronize with the work in these streams before returning them to the pool.
* If `NCCL_BLOCKING_WAIT=1`, mark the `Future` as complete after all streams are synchronized. Otherwise, mark the `Future` as complete when TensorPipe `pipe->read` returns. 


### Discussion

1. Do we need a "power-user" mode to disable stream pool on the caller? We can remember the current stream when sending the requests and use them to receive responses. This can give users more control. 
2. Regarding stream pool, will use `c10::cuda::getStreamFromPool` in the first version, and then will add a new stream pool implement, which instead of using a fixed number of stream, it will create new ones when its queue is depleted. 

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @rohan-varma @xush6528 @jjlilley @osalpekar @jiayisuse @lw @beauby