process_group: added ManagedProcessGroup

Author: d4l3k

torchft/process_group.py

@@ -420,6 +420,34 @@ def allreduce(self, tensors: List[torch.Tensor], opts: object) -> Work:
             return _DummyWork(tensors)
 
 
+class ManagedProcessGroup(ErrorSwallowingProcessGroupWrapper):
+    """
+    This is a wrapper around any ProcessGroup that is managed by a torchft
+    Manager.
+    """
+
+    def __init__(self, manager: "Manager") -> None:
+        super().__init__(manager._pg)
+
+        self._manager = manager
+
+    def report_error(self, e: Exception) -> None:
+        """
+        Report an error to this process group. This will cause all future
+        operations to be skipped until the process group is reconfigured via
+        ``configure``.
+
+        Args:
+            e: exception to report
+        """
+        super().report_error(e)
+
+        self._manager.report_error()
+
+    def size(self) -> int:
+        return self._manager.num_participants()
+
+
 class _BabyWork(Work):
     def __init__(
         self,

torchft/process_group_test.py

@@ -15,12 +15,14 @@
 from torch._C._distributed_c10d import _resolve_process_group
 from torch.distributed import _functional_collectives, ReduceOp, TCPStore
 from torch.distributed.device_mesh import init_device_mesh
+from torchft.manager import Manager
 
 from torchft.process_group import (
     _DummyWork,
     _ErrorSwallowingWork,
     ErrorSwallowingProcessGroupWrapper,
     extend_device_mesh,
+    ManagedProcessGroup,
     ProcessGroup,
     ProcessGroupBabyGloo,
     ProcessGroupBabyNCCL,
@@ -231,3 +233,16 @@ def test_error_swallowing_process_group_wrapper(self) -> None:
         work.wait()
         fut = work.get_future()
         fut.wait()
+
+    def test_managed_process_group(self) -> None:
+        manager = Mock(spec=Manager)
+        manager._pg = ProcessGroupDummy(0, 1)
+        pg = ManagedProcessGroup(manager)
+        manager.num_participants.return_value = 123
+
+        self.assertEqual(pg.size(), 123)
+
+        err = RuntimeError("test")
+        pg.report_error(err)
+        self.assertEqual(pg.error(), err)
+        self.assertEqual(manager.report_error.call_count, 1)