Unify invoking JIT operands (#57850)

lw · facebook-github-bot · commit 77428159f529 · 2021-05-21T13:15:00.000-07:00
Summary: Pull Request resolved: #57850 What I envision is a modular decomposed code, with separate steps which each consume/produce Futures, and which can be chained together to obtain the desired results. One common "starting point" for these chains is the execution of a remote function (Python or JIT or otherwise). I'm thus creating a helper function for one of these, the JIT operators (by deduplicating the places where we used to run them). More will follow. This deduplication will also help to add CUDA support to JIT RPC, since the execution of the JIT function/operators is where we need to set our custom streams. ghstack-source-id: 129567058 Test Plan: CI Reviewed By: mrshenli Differential Revision: D28253847 fbshipit-source-id: 24ab67ad89c8796861e9bbcb78878b26704c0c48
diff --git a/torch/csrc/distributed/rpc/request_callback_no_python.cpp b/torch/csrc/distributed/rpc/request_callback_no_python.cpp
@@ -152,13 +152,10 @@ void RequestCallbackNoPython::processScriptCallOp(
     const std::function<void(Message)>& markComplete,
     std::vector<at::IValue>& stack) const {
   TORCH_INTERNAL_ASSERT(scriptCall.hasOp());
-  scriptCall.op()->getOperation()(&stack);
-  TORCH_INTERNAL_ASSERT(
-      stack.size() == 1,
-      "Return value of a builtin operator or a TorchScript function should be "
-      "a single IValue, got a vector of size ",
-      stack.size());
-  markComplete(std::move(ScriptResp(std::move(stack.front()))).toMessage());
+  auto future = runJitOperator(*scriptCall.op(), stack);
+  future->addCallback([markComplete](JitFuture& future) {
+    markComplete(ScriptResp(future.value()).toMessage());
+  });
 }
 
 TypePtr RequestCallbackNoPython::getScriptRemoteCallType(
@@ -246,23 +243,15 @@ void RequestCallbackNoPython::processScriptRemoteCallOp(
     std::vector<at::IValue>& stack,
     const c10::intrusive_ptr<OwnerRRef>& ownerRRef) const {
   TORCH_INTERNAL_ASSERT(scriptRemoteCall.hasOp());
-  try {
-    scriptRemoteCall.op()->getOperation()(&stack);
-  } catch (const std::exception& e) {
-    // Don't throw in this call, but rather transfer the exception
-    // to the rref.
-    ownerRRef->setError(std::current_exception());
+  auto future = runJitOperator(*scriptRemoteCall.op(), stack);
+  future->addCallback([ownerRRef, postProcessing](JitFuture& future) {
+    if (future.hasError()) {
+      ownerRRef->setError(future.exception_ptr());
+    } else {
+      ownerRRef->setValue(future.value());
+    }
     postProcessing();
-    return;
-  }
-  TORCH_INTERNAL_ASSERT(
-      stack.size() == 1,
-      "Return value of a builtin operator or a "
-      "TorchScript function should be a single IValue, got a vector of "
-      "size ",
-      stack.size());
-  ownerRRef->setValue(std::move(stack.front()));
-  postProcessing();
+  });
 }
 
 void RequestCallbackNoPython::processScriptRRefFetchCall(
@@ -660,6 +649,23 @@ bool RequestCallbackNoPython::cudaAvailable() const {
 #endif
 }
 
+c10::intrusive_ptr<JitFuture> RequestCallbackNoPython::runJitOperator(
+    const jit::Operator& op,
+    std::vector<at::IValue>& stack) const {
+  try {
+    op.getOperation()(&stack);
+  } catch (const std::exception&) {
+    return asFuture(std::current_exception());
+  }
+  TORCH_INTERNAL_ASSERT(
+      stack.size() == 1,
+      "Return value of a builtin operator or a TorchScript function should be "
+      "a single IValue, got a vector of size ",
+      stack.size());
+  TypePtr type = stack.front().type();
+  return asFuture(std::move(stack.front()), std::move(type));
+}
+
 c10::intrusive_ptr<JitFuture> RequestCallbackNoPython::asFuture(
     IValue value,
     TypePtr type) const {
diff --git a/torch/csrc/distributed/rpc/request_callback_no_python.h b/torch/csrc/distributed/rpc/request_callback_no_python.h
@@ -124,6 +124,12 @@ class TORCH_API RequestCallbackNoPython : public RequestCallback {
       RpcCommandBase& rpc,
       const c10::intrusive_ptr<JitFuture>& responseFuture) const;
 
+  // Helpers to run user-defined functions, operators and other computations.
+
+  c10::intrusive_ptr<JitFuture> runJitOperator(
+      const jit::Operator& op,
+      std::vector<at::IValue>& stack) const;
+
   // Helpers to convert various kinds of objects into already-completed futures.
 
   c10::intrusive_ptr<JitFuture> asFuture(IValue value, TypePtr type) const;
