chore: import upstream snapshot with attribution

python/ray/serve/tests/test_proxy_response_generator.py

@@ -0,0 +1,416 @@
+import asyncio
+from typing import AsyncIterator, Tuple
+
+import pytest
+
+from ray import serve
+from ray._common.test_utils import SignalActor, async_wait_for_condition
+from ray.serve._private.proxy_response_generator import ProxyResponseGenerator
+
+
+def disconnect_task_and_event() -> Tuple[asyncio.Event, asyncio.Task]:
+    """Return an event and a task waiting on it for testing disconnect logic."""
+    event = asyncio.Event()
+
+    async def wait_for_event():
+        return await event.wait()
+
+    return event, asyncio.ensure_future(wait_for_event())
+
+
+@pytest.mark.asyncio
+class TestUnary:
+    async def test_basic(self, serve_instance):
+        @serve.deployment
+        class D:
+            def __call__(self, name: str) -> str:
+                return f"Hello {name}!"
+
+            def error(self):
+                raise RuntimeError("oopsies")
+
+        h = serve.run(D.bind()).options(
+            stream=False,
+        )
+        gen = ProxyResponseGenerator(h.remote("Alice"))
+
+        # Test simple response.
+        responses = [r async for r in gen]
+        assert len(responses) == 1
+        assert set(responses) == {"Hello Alice!"}
+
+        with pytest.raises(StopAsyncIteration):
+            await gen.__anext__()
+
+        # Test response raising exception.
+        gen = ProxyResponseGenerator(h.error.remote())
+        with pytest.raises(RuntimeError, match="oopsies"):
+            await gen.__anext__()
+
+        with pytest.raises(StopAsyncIteration):
+            await gen.__anext__()
+
+    async def test_result_callback(self, serve_instance):
+        @serve.deployment
+        class D:
+            def __call__(self, name: str) -> str:
+                return f"Hello {name}!"
+
+        h = serve.run(D.bind()).options(
+            stream=False,
+        )
+
+        def result_callback(result: str) -> str:
+            return f"Callback called on: {result}"
+
+        gen = ProxyResponseGenerator(h.remote("Alice"), result_callback=result_callback)
+
+        responses = [r async for r in gen]
+        assert len(responses) == 1
+        assert set(responses) == {"Callback called on: Hello Alice!"}
+
+    async def test_timeout_while_assigning(self, serve_instance):
+        signal_actor = SignalActor.remote()
+
+        @serve.deployment(max_ongoing_requests=1)
+        class D:
+            async def __call__(self):
+                await signal_actor.wait.remote()
+
+        h = serve.run(D.bind()).options(
+            stream=False,
+        )
+        h.remote()
+
+        async def one_waiter():
+            return await signal_actor.cur_num_waiters.remote() == 1
+
+        await async_wait_for_condition(one_waiter)
+
+        gen = ProxyResponseGenerator(h.remote(), timeout_s=0.1)
+
+        with pytest.raises(TimeoutError):
+            await gen.__anext__()
+
+        assert gen.cancelled()
+
+        assert await signal_actor.cur_num_waiters.remote() == 1
+        await signal_actor.send.remote()
+
+    async def test_timeout_while_executing(self, serve_instance):
+        signal_actor = SignalActor.remote()
+
+        @serve.deployment
+        class D:
+            async def __call__(self):
+                await signal_actor.wait.remote()
+
+        h = serve.run(D.bind()).options(
+            stream=False,
+        )
+
+        gen = ProxyResponseGenerator(h.remote(), timeout_s=0.1)
+
+        with pytest.raises(TimeoutError):
+            await gen.__anext__()
+
+        assert gen.cancelled()
+
+        assert await signal_actor.cur_num_waiters.remote() == 1
+        await signal_actor.send.remote()
+
+    async def test_disconnect_while_assigning(self, serve_instance):
+        signal_actor = SignalActor.remote()
+        disconnect_event, disconnect_task = disconnect_task_and_event()
+
+        @serve.deployment(max_ongoing_requests=1)
+        class D:
+            async def __call__(self):
+                await signal_actor.wait.remote()
+
+        h = serve.run(D.bind()).options(
+            stream=False,
+        )
+        h.remote()
+
+        async def one_waiter():
+            return await signal_actor.cur_num_waiters.remote() == 1
+
+        await async_wait_for_condition(one_waiter)
+
+        gen = ProxyResponseGenerator(h.remote(), disconnected_task=disconnect_task)
+
+        async def get_next():
+            return await gen.__anext__()
+
+        gen_next = asyncio.ensure_future(get_next())
+        done, _ = await asyncio.wait([gen_next], timeout=0.1)
+        assert len(done) == 0
+
+        # Set the disconnect event, causing the disconnect task to finish.
+        disconnect_event.set()
+        with pytest.raises(asyncio.CancelledError):
+            await gen_next
+
+        assert gen.cancelled()
+
+        assert await signal_actor.cur_num_waiters.remote() == 1
+        await signal_actor.send.remote()
+
+    async def test_disconnect_while_executing(self, serve_instance):
+        signal_actor = SignalActor.remote()
+        disconnect_event, disconnect_task = disconnect_task_and_event()
+
+        @serve.deployment
+        class D:
+            async def __call__(self):
+                await signal_actor.wait.remote()
+
+        h = serve.run(D.bind()).options(
+            stream=False,
+        )
+
+        gen = ProxyResponseGenerator(h.remote(), disconnected_task=disconnect_task)
+
+        async def get_next():
+            return await gen.__anext__()
+
+        gen_next = asyncio.ensure_future(get_next())
+        done, _ = await asyncio.wait([gen_next], timeout=0.1)
+        assert len(done) == 0
+
+        # Set the disconnect event, causing the disconnect task to finish.
+        disconnect_event.set()
+        with pytest.raises(asyncio.CancelledError):
+            await gen_next
+
+        assert gen.cancelled()
+
+        assert await signal_actor.cur_num_waiters.remote() == 1
+        await signal_actor.send.remote()
+
+
+@pytest.mark.asyncio
+class TestStreaming:
+    async def test_basic(self, serve_instance):
+        @serve.deployment
+        class D:
+            def __call__(self, name: str) -> AsyncIterator[str]:
+                for _ in range(5):
+                    yield f"Hello {name}!"
+
+            def error(self, name: str) -> AsyncIterator[str]:
+                yield f"Hello {name}!"
+                raise RuntimeError("oopsies")
+
+        h = serve.run(D.bind()).options(
+            stream=True,
+        )
+        gen = ProxyResponseGenerator(h.remote("Alice"))
+
+        # Test simple response.
+        responses = [r async for r in gen]
+        assert len(responses) == 5
+        assert set(responses) == {"Hello Alice!"}
+        with pytest.raises(StopAsyncIteration):
+            await gen.__anext__()
+
+        # Test exception in the middle of response stream.
+        gen = ProxyResponseGenerator(h.error.remote("Alice"))
+        assert await gen.__anext__() == "Hello Alice!"
+        with pytest.raises(RuntimeError, match="oopsies"):
+            await gen.__anext__()
+        with pytest.raises(StopAsyncIteration):
+            await gen.__anext__()
+
+    async def test_result_callback(self, serve_instance):
+        @serve.deployment
+        class D:
+            def __call__(self, name: str) -> AsyncIterator[str]:
+                for _ in range(5):
+                    yield f"Hello {name}!"
+
+        h = serve.run(D.bind()).options(
+            stream=True,
+        )
+
+        def result_callback(result: str) -> str:
+            return f"Callback called on: {result}"
+
+        gen = ProxyResponseGenerator(h.remote("Alice"), result_callback=result_callback)
+
+        responses = [r async for r in gen]
+        assert len(responses) == 5
+        assert set(responses) == {"Callback called on: Hello Alice!"}
+
+    async def test_timeout_while_assigning(self, serve_instance):
+        signal_actor = SignalActor.remote()
+
+        @serve.deployment(max_ongoing_requests=1)
+        class D:
+            async def __call__(self) -> AsyncIterator[str]:
+                yield "hi"
+                await signal_actor.wait.remote()
+
+        h = serve.run(D.bind()).options(
+            stream=True,
+        )
+        h.remote()
+
+        async def one_waiter():
+            return await signal_actor.cur_num_waiters.remote() == 1
+
+        await async_wait_for_condition(one_waiter)
+
+        gen = ProxyResponseGenerator(h.remote(), timeout_s=0.1)
+
+        with pytest.raises(TimeoutError):
+            await gen.__anext__()
+
+        assert gen.cancelled()
+
+        assert await signal_actor.cur_num_waiters.remote() == 1
+        await signal_actor.send.remote()
+
+    async def test_timeout_while_executing(self, serve_instance):
+        signal_actor = SignalActor.remote()
+
+        @serve.deployment
+        class D:
+            async def __call__(self) -> AsyncIterator[str]:
+                yield "hi"
+                await signal_actor.wait.remote()
+
+        h = serve.run(D.bind()).options(
+            stream=True,
+        )
+
+        gen = ProxyResponseGenerator(h.remote(), timeout_s=0.1)
+        assert (await gen.__anext__()) == "hi"
+        with pytest.raises(TimeoutError):
+            await gen.__anext__()
+
+        assert gen.cancelled()
+
+        assert await signal_actor.cur_num_waiters.remote() == 1
+        await signal_actor.send.remote()
+
+    async def test_disconnect_while_assigning(self, serve_instance):
+        signal_actor = SignalActor.remote()
+        disconnect_event, disconnect_task = disconnect_task_and_event()
+
+        @serve.deployment(max_ongoing_requests=1)
+        class D:
+            async def __call__(self) -> AsyncIterator[str]:
+                yield "hi"
+                await signal_actor.wait.remote()
+
+        h = serve.run(D.bind()).options(
+            stream=True,
+        )
+        h.remote()
+
+        async def one_waiter():
+            return await signal_actor.cur_num_waiters.remote() == 1
+
+        await async_wait_for_condition(one_waiter)
+
+        gen = ProxyResponseGenerator(h.remote(), disconnected_task=disconnect_task)
+
+        async def get_next():
+            return await gen.__anext__()
+
+        gen_next = asyncio.ensure_future(get_next())
+        done, _ = await asyncio.wait([gen_next], timeout=0.1)
+        assert len(done) == 0
+
+        # Set the disconnect event, causing the disconnect task to finish.
+        disconnect_event.set()
+        with pytest.raises(asyncio.CancelledError):
+            await gen_next
+
+        assert gen.cancelled()
+
+        assert await signal_actor.cur_num_waiters.remote() == 1
+        await signal_actor.send.remote()
+
+    async def test_disconnect_while_executing(self, serve_instance):
+        signal_actor = SignalActor.remote()
+        disconnect_event, disconnect_task = disconnect_task_and_event()
+
+        @serve.deployment
+        class D:
+            async def __call__(self) -> AsyncIterator[str]:
+                yield "hi"
+                await signal_actor.wait.remote()
+
+        h = serve.run(D.bind()).options(
+            stream=True,
+        )
+
+        gen = ProxyResponseGenerator(h.remote(), disconnected_task=disconnect_task)
+        assert (await gen.__anext__()) == "hi"
+
+        async def get_next():
+            return await gen.__anext__()
+
+        gen_next = asyncio.ensure_future(get_next())
+        done, _ = await asyncio.wait([gen_next], timeout=0.1)
+        assert len(done) == 0
+
+        # Set the disconnect event, causing the disconnect task to finish.
+        disconnect_event.set()
+        with pytest.raises(asyncio.CancelledError):
+            await gen_next
+
+        assert gen.cancelled()
+
+        assert await signal_actor.cur_num_waiters.remote() == 1
+        await signal_actor.send.remote()
+
+    async def test_stop_checking_for_disconnect(self, serve_instance):
+        signal_actor = SignalActor.remote()
+        disconnect_event, disconnect_task = disconnect_task_and_event()
+
+        @serve.deployment
+        class D:
+            async def __call__(self) -> AsyncIterator[str]:
+                yield "hi"
+                await signal_actor.wait.remote()
+
+        h = serve.run(D.bind()).options(
+            stream=True,
+        )
+
+        gen = ProxyResponseGenerator(
+            h.remote(),
+            disconnected_task=disconnect_task,
+        )
+        assert (await gen.__anext__()) == "hi"
+
+        gen.stop_checking_for_disconnect()
+
+        async def get_next():
+            return await gen.__anext__()
+
+        gen_next = asyncio.ensure_future(get_next())
+        done, _ = await asyncio.wait([gen_next], timeout=0.1)
+        assert len(done) == 0
+
+        # Set the disconnect event, causing the disconnect task to finish.
+        # However, because stop_checking_for_disconnect was called this
+        # should still time out.
+        disconnect_event.set()
+        done, _ = await asyncio.wait([gen_next], timeout=0.1)
+        assert len(done) == 0
+
+        assert not gen.cancelled()
+
+        assert await signal_actor.cur_num_waiters.remote() == 1
+        await signal_actor.send.remote()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-s", __file__]))