wehub-resource-sync
80 lines
2.5 KiB
Python
80 lines
2.5 KiB
Python
import asyncio
|
|
import sys
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from ray.serve._private.common import RequestMetadata
|
|
from ray.serve._private.replica import Replica
|
|
from ray.serve._private.utils import Semaphore
|
|
|
|
|
|
def _make_metadata() -> RequestMetadata:
|
|
return RequestMetadata(
|
|
request_id="test-request",
|
|
internal_request_id="test-internal-request",
|
|
is_direct_ingress=True,
|
|
)
|
|
|
|
|
|
def _make_fake_replica(max_ongoing_requests: int):
|
|
"""Minimal stand-in exposing the real slot/queue accounting methods."""
|
|
fake = MagicMock()
|
|
fake._num_queued_requests = 0
|
|
fake._semaphore = Semaphore(lambda: max_ongoing_requests)
|
|
fake._start_request = Replica._start_request.__get__(fake, Replica)
|
|
fake._track_queued_request = Replica._track_queued_request.__get__(fake, Replica)
|
|
return fake
|
|
|
|
|
|
async def _wait_for_semaphore_waiter(semaphore: Semaphore):
|
|
for _ in range(1000):
|
|
await asyncio.sleep(0)
|
|
if semaphore._waiters:
|
|
return
|
|
raise AssertionError("no coroutine blocked on the semaphore")
|
|
|
|
|
|
class TestTrackQueuedRequest:
|
|
def test_release_decrements_exactly_once(self):
|
|
fake = _make_fake_replica(max_ongoing_requests=1)
|
|
|
|
with fake._track_queued_request() as release:
|
|
assert fake._num_queued_requests == 1
|
|
release()
|
|
assert fake._num_queued_requests == 0
|
|
# Extra calls are a no-op.
|
|
release()
|
|
assert fake._num_queued_requests == 0
|
|
|
|
# Exiting the block must not decrement below zero.
|
|
assert fake._num_queued_requests == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_count_released_when_cancelled_while_waiting(self):
|
|
"""A request cancelled while blocked on the slot is released on block
|
|
exit, as the direct-ingress handlers wire it up."""
|
|
fake = _make_fake_replica(max_ongoing_requests=1)
|
|
|
|
# Hold the only slot so the request blocks while queued.
|
|
await fake._semaphore.acquire()
|
|
|
|
async def handler():
|
|
with fake._track_queued_request() as release:
|
|
async with fake._start_request(_make_metadata()):
|
|
release()
|
|
|
|
task = asyncio.ensure_future(handler())
|
|
await _wait_for_semaphore_waiter(fake._semaphore)
|
|
assert fake._num_queued_requests == 1
|
|
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
|
|
assert fake._num_queued_requests == 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", "-s", __file__]))
|