Fix CUDA RPC Stream Synchronization

When converting RPC Message into Python objects, we were not using
a CUDAFuture for the chained Future. As a result, the streams are
not synchronized when calling `rpc_async(...).wait()`. This commit
uses `Future::then` API to create the chained Future, which will
be creating a CUDAFuture if the existing Future is a CUDA one.

fixes #50881
fixes #50839

ghstack-source-id: 56c79004a6250bb608d473300260d181a3b11cc9
Pull Request resolved: #50949

torch/csrc/distributed/rpc/python_functions.cpp

@@ -138,36 +138,29 @@ c10::intrusive_ptr<JitFuture> toPyJitFuture(
     const std::shared_ptr<JitFuture>& messageJitFuture,
     bool hasValue) {
   if (hasValue) {
-    c10::intrusive_ptr<JitFuture> pyJitFuture =
-        c10::make_intrusive<JitFuture>(PyObjectType::get());
     std::weak_ptr<JitFuture> wp = messageJitFuture;
-    messageJitFuture->addCallback(
-        at::wrapPropagateTLSState<void>([pyJitFuture, wp]() {
+    return messageJitFuture->then(
+        at::wrapPropagateTLSState<IValue>([wp]() {
           auto future = wp.lock();
           if (future->hasError()) {
-            pyJitFuture->setError(future->exception_ptr());
+            std::rethrow_exception(future->exception_ptr());
           } else {
-            pyJitFuture->markCompleted(
-                toPyIValue(*future->value().toCustomClass<Message>()));
+            return toPyIValue(*future->value().toCustomClass<Message>());
           }
-        }));
-
-    return pyJitFuture;
+        }),
+        PyObjectType::get());
   } else {
-    c10::intrusive_ptr<JitFuture> pyJitFuture =
-        c10::make_intrusive<JitFuture>(NoneType::get());
     std::weak_ptr<JitFuture> wp = messageJitFuture;
-    messageJitFuture->addCallback(
-        at::wrapPropagateTLSState<void>([wp, pyJitFuture]() {
+    return messageJitFuture->then(
+        at::wrapPropagateTLSState<IValue>([wp]() {
           auto future = wp.lock();
           if (future->hasError()) {
-            pyJitFuture->setError(future->exception_ptr());
+            std::rethrow_exception(future->exception_ptr());
           } else {
-            pyJitFuture->markCompleted(IValue());
+            return IValue();
           }
-        }));
-
-    return pyJitFuture;
+        }),
+       NoneType::get());
   }
 }