Update on "Add autograd hook for python rpc call"

zhaojuanmao · zhaojuanmao · commit c0cd280289ce · 2019-10-17T12:07:16.000-07:00
1. currently if autograd context is valid, even tensors do not require grads and grads function are not attached. it still send rpc with autograd meta. This is not ideal. This diff makes some change to make sure rpc with autograd meta is sent only if autograd context is valid and tensors require grads 2. meanwhile create a utiliy to attach autograd info and functions as needed 3. add autograd send/recv functions for python rpc call 4. make changes to support nested python rpc calls 5. disallow nested dist autograd context (was landed in #27022) Differential Revision: [D17819153](https://our.internmc.facebook.com/intern/diff/D17819153/) [ghstack-poisoned]
diff --git a/torch/csrc/distributed/autograd/utils.h b/torch/csrc/distributed/autograd/utils.h
@@ -36,13 +36,13 @@ TORCH_API DistAutogradContext* addRecvRpcBackward(
 // and attach autograd function for each type of rpc call if it has valid
 // context and tensors require grads, in this case, return RpcWithAutograd
 // message; otherwise return original rpc message.
-rpc::Message getMessageWithAutogradCheck(
+TORCH_API rpc::Message getMessageWithAutogradCheck(
     const rpc::worker_id_t dstId,
     rpc::Message&& wrappedRpcMsg,
     rpc::MessageType msgType);
 
 // Send message after autograd checking
-std::shared_ptr<torch::distributed::rpc::FutureMessage> sendMessage(
+TORCH_API std::shared_ptr<torch::distributed::rpc::FutureMessage> sendMessage(
     rpc::RpcAgent& agent,
     const rpc::WorkerInfo& dst,
     rpc::Message&& wrappedRpcMsg,
