Update on "sync and async torch.distributed.rpc for builtin operators"

mrshenli · mrshenli · commit f0d6fa304675 · 2019-08-05T15:44:44.000-07:00
Features: * sync and async RPC for builtin operators * RpcAgent API * ProcessGroupAgent implementation Goal: This is the first PR for #23110, and there will be many followup ones. So let's focus on the overall API and code structure. Details like efficiency and error handling can be improved in future PRs. * have a minimum working and testable RPC implementation. * make sure the RpcAgent API is sufficient for future ThriftAgent and TensorPipeAgent implementation * For tensor pipe implementation, it might allocate multiple underlying communication channels with different types, and might also use streaming serialization/deserialization for large tensors. To support this requirement, the current implementation only convert a BuiltinOp into a Message which contains a byte vector and a tensor table. It is up to the RpcAgent implementation to determine how it would like to serialize a Message object. * For ThriftAgent, as Thrift has it own request/response matching solution, the Message.id is no longer necessary. Hence the id can be dropped during serialization. All it needs to do is to pass the response Message object to the Future returned by send(...). * support blocking and non-blocking RequestCallback * blocking means the callback won't return before sending out the response * non-blocking can be achieved by enqueue the `(from, request, RpcAgent&)` tuple and use a different thread to process them. That is why there is an `RpcAgent&` arg in the param list. Differential Revision: [D15194693](https://our.internmc.facebook.com/intern/diff/D15194693/)
diff --git a/torch/csrc/distributed/rpc/ScriptCall.cpp b/torch/csrc/distributed/rpc/ScriptCall.cpp
@@ -11,6 +11,10 @@ using torch::jit::Unpickler;
 
 } // namespace
 
+static constexpr char BUILTIN_OP_NAMESPACE_[] = "torch.ops.aten.";
+static constexpr char ATEN_PREFIX_[] = "aten::";
+static constexpr int ATEN_PREFIX_LEN_ = 6;
+
 ScriptCall::ScriptCall(
     std::shared_ptr<Operator> op, std::vector<at::IValue>&& args)
     : op_(std::move(op)), stack_(args) {}
diff --git a/torch/csrc/distributed/rpc/ScriptCall.h b/torch/csrc/distributed/rpc/ScriptCall.h
@@ -32,10 +32,6 @@ class TORCH_API ScriptCall final {
   static std::shared_ptr<Operator> matchOperator(
       at::Symbol& symbol, const std::string& str_schema);
 
-  static constexpr char BUILTIN_OP_NAMESPACE_[] = "torch.ops.aten.";
-  static constexpr char ATEN_PREFIX_[] = "aten::";
-  static constexpr int ATEN_PREFIX_LEN_ = 6;
-
   // This field has value if this ScriptCall represents invocation of a builtin
   // operator.
   c10::optional<std::shared_ptr<Operator>> op_;
