Repository: pytorch
Issue: Add support to profile all RPC worker threads that are run on a process
Labels: oncall: distributed

## 🚀 Feature

Add support to the profiler to profile all RPC worker threads that are run on a process.

## Motivation

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->

It is hard to gain visibility into parameter server implementations with the current state of the profiler. A vanilla parameter server implementation using RPC will have a @rpc.functions.async_execution decorated function. This function will be executed by worker threads created during the RPC initialization. The code below is an example of a parameter server implementation.

```python
import os
from threading import Lock

import torch
import torch.distributed.rpc as rpc
import torch.multiprocessing as mp
import torch.nn as nn
from torch.distributed.rpc import RRef, TensorPipeRpcBackendOptions
from torch.profiler import ProfilerActivity, profile


class MyParameterServer:
    def __init__(self, trainers):
        self.lock = Lock()
        self.trainers = trainers
        self.iteration = 0
        self.updates = 0
        self.futures = []
        self.total = None
        self.p = profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA])
        self.p.start()

    @staticmethod
    def stop(rref):
        self = rref.local_value()
        self.p.stop()

    @staticmethod
    @rpc.functions.async_execution
    def average(rref, riteration, tensor):
        self = rref.local_value()
        fut = torch.futures.Future()
        with self.lock:
            if riteration > self.iteration:
                self.iteration = riteration
                self.updates = 0
                self.futures.clear()
            self.futures.append(fut)
            if self.total is None:
                self.total = tensor
            else:
                self.total += tensor
            self.updates += 1
            if self.trainers == self.updates:
                for fut in self.futures:
                    fut.set_result(self.total / float(self.trainers))
        return fut


class MyEmbeddingBagModel(torch.nn.Module):
    def __init__(self, sparse):
        super().__init__()
        self.eb = torch.nn.EmbeddingBag(
            10,
            10,
            sparse=sparse
        )

    def forward(self, x):
        return self.eb(x)


def trainer_func(rref, sparse=False):
    m = MyEmbeddingBagModel(sparse=sparse)
    loss_fn = nn.MSELoss()
    for i in range(10):
        outputs = m(torch.rand(10, 10).long())
        loss_fn(outputs, torch.rand(10, 10)).backward()
        gradient = list(m.parameters())[0].grad
        fut = rref.rpc_async().average(rref, i, gradient.to_sparse())
        fut.wait()


def basic_ps(rank, world_size):
    rpc_backend_options = TensorPipeRpcBackendOptions(
        num_worker_threads=2
    )
    os.environ['MASTER_ADDR'] = "localhost"
    os.environ['MASTER_PORT'] = "12345"
    if rank == 0:
        rpc.init_rpc(
            f"server{rank}",
            rank=rank,
            world_size=world_size,
            rpc_backend_options=rpc_backend_options
        )
        ps_rref = RRef(MyParameterServer(2))
        futures = []
        for index in range(1, world_size):
            futures.append(
                rpc.rpc_async(
                    f"trainer{index}",
                    trainer_func,
                    args=(
                        ps_rref,
                    ),
                )
            )
        torch.futures.wait_all(futures)
        ps_rref.rpc_sync().stop(ps_rref)
    else:
        rpc.init_rpc(
            f"trainer{rank}",
            rank=rank,
            world_size=world_size,
            rpc_backend_options=rpc_backend_options
        )

    rpc.shutdown()

if __name__ == "__main__":
    mp.spawn(basic_ps, nprocs=3, args=(3,))

```


## Pitch

The profiler is able to profile all RPC worker threads that are run on a process. 

## Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

## Additional context