manager: expose participating_rank (#94)

Author: d4l3k

torchft/manager.py

@@ -652,15 +652,35 @@ def batches_committed(self) -> int:
         """
         return self._batches_committed
 
+    def participating_rank(self) -> Optional[int]:
+        """
+        Get the replica group rank of the current quorum. This will be the same on all
+        ranks within the replica group.
+
+        If this replica group is not participating in the current quorum, this will be None.
+
+        This will block on the async quorum if it is not yet ready.
+
+        Returns:
+            the rank of the current quorum
+        """
+        self.wait_quorum()
+
+        return self._participating_rank
+
     def num_participants(self) -> int:
         """
         Get the number of participants in the current quorum.
 
         This is the number of replicas participating in the current step.
 
+        This will block on the async quorum if it is not yet ready.
+
         Returns:
             the number of participants in the current quorum
         """
+        self.wait_quorum()
+
         assert self._participating_world_size >= 0, "internal error"
         return self._participating_world_size
 

torchft/manager_test.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import concurrent
 from datetime import timedelta
 from typing import Optional
 from unittest import TestCase
@@ -521,10 +522,16 @@ def test_manager_wrap_future_timeout(self, client_mock: MagicMock) -> None:
     def test_manager_numerics(self, client_mock: MagicMock) -> None:
         manager = self._create_manager()
 
-        manager._quorum_future = MagicMock()
+        manager._quorum_future = quorum_future = MagicMock(
+            spec=concurrent.futures.Future
+        )
         manager._participating_rank = 1
         manager._participating_world_size = 5
         self.assertEqual(manager.num_participants(), 5)
+        self.assertEqual(quorum_future.result.call_count, 1)
+        self.assertEqual(manager.participating_rank(), 1)
+        self.assertEqual(quorum_future.result.call_count, 2)
+
         # pyre-ignore[16]: _pg is mocked
         manager._pg.allreduce.return_value = _DummyWork(None)