fix

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

Author: 0oshowero0

tests/test_samplers.py

@@ -480,9 +480,9 @@ def test_rank_aware_sampler_second_rank_gets_cached(self):
         )
 
         assert sampled1 == sampled2 == [0, 1, 2]
-        # First rank returns empty consumed (not all ranks have fetched yet)
+        # First rank already returns consumed indexes
         assert consumed1 == [0, 1, 2]
-        # Last rank returns consumed when all ranks have fetched
+        # Second rank also sees the same consumed indexes; state is then cleaned up
         assert consumed2 == [0, 1, 2]
         # State should be cleaned up
         assert sampler._states == {}

transfer_queue/sampler/rank_aware_sampler.py

@@ -75,20 +75,22 @@ def sample(
                 ready samples, all available samples will be returned.
             dp_group: The group id of current data parallel group. Used to
                 identify which DP group this rank belongs to.
-            dp_world_size: Number of ranks in the data parallel group. Used to
+            dp_world_size: Number of ranks in the data parallelism group. Used to
                 determine when all ranks have fetched their samples.
             world_size: Total number of ranks across all parallelism dimensions.
                 Used to determine when all ranks have fetched their samples.
             *args: Additional positional arguments (ignored).
             **kwargs: Additional keyword arguments (ignored).
 
         Returns:
-            List of sampled global indices of length batch_size
+            List of sampled global indices. The length is
+            min(batch_size, len(ready_indexes)), and may be smaller than
+            batch_size if fewer ready samples are available.
 
-            List of global indices of length batch_size that should be labeled as consumed
-            (will never be retrieved in the future)
+            List of global indices that should be labeled as consumed
+            (will never be retrieved by other dp_groups in the future).
 
-        Raise:
+        Raises:
             RuntimeError: If ``world_size`` is not divisible by ``dp_world_size``.
 
         Note: