fix managed pg allreduce

Summary:
managed pg allreduce should just call manager's allreduce

Author: tushar00jain

torchft/process_group.py

@@ -1059,30 +1059,6 @@ def callback(
         return work
 
 
-class _ManagedWork(Work):
-    def __init__(self, manager: "Manager", work: Work, default_result: object) -> None:
-        super().__init__()
-
-        self._manager = manager
-        self._work = work
-        self._default_result = default_result
-
-    def wait(self, timeout: Optional[timedelta] = None) -> bool:
-        try:
-            if self._work is not None:
-                if timeout is not None:
-                    self._work.wait(timeout)
-                else:
-                    self._work.wait()
-        except Exception as e:
-            self._manager.report_error(e)
-
-        return True
-
-    def get_future(self) -> Future[object]:
-        return self._manager.wrap_future(self._work.get_future(), self._default_result)
-
-
 class ManagedProcessGroup(ProcessGroupWrapper):
     """
     This is a wrapper around any ProcessGroup that is managed by a torchft
@@ -1102,23 +1078,7 @@ def __init__(self, manager: "Manager") -> None:
         self._manager = manager
 
     def allreduce(self, tensors: List[torch.Tensor], opts: object) -> Work:
-        # Ensure we have a valid quorum and are configured before trying to do
-        # any work.
-        self._manager.wait_quorum()
-
-        if self._manager.errored() is not None:
-            return _DummyWork(tensors)
-        try:
-            work = super().allreduce(tensors, opts)
-        except Exception as e:
-            self._manager.report_error(e)
-            return _DummyWork(tensors)
-
-        return _ManagedWork(
-            self._manager,
-            work,
-            tensors,
-        )
+        return self._manager.allreduce(tensors)
 
     def size(self) -> int:
         return self._manager.num_participants()

torchft/process_group_test.py

@@ -48,7 +48,6 @@
     ProcessGroupNCCL,
     ProcessGroupWrapper,
     _ErrorSwallowingWork,
-    _ManagedWork,
     extend_device_mesh,
     ft_init_device_mesh,
 )
@@ -810,11 +809,8 @@ def test_managed_process_group(self) -> None:
         self.assertEqual(pg.size(), 123)
 
         works = _test_pg(pg)
-        self.assertIsInstance(list(works.values())[0], _ManagedWork)
 
-        self.assertEqual(manager.report_error.call_count, 0)
-        self.assertEqual(manager.wrap_future.call_count, 2)
-        self.assertEqual(manager.wait_quorum.call_count, 2)
+        self.assertEqual(manager.allreduce.call_count, 2)
 
 
 class DeviceMeshTest(TestCase):