wehub-resource-sync
353 lines
11 KiB
Python
353 lines
11 KiB
Python
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__]))
|