feat: test_assign_receiver_ranks.py: add unit tests for _assign_receiver_ranks function

Author: specture724

checkpoint_engine/ps.py

@@ -11,7 +11,7 @@
 from collections.abc import Callable
 from datetime import timedelta
 from functools import lru_cache
-from typing import TYPE_CHECKING, Annotated, Any, BinaryIO, NamedTuple
+from typing import TYPE_CHECKING, Annotated, Any, BinaryIO, NamedTuple, TypeVar
 
 import httpx
 import numpy as np
@@ -531,17 +531,25 @@ def _gen_h2d_buckets(
         return _assign_receiver_ranks(buckets, actual_local_topo, remote_topo)
 
 
+T = TypeVar("T")
+
+
 def _assign_receiver_ranks(
-    buckets: list[tuple[int, H2DBucket]],
+    buckets: list[tuple[int, "T"]],
     local_topo: dict[str, set[int]],
     remote_topo: dict[str, set[int]],
-) -> list[tuple[int, int, H2DBucket]]:
+) -> list[tuple[int, int, "T"]]:
     """
     (owner_rank, bucket) -> (receiver_rank, owner_rank, bucket)
 
     Assign receiver ranks to buckets. If ranks is empty, assign the owner_rank as receiver_rank.
     GPU-rdma_device topology will be considered to make full use of the bandwidth.
     """
+    assert local_topo, "local_topo should not be empty"
+    assert remote_topo, "remote_topo should not be empty"
+    if not buckets:
+        logger.warning("bucket list is empty, no need to assign receiver ranks")
+        return []
     rank_to_rdma_device = {
         rank: rdma_device for rdma_device, ranks in remote_topo.items() for rank in ranks
     }

tests/test_assign_receiver_ranks.py

@@ -0,0 +1,36 @@
+import os
+import sys
+
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from checkpoint_engine.ps import _assign_receiver_ranks
+
+
+class TestAssignReceiverRanks:
+    def test_basic_functionality(self):
+        buckets = [(i % 8, f"bucket{i}") for i in range(100)]
+        local_topo = {f"rdma{i}": {i} for i in range(8)}
+        remote_topo = {f"rdma{i}": {i} for i in range(8)}
+
+        result = _assign_receiver_ranks(buckets, local_topo, remote_topo)
+
+        assert len(result) == 100
+        for item in result:
+            assert len(item) == 3
+            assert isinstance(item[0], int)  # receiver_rank
+            assert isinstance(item[1], int)  # owner_rank
+            assert isinstance(item[2], str)  # bucket
+
+        for receiver_rank, owner_rank, bucket in result:
+            assert receiver_rank in range(8)
+            assert owner_rank % 8 == receiver_rank
+            assert bucket in {f"bucket{i}" for i in range(100)}
+
+    def test_empty_buckets(self):
+        buckets = []
+        local_topo = {"rdma0": {0}}
+        remote_topo = {"rdma0": {0}}
+
+        result = _assign_receiver_ranks(buckets, local_topo, remote_topo)
+
+        assert result == []