Revert D5689636: Add RpcAgentTestFixture to extract duplicate code

Test Plan: revert-hammer

Differential Revision:
D5689636

Original commit changeset: f35eea1359ad

fbshipit-source-id: 31928fce5e96b3beceefbc9a03f54769f10b7e1a

Author: ezyang

test/dist_autograd_test.py

@@ -7,9 +7,7 @@
 import torch.distributed as dist
 import torch.distributed.autograd as dist_autograd
 import torch.distributed.rpc as rpc
-import dist_utils
-from dist_utils import dist_init
-from rpc_agent_test_fixture import RpcAgentTestFixture
+from dist_utils import INIT_METHOD_TEMPLATE, dist_init, TEST_CONFIG
 
 import threading
 
@@ -161,7 +159,7 @@ class ExecMode(Enum):
 @unittest.skipIf(
     not torch._six.PY3, "Pytorch distributed autograd package " "does not support python2"
 )
-class DistAutogradTest(RpcAgentTestFixture):
+class DistAutogradTest(object):
 
     def _initialize_pg(self):
         # This is for tests using `dist.barrier`.
@@ -202,6 +200,14 @@ def _next_rank(self):
     def _check_rpc_done(self, rank_distance):
         _check_rpc_done(rank_distance)
 
+    @property
+    def world_size(self):
+        return 4
+
+    @property
+    def init_method(self):
+        return INIT_METHOD_TEMPLATE.format(file_name=self.file_name)
+
     @dist_init
     def test_autograd_context(self):
         # Verify max possible id.
@@ -1049,7 +1055,7 @@ def test_backward_autograd_engine_error(self):
                 # Run backwards, and validate we receive an error.
                 dist_autograd.backward([val.sum()])
 
-    @unittest.skipIf(dist_utils.TEST_CONFIG.rpc_backend_name == "PROCESS_GROUP",
+    @unittest.skipIf(TEST_CONFIG.rpc_backend_name == "PROCESS_GROUP",
                      "Skipping this test temporarily since ProcessGroupAgent does not report errors on node failures")
     @dist_init(clean_shutdown=False)
     def test_backward_node_failure(self):
@@ -1220,7 +1226,7 @@ def _wait_backward_done():
         while not DistAutogradTest._backward_done:
             time.sleep(0.1)
 
-    @unittest.skipIf(dist_utils.TEST_CONFIG.rpc_backend_name == "PROCESS_GROUP",
+    @unittest.skipIf(TEST_CONFIG.rpc_backend_name == "PROCESS_GROUP",
                      "Skipping this test temporarily since ProcessGroupAgent " +
                      "does not report errors on node failures")
     @dist_init(clean_shutdown=False)

test/dist_optimizer_test.py

@@ -2,14 +2,13 @@
 
 import unittest
 
-from dist_utils import dist_init
+from dist_utils import INIT_METHOD_TEMPLATE, dist_init
 from torch import optim
 from torch.distributed.optim import DistributedOptimizer
 import torch
 import torch.distributed.autograd as dist_autograd
 import torch.distributed.rpc as rpc
 import threading
-from rpc_agent_test_fixture import RpcAgentTestFixture
 
 
 class MyModule:
@@ -84,7 +83,17 @@ def rpc_async_method(method, obj_rref, *args, **kwargs):
 @unittest.skipIf(
     not torch._six.PY3, "Pytorch distributed optim does not support python2"
 )
-class DistOptimizerTest(RpcAgentTestFixture):
+class DistOptimizerTest(object):
+
+    @property
+    def world_size(self):
+        return 4
+
+    @property
+    def init_method(self):
+        return INIT_METHOD_TEMPLATE.format(
+            file_name=self.file_name, rank=self.rank, world_size=self.world_size
+        )
 
     @dist_init()
     def test_dist_optim_exception(self):

test/dist_utils.py

@@ -85,7 +85,7 @@ def new_test_method(self, *arg, **kwargs):
             # Use enough 'num_send_recv_threads' until we fix https://github.com/pytorch/pytorch/issues/26359
             rpc.init_rpc(
                 self_name="worker%d" % self.rank,
-                backend=self.rpc_backend,
+                backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name],
                 init_method=self.init_method,
                 self_rank=self.rank,
                 worker_name_to_id=self.worker_name_to_id,

test/rpc_agent_test_fixture.py

@@ -12,16 +12,14 @@
 import torch.distributed.rpc as rpc
 from torch.distributed.rpc import RRef
 from common_utils import load_tests
-import dist_utils
-from dist_utils import dist_init
+from dist_utils import INIT_METHOD_TEMPLATE, TEST_CONFIG, dist_init
 from torch.distributed.rpc.internal import PythonUDF, _internal_rpc_pickler
-from rpc_agent_test_fixture import RpcAgentTestFixture
 
 
 def requires_process_group_agent(message=""):
     def decorator(old_func):
         return unittest.skipUnless(
-            dist_utils.TEST_CONFIG.rpc_backend_name == "PROCESS_GROUP", message
+            TEST_CONFIG.rpc_backend_name == "PROCESS_GROUP", message
         )(old_func)
 
     return decorator
@@ -212,7 +210,15 @@ def raise_func():
     sys.version_info < (3, 0),
     "Pytorch distributed rpc package " "does not support python2",
 )
-class RpcTest(RpcAgentTestFixture):
+class RpcTest(object):
+    @property
+    def world_size(self):
+        return 4
+
+    @property
+    def init_method(self):
+        return INIT_METHOD_TEMPLATE.format(file_name=self.file_name)
+
     @dist_init
     def test_worker_id(self):
         n = self.rank + 1
@@ -308,7 +314,7 @@ def test_duplicate_name(self):
                 self.init_method, rank=self.rank, world_size=self.world_size
             ))
             rpc._init_rpc_backend(
-                backend=self.rpc_backend,
+                backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name],
                 store=store,
                 self_name="duplicate_name",
                 self_rank=self.rank,
@@ -320,7 +326,7 @@ def test_duplicate_name(self):
     def test_reinit(self):
         rpc.init_rpc(
             self_name="worker{}".format(self.rank),
-            backend=self.rpc_backend,
+            backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name],
             init_method=self.init_method,
             self_rank=self.rank,
             worker_name_to_id=self.worker_name_to_id,
@@ -342,7 +348,7 @@ def test_reinit(self):
         with self.assertRaisesRegex(RuntimeError, "is already initialized"):
             rpc.init_rpc(
                 self_name="worker{}".format(self.rank),
-                backend=self.rpc_backend,
+                backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name],
                 init_method=self.init_method,
                 self_rank=self.rank,
                 worker_name_to_id=self.worker_name_to_id,
@@ -356,7 +362,7 @@ def test_invalid_names(self):
                 self.init_method, rank=self.rank, world_size=self.world_size
             ))
             rpc._init_rpc_backend(
-                backend=self.rpc_backend,
+                backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name],
                 store=store,
                 self_name="abc*",
                 self_rank=self.rank,
@@ -373,7 +379,7 @@ def test_invalid_names(self):
                 self.init_method, rank=self.rank, world_size=self.world_size
             ))
             rpc._init_rpc_backend(
-                backend=self.rpc_backend,
+                backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name],
                 store=store,
                 self_name=" ",
                 self_rank=self.rank,
@@ -388,7 +394,7 @@ def test_invalid_names(self):
                 self.init_method, rank=self.rank, world_size=self.world_size
             ))
             rpc._init_rpc_backend(
-                backend=self.rpc_backend,
+                backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name],
                 store=store,
                 self_name="",
                 self_rank=self.rank,
@@ -405,7 +411,7 @@ def test_invalid_names(self):
                 self.init_method, rank=self.rank, world_size=self.world_size
             ))
             rpc._init_rpc_backend(
-                backend=self.rpc_backend,
+                backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name],
                 store=store,
                 self_name="".join(["a" for i in range(500)]),
                 self_rank=self.rank,
@@ -497,7 +503,7 @@ def test_join_rpc(self):
         # Initialize RPC.
         rpc.init_rpc(
             self_name="worker%d" % self.rank,
-            backend=self.rpc_backend,
+            backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name],
             init_method=self.init_method,
             self_rank=self.rank,
             worker_name_to_id=self.worker_name_to_id,
@@ -1078,7 +1084,7 @@ def test_get_rpc_timeout(self):
         timeout = timedelta(seconds=1)
         rpc.init_rpc(
             self_name="worker{}".format(self.rank),
-            backend=self.rpc_backend,
+            backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name],
             init_method=self.init_method,
             self_rank=self.rank,
             worker_name_to_id=self.worker_name_to_id,
@@ -1123,7 +1129,7 @@ def test_requires_process_group_agent_decorator(self):
         def test_func():
             return "expected result"
 
-        if dist_utils.TEST_CONFIG.rpc_backend_name == "PROCESS_GROUP":
+        if TEST_CONFIG.rpc_backend_name == "PROCESS_GROUP":
             self.assertEqual(test_func(), "expected result")
 
     def test_dist_init_decorator(self):

test/rpc_test.py

@@ -70,7 +70,6 @@
     TESTS.extend([
         'rpc_spawn',
         'dist_autograd_spawn',
-        'dist_optimizer_spawn',
     ])
 
 # skip < 3.6 b/c fstrings added in 3.6
@@ -81,18 +80,20 @@
 
 WINDOWS_BLACKLIST = [
     'distributed',
+    'rpc_fork',
     'rpc_spawn',
+    'dist_autograd_fork',
     'dist_autograd_spawn',
-    'dist_optimizer_spawn',
 ]
 
 ROCM_BLACKLIST = [
     'cpp_extensions',
     'distributed',
     'multiprocessing',
+    'rpc_fork',
     'rpc_spawn',
+    'dist_autograd_fork',
     'dist_autograd_spawn',
-    'dist_optimizer_spawn',
 ]
 
 DISTRIBUTED_TESTS_CONFIG = {}