[reland] Create CUDA-aware futures in RequestCallback (#59209)

lw · facebook-github-bot · commit f8a3fd4e34eb · 2021-06-02T05:46:00.000-07:00
Summary: Pull Request resolved: #59209 Reland of #58426 The operations in RequestCallback can return CUDA tensors, thus the futures used to hold them must be CUDA-aware. ghstack-source-id: 130202844 Test Plan: CI Reviewed By: mrshenli Differential Revision: D28623887 fbshipit-source-id: 53561b8ae011458d8f848f0a03830925aff2f0c2
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;
 }
