[TensorPipe] Allow passing args to agent options constructor (#37918)

lw · facebook-github-bot · commit 978ad162905b · 2020-05-07T02:50:47.000-07:00
Summary: Pull Request resolved: #37918 ghstack-source-id: 103569096 Test Plan: Tested top of stack Reviewed By: jiayisuse Differential Revision: D21425537 fbshipit-source-id: 2e78d700ea774944c7fd8b22e152d8e459dd422a
diff --git a/torch/csrc/distributed/rpc/init.cpp b/torch/csrc/distributed/rpc/init.cpp
@@ -448,7 +448,11 @@ If the future completes with an error, an exception is thrown.
   // Base class: torch.distributed.rpc.RpcBackendOptions.
   py::class_<TensorPipeRpcBackendOptions>(
       module, "TensorPipeRpcBackendOptions", rpcBackendOptions)
-      .def(py::init<>())
+      .def(
+          py::init<std::map<std::string, worker_id_t>, float, std::string>(),
+          py::arg("worker_name_to_id"),
+          py::arg("rpc_timeout") = kDefaultRpcTimeoutSeconds,
+          py::arg("init_method") = kDefaultInitMethod)
       .def_readwrite(
           "worker_name_to_id", &TensorPipeRpcBackendOptions::workerNameToId);
 
diff --git a/torch/csrc/distributed/rpc/tensorpipe_agent.h b/torch/csrc/distributed/rpc/tensorpipe_agent.h
@@ -15,6 +15,13 @@ namespace distributed {
 namespace rpc {
 
 struct TensorPipeRpcBackendOptions : public RpcBackendOptions {
+  TensorPipeRpcBackendOptions(
+      std::map<std::string, worker_id_t> worker_name_to_id,
+      float rpc_timeout,
+      std::string init_method)
+      : RpcBackendOptions(rpc_timeout, init_method),
+        workerNameToId(std::move(worker_name_to_id)) {}
+
   std::map<std::string, worker_id_t> workerNameToId;
 };
 
