chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,959 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from typing import List
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray import serve
|
||||
from ray._common.utils import get_or_create_event_loop
|
||||
from ray.serve._private.common import DeploymentID, ReplicaID
|
||||
from ray.serve._private.config import DeploymentConfig
|
||||
from ray.serve._private.constants import SERVE_LOGGER_NAME
|
||||
from ray.serve.batching import _BatchQueue
|
||||
from ray.serve.exceptions import RayServeException
|
||||
|
||||
# Setup the global replica context for the test.
|
||||
default_deployment_config = DeploymentConfig()
|
||||
ray.serve.context._set_internal_replica_context(
|
||||
replica_id=ReplicaID(unique_id="test", deployment_id=DeploymentID(name="test")),
|
||||
servable_object=None,
|
||||
_deployment_config=default_deployment_config,
|
||||
rank=0,
|
||||
world_size=1,
|
||||
)
|
||||
|
||||
|
||||
class FakeStream:
|
||||
def __init__(self):
|
||||
self.messages = []
|
||||
|
||||
def write(self, buf):
|
||||
self.messages.append(buf)
|
||||
|
||||
def reset_message(self):
|
||||
self.messages = []
|
||||
|
||||
|
||||
# We use a single event loop for the entire test session. Without this
|
||||
# fixture, the event loop is sometimes prematurely terminated by pytest.
|
||||
@pytest.fixture(scope="session")
|
||||
def event_loop():
|
||||
loop = get_or_create_event_loop()
|
||||
yield loop
|
||||
loop.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_decorator_validation():
|
||||
@serve.batch
|
||||
async def function():
|
||||
pass
|
||||
|
||||
@serve.batch(max_batch_size=10, batch_wait_timeout_s=1.5)
|
||||
async def function2():
|
||||
pass
|
||||
|
||||
class Class:
|
||||
@serve.batch
|
||||
async def method(self):
|
||||
pass
|
||||
|
||||
class Class2:
|
||||
@serve.batch(max_batch_size=10, batch_wait_timeout_s=1.5)
|
||||
async def method(self):
|
||||
pass
|
||||
|
||||
with pytest.raises(TypeError, match="async def"):
|
||||
|
||||
@serve.batch
|
||||
def non_async_function():
|
||||
pass
|
||||
|
||||
with pytest.raises(TypeError, match="async def"):
|
||||
|
||||
class NotAsync:
|
||||
@serve.batch
|
||||
def method(self, requests):
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
|
||||
class ZeroBatch:
|
||||
@serve.batch(max_batch_size=0)
|
||||
async def method(self, requests):
|
||||
pass
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
|
||||
class FloatNonIntBatch:
|
||||
@serve.batch(max_batch_size=1.1)
|
||||
async def method(self, requests):
|
||||
pass
|
||||
|
||||
class FloatIntegerBatch:
|
||||
@serve.batch(max_batch_size=1.0)
|
||||
async def method(self, requests):
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
|
||||
class NegativeTimeout:
|
||||
@serve.batch(batch_wait_timeout_s=-0.1)
|
||||
async def method(self, requests):
|
||||
pass
|
||||
|
||||
class FloatZeroTimeout:
|
||||
@serve.batch(batch_wait_timeout_s=0.0)
|
||||
async def method(self, requests):
|
||||
pass
|
||||
|
||||
class IntZeroTimeout:
|
||||
@serve.batch(batch_wait_timeout_s=0)
|
||||
async def method(self, requests):
|
||||
pass
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
|
||||
class NonTimeout:
|
||||
@serve.batch(batch_wait_timeout_s="a")
|
||||
async def method(self, requests):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("use_class", [True, False])
|
||||
async def test_batch_size_one_long_timeout(use_class):
|
||||
@serve.batch(max_batch_size=1, batch_wait_timeout_s=1000)
|
||||
async def long_timeout(requests):
|
||||
if "raise" in requests:
|
||||
_ = 1 / 0
|
||||
return requests
|
||||
|
||||
class LongTimeout:
|
||||
@serve.batch(max_batch_size=1, batch_wait_timeout_s=1000)
|
||||
async def long_timeout(self, requests):
|
||||
if "raise" in requests:
|
||||
_ = 1 / 0
|
||||
return requests
|
||||
|
||||
cls = LongTimeout()
|
||||
|
||||
async def call(arg):
|
||||
if use_class:
|
||||
return await cls.long_timeout(arg)
|
||||
else:
|
||||
return await long_timeout(arg)
|
||||
|
||||
assert await call("hi") == "hi"
|
||||
with pytest.raises(ZeroDivisionError):
|
||||
await call("raise")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("use_class", [True, False])
|
||||
async def test_batch_size_multiple_zero_timeout(use_class):
|
||||
block_execution_event = asyncio.Event()
|
||||
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=0)
|
||||
async def zero_timeout(requests):
|
||||
await block_execution_event.wait()
|
||||
if "raise" in requests:
|
||||
_ = 1 / 0
|
||||
return requests
|
||||
|
||||
class ZeroTimeout:
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=0)
|
||||
async def zero_timeout(self, requests):
|
||||
await block_execution_event.wait()
|
||||
if "raise" in requests:
|
||||
_ = 1 / 0
|
||||
return requests
|
||||
|
||||
cls = ZeroTimeout()
|
||||
|
||||
async def call(arg):
|
||||
if use_class:
|
||||
return await cls.zero_timeout(arg)
|
||||
else:
|
||||
return await zero_timeout(arg)
|
||||
|
||||
block_execution_event.set()
|
||||
assert await call("hi") == "hi"
|
||||
with pytest.raises(ZeroDivisionError):
|
||||
await call("raise")
|
||||
block_execution_event.clear()
|
||||
|
||||
# Check that 2 requests will be executed together if available.
|
||||
# The first should cause a size-one batch to be executed, then
|
||||
# the next two should be executed together (signaled by both
|
||||
# having the exception).
|
||||
t1 = get_or_create_event_loop().create_task(call("hi1"))
|
||||
with pytest.raises(asyncio.TimeoutError):
|
||||
await asyncio.wait_for(asyncio.shield(t1), timeout=0.001)
|
||||
|
||||
t2 = get_or_create_event_loop().create_task(call("hi2"))
|
||||
t3 = get_or_create_event_loop().create_task(call("raise"))
|
||||
|
||||
block_execution_event.set()
|
||||
assert await t1 == "hi1"
|
||||
|
||||
with pytest.raises(ZeroDivisionError):
|
||||
await t2
|
||||
with pytest.raises(ZeroDivisionError):
|
||||
await t3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_timeout_empty_queue():
|
||||
"""Check that Serve waits when creating batches.
|
||||
|
||||
Serve should wait a full batch_wait_timeout_s after receiving the first
|
||||
request in the next batch before processing the batch.
|
||||
"""
|
||||
|
||||
@serve.batch(max_batch_size=10, batch_wait_timeout_s=0.25)
|
||||
async def no_op(requests):
|
||||
return ["No-op"] * len(requests)
|
||||
|
||||
num_iterations = 2
|
||||
for iteration in range(num_iterations):
|
||||
tasks = [get_or_create_event_loop().create_task(no_op(None)) for _ in range(9)]
|
||||
done, _ = await asyncio.wait(tasks, timeout=0.05)
|
||||
|
||||
# Due to the long timeout, none of the tasks should finish until a tenth
|
||||
# request is submitted
|
||||
assert len(done) == 0
|
||||
|
||||
tasks.append(get_or_create_event_loop().create_task(no_op(None)))
|
||||
done, _ = await asyncio.wait(tasks, timeout=0.05)
|
||||
|
||||
# All the timeout tasks should be finished
|
||||
assert set(tasks) == set(done)
|
||||
assert all(t.result() == "No-op" for t in tasks)
|
||||
|
||||
if iteration < num_iterations - 1:
|
||||
# Leave queue empty for batch_wait_timeout_s between batches
|
||||
time.sleep(0.25)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_wait_queue_exceeds_batch_size_race_condition():
|
||||
"""Check that the wait queue can exceed the batch size.
|
||||
|
||||
This test was added to guard against a race condition documented in
|
||||
https://github.com/ray-project/ray/pull/42705#discussion_r1466653910.
|
||||
"""
|
||||
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=10000)
|
||||
async def no_op(requests):
|
||||
return ["No-op"] * len(requests)
|
||||
|
||||
tasks = [get_or_create_event_loop().create_task(no_op(None)) for _ in range(10)]
|
||||
|
||||
# All the tasks should finish.
|
||||
done, pending = await asyncio.wait(tasks, timeout=0.5)
|
||||
assert len(done) == len(tasks)
|
||||
assert len(pending) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("use_class", [True, False])
|
||||
async def test_batch_size_multiple_long_timeout(use_class):
|
||||
@serve.batch(max_batch_size=3, batch_wait_timeout_s=1000)
|
||||
async def long_timeout(requests):
|
||||
if "raise" in requests:
|
||||
_ = 1 / 0
|
||||
return requests
|
||||
|
||||
class LongTimeout:
|
||||
@serve.batch(max_batch_size=3, batch_wait_timeout_s=1000)
|
||||
async def long_timeout(self, requests):
|
||||
if "raise" in requests:
|
||||
_ = 1 / 0
|
||||
return requests
|
||||
|
||||
cls = LongTimeout()
|
||||
|
||||
async def call(arg):
|
||||
if use_class:
|
||||
return await cls.long_timeout(arg)
|
||||
else:
|
||||
return await long_timeout(arg)
|
||||
|
||||
t1 = get_or_create_event_loop().create_task(call("hi1"))
|
||||
t2 = get_or_create_event_loop().create_task(call("hi2"))
|
||||
done, pending = await asyncio.wait([t1, t2], timeout=0.1)
|
||||
assert len(done) == 0
|
||||
t3 = get_or_create_event_loop().create_task(call("hi3"))
|
||||
done, pending = await asyncio.wait([t1, t2, t3], timeout=100)
|
||||
assert set(done) == {t1, t2, t3}
|
||||
assert [t1.result(), t2.result(), t3.result()] == ["hi1", "hi2", "hi3"]
|
||||
|
||||
t1 = get_or_create_event_loop().create_task(call("hi1"))
|
||||
t2 = get_or_create_event_loop().create_task(call("raise"))
|
||||
done, pending = await asyncio.wait([t1, t2], timeout=0.1)
|
||||
assert len(done) == 0
|
||||
t3 = get_or_create_event_loop().create_task(call("hi3"))
|
||||
done, pending = await asyncio.wait([t1, t2, t3], timeout=100)
|
||||
assert set(done) == {t1, t2, t3}
|
||||
assert all(isinstance(t.exception(), ZeroDivisionError) for t in done)
|
||||
with pytest.raises(ZeroDivisionError):
|
||||
t1.result()
|
||||
with pytest.raises(ZeroDivisionError):
|
||||
t2.result()
|
||||
with pytest.raises(ZeroDivisionError):
|
||||
t3.result()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mode", ["args", "kwargs", "mixed", "out-of-order"])
|
||||
@pytest.mark.parametrize("use_class", [True, False])
|
||||
async def test_batch_args_kwargs(mode, use_class):
|
||||
if use_class:
|
||||
|
||||
class MultipleArgs:
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
|
||||
async def method(self, key1, key2):
|
||||
return [(key1[i], key2[i]) for i in range(len(key1))]
|
||||
|
||||
instance = MultipleArgs()
|
||||
func = instance.method
|
||||
|
||||
else:
|
||||
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
|
||||
async def func(key1, key2):
|
||||
return [(key1[i], key2[i]) for i in range(len(key1))]
|
||||
|
||||
if mode == "args":
|
||||
coros = [func("hi1", "hi2"), func("hi3", "hi4")]
|
||||
elif mode == "kwargs":
|
||||
coros = [func(key1="hi1", key2="hi2"), func(key1="hi3", key2="hi4")]
|
||||
elif mode == "mixed":
|
||||
coros = [func("hi1", key2="hi2"), func("hi3", key2="hi4")]
|
||||
elif mode == "out-of-order":
|
||||
coros = [func(key2="hi2", key1="hi1"), func(key2="hi4", key1="hi3")]
|
||||
|
||||
result = await asyncio.gather(*coros)
|
||||
assert result == [("hi1", "hi2"), ("hi3", "hi4")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("use_class", [True, False])
|
||||
@pytest.mark.parametrize("use_gen", [True, False])
|
||||
async def test_batch_cancellation(use_class, use_gen):
|
||||
block_requests = asyncio.Event()
|
||||
request_was_cancelled = True
|
||||
|
||||
async def unary_implementation(key1, key2):
|
||||
nonlocal block_requests, request_was_cancelled
|
||||
await block_requests.wait()
|
||||
request_was_cancelled = False
|
||||
return [(key1[i], key2[i]) for i in range(len(key1))]
|
||||
|
||||
async def streaming_implementation(key1, key2):
|
||||
nonlocal block_requests, request_was_cancelled
|
||||
await block_requests.wait()
|
||||
request_was_cancelled = False
|
||||
yield [(key1[i], key2[i]) for i in range(len(key1))]
|
||||
|
||||
if use_class:
|
||||
|
||||
class MultipleArgs:
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
|
||||
async def unary_method(self, key1, key2):
|
||||
return await unary_implementation(key1, key2)
|
||||
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
|
||||
async def streaming_method(self, key1, key2):
|
||||
async for value in streaming_implementation(key1, key2):
|
||||
yield value
|
||||
|
||||
instance = MultipleArgs()
|
||||
|
||||
if use_gen:
|
||||
func = instance.streaming_method
|
||||
else:
|
||||
func = instance.unary_method
|
||||
|
||||
else:
|
||||
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
|
||||
async def unary_func(key1, key2):
|
||||
return await unary_implementation(key1, key2)
|
||||
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
|
||||
async def streaming_func(key1, key2):
|
||||
async for value in streaming_implementation(key1, key2):
|
||||
yield value
|
||||
|
||||
if use_gen:
|
||||
func = streaming_func
|
||||
else:
|
||||
func = unary_func
|
||||
|
||||
if use_gen:
|
||||
gens = [func("hi1", "hi2"), func("hi3", "hi4")]
|
||||
tasks = [asyncio.create_task(gen.__anext__()) for gen in gens]
|
||||
else:
|
||||
tasks = [
|
||||
asyncio.create_task(func("hi1", "hi2")),
|
||||
asyncio.create_task(func("hi3", "hi4")),
|
||||
]
|
||||
|
||||
print("Submitted requests.")
|
||||
|
||||
# The requests should be blocked on the long request_timeout
|
||||
done, pending = await asyncio.wait(tasks, timeout=0.01)
|
||||
assert len(done) == 0
|
||||
assert len(pending) == 2
|
||||
|
||||
print("Requests are blocked, as expected.")
|
||||
|
||||
# Cancel the first request. The second request should still be blocked on
|
||||
# the long request_timeout
|
||||
tasks[0].cancel()
|
||||
pending, done = await asyncio.wait(tasks, timeout=0.01)
|
||||
assert len(done) == 1
|
||||
assert len(pending) == 1
|
||||
|
||||
print("Cancelled first request.")
|
||||
|
||||
# Cancel the second request. Both requests should be done.
|
||||
tasks[1].cancel()
|
||||
done, pending = await asyncio.wait(tasks, timeout=0.01)
|
||||
assert len(done) == 2
|
||||
assert len(pending) == 0
|
||||
|
||||
print("Cancelled second request. Sending new requests with no timeout.")
|
||||
|
||||
# Sanity check that the request was actually cancelled.
|
||||
assert request_was_cancelled
|
||||
|
||||
# Unblock requests. The requests should succeed.
|
||||
block_requests.set()
|
||||
|
||||
if use_gen:
|
||||
gens = [func("hi1", "hi2"), func("hi3", "hi4")]
|
||||
tasks = [asyncio.create_task(gen.__anext__()) for gen in gens]
|
||||
else:
|
||||
tasks = [
|
||||
asyncio.create_task(func("hi1", "hi2")),
|
||||
asyncio.create_task(func("hi3", "hi4")),
|
||||
]
|
||||
|
||||
result = await asyncio.gather(*tasks)
|
||||
assert result == [("hi1", "hi2"), ("hi3", "hi4")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("use_class", [True, False])
|
||||
@pytest.mark.parametrize("use_gen", [True, False])
|
||||
async def test_cancellation_after_error(use_class, use_gen):
|
||||
"""Cancelling a request after it errors should be supported."""
|
||||
|
||||
raise_error = asyncio.Event()
|
||||
|
||||
async def unary_implementation(key1, key2):
|
||||
if not raise_error.is_set():
|
||||
raise ValueError()
|
||||
return [(key1[i], key2[i]) for i in range(len(key1))]
|
||||
|
||||
async def streaming_implementation(key1, key2):
|
||||
if not raise_error.is_set():
|
||||
raise ValueError()
|
||||
yield [(key1[i], key2[i]) for i in range(len(key1))]
|
||||
|
||||
if use_class:
|
||||
|
||||
class MultipleArgs:
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
|
||||
async def unary_method(self, key1, key2):
|
||||
return await unary_implementation(key1, key2)
|
||||
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
|
||||
async def streaming_method(self, key1, key2):
|
||||
async for value in streaming_implementation(key1, key2):
|
||||
yield value
|
||||
|
||||
instance = MultipleArgs()
|
||||
|
||||
if use_gen:
|
||||
func = instance.streaming_method
|
||||
else:
|
||||
func = instance.unary_method
|
||||
|
||||
else:
|
||||
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
|
||||
async def unary_func(key1, key2):
|
||||
return await unary_implementation(key1, key2)
|
||||
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
|
||||
async def streaming_func(key1, key2):
|
||||
async for value in streaming_implementation(key1, key2):
|
||||
yield value
|
||||
|
||||
if use_gen:
|
||||
func = streaming_func
|
||||
else:
|
||||
func = unary_func
|
||||
|
||||
# Submit requests and then cancel them.
|
||||
if use_gen:
|
||||
gens = [func("hi1", "hi2"), func("hi3", "hi4")]
|
||||
tasks = [asyncio.create_task(gen.__anext__()) for gen in gens]
|
||||
else:
|
||||
tasks = [
|
||||
asyncio.create_task(func("hi1", "hi2")),
|
||||
asyncio.create_task(func("hi3", "hi4")),
|
||||
]
|
||||
|
||||
print("Submitted initial batch of requests.")
|
||||
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
|
||||
print("Closed initial batch of requests.")
|
||||
|
||||
raise_error.set()
|
||||
|
||||
# Submit requests and check that they still work.
|
||||
if use_gen:
|
||||
gens = [func("hi1", "hi2"), func("hi3", "hi4")]
|
||||
tasks = [asyncio.create_task(gen.__anext__()) for gen in gens]
|
||||
else:
|
||||
tasks = [
|
||||
asyncio.create_task(func("hi1", "hi2")),
|
||||
asyncio.create_task(func("hi3", "hi4")),
|
||||
]
|
||||
|
||||
print("Submitted new batch of requests.")
|
||||
|
||||
result = await asyncio.gather(*tasks)
|
||||
assert result == [("hi1", "hi2"), ("hi3", "hi4")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("use_class", [True, False])
|
||||
async def test_batch_setters(use_class):
|
||||
if use_class:
|
||||
|
||||
class C:
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
|
||||
async def method(self, key1, key2):
|
||||
return [(key1[i], key2[i]) for i in range(len(key1))]
|
||||
|
||||
instance = C()
|
||||
func = instance.method
|
||||
|
||||
else:
|
||||
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
|
||||
async def func(key1, key2):
|
||||
return [(key1[i], key2[i]) for i in range(len(key1))]
|
||||
|
||||
assert func._get_max_batch_size() == 2
|
||||
assert func._get_batch_wait_timeout_s() == 1000
|
||||
|
||||
# @serve.batch should create batches of size 2
|
||||
tasks = [
|
||||
get_or_create_event_loop().create_task(func("hi1", "hi2")),
|
||||
get_or_create_event_loop().create_task(func("hi3", "hi4")),
|
||||
]
|
||||
done, pending = await asyncio.wait(tasks, timeout=0.1)
|
||||
assert len(pending) == 0
|
||||
assert {task.result() for task in done} == {("hi1", "hi2"), ("hi3", "hi4")}
|
||||
|
||||
# Set new values
|
||||
func.set_max_batch_size(3)
|
||||
func.set_batch_wait_timeout_s(15000)
|
||||
|
||||
assert func._get_max_batch_size() == 3
|
||||
assert func._get_batch_wait_timeout_s() == 15000
|
||||
|
||||
# @serve.batch should create batches of size 3
|
||||
tasks = [
|
||||
get_or_create_event_loop().create_task(func("hi1", "hi2")),
|
||||
get_or_create_event_loop().create_task(func("hi3", "hi4")),
|
||||
get_or_create_event_loop().create_task(func("hi5", "hi6")),
|
||||
]
|
||||
done, pending = await asyncio.wait(tasks, timeout=0.1)
|
||||
assert len(pending) == 0
|
||||
assert {task.result() for task in done} == {
|
||||
("hi1", "hi2"),
|
||||
("hi3", "hi4"),
|
||||
("hi5", "hi6"),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_use_earliest_setters():
|
||||
"""@serve.batch should use the right settings when constructing a batch.
|
||||
|
||||
When the @serve.batch setters get called before a batch has started
|
||||
accumulating, the next batch should use the setters' values. When they
|
||||
get called while a batch is accumulating, the previous values should be
|
||||
used.
|
||||
"""
|
||||
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
|
||||
async def func(key1, key2):
|
||||
return [(key1[i], key2[i]) for i in range(len(key1))]
|
||||
|
||||
assert func._get_max_batch_size() == 2
|
||||
assert func._get_batch_wait_timeout_s() == 1000
|
||||
|
||||
# Set new values
|
||||
func.set_max_batch_size(3)
|
||||
func.set_batch_wait_timeout_s(15000)
|
||||
|
||||
assert func._get_max_batch_size() == 3
|
||||
assert func._get_batch_wait_timeout_s() == 15000
|
||||
|
||||
# Should create batches of size 3, even if setters are called while
|
||||
# batch is accumulated
|
||||
tasks = [
|
||||
get_or_create_event_loop().create_task(func("hi1", "hi2")),
|
||||
get_or_create_event_loop().create_task(func("hi3", "hi4")),
|
||||
]
|
||||
|
||||
# Batch should be waiting for last request
|
||||
done, pending = await asyncio.wait(tasks, timeout=0.1)
|
||||
assert len(done) == 0 and len(pending) == 2
|
||||
|
||||
func.set_max_batch_size(1)
|
||||
func.set_batch_wait_timeout_s(0)
|
||||
|
||||
assert func._get_max_batch_size() == 1
|
||||
assert func._get_batch_wait_timeout_s() == 0
|
||||
|
||||
# Batch should still be waiting for last request
|
||||
done, pending = await asyncio.wait(pending, timeout=0.1)
|
||||
assert len(done) == 0 and len(pending) == 2
|
||||
|
||||
# Batch should execute after last request
|
||||
pending.add(get_or_create_event_loop().create_task(func("hi5", "hi6")))
|
||||
done, pending = await asyncio.wait(pending, timeout=0.1)
|
||||
assert len(pending) == 0
|
||||
assert {task.result() for task in done} == {
|
||||
("hi1", "hi2"),
|
||||
("hi3", "hi4"),
|
||||
("hi5", "hi6"),
|
||||
}
|
||||
|
||||
# Next batch should use updated values
|
||||
tasks = [get_or_create_event_loop().create_task(func("hi1", "hi2"))]
|
||||
done, pending = await asyncio.wait(tasks, timeout=0.1)
|
||||
assert len(done) == 1 and len(pending) == 0
|
||||
assert done.pop().result() == ("hi1", "hi2")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mode", ["args", "kwargs", "mixed", "out-of-order"])
|
||||
@pytest.mark.parametrize("use_class", [True, False])
|
||||
@pytest.mark.parametrize("generator_length", [0, 2, 5])
|
||||
async def test_batch_generator_basic(mode, use_class, generator_length):
|
||||
if use_class:
|
||||
|
||||
class MultipleArgs:
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
|
||||
async def method(self, key1, key2):
|
||||
for gen_idx in range(generator_length):
|
||||
yield [(gen_idx, key1[i], key2[i]) for i in range(len(key1))]
|
||||
|
||||
instance = MultipleArgs()
|
||||
func = instance.method
|
||||
|
||||
else:
|
||||
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
|
||||
async def func(key1, key2):
|
||||
for gen_idx in range(generator_length):
|
||||
yield [(gen_idx, key1[i], key2[i]) for i in range(len(key1))]
|
||||
|
||||
if mode == "args":
|
||||
generators = [func("hi1", "hi2"), func("hi3", "hi4")]
|
||||
elif mode == "kwargs":
|
||||
generators = [func(key1="hi1", key2="hi2"), func(key1="hi3", key2="hi4")]
|
||||
elif mode == "mixed":
|
||||
generators = [func("hi1", key2="hi2"), func("hi3", key2="hi4")]
|
||||
elif mode == "out-of-order":
|
||||
generators = [func(key2="hi2", key1="hi1"), func(key2="hi4", key1="hi3")]
|
||||
|
||||
results = [
|
||||
[result async for result in generators[0]],
|
||||
[result async for result in generators[1]],
|
||||
]
|
||||
|
||||
assert results == [
|
||||
[(gen_idx, "hi1", "hi2") for gen_idx in range(generator_length)],
|
||||
[(gen_idx, "hi3", "hi4") for gen_idx in range(generator_length)],
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("error_type", ["runtime_error", "mismatched_lengths"])
|
||||
async def test_batch_generator_exceptions(error_type):
|
||||
GENERATOR_LENGTH = 5
|
||||
ERROR_IDX = 2
|
||||
ERROR_MSG = "Testing error"
|
||||
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
|
||||
async def func(key1, key2):
|
||||
for gen_idx in range(GENERATOR_LENGTH):
|
||||
results = [(gen_idx, key1[i], key2[i]) for i in range(len(key1))]
|
||||
if gen_idx == ERROR_IDX:
|
||||
if error_type == "runtime_error":
|
||||
raise RuntimeError(ERROR_MSG)
|
||||
elif error_type == "mismatched_lengths":
|
||||
yield results * 2
|
||||
yield results
|
||||
|
||||
generators = [func("hi1", "hi2"), func("hi3", "hi4")]
|
||||
|
||||
for generator in generators:
|
||||
for _ in range(ERROR_IDX):
|
||||
await generator.__anext__()
|
||||
|
||||
if error_type == "runtime_error":
|
||||
with pytest.raises(RuntimeError, match=ERROR_MSG):
|
||||
await generator.__anext__()
|
||||
elif error_type == "mismatched_lengths":
|
||||
with pytest.raises(RayServeException):
|
||||
await generator.__anext__()
|
||||
|
||||
with pytest.raises(StopAsyncIteration):
|
||||
await generator.__anext__()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stop_token", [StopAsyncIteration, StopIteration])
|
||||
async def test_batch_generator_early_termination(stop_token):
|
||||
NUM_CALLERS = 4
|
||||
|
||||
event = asyncio.Event()
|
||||
|
||||
@serve.batch(max_batch_size=NUM_CALLERS, batch_wait_timeout_s=1000)
|
||||
async def sequential_terminator(ids: List[int]):
|
||||
"""Terminates callers one-after-another in order of call."""
|
||||
|
||||
for num_finished_callers in range(1, NUM_CALLERS + 1):
|
||||
event.clear()
|
||||
responses = [stop_token for _ in range(num_finished_callers)]
|
||||
responses += [ids[idx] for idx in range(num_finished_callers, NUM_CALLERS)]
|
||||
yield responses
|
||||
await event.wait()
|
||||
|
||||
ids = list(range(NUM_CALLERS))
|
||||
generators = [sequential_terminator(id) for id in ids]
|
||||
for id, generator in zip(ids, generators):
|
||||
async for result in generator:
|
||||
assert result == id
|
||||
|
||||
# Each terminated caller frees the sequential_terminator to process
|
||||
# another iteration.
|
||||
event.set()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_generator_setters():
|
||||
"""@serve.batch setters should succeed while the current batch streams."""
|
||||
|
||||
@serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
|
||||
async def yield_three_times(key1, key2):
|
||||
for _ in range(3):
|
||||
yield [(key1[i], key2[i]) for i in range(len(key1))]
|
||||
|
||||
assert yield_three_times._get_max_batch_size() == 2
|
||||
assert yield_three_times._get_batch_wait_timeout_s() == 1000
|
||||
|
||||
args_list = [("hi1", "hi2"), ("hi3", "hi4")]
|
||||
coros = [yield_three_times(*args) for args in args_list]
|
||||
|
||||
# Partially consume generators
|
||||
for coro, expected_result in zip(coros, args_list):
|
||||
for _ in range(2):
|
||||
await coro.__anext__() == expected_result
|
||||
|
||||
# Set new values
|
||||
yield_three_times.set_max_batch_size(3)
|
||||
yield_three_times.set_batch_wait_timeout_s(15000)
|
||||
|
||||
assert yield_three_times._get_max_batch_size() == 3
|
||||
assert yield_three_times._get_batch_wait_timeout_s() == 15000
|
||||
|
||||
# Execute three more requests
|
||||
args_list_2 = [("hi1", "hi2"), ("hi3", "hi4"), ("hi5", "hi6")]
|
||||
coros_2 = [yield_three_times(*args) for args in args_list_2]
|
||||
|
||||
# Finish consuming original requests
|
||||
for coro, expected_result in zip(coros, args_list):
|
||||
await coro.__anext__() == expected_result
|
||||
with pytest.raises(StopAsyncIteration):
|
||||
await coro.__anext__()
|
||||
|
||||
# Consume new requests
|
||||
for coro, expected_result in zip(coros_2, args_list_2):
|
||||
for _ in range(3):
|
||||
await coro.__anext__() == expected_result
|
||||
with pytest.raises(StopAsyncIteration):
|
||||
await coro.__anext__()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_size_fn_deferred_item_early_break():
|
||||
batches_processed = []
|
||||
|
||||
@serve.batch(
|
||||
max_batch_size=10,
|
||||
batch_wait_timeout_s=0.05,
|
||||
batch_size_fn=lambda items: sum(item["size"] for item in items),
|
||||
)
|
||||
async def batch_handler(requests):
|
||||
batches_processed.append([req["value"] for req in requests])
|
||||
return [req["value"] for req in requests]
|
||||
|
||||
# Request 1: size=6 (fits in batch)
|
||||
# Request 2: size=6 (would make total 12 > 10, should be deferred)
|
||||
# Each should be processed in its own batch
|
||||
t1 = get_or_create_event_loop().create_task(
|
||||
batch_handler({"size": 6, "value": "first"})
|
||||
)
|
||||
t2 = get_or_create_event_loop().create_task(
|
||||
batch_handler({"size": 6, "value": "second"})
|
||||
)
|
||||
|
||||
done, pending = await asyncio.wait([t1, t2], timeout=1.0)
|
||||
|
||||
assert len(done) == 2, "Both tasks should complete"
|
||||
assert len(pending) == 0
|
||||
|
||||
results = {t1.result(), t2.result()}
|
||||
assert results == {"first", "second"}
|
||||
|
||||
# Verify they were processed in separate batches due to size constraint
|
||||
assert (
|
||||
len(batches_processed) == 2
|
||||
), f"Expected 2 separate batches, got {batches_processed}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_size_fn_fail_to_fit():
|
||||
batches_processed = []
|
||||
|
||||
@serve.batch(
|
||||
max_batch_size=10,
|
||||
batch_wait_timeout_s=0.05,
|
||||
batch_size_fn=lambda items: sum(item["size"] for item in items),
|
||||
)
|
||||
async def batch_handler(requests):
|
||||
batches_processed.append([req["value"] for req in requests])
|
||||
return [req["value"] for req in requests]
|
||||
|
||||
# Request 1: size=6 (fits in batch)
|
||||
# Request 2: size=6 (would make total 12 > 10, should be deferred)
|
||||
# Each should be processed in its own batch
|
||||
t1 = get_or_create_event_loop().create_task(
|
||||
batch_handler({"size": 6, "value": "first"})
|
||||
)
|
||||
t2 = get_or_create_event_loop().create_task(
|
||||
batch_handler({"size": 12, "value": "second"})
|
||||
)
|
||||
|
||||
t1_result = await t1
|
||||
assert t1_result == "first"
|
||||
with pytest.raises(RuntimeError):
|
||||
await t2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_size_fn_multiple_items_fit():
|
||||
"""Test that multiple items are batched together when they fit within max_batch_size."""
|
||||
batches_seen = []
|
||||
|
||||
@serve.batch(
|
||||
max_batch_size=20,
|
||||
batch_wait_timeout_s=0.1,
|
||||
batch_size_fn=lambda items: sum(item["size"] for item in items),
|
||||
)
|
||||
async def batch_handler(requests):
|
||||
batch_values = [req["value"] for req in requests]
|
||||
batches_seen.append(batch_values)
|
||||
return batch_values
|
||||
|
||||
# All three requests fit: 5 + 6 + 7 = 18 <= 20
|
||||
t1 = get_or_create_event_loop().create_task(
|
||||
batch_handler({"size": 5, "value": "a"})
|
||||
)
|
||||
t2 = get_or_create_event_loop().create_task(
|
||||
batch_handler({"size": 6, "value": "b"})
|
||||
)
|
||||
t3 = get_or_create_event_loop().create_task(
|
||||
batch_handler({"size": 7, "value": "c"})
|
||||
)
|
||||
|
||||
done, pending = await asyncio.wait([t1, t2, t3], timeout=1.0)
|
||||
assert len(done) == 3
|
||||
assert len(pending) == 0
|
||||
|
||||
# All three should be in the same batch since they fit
|
||||
assert len(batches_seen) == 1, f"Expected 1 batch, got {len(batches_seen)}"
|
||||
assert set(batches_seen[0]) == {"a", "b", "c"}
|
||||
|
||||
|
||||
def test_warn_if_max_batch_size_exceeds_max_ongoing_requests():
|
||||
"""Test warn_if_max_batch_size_exceeds_max_ongoing_requests() logged the warning
|
||||
message correctly.
|
||||
|
||||
When the queue starts with or updated `max_batch_size` to be larger than
|
||||
max_ongoing_requests, log the warning to suggest configuring `max_ongoing_requests`.
|
||||
When the queue starts with or updated `max_batch_size` to be smaller or equal than
|
||||
max_ongoing_requests, no warning should be logged.
|
||||
"""
|
||||
logger = logging.getLogger(SERVE_LOGGER_NAME)
|
||||
stream = FakeStream()
|
||||
stream_handler = logging.StreamHandler(stream)
|
||||
logger.addHandler(stream_handler)
|
||||
bound = default_deployment_config.max_ongoing_requests
|
||||
over_bound = bound + 1
|
||||
under_bound = bound - 1
|
||||
over_bound_warning_message = (
|
||||
f"`max_batch_size` ({over_bound}) * `max_concurrent_batches` "
|
||||
f"({1}) is larger than `max_ongoing_requests` "
|
||||
f"({bound}). This means the replica will never achieve "
|
||||
"the configured `max_batch_size` concurrently. Please update "
|
||||
"`max_ongoing_requests` to be >= `max_batch_size` * `max_concurrent_batches`.\n"
|
||||
)
|
||||
|
||||
# Start queue above the bound will log warning. Start at under or at the bound will
|
||||
# not log warning
|
||||
for max_batch_size in [over_bound, under_bound, bound]:
|
||||
queue = _BatchQueue(
|
||||
max_batch_size=max_batch_size,
|
||||
batch_wait_timeout_s=1000,
|
||||
max_concurrent_batches=1,
|
||||
)
|
||||
if max_batch_size > bound:
|
||||
assert over_bound_warning_message in stream.messages
|
||||
else:
|
||||
assert over_bound_warning_message not in stream.messages
|
||||
stream.reset_message()
|
||||
|
||||
# Update queue above the bound will log warning. Update at under or at the bound
|
||||
# will not log warning
|
||||
for max_batch_size in [over_bound, under_bound, bound]:
|
||||
queue.set_max_batch_size(max_batch_size)
|
||||
if max_batch_size > bound:
|
||||
assert over_bound_warning_message in stream.messages
|
||||
else:
|
||||
assert over_bound_warning_message not in stream.messages
|
||||
stream.reset_message()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", "-s", __file__]))
|
||||
Reference in New Issue
Block a user