chore: import upstream snapshot with attribution

python/ray/serve/tests/unit/test_grpc_replica_result.py

@@ -0,0 +1,391 @@
+import asyncio
+import sys
+import threading
+
+import pytest
+
+from ray import ActorID, cloudpickle
+from ray._common.test_utils import wait_for_condition
+from ray.serve._private.common import RequestMetadata
+from ray.serve._private.replica_result import gRPCReplicaResult
+from ray.serve.generated import serve_pb2
+
+
+class FakegRPCUnaryCall:
+    def __init__(self, item, is_error: bool = False):
+        self._loop = asyncio.get_running_loop()
+        self._item = item
+        self._is_error = is_error
+
+    def __await__(self):
+        if asyncio.get_running_loop() != self._loop:
+            raise RuntimeError("Tried to fetch from a different loop!")
+
+        yield
+        return serve_pb2.ASGIResponse(
+            serialized_message=cloudpickle.dumps(self._item), is_error=self._is_error
+        )
+
+    def add_done_callback(self, cb):
+        pass
+
+
+class FakegRPCStreamCall:
+    def __init__(self, items, *, event: threading.Event = None):
+        self._loop = asyncio.get_running_loop()
+        self._items = items
+        self._event = event
+
+    def is_empty(self) -> bool:
+        assert len(self._items) == 0
+        return True
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        if asyncio.get_running_loop() != self._loop:
+            raise RuntimeError("Tried to fetch from a different loop!")
+
+        if not self._items:
+            raise StopAsyncIteration
+
+        if self._event:
+            await self._loop.run_in_executor(None, self._event.wait)
+
+        item, is_error = self._items.pop(0)
+        return serve_pb2.ASGIResponse(
+            serialized_message=cloudpickle.dumps(item),
+            is_error=is_error,
+        )
+
+    def add_done_callback(self, cb):
+        pass
+
+
+@pytest.fixture
+def create_asyncio_event_loop_in_thread():
+    async_loop = asyncio.new_event_loop()
+    thread = threading.Thread(daemon=True, target=async_loop.run_forever)
+    thread.start()
+    event = threading.Event()
+
+    yield async_loop, event
+
+    # Unblock event in case it's blocking shutdown
+    event.set()
+
+
+@pytest.mark.asyncio
+class TestSameLoop:
+    def make_fake_call(self, is_streaming: bool, *, data=None, error=None):
+        if is_streaming:
+            fake_call = FakegRPCStreamCall(data)
+        else:
+            if error:
+                fake_call = FakegRPCUnaryCall(error, is_error=True)
+            else:
+                fake_call = FakegRPCUnaryCall(data, is_error=False)
+
+        return gRPCReplicaResult(
+            fake_call,
+            metadata=RequestMetadata(
+                request_id="",
+                internal_request_id="",
+                is_streaming=False,
+                _on_separate_loop=False,
+            ),
+            actor_id=ActorID(b"2" * 16),
+            loop=asyncio.get_running_loop(),
+        )
+
+    async def test_unary(self):
+        replica_result = self.make_fake_call(is_streaming=False, data="hello")
+        assert await replica_result.get_async() == "hello"
+
+    async def test_streaming(self):
+        replica_result = self.make_fake_call(
+            is_streaming=True, data=[(1, False), (2, False), (3, False), (4, False)]
+        )
+        assert [r async for r in replica_result] == [1, 2, 3, 4]
+
+    async def test_unary_with_gen(self):
+        replica_result = self.make_fake_call(is_streaming=True, data=[("hello", False)])
+        assert await replica_result.get_async() == "hello"
+
+    async def test_unary_error(self):
+        """Test error is raised correctly."""
+
+        replica_result = self.make_fake_call(
+            is_streaming=False, error=RuntimeError("oh no!")
+        )
+        with pytest.raises(RuntimeError, match="oh no!"):
+            await replica_result.get_async()
+
+    async def test_streaming_error(self):
+        """Test error is raised correctly."""
+
+        replica_result = self.make_fake_call(
+            is_streaming=True, data=[(RuntimeError("oh no!"), True)]
+        )
+        with pytest.raises(RuntimeError, match="oh no!"):
+            await replica_result.__anext__()
+
+
+class TestSeparateLoop:
+    async def make_fake_unary_request(self, data, loop: asyncio.AbstractEventLoop):
+        fake_call = FakegRPCUnaryCall(data)
+        replica_result = gRPCReplicaResult(
+            fake_call,
+            metadata=RequestMetadata(
+                request_id="",
+                internal_request_id="",
+                is_streaming=False,
+                _on_separate_loop=True,
+            ),
+            actor_id=ActorID(b"2" * 16),
+            loop=loop,
+        )
+        return replica_result
+
+    async def make_fake_streaming_request(
+        self,
+        data,
+        loop: asyncio.AbstractEventLoop,
+        on_separate_loop: bool,
+        *,
+        is_streaming: bool = True,
+        event: threading.Event = None,
+        error=None,
+    ):
+        if error:
+            fake_call = FakegRPCStreamCall([(error, True)], event=event)
+        else:
+            fake_call = FakegRPCStreamCall([(d, False) for d in data], event=event)
+        return gRPCReplicaResult(
+            fake_call,
+            metadata=RequestMetadata(
+                request_id="",
+                internal_request_id="",
+                is_streaming=is_streaming,
+                _on_separate_loop=on_separate_loop,
+            ),
+            actor_id=ActorID(b"2" * 16),
+            loop=loop,
+        )
+
+    def test_unary_sync(self, create_asyncio_event_loop_in_thread):
+        loop, _ = create_asyncio_event_loop_in_thread
+
+        fut = asyncio.run_coroutine_threadsafe(
+            self.make_fake_unary_request("hello", loop), loop=loop
+        )
+        replica_result = fut.result()
+
+        assert replica_result.get(None) == "hello"
+
+    @pytest.mark.asyncio
+    async def test_unary_async(self, create_asyncio_event_loop_in_thread):
+        loop, _ = create_asyncio_event_loop_in_thread
+
+        fut = asyncio.run_coroutine_threadsafe(
+            self.make_fake_unary_request("hello", loop), loop=loop
+        )
+        replica_result = fut.result()
+
+        assert await replica_result.get_async() == "hello"
+
+    def test_streaming_sync(self, create_asyncio_event_loop_in_thread):
+        loop, _ = create_asyncio_event_loop_in_thread
+
+        # Instantiate gRPCReplicaResult with FakegRPCStreamCall. This needs
+        # to be run on the "other loop"
+        fut = asyncio.run_coroutine_threadsafe(
+            self.make_fake_streaming_request([1, 2, 3, 4], loop, on_separate_loop=True),
+            loop=loop,
+        )
+        replica_result = fut.result()
+
+        # The async generator should be consumed even if we don't fetch
+        # the items explicitly through the ReplicaResult object
+        wait_for_condition(replica_result._call.is_empty, retry_interval_ms=10)
+
+        # Finally, check results given by gRPCReplicaResult fetched from
+        # the queue are correct
+        assert list(replica_result) == [1, 2, 3, 4]
+
+    @pytest.mark.asyncio
+    async def test_streaming_async(self, create_asyncio_event_loop_in_thread):
+        loop, _ = create_asyncio_event_loop_in_thread
+
+        fut = asyncio.run_coroutine_threadsafe(
+            self.make_fake_streaming_request([1, 2, 3, 4], loop, on_separate_loop=True),
+            loop=loop,
+        )
+        replica_result = fut.result()
+
+        # Check async generator is consumed on its own
+        wait_for_condition(replica_result._call.is_empty, retry_interval_ms=10)
+        assert [r async for r in replica_result] == [1, 2, 3, 4]
+
+    @pytest.mark.asyncio
+    async def test_streaming_blocked(self, create_asyncio_event_loop_in_thread):
+        """Use threading event to block async generator, check everything works"""
+
+        loop, event = create_asyncio_event_loop_in_thread
+
+        fut = asyncio.run_coroutine_threadsafe(
+            self.make_fake_streaming_request(
+                [1, 2, 3, 4], loop, on_separate_loop=True, event=event
+            ),
+            loop=loop,
+        )
+        replica_result = fut.result()
+
+        async def fetch():
+            return [r async for r in replica_result]
+
+        t = asyncio.create_task(fetch())
+
+        with pytest.raises(asyncio.TimeoutError):
+            await asyncio.wait_for(asyncio.shield(t), 0.01)
+
+        event.set()
+        assert await t == [1, 2, 3, 4]
+
+    def test_unary_with_gen_sync(self, create_asyncio_event_loop_in_thread):
+        loop, _ = create_asyncio_event_loop_in_thread
+
+        fut = asyncio.run_coroutine_threadsafe(
+            self.make_fake_streaming_request(
+                ["hello"], loop, on_separate_loop=True, is_streaming=False
+            ),
+            loop=loop,
+        )
+        replica_result = fut.result()
+
+        # Check async generator is consumed on its own
+        wait_for_condition(replica_result._call.is_empty, retry_interval_ms=10)
+        assert replica_result.get(None) == "hello"
+
+    @pytest.mark.asyncio
+    async def test_unary_with_gen_async(self, create_asyncio_event_loop_in_thread):
+        loop, _ = create_asyncio_event_loop_in_thread
+
+        fut = asyncio.run_coroutine_threadsafe(
+            self.make_fake_streaming_request(
+                ["hello"], loop, on_separate_loop=True, is_streaming=False
+            ),
+            loop=loop,
+        )
+        replica_result = fut.result()
+
+        # Check async generator is consumed on its own
+        wait_for_condition(replica_result._call.is_empty, retry_interval_ms=10)
+        assert await replica_result.get_async() == "hello"
+
+    @pytest.mark.asyncio
+    async def test_unary_with_gen_blocked(self, create_asyncio_event_loop_in_thread):
+        """Use threading event to block async generator, check everything works"""
+
+        loop, event = create_asyncio_event_loop_in_thread
+
+        fut = asyncio.run_coroutine_threadsafe(
+            self.make_fake_streaming_request(
+                ["hello"], loop, on_separate_loop=True, event=event
+            ),
+            loop=loop,
+        )
+        replica_result = fut.result()
+
+        t = asyncio.create_task(replica_result.get_async())
+        with pytest.raises(asyncio.TimeoutError):
+            await asyncio.wait_for(asyncio.shield(t), 0.01)
+
+        event.set()
+        assert await t == "hello"
+
+    def test_unary_with_timeout(self, create_asyncio_event_loop_in_thread):
+        """Test get() with timeout."""
+
+        loop, event = create_asyncio_event_loop_in_thread
+
+        fut = asyncio.run_coroutine_threadsafe(
+            self.make_fake_streaming_request(
+                ["hello"], loop, on_separate_loop=True, event=event
+            ),
+            loop=loop,
+        )
+        replica_result = fut.result()
+
+        with pytest.raises(TimeoutError):
+            replica_result.get(timeout_s=0.01)
+
+        event.set()
+        assert replica_result.get(timeout_s=0.01) == "hello"
+
+    def test_unary_error_sync(self, create_asyncio_event_loop_in_thread):
+        """Test error is raised correctly."""
+
+        loop, _ = create_asyncio_event_loop_in_thread
+        fut = asyncio.run_coroutine_threadsafe(
+            self.make_fake_streaming_request(
+                None, loop, on_separate_loop=True, error=RuntimeError("oh no!")
+            ),
+            loop=loop,
+        )
+        replica_result = fut.result()
+
+        with pytest.raises(RuntimeError, match="oh no!"):
+            replica_result.get(None)
+
+    @pytest.mark.asyncio
+    async def test_unary_error_async(self, create_asyncio_event_loop_in_thread):
+        """Test error is raised correctly."""
+
+        loop, _ = create_asyncio_event_loop_in_thread
+        fut = asyncio.run_coroutine_threadsafe(
+            self.make_fake_streaming_request(
+                None, loop, on_separate_loop=True, error=RuntimeError("oh no!")
+            ),
+            loop=loop,
+        )
+        replica_result = fut.result()
+
+        with pytest.raises(RuntimeError, match="oh no!"):
+            await replica_result.get_async()
+
+    def test_streaming_error_sync(self, create_asyncio_event_loop_in_thread):
+        """Test error is raised correctly."""
+
+        loop, _ = create_asyncio_event_loop_in_thread
+        fut = asyncio.run_coroutine_threadsafe(
+            self.make_fake_streaming_request(
+                None, loop, on_separate_loop=True, error=RuntimeError("oh no!")
+            ),
+            loop=loop,
+        )
+        replica_result = fut.result()
+
+        with pytest.raises(RuntimeError, match="oh no!"):
+            replica_result.__next__()
+
+    @pytest.mark.asyncio
+    async def test_streaming_error_async(self, create_asyncio_event_loop_in_thread):
+        """Test error is raised correctly."""
+
+        loop, _ = create_asyncio_event_loop_in_thread
+        fut = asyncio.run_coroutine_threadsafe(
+            self.make_fake_streaming_request(
+                None, loop, on_separate_loop=True, error=RuntimeError("oh no!")
+            ),
+            loop=loop,
+        )
+        replica_result = fut.result()
+
+        with pytest.raises(RuntimeError, match="oh no!"):
+            await replica_result.__anext__()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))