Unify fetching RRefs (#57859)

lw · facebook-github-bot · commit 797dff55b520 · 2021-05-21T13:15:15.000-07:00
Summary: Pull Request resolved: #57859 Just like with assigning OwnerRRefs, we can also deduplicate the code paths for fetching their values. In fact this was duplicated three times, with different ways of post-processing the value (once for JIT, once for Python, once for autograd). Thanks to future, we can have that logic once, and then connect it to different follow-up steps. ghstack-source-id: 129567050 Test Plan: CI Reviewed By: mrshenli Differential Revision: D28286172 fbshipit-source-id: e0742a99cf555755e848057ab6fee5285ff0df2a
diff --git a/torch/csrc/distributed/rpc/py_rref.cpp b/torch/csrc/distributed/rpc/py_rref.cpp
@@ -307,37 +307,41 @@ void PyRRef::backward(int64_t autogradContextId, bool retainGraph) {
   backward(autogradContextId, retainGraph, rref_);
 }
 
-void PyRRef::backward(
+void PyRRef::backwardOwnerRRef(
     int64_t autogradContextId,
     bool retainGraph,
-    const c10::intrusive_ptr<RRef>& rref) {
-  if (rref->isOwner()) {
-    auto value =
-        c10::static_intrusive_pointer_cast<const OwnerRRef>(rref)->getValue();
-
-    // If we have a PyObj, retrieve the underlying tensor.
-    if (rref->isPyObj()) {
-      py::gil_scoped_acquire gil;
-      py::object obj = torch::jit::toPyObject(value);
-      try {
-        value = torch::jit::toIValue(obj, c10::TensorType::get());
-      } catch (py::cast_error& e) {
-        throw std::runtime_error(
-            "RRef should contain a tensor for .backward()");
-      }
+    IValue value) {
+  // If we have a PyObj, retrieve the underlying tensor.
+  if (value.isPyObject()) {
+    py::gil_scoped_acquire gil;
+    py::object obj = torch::jit::toPyObject(value);
+    try {
+      value = torch::jit::toIValue(obj, c10::TensorType::get());
+    } catch (py::cast_error& e) {
+      throw std::runtime_error("RRef should contain a tensor for .backward()");
     }
+  }
 
-    TORCH_CHECK(
-        value.isTensor(), "RRef should contain a tensor for .backward()");
-    auto root = value.toTensor();
+  TORCH_CHECK(value.isTensor(), "RRef should contain a tensor for .backward()");
+  auto root = value.toTensor();
 
-    if (autogradContextId == -1) {
-      torch::autograd::backward({root});
-    } else {
-      torch::distributed::autograd::backward(
-          autogradContextId, {root}, retainGraph);
-    }
+  if (autogradContextId == -1) {
+    torch::autograd::backward({root});
+  } else {
+    torch::distributed::autograd::backward(
+        autogradContextId, {root}, retainGraph);
+  }
+}
 
+void PyRRef::backward(
+    int64_t autogradContextId,
+    bool retainGraph,
+    const c10::intrusive_ptr<RRef>& rref) {
+  if (rref->isOwner()) {
+    backwardOwnerRRef(
+        autogradContextId,
+        retainGraph,
+        c10::static_intrusive_pointer_cast<const OwnerRRef>(rref)->getValue());
   } else {
     TORCH_CHECK(
         autogradContextId != -1,
diff --git a/torch/csrc/distributed/rpc/py_rref.h b/torch/csrc/distributed/rpc/py_rref.h
@@ -66,6 +66,12 @@ class PYBIND11_EXPORT PyRRef {
       bool retainGraph,
       const c10::intrusive_ptr<RRef>& rref);
 
+  // Specialization of backward if the rref is an OwnerRRef.
+  static void backwardOwnerRRef(
+      int64_t autogradContextId,
+      bool retainGraph,
+      IValue value);
+
  private:
   c10::intrusive_ptr<RRef> rref_;
   c10::optional<c10::intrusive_ptr<JitFuture>> profilingFuture_;
diff --git a/torch/csrc/distributed/rpc/request_callback_impl.cpp b/torch/csrc/distributed/rpc/request_callback_impl.cpp
@@ -220,52 +220,20 @@ c10::intrusive_ptr<JitFuture> RequestCallbackImpl::processPythonRemoteCall(
       uprc.rrefId(), uprc.forkId(), std::move(future), std::move(lsctx));
 }
 
-void RequestCallbackImpl::processPythonRRefFetchCall(
+c10::intrusive_ptr<JitFuture> RequestCallbackImpl::processPythonRRefFetchCall(
     RpcCommandBase& rpc,
-    const c10::intrusive_ptr<JitFuture>& responseFuture,
     std::shared_ptr<LazyStreamContext> lsctx) const {
-  // Making this lambda mutable to allow move-capture it in callbacks
-  auto postProcessing = [responseFuture, lsctx = std::move(lsctx)](
-                            const c10::intrusive_ptr<OwnerRRef>& rref) mutable {
-    auto whenValueSet = rref->getFuture();
-    if (whenValueSet->hasError()) {
-      responseFuture->setError(whenValueSet->exception_ptr());
-      return;
-    }
-    try {
-      SerializedPyObj result = serializePyObject(rref->getValue());
-      Message m = PythonRRefFetchRet(std::move(result).toIValues()).toMessage();
-      rref->blockAllStreams(lsctx);
-      responseFuture->markCompleted(
-          IValue(c10::make_intrusive<Message>(std::move(m))));
-    } catch (const std::exception& /* unused */) {
-      responseFuture->setError(std::current_exception());
-    }
-  };
-
   auto& prf = static_cast<PythonRRefFetchCall&>(rpc);
-  auto& ctx = RRefContext::getInstance();
-
-  auto futureOwner = ctx.getOwnerRRef(prf.rrefId());
-  if (futureOwner->completed()) {
-    auto rref = fromRRefInterface(futureOwner->constValue().toRRef());
-    if (rref->hasValue()) {
-      // optional fast-path, the OwnerRRef has been created
-      postProcessing(rref);
-      return;
-    }
-  }
 
-  futureOwner->addCallback([postProcessing{std::move(postProcessing)}](
-                               JitFuture& futureOwner) mutable {
-    const auto& rref = fromRRefInterface(futureOwner.constValue().toRRef());
+  auto future = retrieveOwnerRRef(prf.rrefId(), std::move(lsctx));
 
-    // Our response is satisfied when the the rpc.remote() request
-    // finishes executing on the owner.
-    rref->getFuture()->addCallback(
-        [rref, postProcessing{std::move(postProcessing)}](
-            JitFuture& /* unused */) mutable { postProcessing(rref); });
-  });
+  return future->then(
+      [](JitFuture& future) {
+        SerializedPyObj result = serializePyObject(future.value());
+        return c10::make_intrusive<Message>(
+            PythonRRefFetchRet(std::move(result).toIValues()).toMessage());
+      },
+      c10::getCustomClassType<c10::intrusive_ptr<Message>>());
 }
 
 void RequestCallbackImpl::handleRRefDelete(
@@ -308,43 +276,23 @@ bool RequestCallbackImpl::cudaAvailable() const {
 #endif
 }
 
-void RequestCallbackImpl::processRRefBackward(
-    RpcCommandBase& rpc,
-    const c10::intrusive_ptr<JitFuture>& responseFuture) const {
+c10::intrusive_ptr<JitFuture> RequestCallbackImpl::processRRefBackward(
+    RpcCommandBase& rpc) const {
   auto& rrefBackwardReq = static_cast<RRefBackwardReq&>(rpc);
 
-  // Get all fields
-  const auto& rrefId = rrefBackwardReq.getRRefId();
-  const auto& autogradContextId = rrefBackwardReq.getAutogradContextId();
-  const auto& retainGraph = rrefBackwardReq.retainGraph();
-
-  auto futureOwner = RRefContext::getInstance().getOwnerRRef(rrefId);
-  futureOwner->addCallback(
-      [responseFuture, autogradContextId, retainGraph](JitFuture& futureOwner) {
-        const auto& rref = fromRRefInterface(futureOwner.constValue().toRRef());
-        auto whenValueSet = rref->getFuture();
-
-        whenValueSet->addCallback(
-            [responseFuture, rref, autogradContextId, retainGraph](
-                JitFuture& whenValueSet) {
-              if (whenValueSet.hasError()) {
-                responseFuture->setError(whenValueSet.exception_ptr());
-                return;
-              }
-
-              try {
-                // Run backward (TODO: make this async?).
-                PyRRef::backward(autogradContextId, retainGraph, rref);
-
-                // Return the response.
-                Message m = RRefBackwardResp().toMessage();
-                responseFuture->markCompleted(
-                    IValue(c10::make_intrusive<Message>(std::move(m))));
-              } catch (const std::exception& /* unused */) {
-                responseFuture->setError(std::current_exception());
-              }
-            });
-      });
+  auto future =
+      retrieveOwnerRRef(rrefBackwardReq.getRRefId(), /*lsctx=*/nullptr);
+
+  return future->then(
+      [autogradContextId = rrefBackwardReq.getAutogradContextId(),
+       retainGraph = rrefBackwardReq.retainGraph()](JitFuture& future) {
+        // Run backward (TODO: make this async?).
+        PyRRef::backwardOwnerRRef(
+            autogradContextId, retainGraph, future.value());
+
+        return c10::make_intrusive<Message>(RRefBackwardResp().toMessage());
+      },
+      c10::getCustomClassType<c10::intrusive_ptr<Message>>());
 }
 
 c10::intrusive_ptr<JitFuture> RequestCallbackImpl::runJitFunction(
diff --git a/torch/csrc/distributed/rpc/request_callback_impl.h b/torch/csrc/distributed/rpc/request_callback_impl.h
@@ -28,9 +28,8 @@ class TORCH_API RequestCallbackImpl : public RequestCallbackNoPython {
       RpcCommandBase& rpc,
       std::shared_ptr<LazyStreamContext> ctx) const override;
 
-  void processPythonRRefFetchCall(
+  c10::intrusive_ptr<JitFuture> processPythonRRefFetchCall(
       RpcCommandBase& rpc,
-      const c10::intrusive_ptr<JitFuture>& responseFuture,
       std::shared_ptr<LazyStreamContext> ctx) const override;
 
   void handleRRefDelete(c10::intrusive_ptr<RRef>& rref) const override;
@@ -42,9 +41,8 @@ class TORCH_API RequestCallbackImpl : public RequestCallbackNoPython {
 
   bool cudaAvailable() const override;
 
-  void processRRefBackward(
-      RpcCommandBase& rpc,
-      const c10::intrusive_ptr<JitFuture>& responseFuture) const override;
+  c10::intrusive_ptr<JitFuture> processRRefBackward(
+      RpcCommandBase& rpc) const override;
 
   // Helpers to run user-defined functions, operators and other computations.
 
diff --git a/torch/csrc/distributed/rpc/request_callback_no_python.cpp b/torch/csrc/distributed/rpc/request_callback_no_python.cpp
@@ -220,50 +220,51 @@ c10::intrusive_ptr<JitFuture> RequestCallbackNoPython::processScriptRemoteCall(
       /*lsctx=*/nullptr);
 }
 
-void RequestCallbackNoPython::processScriptRRefFetchCall(
-    RpcCommandBase& rpc,
-    const std::function<void(Message)>& markComplete,
-    const c10::intrusive_ptr<JitFuture>& responseFuture) const {
-  auto& srf = static_cast<ScriptRRefFetchCall&>(rpc);
+c10::intrusive_ptr<JitFuture> RequestCallbackNoPython::retrieveOwnerRRef(
+    const RRefId& rrefId,
+    std::shared_ptr<LazyStreamContext> lsctx) const {
   auto& ctx = RRefContext::getInstance();
 
-  auto futureOwner = ctx.getOwnerRRef(srf.rrefId());
+  auto rrefFuture = ctx.getOwnerRRef(rrefId);
+
+  at::TypePtr type = rrefFuture->elementType();
+  TORCH_INTERNAL_ASSERT(type->kind() == at::RRefType::Kind);
+  return rrefFuture->thenAsync(
+      [lsctx](JitFuture& rrefFuture) {
+        c10::intrusive_ptr<OwnerRRef> rref =
+            fromRRefInterface(rrefFuture.value().toRRef());
+        auto valueFuture = rref->getFuture();
+        // FIXME This is a temporary fix to synchronize CUDA streams. Once the
+        // OwnerRRef's internal Future becomes CUDA-aware this will be automatic
+        // and we can remove this hack.
+        return valueFuture->then(
+            [rref, lsctx](JitFuture& future) {
+              rref->blockAllStreams(lsctx);
+              return future.value();
+            },
+            valueFuture->elementType());
+      },
+      type->cast<at::RRefType>()->getElementType());
+}
+
+c10::intrusive_ptr<JitFuture> RequestCallbackNoPython::
+    processScriptRRefFetchCall(RpcCommandBase& rpc) const {
+  auto& srf = static_cast<ScriptRRefFetchCall&>(rpc);
 
-  if (futureOwner->completed()) { // optional fast-path
-    // the OwnerRRef has been created
-    const auto& rref = fromRRefInterface(futureOwner->constValue().toRRef());
-    if (rref->hasValue()) {
-      markComplete(ScriptRRefFetchRet({rref->getValue()}).toMessage());
-      return;
-    }
-  }
+  auto future = retrieveOwnerRRef(srf.rrefId(), /*lsctx=*/nullptr);
 
-  futureOwner->addCallback([responseFuture](JitFuture& futureOwner) {
-    const auto& rref = fromRRefInterface(futureOwner.constValue().toRRef());
-    auto whenValueSet = rref->getFuture();
-
-    // Our response is satisfied when the rpc.remote() request
-    // finishes executing on the owner.
-    whenValueSet->addCallback([responseFuture, rref](JitFuture& whenValueSet) {
-      if (whenValueSet.hasError()) {
-        responseFuture->setError(whenValueSet.exception_ptr());
-        return;
-      }
-      try {
-        Message m = ScriptRRefFetchRet({rref->getValue()}).toMessage();
-        responseFuture->markCompleted(
-            IValue(c10::make_intrusive<Message>(std::move(m))));
-      } catch (const std::exception& /* unused */) {
-        responseFuture->setError(std::current_exception());
-      }
-    });
-  });
+  return future->then(
+      [](JitFuture& future) {
+        return c10::make_intrusive<Message>(
+            ScriptRRefFetchRet({future.value()}).toMessage());
+      },
+      c10::getCustomClassType<c10::intrusive_ptr<Message>>());
 }
 
-void RequestCallbackNoPython::processPythonRRefFetchCall(
-    RpcCommandBase& rpc,
-    const c10::intrusive_ptr<JitFuture>& /* unused */,
-    std::shared_ptr<LazyStreamContext> /* unused */) const {
+c10::intrusive_ptr<JitFuture> RequestCallbackNoPython::
+    processPythonRRefFetchCall(
+        RpcCommandBase& rpc,
+        std::shared_ptr<LazyStreamContext> /* unused */) const {
   C10_THROW_ERROR(Error, "Python call not supported!");
 }
 
@@ -503,9 +504,8 @@ c10::intrusive_ptr<JitFuture> RequestCallbackNoPython::
   }
 }
 
-void RequestCallbackNoPython::processRRefBackward(
-    RpcCommandBase& rpc,
-    const c10::intrusive_ptr<JitFuture>& /* unused */) const {
+c10::intrusive_ptr<JitFuture> RequestCallbackNoPython::processRRefBackward(
+    RpcCommandBase& rpc) const {
   C10_THROW_ERROR(Error, "Python call not supported!");
 }
 
@@ -541,12 +541,10 @@ c10::intrusive_ptr<JitFuture> RequestCallbackNoPython::processRpc(
       return processPythonRemoteCall(rpc, std::move(ctx));
     }
     case MessageType::SCRIPT_RREF_FETCH_CALL: {
-      processScriptRRefFetchCall(rpc, markComplete, responseFuture);
-      return responseFuture;
+      return processScriptRRefFetchCall(rpc);
     }
     case MessageType::PYTHON_RREF_FETCH_CALL: {
-      processPythonRRefFetchCall(rpc, responseFuture, std::move(ctx));
-      return responseFuture;
+      return processPythonRRefFetchCall(rpc, std::move(ctx));
     }
     case MessageType::RREF_USER_DELETE: {
       processRRefUserDelete(rpc, markComplete);
@@ -575,8 +573,7 @@ c10::intrusive_ptr<JitFuture> RequestCallbackNoPython::processRpc(
       return processRunWithProfilingReq(rpc);
     }
     case MessageType::RREF_BACKWARD_REQ: {
-      processRRefBackward(rpc, responseFuture);
-      return responseFuture;
+      return processRRefBackward(rpc);
     }
     default: {
       TORCH_INTERNAL_ASSERT(
diff --git a/torch/csrc/distributed/rpc/request_callback_no_python.h b/torch/csrc/distributed/rpc/request_callback_no_python.h
@@ -42,14 +42,15 @@ class TORCH_API RequestCallbackNoPython : public RequestCallback {
       RpcCommandBase& rpc,
       std::shared_ptr<LazyStreamContext> ctx) const;
 
-  void processScriptRRefFetchCall(
-      RpcCommandBase& rpc,
-      const std::function<void(Message)>& markComplete,
-      const c10::intrusive_ptr<JitFuture>& responseFuture) const;
+  c10::intrusive_ptr<JitFuture> retrieveOwnerRRef(
+      const RRefId& rrefId,
+      std::shared_ptr<LazyStreamContext> lsctx) const;
 
-  virtual void processPythonRRefFetchCall(
+  c10::intrusive_ptr<JitFuture> processScriptRRefFetchCall(
+      RpcCommandBase& rpc) const;
+
+  virtual c10::intrusive_ptr<JitFuture> processPythonRRefFetchCall(
       RpcCommandBase& rpc,
-      const c10::intrusive_ptr<JitFuture>& responseFuture,
       std::shared_ptr<LazyStreamContext> ctx) const;
 
   void processRRefUserDelete(
@@ -98,9 +99,8 @@ class TORCH_API RequestCallbackNoPython : public RequestCallback {
 
   virtual bool cudaAvailable() const;
 
-  virtual void processRRefBackward(
-      RpcCommandBase& rpc,
-      const c10::intrusive_ptr<JitFuture>& responseFuture) const;
+  virtual c10::intrusive_ptr<JitFuture> processRRefBackward(
+      RpcCommandBase& rpc) const;
 
   // Helpers to run user-defined functions, operators and other computations.
 
diff --git a/torch/csrc/distributed/rpc/rref_impl.cpp b/torch/csrc/distributed/rpc/rref_impl.cpp
@@ -306,7 +306,7 @@ void OwnerRRef::recordAllStreams(
   }
 }
 
-void OwnerRRef::blockAllStreams(std::shared_ptr<LazyStreamContext>& ctx) {
+void OwnerRRef::blockAllStreams(const std::shared_ptr<LazyStreamContext>& ctx) {
   if (ctx) {
     for (c10::Event& event : events_) {
       event.block(ctx->getStream(event.device()));
diff --git a/torch/csrc/distributed/rpc/rref_impl.h b/torch/csrc/distributed/rpc/rref_impl.h
