Update ProcessGroupRoundRobin (#91172)

H-Huang · pytorchmergebot · commit f8b348c1fc96 · 2022-12-20T19:53:34.000Z
Summary: Temporary fix to unblock jobs in https://fb.workplace.com/groups/300451907202972/permalink/906337097050850/ Real fix would be to remove use of _round_robin_process_group API and update corresponding references (e.g. PyText) Test Plan: sandcastle Differential Revision: D42169592 Pull Request resolved: #91172 Approved by: https://github.com/awgu
diff --git a/test/distributed/test_c10d_gloo.py b/test/distributed/test_c10d_gloo.py
@@ -1415,11 +1415,10 @@ def test_barrier_implies_wait(self):
     def test_round_robin(self):
         num_process_groups = 2
         store = c10d.FileStore(self.file_name, self.world_size)
+        c10d.init_process_group(backend="gloo", store=store, rank=self.rank, world_size=self.world_size)
         pg = c10d._round_robin_process_groups(
             [
-                self._create_process_group_gloo(
-                    c10d.PrefixStore(str(i), store), self.rank, self.world_size, self.opts()
-                )
+                c10d.new_group(pg_options=self.opts())
                 for i in range(num_process_groups)
             ]
         )
@@ -1434,13 +1433,12 @@ def test_round_robin(self):
     @requires_gloo()
     def test_round_robin_create_destroy(self):
         store = c10d.FileStore(self.file_name, self.world_size)
+        c10d.init_process_group(backend="gloo", store=store, rank=self.rank, world_size=self.world_size)
 
         def create(num, prefix):
             return c10d._round_robin_process_groups(
                 [
-                    self._create_process_group_gloo(
-                        c10d.PrefixStore("%s/%d" % (prefix, i), store), self.rank, self.world_size, self.opts()
-                    )
+                    c10d.new_group(pg_options=self.opts())
                     for i in range(num)
                 ]
             )
diff --git a/torch/csrc/distributed/c10d/ProcessGroupRoundRobin.cpp b/torch/csrc/distributed/c10d/ProcessGroupRoundRobin.cpp
@@ -5,7 +5,7 @@ namespace c10d {
 ProcessGroupRoundRobin::ProcessGroupRoundRobin(
     int rank,
     int size,
-    std::vector<c10::intrusive_ptr<Backend>> processGroups)
+    std::vector<c10::intrusive_ptr<ProcessGroup>> processGroups)
     : ProcessGroup(rank, size), processGroups_(std::move(processGroups)) {
   TORCH_WARN(
       "ProcessGroupRoundRobin is deprecated and scheduled to be removed after this current release (1.13). ",
@@ -114,7 +114,7 @@ c10::intrusive_ptr<Work> ProcessGroupRoundRobin::barrier(
   TORCH_CHECK(false, "ProcessGroupRoundRobin does not support barrier");
 };
 
-const c10::intrusive_ptr<Backend>& ProcessGroupRoundRobin::next() {
+const c10::intrusive_ptr<ProcessGroup>& ProcessGroupRoundRobin::next() {
   auto& processGroup = *iterator_;
   iterator_++;
   if (iterator_ == processGroups_.end()) {
diff --git a/torch/csrc/distributed/c10d/ProcessGroupRoundRobin.hpp b/torch/csrc/distributed/c10d/ProcessGroupRoundRobin.hpp
@@ -23,7 +23,7 @@ class TORCH_API ProcessGroupRoundRobin final : public ProcessGroup {
   explicit ProcessGroupRoundRobin(
       int rank,
       int size,
-      std::vector<c10::intrusive_ptr<Backend>> processGroups);
+      std::vector<c10::intrusive_ptr<ProcessGroup>> processGroups);
 
   ~ProcessGroupRoundRobin() override;
 
@@ -103,11 +103,11 @@ class TORCH_API ProcessGroupRoundRobin final : public ProcessGroup {
       const BarrierOptions& opts = BarrierOptions()) override;
 
  private:
-  std::vector<c10::intrusive_ptr<Backend>> processGroups_;
-  std::vector<c10::intrusive_ptr<Backend>>::const_iterator iterator_;
+  std::vector<c10::intrusive_ptr<ProcessGroup>> processGroups_;
+  std::vector<c10::intrusive_ptr<ProcessGroup>>::const_iterator iterator_;
 
   // Returns the next ProcessGroup to use.
-  const c10::intrusive_ptr<Backend>& next();
+  const c10::intrusive_ptr<ProcessGroup>& next();
 };
 
 } // namespace c10d
diff --git a/torch/csrc/distributed/c10d/init.cpp b/torch/csrc/distributed/c10d/init.cpp
@@ -1627,7 +1627,7 @@ options :class:`~torch.distributed.ProcessGroupNCCL.Options`).
 #ifndef _WIN32
   module.def(
       "_round_robin_process_groups",
-      [](std::vector<c10::intrusive_ptr<::c10d::Backend>> processGroups)
+      [](std::vector<c10::intrusive_ptr<::c10d::ProcessGroup>> processGroups)
           -> c10::intrusive_ptr<::c10d::ProcessGroup> {
         if (processGroups.size() == 0) {
           throw std::invalid_argument("Specify at least 1 process group");

Original file line number	Diff line number	Diff line change
`@@ -1415,11 +1415,10 @@ def test_barrier_implies_wait(self):`
`1415`	`1415`	`def test_round_robin(self):`
`1416`	`1416`	`num_process_groups = 2`
`1417`	`1417`	`store = c10d.FileStore(self.file_name, self.world_size)`
	`1418`	`+ c10d.init_process_group(backend="gloo", store=store, rank=self.rank, world_size=self.world_size)`
`1418`	`1419`	`pg = c10d._round_robin_process_groups(`
`1419`	`1420`	`[`
`1420`		`- self._create_process_group_gloo(`
`1421`		`- c10d.PrefixStore(str(i), store), self.rank, self.world_size, self.opts()`
`1422`		`- )`
	`1421`	`+ c10d.new_group(pg_options=self.opts())`
`1423`	`1422`	`for i in range(num_process_groups)`
`1424`	`1423`	`]`
`1425`	`1424`	`)`
`@@ -1434,13 +1433,12 @@ def test_round_robin(self):`
`1434`	`1433`	`@requires_gloo()`
`1435`	`1434`	`def test_round_robin_create_destroy(self):`
`1436`	`1435`	`store = c10d.FileStore(self.file_name, self.world_size)`
	`1436`	`+ c10d.init_process_group(backend="gloo", store=store, rank=self.rank, world_size=self.world_size)`
`1437`	`1437`
`1438`	`1438`	`def create(num, prefix):`
`1439`	`1439`	`return c10d._round_robin_process_groups(`
`1440`	`1440`	`[`
`1441`		`- self._create_process_group_gloo(`
`1442`		`- c10d.PrefixStore("%s/%d" % (prefix, i), store), self.rank, self.world_size, self.opts()`
`1443`		`- )`
	`1441`	`+ c10d.new_group(pg_options=self.opts())`
`1444`	`1442`	`for i in range(num)`
`1445`	`1443`	`]`
`1446`	`1444`	`)`