fix test_process_group_debug_info flaky test (pytorch#31533)

Summary: Pull Request resolved: pytorch#31533 Fixes this test that was flaky and has been disabled (see pytorch#31112) ghstack-source-id: 96038999 Test Plan: Run the test 1000 times and ensure that it passes. Differential Revision: D19203366 fbshipit-source-id: 7978cbb8ca0989a0a370a36349cdd4db3bb8345b

Author: rohan-varma

test/rpc_test.py

@@ -1203,11 +1203,17 @@ def test_rref_context_debug_info(self):
  # barrier after check 3
  dist.barrier()
 
- @unittest.skip("Test is flaky, see https://github.com/pytorch/pytorch/issues/31112")
  @dist_init
  @requires_process_group_agent("PROCESS_GROUP rpc backend specific test, skip")
  def test_process_group_debug_info(self):
  from torch.distributed.rpc.api import _agent
+ if not dist.is_initialized():
+ dist.init_process_group(
+ backend="gloo",
+ init_method=self.init_method,
+ rank=self.rank,
+ world_size=self.world_size,
+ )
 
  NUM_THREAD = self.rpc_backend_options.num_send_recv_threads
 
@@ -1218,7 +1224,10 @@ def test_process_group_debug_info(self):
  self.assertEqual(int(info["num_pending_requests"]), 0)
  self.assertEqual(int(info["thread_pool_size"]), NUM_THREAD)
  self.assertEqual(int(info["num_idle_threads"]), NUM_THREAD)
-
+ # for the above check, add a barrier to ensure that another worker
+ # cannot send a request before we check num_idle_threads, since we'd
+ # use up an idle thread if we start processing that request.
+ dist.barrier()
  dst_rank = (self.rank + 1) % self.world_size
  fut = rpc.rpc_async(
  "worker{}".format(dst_rank),
@@ -1239,14 +1248,6 @@ def test_process_group_debug_info(self):
  # might be either 1 or 2 busy threads
  self.assertTrue(num_idle_threads in [NUM_THREAD - 1, NUM_THREAD - 2])
 
- if not dist.is_initialized():
- dist.init_process_group(
- backend="gloo",
- init_method=self.init_method,
- rank=self.rank,
- world_size=self.world_size,
- )
-
  # add a barrier to make sure the request is not finished before checking
  # num_pending_requests
  dist.barrier()