chore: import upstream snapshot with attribution

python/ray/serve/tests/test_broadcast.py

@@ -0,0 +1,206 @@
+import os
+import sys
+
+import pytest
+
+from ray import serve
+
+
+def test_broadcast_basic(serve_instance):
+    """Test that broadcast() calls every replica."""
+
+    @serve.deployment(num_replicas=3)
+    class D:
+        def get_pid(self):
+            return os.getpid()
+
+    serve.run(D.bind())
+    handle = serve.get_deployment_handle("D", "default")
+
+    # broadcast() internally waits for replicas to be available.
+    pids = handle.broadcast("get_pid").results(timeout_s=10)
+
+    assert len(pids) == 3
+    # Each replica should have a unique PID.
+    assert len(set(pids)) == 3
+
+
+def test_broadcast_with_args(serve_instance):
+    """Test broadcast with positional and keyword arguments."""
+
+    @serve.deployment(num_replicas=2)
+    class D:
+        def add(self, a, b=0):
+            return a + b
+
+    serve.run(D.bind())
+    handle = serve.get_deployment_handle("D", "default")
+
+    results = handle.broadcast("add", 1, b=2).results(timeout_s=10)
+
+    assert len(results) == 2
+    assert all(r == 3 for r in results)
+
+
+def test_broadcast_stateful(serve_instance):
+    """Test broadcast for state mutation (the cache-reset use case)."""
+
+    @serve.deployment(num_replicas=2)
+    class D:
+        def __init__(self):
+            self.cache = {"key": "value"}
+
+        def reset_cache(self):
+            self.cache.clear()
+            return "cleared"
+
+        def get_cache_size(self):
+            return len(self.cache)
+
+    serve.run(D.bind())
+    handle = serve.get_deployment_handle("D", "default")
+
+    # All replicas should start with cache size 1.
+    sizes = handle.broadcast("get_cache_size").results(timeout_s=10)
+    assert all(s == 1 for s in sizes)
+
+    # Broadcast cache reset.
+    results = handle.broadcast("reset_cache").results(timeout_s=10)
+    assert all(r == "cleared" for r in results)
+
+    # All replicas should now have empty caches.
+    sizes = handle.broadcast("get_cache_size").results(timeout_s=10)
+    assert all(s == 0 for s in sizes)
+
+
+@pytest.mark.asyncio
+async def test_broadcast_async(serve_instance):
+    """Test the async results path."""
+
+    @serve.deployment(num_replicas=2)
+    class D:
+        def get_pid(self):
+            return os.getpid()
+
+    serve.run(D.bind())
+    handle = serve.get_deployment_handle("D", "default")
+
+    pids = await handle.broadcast("get_pid").results_async()
+
+    assert len(pids) == 2
+    assert len(set(pids)) == 2
+
+
+def test_broadcast_single_replica(serve_instance):
+    """Test broadcast with a single replica."""
+
+    @serve.deployment(num_replicas=1)
+    class D:
+        def ping(self):
+            return "pong"
+
+    serve.run(D.bind())
+    handle = serve.get_deployment_handle("D", "default")
+
+    results = handle.broadcast("ping").results(timeout_s=10)
+    assert results == ["pong"]
+
+
+def test_broadcast_ignores_streaming_handle_option(serve_instance):
+    """Test broadcast on a handle configured with stream=True."""
+
+    @serve.deployment(num_replicas=2)
+    class D:
+        def ping(self):
+            return "pong"
+
+    serve.run(D.bind())
+    handle = serve.get_deployment_handle("D", "default").options(stream=True)
+
+    results = handle.broadcast("ping").results(timeout_s=10)
+
+    assert len(results) == 2
+    assert all(r == "pong" for r in results)
+
+
+def test_broadcast_return_exceptions_sync(serve_instance):
+    """Test best-effort sync result collection with exceptions."""
+
+    @serve.deployment(num_replicas=2)
+    class D:
+        def fail(self):
+            raise RuntimeError("boom")
+
+    serve.run(D.bind())
+    handle = serve.get_deployment_handle("D", "default")
+
+    with pytest.raises(Exception):
+        handle.broadcast("fail").results(timeout_s=10)
+
+    results = handle.broadcast("fail").results(
+        timeout_s=10,
+        return_exceptions=True,
+    )
+    assert len(results) == 2
+    assert all(isinstance(r, Exception) for r in results)
+
+
+@pytest.mark.asyncio
+async def test_broadcast_return_exceptions_async(serve_instance):
+    """Test best-effort async result collection with exceptions."""
+
+    @serve.deployment(num_replicas=2)
+    class D:
+        def fail(self):
+            raise RuntimeError("boom")
+
+    serve.run(D.bind())
+    handle = serve.get_deployment_handle("D", "default")
+
+    with pytest.raises(Exception):
+        await handle.broadcast("fail").results_async()
+
+    results = await handle.broadcast("fail").results_async(return_exceptions=True)
+    assert len(results) == 2
+    assert all(isinstance(r, Exception) for r in results)
+
+
+def test_broadcast_sync_timeout(serve_instance):
+    """Test sync timeout handling for broadcast results collection."""
+
+    @serve.deployment(num_replicas=2)
+    class D:
+        def slow(self):
+            import time
+
+            time.sleep(10)
+            return "done"
+
+    serve.run(D.bind())
+    handle = serve.get_deployment_handle("D", "default")
+
+    with pytest.raises(TimeoutError):
+        handle.broadcast("slow").results(timeout_s=0.5)
+
+
+@pytest.mark.asyncio
+async def test_broadcast_async_timeout(serve_instance):
+    """Test async timeout handling for broadcast results collection."""
+
+    @serve.deployment(num_replicas=2)
+    class D:
+        def slow(self):
+            import time
+
+            time.sleep(10)
+            return "done"
+
+    serve.run(D.bind())
+    handle = serve.get_deployment_handle("D", "default")
+
+    with pytest.raises(TimeoutError):
+        await handle.broadcast("slow").results_async(timeout_s=0.5)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))