Add bf.wait alias for synchronize

bluefog/torch/__init__.py

@@ -48,7 +48,7 @@
 from bluefog.torch.mpi_ops import neighbor_allreduce, neighbor_allreduce_nonblocking
 from bluefog.torch.mpi_ops import hierarchical_neighbor_allreduce
 from bluefog.torch.mpi_ops import hierarchical_neighbor_allreduce_nonblocking
-from bluefog.torch.mpi_ops import poll, synchronize, barrier
+from bluefog.torch.mpi_ops import poll, synchronize, wait, barrier
 
 from bluefog.torch.mpi_ops import win_create, win_free
 from bluefog.torch.mpi_ops import win_update, win_update_then_collect

bluefog/torch/mpi_ops.py

@@ -822,13 +822,12 @@ def pair_gossip_nonblocking(tensor: torch.Tensor, target_rank: int, self_weight:
 
 def poll(handle: int) -> bool:
     """
-    Polls an allreduce, allgather or broadcast handle to determine whether underlying
-    nonblocking operation has completed. After `poll()` returns `True`, `synchronize()`
+    Polls an allreduce, neighbor_allreduce, etc operation handle to determine whether underlying
+    nonblocking operation has completed. After `poll()` returns `True`, `wait()`
     will return without blocking.
 
     Arguments:
-        handle: A handle returned by an allreduce, allgather, broadcast, neighbor_allgather,
-        and neighbro_allreduce nonblocking operation.
+        handle: A handle returned by an allreduce, neighbor_allreduce, etc. nonblocking operation.
 
     Returns:
         A flag indicating whether the operation has completed.
@@ -838,12 +837,12 @@ def poll(handle: int) -> bool:
 
 def synchronize(handle: int) -> torch.Tensor:
     """
-    Synchronizes an nonblocking allreduce, allgather or broadcast operation until
+    Wait an allreduce, neighbor_allreduce, etc operation until
     it's completed. Returns the result of the operation.
+    It is the same function as `wait()`.
 
     Args:
-        handle: A handle returned by an allreduce, allgather or broadcast nonblocking
-                operation.
+        handle: A handle returned by an allreduce, neighbor_allreduce, etc. nonblocking operation.
 
     Returns:
         torch.Tensor: An output tensor of the operation.
@@ -855,6 +854,21 @@ def synchronize(handle: int) -> torch.Tensor:
     return output
 
 
+def wait(handle: int) -> torch.Tensor:
+    """
+    Wait an allreduce, neighbor_allreduce, etc operation until
+    it's completed. Returns the result of the operation.
+    It is just alias of `synchronize()` function.
+
+    Args:
+        handle: A handle returned by an allreduce, neighbor_allreduce, etc. nonblocking operation.
+
+    Returns:
+        torch.Tensor: An output tensor of the operation.
+    """
+    return synchronize(handle)
+
+
 def barrier():
     """Barrier function to sychronize all MPI processes.