Repository: pytorch
Issue: Cuda RPC error when using then()
Labels: oncall: distributed,module: tensorpipe

## 🐛 Bug

The following results in an error on master (add into `rpc/rpc_test.py`):
```
def get_myclass():
    return MyClass(4)

@skip_if_lt_x_gpu(1)
    def test_rpc_cuda_then(self):
        dst = worker_name((self.rank + 1) % self.world_size)
        options = self.rpc_backend_options
        options.set_device_map(dst, {0: 0})

        rpc.init_rpc(
            name=worker_name(self.rank),
            backend=self.rpc_backend,
            rank=self.rank,
            world_size=self.world_size,
            rpc_backend_options=options,
        )

        if self.rank == 1:
            fut = rpc.rpc_async(dst, get_myclass, args=()).then(lambda res: res)
            print(fut.wait())

        rpc.shutdown()
```

Stacktrace:

```
======================================================================
ERROR: test_rpc_cuda_then (test_tensorpipe_agent.TensorPipeTensorPipeAgentCudaRpcTestWithSpawn)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/data/users/rvarm1/fbsource/fbcode/buck-out/dev/gen/caffe2/test/distributed/rpc/cuda/tensorpipe_agent_cuda#binary,link-tree/torch/testing/_internal/common_distributed.py", line 322, in wrapper
    self._join_processes(fn)
  File "/data/users/rvarm1/fbsource/fbcode/buck-out/dev/gen/caffe2/test/distributed/rpc/cuda/tensorpipe_agent_cuda#binary,link-tree/torch/testing/_internal/common_distributed.py", line 515, in _join_processes
    self._check_return_codes(elapsed_time)
  File "/data/users/rvarm1/fbsource/fbcode/buck-out/dev/gen/caffe2/test/distributed/rpc/cuda/tensorpipe_agent_cuda#binary,link-tree/torch/testing/_internal/common_distributed.py", line 558, in _check_return_codes
    raise RuntimeError(error)
RuntimeError: Process 1 exited with error code 10 and exception:
Traceback (most recent call last):
  File "/data/users/rvarm1/fbsource/fbcode/buck-out/dev/gen/caffe2/test/distributed/rpc/cuda/tensorpipe_agent_cuda#binary,link-tree/torch/testing/_internal/common_distributed.py", line 441, in run_test
    getattr(self, test_name)()
  File "/data/users/rvarm1/fbsource/fbcode/buck-out/dev/gen/caffe2/test/distributed/rpc/cuda/tensorpipe_agent_cuda#binary,link-tree/torch/testing/_internal/common_distributed.py", line 324, in wrapper
    fn()
  File "/data/users/rvarm1/fbsource/fbcode/buck-out/dev/gen/caffe2/test/distributed/rpc/cuda/tensorpipe_agent_cuda#binary,link-tree/torch/testing/_internal/common_distributed.py", line 89, in wrapper
    return func(*args, **kwargs)
  File "/data/users/rvarm1/fbsource/fbcode/buck-out/dev/gen/caffe2/test/distributed/rpc/cuda/tensorpipe_agent_cuda#binary,link-tree/torch/testing/_internal/distributed/rpc/rpc_test.py", line 5605, in test_rpc_cuda_then
    print(fut.wait())
RuntimeError: match.success()INTERNAL ASSERT FAILED at "caffe2/aten/src/ATen/core/ivalue.cpp":221, please report a bug to PyTorch. Cannot infer type of <torch._C.Future object at 0x7fa9240fbdf0>
:Only tensors and (possibly nested) tuples of tensors, lists, or dictsare supported as inputs or outputs of traced functions, but instead got value of type Future.
```

From my understanding, the root cause is as follows. I could be mistaken though as I'm not too familiar with CUDA rpc: 
1. In CUDA mode, tensorpipe returns AtomicJitFuture which is a RpcCUDAFuture which is a CUDAFuture. 
2. when processing a response from server, rpc client calls markCompleted on the future for the rpc with a Message object, and then extractDataPtrs runs. Specifically `extractDataPtrs` defined for `RpcCudaFuture` runs. 
3. If the future in (2) has a then() callback attached to it, then the resulting future returned by `then()` is not a c10::ivalue::Future, it's a at::CudaFuture per  https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/core/ivalue_inl.h#L469 (createInstance is overriden by CUDAFuture). Since the future in (2) is of type `RpcCudaFuture` the createInstance should dispatch to `at::cudaFuture`. 
4. When markCompleted is called on this future from `then()`, it seems like extractDataPtrs defined in `CUDAFuture` runs, and errors out with the above error. I was confused as to why extractDataPtrs defined in `RpcCudaFuture` wasn't running, but I think this is because the future in (3) is `at::CudaFuture` and not `RpcCudaFuture`. extractDataPtrs uses `getSubValues` which throws on non-torchscript python objs.

## Additional context
Hit this bug while looking into https://github.com/pytorch/pytorch/issues/55757


cc @osalpekar @jiayisuse @lw @beauby @pritamdamania87 @mrshenli @jjlilley @gqchen @rohan-varma @pietern @zhaojuanmao @satgera @aazzolini @agolynski @SciPioneer @H-Huang @mrzzd @cbalioglu