[reland] Create CUDA-aware futures in RequestCallback

lw · lw · commit 5f6aed6e3e39 · 2021-05-31T05:26:19.000-07:00
Reland of #58426 The operations in RequestCallback can return CUDA tensors, thus the futures used to hold them must be CUDA-aware. Differential Revision: [D28623887](https://our.internmc.facebook.com/intern/diff/D28623887/) [ghstack-poisoned]
diff --git a/torch/csrc/distributed/rpc/request_callback_no_python.cpp b/torch/csrc/distributed/rpc/request_callback_no_python.cpp
@@ -602,15 +602,17 @@ c10::intrusive_ptr<JitFuture> RequestCallbackNoPython::runJitOperator(
 c10::intrusive_ptr<JitFuture> RequestCallbackNoPython::asFuture(
     IValue value,
     TypePtr type) const {
-  auto future = c10::make_intrusive<JitFuture>(std::move(type));
+  auto future = c10::make_intrusive<JitFuture>(
+      std::move(type), RpcAgent::getCurrentRpcAgent()->getDevices());
   future->markCompleted(std::move(value));
   return future;
 }
 
 c10::intrusive_ptr<JitFuture> RequestCallbackNoPython::asFuture(
     c10::intrusive_ptr<Message> message) const {
   auto future = c10::make_intrusive<JitFuture>(
-      at::getCustomClassType<c10::intrusive_ptr<Message>>());
+      at::getCustomClassType<c10::intrusive_ptr<Message>>(),
+      RpcAgent::getCurrentRpcAgent()->getDevices());
   std::vector<std::reference_wrapper<const at::DataPtr>> dataPtrs =
       message->getDataPtrs();
   future->markCompleted(std::move(message), std::move(dataPtrs));
@@ -619,7 +621,8 @@ c10::intrusive_ptr<JitFuture> RequestCallbackNoPython::asFuture(
 
 c10::intrusive_ptr<JitFuture> RequestCallbackNoPython::asFuture(
     std::exception_ptr err) const {
-  auto future = c10::make_intrusive<JitFuture>(at::NoneType::get());
+  auto future = c10::make_intrusive<JitFuture>(
+      at::NoneType::get(), RpcAgent::getCurrentRpcAgent()->getDevices());
   future->setError(err);
   return future;
 }
