wehub-resource-sync
404 lines
12 KiB
Python
404 lines
12 KiB
Python
import sys
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray import serve
|
|
from ray._common.test_utils import (
|
|
SignalActor,
|
|
async_wait_for_condition,
|
|
wait_for_condition,
|
|
)
|
|
from ray.serve._private.constants import RAY_SERVE_FORCE_LOCAL_TESTING_MODE
|
|
from ray.serve._private.test_utils import send_signal_on_cancellation, tlog
|
|
from ray.serve.exceptions import RequestCancelledError
|
|
|
|
|
|
def test_cancel_sync_handle_call_during_execution(serve_instance):
|
|
"""Test cancelling handle request during execution (sync context)."""
|
|
running_signal_actor = SignalActor.remote()
|
|
cancelled_signal_actor = SignalActor.remote()
|
|
|
|
@serve.deployment
|
|
class Ingress:
|
|
async def __call__(self, *args):
|
|
async with send_signal_on_cancellation(cancelled_signal_actor):
|
|
await running_signal_actor.send.remote()
|
|
|
|
h = serve.run(Ingress.bind())
|
|
|
|
# Send a request and wait for it to start executing.
|
|
r = h.remote()
|
|
ray.get(running_signal_actor.wait.remote(), timeout=10)
|
|
|
|
# Cancel it and verify that it is cancelled via signal.
|
|
r.cancel()
|
|
ray.get(cancelled_signal_actor.wait.remote(), timeout=10)
|
|
|
|
with pytest.raises(RequestCancelledError):
|
|
r.result()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
|
|
reason="local_testing_mode doesn't have assignment/execution split",
|
|
)
|
|
def test_cancel_sync_handle_call_during_assignment(serve_instance):
|
|
"""Test cancelling handle request during assignment (sync context)."""
|
|
signal_actor = SignalActor.remote()
|
|
|
|
@serve.deployment(max_ongoing_requests=1)
|
|
class Ingress:
|
|
def __init__(self):
|
|
self._num_requests = 0
|
|
|
|
async def __call__(self, *args):
|
|
self._num_requests += 1
|
|
await signal_actor.wait.remote()
|
|
|
|
return self._num_requests
|
|
|
|
h = serve.run(Ingress.bind())
|
|
|
|
# Send a request and wait for it to be ongoing so we know that further requests
|
|
# will block trying to assign a replica.
|
|
initial_response = h.remote()
|
|
wait_for_condition(lambda: ray.get(signal_actor.cur_num_waiters.remote()) == 1)
|
|
|
|
# Make a second request, cancel it, and verify that it is cancelled.
|
|
second_response = h.remote()
|
|
second_response.cancel()
|
|
with pytest.raises(RequestCancelledError):
|
|
second_response.result()
|
|
|
|
# Now signal the initial request to finish and check that the second request
|
|
# never reached the replica.
|
|
ray.get(signal_actor.send.remote())
|
|
assert initial_response.result() == 1
|
|
for i in range(2, 12):
|
|
assert h.remote().result() == i
|
|
|
|
|
|
def test_cancel_async_handle_call_during_execution(serve_instance):
|
|
"""Test cancelling handle request during execution (async context)."""
|
|
running_signal_actor = SignalActor.remote()
|
|
cancelled_signal_actor = SignalActor.remote()
|
|
|
|
@serve.deployment
|
|
class Downstream:
|
|
async def __call__(self, *args):
|
|
async with send_signal_on_cancellation(cancelled_signal_actor):
|
|
await running_signal_actor.send.remote()
|
|
|
|
@serve.deployment
|
|
class Ingress:
|
|
def __init__(self, handle):
|
|
self._h = handle
|
|
|
|
async def __call__(self, *args):
|
|
# Send a request and wait for it to start executing.
|
|
r = self._h.remote()
|
|
await running_signal_actor.wait.remote()
|
|
|
|
# Cancel it and verify that it is cancelled via signal.
|
|
r.cancel()
|
|
await cancelled_signal_actor.wait.remote()
|
|
|
|
with pytest.raises(RequestCancelledError):
|
|
await r
|
|
|
|
h = serve.run(Ingress.bind(Downstream.bind()))
|
|
h.remote().result() # Would raise if test failed.
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
|
|
reason="local_testing_mode doesn't have assignment/execution split",
|
|
)
|
|
def test_cancel_async_handle_call_during_assignment(serve_instance):
|
|
"""Test cancelling handle request during assignment (async context)."""
|
|
signal_actor = SignalActor.remote()
|
|
|
|
@serve.deployment(max_ongoing_requests=1)
|
|
class Downstream:
|
|
def __init__(self):
|
|
self._num_requests = 0
|
|
|
|
async def __call__(self, *args):
|
|
self._num_requests += 1
|
|
await signal_actor.wait.remote()
|
|
|
|
return self._num_requests
|
|
|
|
@serve.deployment
|
|
class Ingress:
|
|
def __init__(self, handle):
|
|
self._h = handle
|
|
|
|
async def __call__(self, *args):
|
|
# Send a request and wait for it to be ongoing so we know that further
|
|
# requests will block trying to assign a replica.
|
|
initial_response = self._h.remote()
|
|
|
|
async def one_waiter():
|
|
return await signal_actor.cur_num_waiters.remote() == 1
|
|
|
|
await async_wait_for_condition(one_waiter)
|
|
|
|
# Make a second request, cancel it, and verify that it is cancelled.
|
|
second_response = self._h.remote()
|
|
second_response.cancel()
|
|
with pytest.raises(RequestCancelledError):
|
|
await second_response
|
|
|
|
# Now signal the initial request to finish and check that the second request
|
|
# never reached the replica.
|
|
await signal_actor.send.remote()
|
|
assert await initial_response == 1
|
|
for i in range(2, 12):
|
|
assert await self._h.remote() == i
|
|
|
|
h = serve.run(Ingress.bind(Downstream.bind()))
|
|
h.remote().result() # Would raise if test failed.
|
|
|
|
|
|
def test_cancel_generator_sync(serve_instance):
|
|
"""Test cancelling streaming handle request during execution."""
|
|
signal_actor = SignalActor.remote()
|
|
|
|
@serve.deployment
|
|
class Ingress:
|
|
async def __call__(self, *args):
|
|
yield "hi"
|
|
async with send_signal_on_cancellation(signal_actor):
|
|
pass
|
|
|
|
h = serve.run(Ingress.bind()).options(stream=True)
|
|
|
|
# Send a request and wait for it to start executing.
|
|
g = h.remote()
|
|
|
|
assert next(g) == "hi"
|
|
|
|
# Cancel it and verify that it is cancelled via signal.
|
|
g.cancel()
|
|
|
|
with pytest.raises(RequestCancelledError):
|
|
next(g)
|
|
|
|
ray.get(signal_actor.wait.remote(), timeout=10)
|
|
|
|
|
|
def test_cancel_generator_async(serve_instance):
|
|
"""Test cancelling streaming handle request during execution."""
|
|
signal_actor = SignalActor.remote()
|
|
|
|
@serve.deployment
|
|
class Downstream:
|
|
async def __call__(self, *args):
|
|
yield "hi"
|
|
async with send_signal_on_cancellation(signal_actor):
|
|
pass
|
|
|
|
@serve.deployment
|
|
class Ingress:
|
|
def __init__(self, handle):
|
|
self._h = handle.options(stream=True)
|
|
|
|
async def __call__(self, *args):
|
|
# Send a request and wait for it to start executing.
|
|
g = self._h.remote()
|
|
assert await g.__anext__() == "hi"
|
|
|
|
# Cancel it and verify that it is cancelled via signal.
|
|
g.cancel()
|
|
|
|
with pytest.raises(RequestCancelledError):
|
|
await g.__anext__()
|
|
|
|
await signal_actor.wait.remote()
|
|
|
|
h = serve.run(Ingress.bind(Downstream.bind()))
|
|
h.remote().result() # Would raise if test failed.
|
|
|
|
|
|
def test_only_relevant_task_is_cancelled(serve_instance):
|
|
"""Test cancelling one request doesn't affect others."""
|
|
signal_actor = SignalActor.remote()
|
|
|
|
@serve.deployment
|
|
class Ingress:
|
|
async def __call__(self, *args):
|
|
await signal_actor.wait.remote()
|
|
return "ok"
|
|
|
|
h = serve.run(Ingress.bind())
|
|
|
|
r1 = h.remote()
|
|
r2 = h.remote()
|
|
|
|
# Wait for both requests to be executing.
|
|
wait_for_condition(lambda: ray.get(signal_actor.cur_num_waiters.remote()) == 2)
|
|
|
|
r1.cancel()
|
|
with pytest.raises(RequestCancelledError):
|
|
r1.result()
|
|
|
|
# Now signal r2 to run to completion and check that it wasn't cancelled.
|
|
ray.get(signal_actor.send.remote())
|
|
assert r2.result() == "ok"
|
|
|
|
|
|
def test_out_of_band_task_is_not_cancelled(serve_instance):
|
|
"""
|
|
Test cancelling a request doesn't cancel tasks submitted
|
|
outside the request's context.
|
|
"""
|
|
signal_actor = SignalActor.remote()
|
|
|
|
@serve.deployment
|
|
class Downstream:
|
|
async def hi(self):
|
|
await signal_actor.wait.remote()
|
|
return "ok"
|
|
|
|
@serve.deployment
|
|
class Ingress:
|
|
def __init__(self, handle):
|
|
self._h = handle
|
|
self._out_of_band_req = self._h.hi.remote()
|
|
|
|
async def __call__(self, *args):
|
|
await self._h.hi.remote()
|
|
|
|
async def get_out_of_band_response(self):
|
|
return await self._out_of_band_req
|
|
|
|
h = serve.run(Ingress.bind(Downstream.bind()))
|
|
|
|
# Send a request, wait for downstream request to start, and cancel it.
|
|
r1 = h.remote()
|
|
wait_for_condition(lambda: ray.get(signal_actor.cur_num_waiters.remote()) == 2)
|
|
|
|
r1.cancel()
|
|
with pytest.raises(RequestCancelledError):
|
|
r1.result()
|
|
|
|
# Now signal out of band request to run to completion and check that it wasn't
|
|
# cancelled.
|
|
ray.get(signal_actor.send.remote())
|
|
assert h.get_out_of_band_response.remote().result() == "ok"
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
|
|
reason="local_testing_mode doesn't implement recursive cancellation",
|
|
)
|
|
def test_recursive_cancellation_during_execution(serve_instance):
|
|
inner_signal_actor = SignalActor.remote()
|
|
outer_signal_actor = SignalActor.remote()
|
|
|
|
@serve.deployment
|
|
async def inner():
|
|
async with send_signal_on_cancellation(inner_signal_actor):
|
|
pass
|
|
|
|
@serve.deployment
|
|
class Ingress:
|
|
def __init__(self, handle):
|
|
self._handle = handle
|
|
|
|
async def __call__(self):
|
|
_ = self._handle.remote()
|
|
async with send_signal_on_cancellation(outer_signal_actor):
|
|
pass
|
|
|
|
h = serve.run(Ingress.bind(inner.bind()))
|
|
|
|
resp = h.remote()
|
|
with pytest.raises(TimeoutError):
|
|
resp.result(timeout_s=0.5)
|
|
|
|
resp.cancel()
|
|
ray.get(inner_signal_actor.wait.remote(), timeout=10)
|
|
ray.get(outer_signal_actor.wait.remote(), timeout=10)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
|
|
reason="local_testing_mode doesn't implement recursive cancellation",
|
|
)
|
|
def test_recursive_cancellation_during_assignment(serve_instance):
|
|
signal = SignalActor.remote()
|
|
|
|
@serve.deployment(max_ongoing_requests=1)
|
|
class Counter:
|
|
def __init__(self):
|
|
self._count = 0
|
|
|
|
async def __call__(self):
|
|
self._count += 1
|
|
await signal.wait.remote()
|
|
|
|
def get_count(self):
|
|
return self._count
|
|
|
|
@serve.deployment
|
|
class Ingress:
|
|
def __init__(self, handle):
|
|
self._handle = handle
|
|
|
|
async def __call__(self):
|
|
self._handle.remote()
|
|
await signal.wait.remote()
|
|
return "hi"
|
|
|
|
async def get_count(self):
|
|
return await self._handle.get_count.remote()
|
|
|
|
async def check_requests_pending_assignment_cache(self):
|
|
requests_pending_assignment = ray.serve.context._requests_pending_assignment
|
|
return {k: list(v.keys()) for k, v in requests_pending_assignment.items()}
|
|
|
|
h = serve.run(Ingress.bind(Counter.bind()))
|
|
|
|
# Send two requests to Ingress. The second should be queued and
|
|
# pending assignment at Ingress because max ongoing requests for
|
|
# Counter is only 1.
|
|
tlog("Sending two requests to Ingress.")
|
|
resp1 = h.remote()
|
|
with pytest.raises(TimeoutError):
|
|
resp1.result(timeout_s=0.5)
|
|
resp2 = h.remote()
|
|
with pytest.raises(TimeoutError):
|
|
resp2.result(timeout_s=0.5)
|
|
|
|
# Cancel second request, which should be pending assignment.
|
|
tlog("Canceling second request.")
|
|
resp2.cancel()
|
|
|
|
# Release signal so that the first request can complete, and any new
|
|
# requests to Counter can be let through
|
|
tlog("Releasing signal.")
|
|
ray.get(signal.send.remote())
|
|
assert resp1.result() == "hi"
|
|
|
|
# The second request, even though it was pending assignment to a
|
|
# Counter replica, should have been properly canceled. Confirm this
|
|
# by making sure that no more calls to __call__ were made
|
|
for _ in range(10):
|
|
assert h.get_count.remote().result() == 1
|
|
|
|
tlog("Confirmed second request was properly canceled.")
|
|
|
|
# Check that cache was cleared so there are no memory leaks
|
|
requests_pending_assignment = (
|
|
h.check_requests_pending_assignment_cache.remote().result()
|
|
)
|
|
for k, v in requests_pending_assignment.items():
|
|
assert len(v) == 0, f"Request {k} has in flight requests in cache: {v}"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", "-s", __file__]))
|