chore: import upstream snapshot with attribution

python/ray/air/tests/execution/test_event_manager.py

@@ -0,0 +1,224 @@
+import time
+from typing import Any, Type
+
+import pytest
+
+import ray
+from ray.air.execution._internal import Barrier
+from ray.air.execution._internal.event_manager import RayEventManager
+from ray.exceptions import RayTaskError
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@ray.remote
+def succeeding(ret: Any = None) -> Any:
+    return ret
+
+
+@ray.remote
+def failing(exc: Type[Exception], *args) -> None:
+    raise exc(*args)
+
+
+@ray.remote
+def sleeping(seconds: int, result: Any) -> Any:
+    time.sleep(seconds)
+    return result
+
+
+def test_track_future_success(ray_start_4_cpus):
+    """Schedule a future that return successfully.
+
+    Check that the on_result callback was triggered.
+    """
+    event_manager = RayEventManager()
+
+    seen = set()
+
+    def on_result(result: Any):
+        seen.add(result)
+
+    event_manager.track_future(succeeding.remote("a"), on_result=on_result)
+
+    event_manager.wait()
+    assert "a" in seen
+
+    assert not event_manager._tracked_futures
+
+
+def test_track_future_success_no_callback(ray_start_4_cpus):
+    """Schedule a future that return successfully.
+
+    Check that passing no callback still succeeds.
+    """
+    event_manager = RayEventManager()
+
+    event_manager.track_future(succeeding.remote("a"))
+
+    event_manager.wait()
+
+    assert not event_manager._tracked_futures
+
+
+def test_track_future_error(ray_start_4_cpus):
+    """Schedule a future that fails.
+
+    Check that the on_error callback was triggered.
+    """
+    event_manager = RayEventManager()
+
+    seen = set()
+
+    class CustomError(RuntimeError):
+        pass
+
+    def on_error(exception: Exception):
+        seen.add(exception)
+
+    event_manager.track_future(failing.remote(CustomError), on_error=on_error)
+
+    event_manager.wait()
+    assert isinstance(seen.pop(), CustomError)
+
+    assert not event_manager._tracked_futures
+
+
+def test_track_future_error_no_callback(ray_start_4_cpus):
+    """Schedule a future that fails.
+
+    Check that passing no callback raises the original error.
+    """
+    event_manager = RayEventManager()
+
+    event_manager.track_future(failing.remote(RuntimeError))
+
+    with pytest.raises(RuntimeError):
+        event_manager.wait()
+
+    assert not event_manager._tracked_futures
+
+
+@pytest.mark.parametrize("results_per_wait", [None, 1, 5, 10, 100])
+def test_many_futures(ray_start_4_cpus, results_per_wait):
+    """Schedule 500 succeeding and failing futures.
+
+    Check that the callbacks get triggered correctly, independent of the number
+    of results we await per call to RayEventManager.wait().
+    """
+    num_futures = 500
+
+    event_manager = RayEventManager()
+
+    seen_results = set()
+    seen_errors = set()
+
+    def on_result(result: Any):
+        seen_results.add(result)
+
+    def on_error(exception: RayTaskError):
+        seen_errors.add(exception.cause.args[0])
+
+    for i in range(num_futures):
+        event_manager.track_futures(
+            [
+                succeeding.remote("a" + str(i)),
+                failing.remote(RuntimeError, "b" + str(i)),
+            ],
+            on_result=on_result,
+            on_error=on_error,
+        )
+
+    while event_manager.num_futures > 0:
+        event_manager.wait(num_results=results_per_wait)
+
+    for i in range(num_futures):
+        assert "a" + str(i) in seen_results
+        assert "b" + str(i) in seen_errors
+
+
+def test_timeout(ray_start_4_cpus):
+    """Test the timeout parameter.
+
+    Start 4 tasks: Two succeed immediately, two after 1 second.
+
+    After waiting for 0.5 seconds, the first two tasks should have returned.
+    After waiting for up to 5 seconds, the other two tasks should have returned.
+    But because the tasks take only 0.5 seconds to run, we should have waited
+    way less than 5 seconds.
+    """
+    event_manager = RayEventManager()
+
+    seen = set()
+
+    def on_result(result: Any):
+        seen.add(result)
+
+    event_manager.track_futures(
+        [
+            succeeding.remote("a"),
+            succeeding.remote("b"),
+            sleeping.remote(1, "c"),
+            sleeping.remote(1, "d"),
+        ],
+        on_result=on_result,
+    )
+
+    start = time.monotonic()
+    event_manager.wait(num_results=None, timeout=0.5)
+    assert "a" in seen
+    assert "b" in seen
+    assert "c" not in seen
+    assert "d" not in seen
+
+    event_manager.wait(num_results=None, timeout=5)
+    taken = time.monotonic() - start
+
+    assert "c" in seen
+    assert "d" in seen
+
+    # Should have returned much earlier than after 5 seconds
+    assert taken < 3
+
+    assert not event_manager._tracked_futures
+
+
+def test_task_barrier(ray_start_4_cpus):
+    event_manager = RayEventManager()
+
+    seen = set()
+
+    def on_completion(barrier: Barrier):
+        seen.update(barrier.get_results())
+
+    barrier = Barrier(max_results=4, on_completion=on_completion)
+
+    event_manager.track_futures(
+        [
+            succeeding.remote("a"),
+            succeeding.remote("b"),
+            succeeding.remote("c"),
+            succeeding.remote("d"),
+            sleeping.remote(2, "e"),
+        ],
+        on_result=barrier.arrive,
+    )
+
+    event_manager.wait(num_results=4)
+
+    assert "a" in seen
+    assert "b" in seen
+    assert "c" in seen
+    assert "d" in seen
+    assert "e" not in seen
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))