wehub-resource-sync
271 lines
8.0 KiB
Python
271 lines
8.0 KiB
Python
import asyncio
|
|
import concurrent.futures
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
|
|
from ray.exceptions import RayTaskError
|
|
from ray.serve._private.proxy_state import ActorProxyWrapper, wrap_as_future
|
|
from ray.serve.schema import LoggingConfig
|
|
|
|
|
|
def _create_object_ref_mock():
|
|
fut = concurrent.futures.Future()
|
|
|
|
ray_object_ref_mock = Mock(future=Mock(return_value=fut))
|
|
|
|
return ray_object_ref_mock, fut
|
|
|
|
|
|
def _create_mocked_actor_proxy_wrapper(actor_handle_mock):
|
|
return ActorProxyWrapper(
|
|
logging_config=LoggingConfig(),
|
|
actor_handle=actor_handle_mock,
|
|
node_id="some_node_id",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_wrap_as_future_timeout():
|
|
object_ref_mock, fut = _create_object_ref_mock()
|
|
|
|
aio_fut = wrap_as_future(ref=object_ref_mock, timeout_s=0)
|
|
|
|
assert not aio_fut.done()
|
|
# Yield the event-loop
|
|
await asyncio.sleep(0.001)
|
|
|
|
assert aio_fut.done()
|
|
with pytest.raises(TimeoutError) as exc_info:
|
|
aio_fut.result()
|
|
|
|
assert "Future cancelled after timeout 0s" in str(exc_info.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_wrap_as_future_success():
|
|
# Test #1: Validate wrapped asyncio future completes, upon completion of the
|
|
# ObjectRef's one
|
|
|
|
object_ref_mock, fut = _create_object_ref_mock()
|
|
|
|
aio_fut = wrap_as_future(ref=object_ref_mock, timeout_s=3600)
|
|
|
|
assert not aio_fut.done()
|
|
# Complete source future (ObjectRef one)
|
|
fut.set_result("test")
|
|
# Yield the event-loop
|
|
await asyncio.sleep(0.001)
|
|
|
|
assert aio_fut.done()
|
|
assert aio_fut.result() == "test"
|
|
|
|
# Test #2: Wrapped asyncio future completed before time-out expiration,
|
|
# should not be affected by the cancellation callback
|
|
|
|
object_ref_mock, fut = _create_object_ref_mock()
|
|
# Purposefully set timeout to 0, ie future has to be cancelled upon next
|
|
# event-loop iteration
|
|
aio_fut = wrap_as_future(ref=object_ref_mock, timeout_s=0)
|
|
|
|
assert not aio_fut.done()
|
|
# Complete source future (ObjectRef one)
|
|
fut.set_result("test")
|
|
# Yield the event-loop
|
|
await asyncio.sleep(0.001)
|
|
|
|
assert aio_fut.done()
|
|
assert aio_fut.result() == "test"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_wrap_as_future_timeout_does_not_assert_on_late_source_completion(
|
|
monkeypatch,
|
|
):
|
|
object_ref_mock, fut = _create_object_ref_mock()
|
|
loop = asyncio.get_running_loop()
|
|
queued_callbacks = []
|
|
|
|
def queue_callback(callback, *args, context=None):
|
|
queued_callbacks.append((callback, args, context))
|
|
|
|
monkeypatch.setattr(loop, "call_soon_threadsafe", queue_callback)
|
|
|
|
aio_fut = wrap_as_future(ref=object_ref_mock, timeout_s=0)
|
|
|
|
fut.set_result("test")
|
|
assert len(queued_callbacks) == 1
|
|
|
|
await asyncio.sleep(0.001)
|
|
|
|
assert aio_fut.done()
|
|
with pytest.raises(TimeoutError):
|
|
aio_fut.result()
|
|
|
|
callback, args, _ = queued_callbacks.pop()
|
|
callback(*args)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("response", "is_ready"),
|
|
[
|
|
# ProxyActor.ready responds with an tuple/array of 2 strings
|
|
('["foo", "bar"]', True),
|
|
("malformed_json", False),
|
|
(Exception(), False),
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_is_ready_check_success(response, is_ready):
|
|
"""Tests calling is_ready method on ProxyActorWrapper, mocking out underlying
|
|
ActorHandle response"""
|
|
object_ref_mock, fut = _create_object_ref_mock()
|
|
actor_handle_mock = Mock(ready=Mock(remote=Mock(return_value=object_ref_mock)))
|
|
|
|
proxy_wrapper = _create_mocked_actor_proxy_wrapper(actor_handle_mock)
|
|
|
|
for _ in range(10):
|
|
assert proxy_wrapper.is_ready(timeout_s=1) is None
|
|
# Yield loop!
|
|
await asyncio.sleep(0.01)
|
|
|
|
# Complete source future
|
|
if isinstance(response, Exception):
|
|
object_ref_mock.future().set_exception(response)
|
|
else:
|
|
object_ref_mock.future().set_result(response)
|
|
|
|
# Yield loop!
|
|
await asyncio.sleep(0)
|
|
# NOTE: Timeout setting is only relevant, in case there's no pending request
|
|
# and one will be issued
|
|
assert proxy_wrapper.is_ready(timeout_s=1) is is_ready
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_is_ready_check_timeout():
|
|
object_ref_mock, fut = _create_object_ref_mock()
|
|
actor_handle_mock = Mock(ready=Mock(remote=Mock(return_value=object_ref_mock)))
|
|
|
|
proxy_wrapper = _create_mocked_actor_proxy_wrapper(actor_handle_mock)
|
|
|
|
# First call, invokes ProxyActor.ready call
|
|
assert proxy_wrapper.is_ready(timeout_s=0) is None
|
|
# Yield loop!
|
|
await asyncio.sleep(0.001)
|
|
|
|
assert proxy_wrapper.is_ready(timeout_s=0) is False
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("response", "is_healthy"),
|
|
[
|
|
(True, True),
|
|
(RayTaskError("check_health", "<traceback>", "cuz"), False),
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_is_healthy_check_success(response, is_healthy):
|
|
"""Tests calling is_healthy method on ProxyActorWrapper, mocking out underlying
|
|
ActorHandle response"""
|
|
object_ref_mock, fut = _create_object_ref_mock()
|
|
actor_handle_mock = Mock(
|
|
check_health=Mock(remote=Mock(return_value=object_ref_mock))
|
|
)
|
|
|
|
proxy_wrapper = _create_mocked_actor_proxy_wrapper(actor_handle_mock)
|
|
|
|
for _ in range(10):
|
|
assert proxy_wrapper.is_healthy(timeout_s=1) is None
|
|
# Yield loop!
|
|
await asyncio.sleep(0.01)
|
|
|
|
object_ref_mock.future.assert_called_once()
|
|
|
|
# Complete source future
|
|
if isinstance(response, Exception):
|
|
object_ref_mock.future.return_value.set_exception(response)
|
|
else:
|
|
object_ref_mock.future.return_value.set_result(response)
|
|
|
|
# Yield loop!
|
|
await asyncio.sleep(0)
|
|
# NOTE: Timeout setting is only relevant, in case there's no pending request
|
|
# and one will be issued
|
|
assert proxy_wrapper.is_healthy(timeout_s=1) is is_healthy
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_is_healthy_check_timeout():
|
|
object_ref_mock, fut = _create_object_ref_mock()
|
|
actor_handle_mock = Mock(
|
|
check_health=Mock(remote=Mock(return_value=object_ref_mock))
|
|
)
|
|
|
|
proxy_wrapper = _create_mocked_actor_proxy_wrapper(actor_handle_mock)
|
|
|
|
# First call, invokes ProxyActor.ready call
|
|
assert proxy_wrapper.is_healthy(timeout_s=0) is None
|
|
# Yield loop!
|
|
await asyncio.sleep(0.001)
|
|
|
|
assert proxy_wrapper.is_healthy(timeout_s=0) is False
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("response", "is_drained"),
|
|
[
|
|
(True, True),
|
|
(False, False),
|
|
(RayTaskError("is_drained", "<traceback>", "cuz"), False),
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_is_drained_check_success(response, is_drained):
|
|
"""Tests calling is_drained method on ProxyActorWrapper, mocking out underlying
|
|
ActorHandle response"""
|
|
object_ref_mock, fut = _create_object_ref_mock()
|
|
actor_handle_mock = Mock(is_drained=Mock(remote=Mock(return_value=object_ref_mock)))
|
|
|
|
proxy_wrapper = _create_mocked_actor_proxy_wrapper(actor_handle_mock)
|
|
|
|
for _ in range(10):
|
|
assert proxy_wrapper.is_drained(timeout_s=1) is None
|
|
# Yield loop!
|
|
await asyncio.sleep(0.01)
|
|
|
|
object_ref_mock.future.assert_called_once()
|
|
|
|
# Complete source future
|
|
if isinstance(response, Exception):
|
|
object_ref_mock.future.return_value.set_exception(response)
|
|
else:
|
|
object_ref_mock.future.return_value.set_result(response)
|
|
|
|
# Yield loop!
|
|
await asyncio.sleep(0)
|
|
# NOTE: Timeout setting is only relevant, in case there's no pending request
|
|
# and one will be issued
|
|
assert proxy_wrapper.is_drained(timeout_s=1) is is_drained
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_is_drained_check_timeout():
|
|
object_ref_mock, fut = _create_object_ref_mock()
|
|
actor_handle_mock = Mock(is_drained=Mock(remote=Mock(return_value=object_ref_mock)))
|
|
|
|
proxy_wrapper = _create_mocked_actor_proxy_wrapper(actor_handle_mock)
|
|
|
|
# First call, invokes ProxyActor.ready call
|
|
assert proxy_wrapper.is_drained(timeout_s=0) is None
|
|
# Yield loop!
|
|
await asyncio.sleep(0.001)
|
|
|
|
assert proxy_wrapper.is_drained(timeout_s=0) is False
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-s", __file__]))
|