chore: import upstream snapshot with attribution

python/ray/serve/tests/test_streaming_response.py

@@ -0,0 +1,352 @@
+import asyncio
+import os
+from typing import AsyncGenerator, Optional
+
+import httpx
+import pytest
+from fastapi import FastAPI
+from starlette.requests import Request
+from starlette.responses import StreamingResponse
+
+import ray
+from ray import serve
+from ray._common.test_utils import SignalActor
+from ray.serve._private.test_utils import (
+    get_application_url,
+    get_application_urls,
+    send_signal_on_cancellation,
+)
+from ray.serve.handle import DeploymentHandle
+
+
+@ray.remote
+class StreamingRequester:
+    async def make_request(
+        self, url: Optional[str] = None
+    ) -> AsyncGenerator[str, None]:
+        url = url or get_application_url("HTTP")
+        with httpx.stream("GET", url) as r:
+            r.raise_for_status()
+            for chunk in r.iter_text():
+                yield chunk
+                await asyncio.sleep(0.001)
+
+
+@pytest.mark.parametrize("use_fastapi", [False, True])
+@pytest.mark.parametrize("use_async", [False, True])
+def test_basic(serve_instance, use_async: bool, use_fastapi: bool):
+    signal_actor = SignalActor.remote()
+
+    async def hi_gen_async():
+        for i in range(10):
+            yield f"hi_{i}"
+            # to avoid coalescing chunks
+            await signal_actor.wait.remote()
+
+    def hi_gen_sync():
+        for i in range(10):
+            yield f"hi_{i}"
+            # to avoid coalescing chunks
+            ray.get(signal_actor.wait.remote())
+
+    if use_fastapi:
+        app = FastAPI()
+
+        @serve.deployment
+        @serve.ingress(app)
+        class SimpleGenerator:
+            @app.get("/")
+            def stream_hi(self, request: Request) -> StreamingResponse:
+                gen = hi_gen_async() if use_async else hi_gen_sync()
+                return StreamingResponse(gen, media_type="text/plain")
+
+    else:
+
+        @serve.deployment
+        class SimpleGenerator:
+            def __call__(self, request: Request) -> StreamingResponse:
+                gen = hi_gen_async() if use_async else hi_gen_sync()
+                return StreamingResponse(gen, media_type="text/plain")
+
+    serve.run(SimpleGenerator.bind())
+
+    url = get_application_url("HTTP")
+    with httpx.stream("GET", url) as r:
+        r.raise_for_status()
+        for i, chunk in enumerate(r.iter_text()):
+            assert chunk == f"hi_{i}"
+            ray.get(signal_actor.send.remote())
+
+
+@pytest.mark.parametrize("use_fastapi", [False, True])
+@pytest.mark.parametrize("use_async", [False, True])
+@pytest.mark.parametrize("use_multiple_replicas", [False, True])
+def test_responses_actually_streamed(
+    serve_instance, use_fastapi: bool, use_async: bool, use_multiple_replicas: bool
+):
+    """Checks that responses are streamed as they are yielded.
+
+    Also checks that responses can be streamed concurrently from a single replica
+    or from multiple replicas.
+    """
+    signal_actor = SignalActor.remote()
+
+    async def wait_on_signal_async():
+        yield f"{os.getpid()}: before signal"
+        await signal_actor.wait.remote()
+        yield f"{os.getpid()}: after signal"
+
+    def wait_on_signal_sync():
+        yield f"{os.getpid()}: before signal"
+        ray.get(signal_actor.wait.remote())
+        yield f"{os.getpid()}: after signal"
+
+    if use_fastapi:
+        app = FastAPI()
+
+        @serve.deployment
+        @serve.ingress(app)
+        class SimpleGenerator:
+            @app.get("/")
+            def stream(self, request: Request) -> StreamingResponse:
+                gen = wait_on_signal_async() if use_async else wait_on_signal_sync()
+                return StreamingResponse(gen, media_type="text/plain")
+
+    else:
+
+        @serve.deployment
+        class SimpleGenerator:
+            def __call__(self, request: Request) -> StreamingResponse:
+                gen = wait_on_signal_async() if use_async else wait_on_signal_sync()
+                return StreamingResponse(gen, media_type="text/plain")
+
+    serve.run(
+        SimpleGenerator.options(
+            ray_actor_options={"num_cpus": 0},
+            num_replicas=2 if use_multiple_replicas else 1,
+        ).bind()
+    )
+
+    urls = get_application_urls("HTTP")
+
+    requester = StreamingRequester.remote()
+    if len(urls) == 2:
+        gen1 = requester.make_request.options(num_returns="streaming").remote(urls[0])
+        gen2 = requester.make_request.options(num_returns="streaming").remote(urls[1])
+    else:
+        gen1 = requester.make_request.options(num_returns="streaming").remote()
+        gen2 = requester.make_request.options(num_returns="streaming").remote()
+
+    # Check that we get the first responses before the signal is sent
+    # (so the generator is still hanging after the first yield).
+    gen1_result = ray.get(next(gen1))
+    gen2_result = ray.get(next(gen2))
+    assert gen1_result.endswith("before signal")
+    assert gen2_result.endswith("before signal")
+    gen1_pid = gen1_result.split(":")[0]
+    gen2_pid = gen2_result.split(":")[0]
+    if use_multiple_replicas:
+        assert gen1_pid != gen2_pid
+    else:
+        assert gen1_pid == gen2_pid
+
+    # Check that the next obj_ref is not ready yet for both generators.
+    assert gen1._next_sync(timeout_s=0.01).is_nil()
+    assert gen2._next_sync(timeout_s=0.01).is_nil()
+
+    # Now send signal to actor, second yield happens and we should get responses.
+    ray.get(signal_actor.send.remote())
+    gen1_result = ray.get(next(gen1))
+    gen2_result = ray.get(next(gen2))
+    assert gen1_result.startswith(gen1_pid)
+    assert gen2_result.startswith(gen2_pid)
+    assert gen1_result.endswith("after signal")
+    assert gen2_result.endswith("after signal")
+
+    # Client should be done getting messages.
+    with pytest.raises(StopIteration):
+        next(gen1)
+
+    with pytest.raises(StopIteration):
+        next(gen2)
+
+
+@pytest.mark.parametrize("use_fastapi", [False, True])
+def test_metadata_preserved(serve_instance, use_fastapi: bool):
+    """Check that status code, headers, and media type are preserved."""
+
+    def hi_gen():
+        for i in range(10):
+            yield f"hi_{i}"
+
+    if use_fastapi:
+        app = FastAPI()
+
+        @serve.deployment
+        @serve.ingress(app)
+        class SimpleGenerator:
+            @app.get("/")
+            def stream_hi(self, request: Request) -> StreamingResponse:
+                return StreamingResponse(
+                    hi_gen(),
+                    status_code=301,
+                    headers={"hello": "world"},
+                    media_type="foo/bar",
+                )
+
+    else:
+
+        @serve.deployment
+        class SimpleGenerator:
+            def __call__(self, request: Request) -> StreamingResponse:
+                return StreamingResponse(
+                    hi_gen(),
+                    status_code=301,
+                    headers={"hello": "world"},
+                    media_type="foo/bar",
+                )
+
+    serve.run(SimpleGenerator.bind())
+
+    url = get_application_url("HTTP")
+    with httpx.stream("GET", url) as r:
+        assert r.status_code == 301
+        assert r.headers["hello"] == "world"
+        assert r.headers["content-type"] == "foo/bar"
+        for i, chunk in enumerate(r.iter_bytes()):
+            assert chunk == f"hi_{i}".encode("utf-8")
+
+
+@pytest.mark.parametrize("use_fastapi", [False, True])
+@pytest.mark.parametrize("use_async", [False, True])
+def test_exception_in_generator(serve_instance, use_async: bool, use_fastapi: bool):
+    async def hi_gen_async():
+        yield "first result"
+        raise Exception("raised in generator")
+
+    def hi_gen_sync():
+        yield "first result"
+        raise Exception("raised in generator")
+
+    if use_fastapi:
+        app = FastAPI()
+
+        @serve.deployment
+        @serve.ingress(app)
+        class SimpleGenerator:
+            @app.get("/")
+            def stream_hi(self, request: Request) -> StreamingResponse:
+                gen = hi_gen_async() if use_async else hi_gen_sync()
+                return StreamingResponse(gen, media_type="text/plain")
+
+    else:
+
+        @serve.deployment
+        class SimpleGenerator:
+            def __call__(self, request: Request) -> StreamingResponse:
+                gen = hi_gen_async() if use_async else hi_gen_sync()
+                return StreamingResponse(gen, media_type="text/plain")
+
+    serve.run(SimpleGenerator.bind())
+
+    url = get_application_url("HTTP")
+    with httpx.stream("GET", url) as r:
+        r.raise_for_status()
+        stream_iter = r.iter_text()
+        assert next(stream_iter) == "first result"
+        with pytest.raises(httpx.HTTPError):
+            next(stream_iter)
+
+
+@pytest.mark.parametrize("use_fastapi", [False, True])
+@pytest.mark.parametrize("use_async", [False, True])
+def test_proxy_from_streaming_handle(
+    serve_instance, use_async: bool, use_fastapi: bool
+):
+    signal_actor = SignalActor.remote()
+
+    @serve.deployment
+    class Streamer:
+        async def hi_gen_async(self):
+            for i in range(10):
+                yield f"hi_{i}"
+                # to avoid coalescing chunks
+                await signal_actor.wait.remote()
+
+        def hi_gen_sync(self):
+            for i in range(10):
+                yield f"hi_{i}"
+                # to avoid coalescing chunks
+                ray.get(signal_actor.wait.remote())
+
+    if use_fastapi:
+        app = FastAPI()
+
+        @serve.deployment
+        @serve.ingress(app)
+        class SimpleGenerator:
+            def __init__(self, handle: DeploymentHandle):
+                self._h = handle.options(stream=True)
+
+            @app.get("/")
+            def stream_hi(self, request: Request) -> StreamingResponse:
+                if use_async:
+                    gen = self._h.hi_gen_async.remote()
+                else:
+                    gen = self._h.hi_gen_sync.remote()
+
+                return StreamingResponse(gen, media_type="text/plain")
+
+    else:
+
+        @serve.deployment
+        class SimpleGenerator:
+            def __init__(self, handle: DeploymentHandle):
+                self._h = handle.options(stream=True)
+
+            def __call__(self, request: Request) -> StreamingResponse:
+                if use_async:
+                    gen = self._h.hi_gen_async.remote()
+                else:
+                    gen = self._h.hi_gen_sync.remote()
+
+                return StreamingResponse(gen, media_type="text/plain")
+
+    serve.run(SimpleGenerator.bind(Streamer.bind()))
+
+    url = get_application_url("HTTP")
+    with httpx.stream("GET", url) as r:
+        r.raise_for_status()
+        for i, chunk in enumerate(r.iter_text()):
+            assert chunk == f"hi_{i}"
+            ray.get(signal_actor.send.remote())
+
+
+def test_http_disconnect(serve_instance):
+    """Test that response generators are cancelled when the client disconnects."""
+    signal_actor = SignalActor.remote()
+
+    @serve.deployment
+    class SimpleGenerator:
+        def __call__(self, request: Request) -> StreamingResponse:
+            async def wait_for_disconnect():
+                yield "hi"
+                async with send_signal_on_cancellation(signal_actor):
+                    pass
+
+            return StreamingResponse(wait_for_disconnect())
+
+    serve.run(SimpleGenerator.bind())
+
+    url = get_application_url("HTTP")
+    with httpx.stream("GET", url):
+        with pytest.raises(TimeoutError):
+            _ = ray.get(signal_actor.wait.remote(), timeout=1)
+
+    ray.get(signal_actor.wait.remote(), timeout=5)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-s", __file__]))