chore: import upstream snapshot with attribution

python/ray/train/v2/tests/test_sync_actor.py

@@ -0,0 +1,189 @@
+import pytest
+
+import ray
+from ray.exceptions import RayTaskError
+from ray.train.v2._internal.constants import DEFAULT_COLLECTIVE_TIMEOUT_S
+from ray.train.v2._internal.exceptions import BroadcastCollectiveTimeoutError
+from ray.train.v2._internal.execution.checkpoint.sync_actor import (
+    SynchronizationActor,
+)
+
+
+@pytest.fixture(autouse=True, scope="module")
+def ray_start_4_cpus():
+    ray.init(num_cpus=4)
+    yield
+    ray.shutdown()
+
+
+@pytest.mark.parametrize("world_size", [1, 10, 1000])
+def test_broadcast_from_rank_0(world_size):
+    """Check that rank 0 can broadcast data to all other workers.
+    Every worker sends data with a string "data-{rank}" that is unique
+    to the worker. Everyone should receive the data from rank 0, which is "data-0".
+    Also assert that the actor state is reset after the broadcast function returns.
+    """
+    sync_actor = SynchronizationActor.remote()
+    # Test broadcast_from_rank_zero with a world size of 10
+    remote_tasks = []
+    for rank in range(world_size):
+        remote_tasks.append(
+            sync_actor.broadcast_from_rank_zero.remote(
+                world_rank=rank,
+                world_size=world_size,
+                data=f"data-{rank}",
+                caller_method_name="broadcast_from_rank_zero",
+            )
+        )
+    # Ensure that all workers have the same consensus data same as rank 0
+    assert all([each == "data-0" for each in ray.get(remote_tasks)])
+    # Ensure all the states are cleared after the broadcast function returns
+    assert ray.get(sync_actor.get_counter.remote()) == 0
+    assert ray.get(sync_actor.get_world_size.remote()) == 0
+    assert ray.get(sync_actor.get_reduced_data.remote()) is None
+
+
+def test_hang_with_timeout():
+    """The test checks if the workers are blocked and hang when the world size
+    is greater than the number of workers. The workers should block and hang
+    until the barrier is lifted.
+    """
+    sync_actor = SynchronizationActor.remote(timeout_s=1, warn_interval_s=0.2)
+    # Test broadcast_from_rank_zero with a world size of 10. But
+    # only 9 workers data, the workers should block and hang
+    remote_tasks = []
+    for rank in range(9):
+        remote_tasks.append(
+            sync_actor.broadcast_from_rank_zero.remote(
+                world_rank=rank,
+                world_size=10,
+                data=f"data-{rank}",
+                caller_method_name="broadcast_from_rank_zero",
+            )
+        )
+    # Ensure that the workers are blocked and raise BroadcastCollectiveTimeoutError
+    # after 1 second
+    with pytest.raises(BroadcastCollectiveTimeoutError) as excinfo:
+        ray.get(remote_tasks)
+    assert "The following ranks have not joined the collective operation: [9]" in str(
+        excinfo.value
+    )
+
+
+def test_hang_without_timeout():
+    """Test the default behavior of running with no collective timeout."""
+    assert DEFAULT_COLLECTIVE_TIMEOUT_S is None
+
+    sync_actor = SynchronizationActor.remote()
+    remote_tasks = []
+    for rank in range(9):
+        remote_tasks.append(
+            sync_actor.broadcast_from_rank_zero.remote(
+                world_rank=rank,
+                world_size=10,
+                data=f"data-{rank}",
+                caller_method_name="broadcast_from_rank_zero",
+            )
+        )
+
+    # Just check for a short timeout to ensure the test doesn't error out.
+    done, _ = ray.wait(remote_tasks, num_returns=len(remote_tasks), timeout=2)
+    assert not done, "All tasks should be hanging, but some are done."
+
+    # Finish up once the last worker joins.
+    remote_tasks.append(
+        sync_actor.broadcast_from_rank_zero.remote(
+            world_rank=9,
+            world_size=10,
+            data="data-9",
+            caller_method_name="broadcast_from_rank_zero",
+        )
+    )
+    ray.get(remote_tasks)
+
+
+def test_world_size_mismatch():
+    """The test checks if the workers are blocked and raise an value error
+    when the world size is different. The workers should block and raise
+    a ValueError.
+    """
+    sync_actor = SynchronizationActor.remote()
+    remote_tasks = []
+
+    # All workers pass use a world size of 10, except for one.
+    for rank in range(9):
+        remote_tasks.append(
+            sync_actor.broadcast_from_rank_zero.remote(
+                world_rank=rank,
+                world_size=10,
+                data=f"data-{rank}",
+                caller_method_name="broadcast_from_rank_zero",
+            )
+        )
+
+    # The last worker calls broadcast with a different world size.
+    # This task should raise an error immediately.
+    mismatch_task = sync_actor.broadcast_from_rank_zero.remote(
+        world_rank=9,
+        world_size=11,
+        data="data-9",
+        caller_method_name="broadcast_from_rank_zero",
+    )
+    with pytest.raises(ValueError, match="same world size"):
+        ray.get(mismatch_task)
+
+
+def test_reset():
+    """Test that calling reset() unblocks all waiting workers with
+    SynchronizationBarrierResetError and leaves the actor usable for
+    subsequent barriers.
+    """
+    sync_actor = SynchronizationActor.remote()
+
+    # 9 out of 10 workers enter the barrier — they should all block.
+    remote_tasks = []
+    for rank in range(9):
+        remote_tasks.append(
+            sync_actor.broadcast_from_rank_zero.remote(
+                world_rank=rank,
+                world_size=10,
+                data=f"data-{rank}",
+                caller_method_name="broadcast_from_rank_zero",
+            )
+        )
+
+    # Verify all tasks are blocking.
+    done, _ = ray.wait(remote_tasks, num_returns=len(remote_tasks), timeout=2)
+    assert not done, "All tasks should be hanging, but some are done."
+
+    # Reset the actor — this should unblock all 9 workers.
+    ray.get(sync_actor.reset.remote())
+
+    # All 9 tasks should raise SynchronizationBarrierResetError.
+    with pytest.raises(RayTaskError) as excinfo:
+        ray.get(remote_tasks)
+    assert "SynchronizationBarrierResetError" in str(excinfo.value)
+
+    # Actor state should be fully cleaned up.
+    assert ray.get(sync_actor.get_counter.remote()) == 0
+    assert ray.get(sync_actor.get_world_size.remote()) == 0
+    assert ray.get(sync_actor.get_reduced_data.remote()) is None
+
+    # The actor should still be usable for a subsequent barrier.
+    remote_tasks = []
+    for rank in range(10):
+        remote_tasks.append(
+            sync_actor.broadcast_from_rank_zero.remote(
+                world_rank=rank,
+                world_size=10,
+                data=f"data-{rank}",
+                caller_method_name="broadcast_from_rank_zero",
+            )
+        )
+    assert all([each == "data-0" for each in ray.get(remote_tasks)])
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))