Revert D18549919: Add RpcAgentOptions struct type, which bundles different required arguments for different RpcAgents

Test Plan: revert-hammer

Differential Revision:
D18549919

Original commit changeset: b9f3f1a41d1f

fbshipit-source-id: 2d5e578d18c0725b59eb99a0e942fbf7fe3341ee

Author: ezyang

test/dist_utils.py

@@ -2,6 +2,7 @@
 
 import threading
 from functools import partial, wraps
+from os import getenv
 
 import torch.distributed as dist
 import torch.distributed.rpc as rpc
@@ -13,15 +14,15 @@
 
 
 class TestConfig:
-    __slots__ = ["rpc_backend_name", "build_rpc_agent_options"]
+    __slots__ = ["rpc_backend_name"]
 
     def __init__(self, *args, **kwargs):
         assert len(args) == 0, "TestConfig only takes kwargs."
         for k, v in kwargs.items():
             setattr(self, k, v)
 
 
-TEST_CONFIG = TestConfig()
+TEST_CONFIG = TestConfig(rpc_backend_name=getenv("RPC_BACKEND_NAME", "PROCESS_GROUP"))
 INIT_METHOD_TEMPLATE = "file://{file_name}"
 
 
@@ -73,20 +74,22 @@ def dist_init(old_test_method=None, setup_rpc=True, clean_shutdown=True):
     @wraps(old_test_method)
     def new_test_method(self, *arg, **kwargs):
         self.worker_id = self.rank
+        self.worker_name_to_id = {
+            "worker{}".format(rank): rank for rank in range(self.world_size)
+        }
 
         if setup_rpc:
             global _ALL_NODE_NAMES
-            _ALL_NODE_NAMES = {
-                "worker{}".format(rank) for rank in range(self.world_size)
-            }
+            _ALL_NODE_NAMES = self.worker_name_to_id.keys()
 
+            # Use enough 'num_send_recv_threads' until we fix https://github.com/pytorch/pytorch/issues/26359
             rpc.init_rpc(
-                name="worker%d" % self.rank,
+                self_name="worker%d" % self.rank,
                 backend=self.rpc_backend,
                 init_method=self.init_method,
-                rank=self.rank,
-                world_size=self.world_size,
-                rpc_agent_options=self.rpc_agent_options,
+                self_rank=self.rank,
+                worker_name_to_id=self.worker_name_to_id,
+                num_send_recv_threads=16,
             )
 
         return_value = old_test_method(self, *arg, **kwargs)
@@ -128,12 +131,3 @@ def new_test_method(self, *arg, **kwargs):
         return return_value
 
     return new_test_method
-
-
-# Set PROCESS_GROUP as the default RPC backend.
-TEST_CONFIG.rpc_backend_name = "PROCESS_GROUP"
-TEST_CONFIG.build_rpc_agent_options = lambda test_object: rpc.backend_registry.construct_rpc_agent_options(
-    test_object.rpc_backend,
-    # Use enough 'num_send_recv_threads' until we fix https://github.com/pytorch/pytorch/issues/26359
-    num_send_recv_threads=16,
-)

test/rpc_test.py

@@ -30,14 +30,8 @@ def decorator(old_func):
 VALUE_FUTURE = concurrent.futures.Future()
 
 
-def _stub_construct_rpc_agent_options_handler(
-    **kwargs
-):
-    return mock.Mock()  # RpcAgentOptions.
-
-
-def _stub_start_rpc_backend_handler(
-    store, name, rank, world_size, rpc_agent_options
+def stub_start_rpc_backend_handler(
+    store, self_name, self_rank, worker_name_to_id, *args, **kwargs
 ):
     return mock.Mock()  # RpcAgent.
 
@@ -288,27 +282,22 @@ def test_register_rpc_backend_and_start_rpc_backend(
         backend_name = "stub_backend"
 
         backend = rpc.backend_registry.register_backend(
-            backend_name,
-            _stub_construct_rpc_agent_options_handler,
-            _stub_start_rpc_backend_handler,
+            backend_name, stub_start_rpc_backend_handler
         )
 
         with self.assertRaisesRegex(
             RuntimeError, "^RPC backend .+: already registered$"
         ):
-            backend = rpc.backend_registry.register_backend(
-                backend_name,
-                _stub_construct_rpc_agent_options_handler,
-                _stub_start_rpc_backend_handler,
+            rpc.backend_registry.register_backend(
+                backend_name, stub_start_rpc_backend_handler
             )
 
         rpc.init_rpc(
-            name="worker1",
+            self_name="worker1",
             backend=backend,
             init_method=self.init_method,
-            rank=self.rank,
-            world_size=self.world_size,
-            rpc_agent_options=self.rpc_agent_options,
+            self_rank=self.rank,
+            worker_name_to_id=self.worker_name_to_id,
         )
 
     @requires_process_group_agent("PROCESS_GROUP rpc backend specific test, skip")
@@ -321,22 +310,20 @@ def test_duplicate_name(self):
             rpc._init_rpc_backend(
                 backend=self.rpc_backend,
                 store=store,
-                name="duplicate_name",
-                rank=self.rank,
-                world_size=self.world_size,
-                rpc_agent_options=self.rpc_agent_options,
+                self_name="duplicate_name",
+                self_rank=self.rank,
+                worker_name_to_id=self.worker_name_to_id,
             )
         rpc.join_rpc()
 
     @dist_init(setup_rpc=False)
     def test_reinit(self):
         rpc.init_rpc(
-            name="worker{}".format(self.rank),
+            self_name="worker{}".format(self.rank),
             backend=self.rpc_backend,
             init_method=self.init_method,
-            rank=self.rank,
-            world_size=self.world_size,
-            rpc_agent_options=self.rpc_agent_options,
+            self_rank=self.rank,
+            worker_name_to_id=self.worker_name_to_id,
         )
 
         # This is for the below `dist.barrier`.
@@ -354,12 +341,11 @@ def test_reinit(self):
 
         with self.assertRaisesRegex(RuntimeError, "is already initialized"):
             rpc.init_rpc(
-                name="worker{}".format(self.rank),
+                self_name="worker{}".format(self.rank),
                 backend=self.rpc_backend,
                 init_method=self.init_method,
-                rank=self.rank,
-                world_size=self.world_size,
-                rpc_agent_options=self.rpc_agent_options,
+                self_rank=self.rank,
+                worker_name_to_id=self.worker_name_to_id,
             )
         rpc.join_rpc()
 
@@ -372,10 +358,10 @@ def test_invalid_names(self):
             rpc._init_rpc_backend(
                 backend=self.rpc_backend,
                 store=store,
-                name="abc*",
-                rank=self.rank,
-                world_size=self.world_size,
-                rpc_agent_options=self.rpc_agent_options,
+                self_name="abc*",
+                self_rank=self.rank,
+                worker_name_to_id=self.worker_name_to_id,
+                num_send_recv_threads=16,
             )
 
         base_file_name = self.file_name
@@ -389,10 +375,10 @@ def test_invalid_names(self):
             rpc._init_rpc_backend(
                 backend=self.rpc_backend,
                 store=store,
-                name=" ",
-                rank=self.rank,
-                world_size=self.world_size,
-                rpc_agent_options=self.rpc_agent_options,
+                self_name=" ",
+                self_rank=self.rank,
+                worker_name_to_id=self.worker_name_to_id,
+                num_send_recv_threads=16,
             )
 
         # Use a different file path for FileStore to avoid rendezvous mismatch.
@@ -404,10 +390,10 @@ def test_invalid_names(self):
             rpc._init_rpc_backend(
                 backend=self.rpc_backend,
                 store=store,
-                name="",
-                rank=self.rank,
-                world_size=self.world_size,
-                rpc_agent_options=self.rpc_agent_options,
+                self_name="",
+                self_rank=self.rank,
+                worker_name_to_id=self.worker_name_to_id,
+                num_send_recv_threads=16,
             )
 
         # Use a different file path for FileStore to avoid rendezvous mismatch.
@@ -421,10 +407,10 @@ def test_invalid_names(self):
             rpc._init_rpc_backend(
                 backend=self.rpc_backend,
                 store=store,
-                name="".join(["a" for i in range(500)]),
-                rank=self.rank,
-                world_size=self.world_size,
-                rpc_agent_options=self.rpc_agent_options,
+                self_name="".join(["a" for i in range(500)]),
+                self_rank=self.rank,
+                worker_name_to_id=self.worker_name_to_id,
+                num_send_recv_threads=16,
             )
 
         from torch.distributed.rpc.api import _agent
@@ -510,12 +496,11 @@ def test_multi_rpc(self):
     def test_join_rpc(self):
         # Initialize RPC.
         rpc.init_rpc(
-            name="worker%d" % self.rank,
+            self_name="worker%d" % self.rank,
             backend=self.rpc_backend,
             init_method=self.init_method,
-            rank=self.rank,
-            world_size=self.world_size,
-            rpc_agent_options=self.rpc_agent_options,
+            self_rank=self.rank,
+            worker_name_to_id=self.worker_name_to_id,
         )
 
         n = self.rank + 1
@@ -1091,19 +1076,13 @@ def test_call_method_on_rref(self):
     @dist_init(setup_rpc=False)
     def test_get_rpc_timeout(self):
         timeout = timedelta(seconds=1)
-
-        # A new `RpcAgentOptions` is constructed
-        # when accessing `self.rpc_agent_options`.
-        rpc_agent_options = self.rpc_agent_options
-        rpc_agent_options.rpc_timeout = timeout
-
         rpc.init_rpc(
-            name="worker{}".format(self.rank),
+            self_name="worker{}".format(self.rank),
             backend=self.rpc_backend,
             init_method=self.init_method,
-            rank=self.rank,
-            world_size=self.world_size,
-            rpc_agent_options=rpc_agent_options,
+            self_rank=self.rank,
+            worker_name_to_id=self.worker_name_to_id,
+            rpc_timeout=timeout
         )
         set_timeout = rpc.get_rpc_timeout()
         self.assertEqual(timeout, set_timeout)

torch/csrc/distributed/rpc/init.cpp

@@ -34,10 +34,6 @@ PyObject* rpc_init(PyObject* /* unused */) {
 
   auto module = py::handle(rpc_module).cast<py::module>();
 
-  auto rpcAgentOptions =
-      shared_ptr_class_<RpcAgentOptions>(module, "RpcAgentOptions")
-          .def_readwrite("rpc_timeout", &RpcAgentOptions::rpcTimeout);
-
   auto workerInfo =
       shared_ptr_class_<WorkerInfo>(module, "WorkerInfo")
           .def_readonly("name", &WorkerInfo::name_)
@@ -102,13 +98,6 @@ PyObject* rpc_init(PyObject* /* unused */) {
               [&](FutureMessage& fut) { return toPyObj(fut.wait()); },
               py::call_guard<py::gil_scoped_release>());
 
-  shared_ptr_class_<ProcessGroupRpcAgentOptions>(
-      module, "ProcessGroupRpcAgentOptions", rpcAgentOptions)
-      .def(py::init<>())
-      .def_readwrite(
-          "num_send_recv_threads",
-          &ProcessGroupRpcAgentOptions::numSendRecvThreads);
-
   shared_ptr_class_<ProcessGroupAgent>(module, "ProcessGroupAgent", rpcAgent)
       .def(
           py::init<

torch/csrc/distributed/rpc/process_group_agent.h

@@ -13,11 +13,6 @@ namespace torch {
 namespace distributed {
 namespace rpc {
 
-struct ProcessGroupRpcAgentOptions : public RpcAgentOptions {
-  ProcessGroupRpcAgentOptions() noexcept = default;
-  int numSendRecvThreads;
-};
-
 // SendWork and RecvWork will be put into a task queue, and later picked up by
 // worker threads from the same ThreadPool.
 struct SendWork {

torch/csrc/distributed/rpc/rpc_agent.h

@@ -12,11 +12,6 @@ namespace torch {
 namespace distributed {
 namespace rpc {
 
-struct RpcAgentOptions {
-  RpcAgentOptions() noexcept = default;
-  std::chrono::milliseconds rpcTimeout;
-};
-
 // A globally unique ID to identify an RpcAgent
 struct TORCH_API WorkerInfo {
   WorkerInfo(std::string name, int id)

torch/distributed/rendezvous.py

@@ -3,8 +3,6 @@
 except ImportError:
     from urlparse import urlparse
 
-import torch._six as six
-import numbers
 import os
 from . import FileStore, TCPStore
 
@@ -44,15 +42,6 @@ def register_rendezvous_handler(scheme, handler):
 
 
 def rendezvous(url, rank=-1, world_size=-1, **kwargs):
-    if not isinstance(url, six.string_classes):
-        raise RuntimeError("`url` must be a string. {}: {}".format(type(url), url))
-
-    if not isinstance(rank, numbers.Integral):
-        raise RuntimeError("`rank` must be an integer. {}".format(rank))
-
-    if not isinstance(world_size, numbers.Integral):
-        raise RuntimeError("`world_size` must be an integer. {}".format(world_size))
-
     # Append node-specific arguments.
     if rank != -1 or world_size != -1:
         assert (