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