File size: 6,278 Bytes
299567b 3c0d20e 299567b 3c0d20e 299567b 5e96f99 299567b 5e96f99 299567b 5e96f99 299567b dfdc75f 299567b 14b0898 299567b 3c0d20e 14b0898 3c0d20e 5e96f99 3c0d20e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | """Unit tests for DynamicBatcher (app/batching.py)."""
import asyncio
import contextlib
import threading
import pytest
from app.batching import DynamicBatcher, QueueFullError, RequestTimeoutError
async def test_size_trigger_flush(make_model):
"""A full batch flushes on the size trigger without waiting out max_wait_ms."""
model = make_model()
batcher = DynamicBatcher(model, max_batch_size=2, max_wait_ms=10_000)
batcher.start()
try:
results = await asyncio.wait_for(
asyncio.gather(batcher.submit(["a"]), batcher.submit(["b"])), timeout=5
)
finally:
await batcher.stop()
assert results[0].embeddings == make_model().encode(["a"]).tolist()
assert results[1].embeddings == make_model().encode(["b"]).tolist()
for result in results:
assert result.ttfr_ms >= result.ttft_ms >= 0
async def test_time_trigger_flush(make_model):
"""A partial batch flushes once max_wait_ms elapses."""
batcher = DynamicBatcher(make_model(), max_batch_size=10, max_wait_ms=50)
batcher.start()
try:
result = await asyncio.wait_for(batcher.submit(["solo"]), timeout=5)
finally:
await batcher.stop()
assert result.embeddings == make_model().encode(["solo"]).tolist()
assert result.ttfr_ms >= result.ttft_ms >= 0
# time trigger waits ~max_wait_ms before encode starts
assert result.ttft_ms >= 20
async def test_request_response_mapping(make_model):
"""Mixed-size requests batched together each get back exactly their own rows."""
batcher = DynamicBatcher(make_model(), max_batch_size=8, max_wait_ms=200)
batcher.start()
requests = [["a"], ["b", "c"], ["d", "e", "f"]]
try:
results = await asyncio.wait_for(
asyncio.gather(*(batcher.submit(req) for req in requests)), timeout=5
)
finally:
await batcher.stop()
model = make_model()
for req, res in zip(requests, results):
assert res.embeddings == model.encode(req).tolist()
assert res.ttfr_ms >= res.ttft_ms >= 0
async def test_encode_uses_max_batch_size(make_model):
"""The worker encodes the aggregated batch in one pass sized by max_batch_size,
not model.encode's hidden default of 32."""
model = make_model()
batcher = DynamicBatcher(model, max_batch_size=7, max_wait_ms=50)
batcher.start()
try:
await asyncio.wait_for(batcher.submit(["a"]), timeout=5)
finally:
await batcher.stop()
assert model.last_batch_size == 7
async def test_exception_propagation(make_model):
"""An inference error surfaces to the awaiting submit()."""
batcher = DynamicBatcher(
make_model(error=RuntimeError("model failed")), max_batch_size=1, max_wait_ms=50
)
batcher.start()
try:
with pytest.raises(RuntimeError, match="model failed"):
await asyncio.wait_for(batcher.submit(["x"]), timeout=5)
finally:
await batcher.stop()
async def test_start_stop_lifecycle(make_model):
"""start() launches the worker; stop() ends it cleanly."""
batcher = DynamicBatcher(make_model(), max_batch_size=4, max_wait_ms=50)
batcher.start()
assert batcher._worker_task is not None
await batcher.stop()
assert batcher._worker_task.done()
async def test_shutdown_drains_queued_requests(make_model):
"""Requests still queued at shutdown are failed with CancelledError."""
gate = threading.Event()
# batch_size=1 so the worker takes request A and immediately enters encode(), which blocks.
batcher = DynamicBatcher(make_model(gate=gate), max_batch_size=1, max_wait_ms=50)
batcher.start()
a = asyncio.create_task(batcher.submit(["A"]))
try:
await asyncio.sleep(0.1) # let the worker pull A and block inside the executor
b = asyncio.create_task(batcher.submit(["B"]))
await asyncio.sleep(0.05) # B is now sitting in the queue
await batcher.stop()
with pytest.raises(asyncio.CancelledError):
await b
finally:
gate.set() # unblock the executor thread so it can exit
a.cancel()
with pytest.raises(asyncio.CancelledError):
await a
async def test_queue_full_rejects(make_model):
"""A submit beyond the queue's capacity is rejected immediately with QueueFullError."""
gate = threading.Event()
# batch_size=1: the worker pulls request A and blocks inside encode(); max_queue_size=1
# leaves room for exactly one queued request (B) before the next submit is rejected.
batcher = DynamicBatcher(
make_model(gate=gate), max_batch_size=1, max_wait_ms=50, max_queue_size=1
)
batcher.start()
a = asyncio.create_task(batcher.submit(["A"]))
await asyncio.sleep(0.1) # worker pulls A and blocks in the executor; queue is now empty
b = asyncio.create_task(batcher.submit(["B"]))
await asyncio.sleep(0.05) # B is now sitting in the queue (at capacity)
try:
with pytest.raises(QueueFullError):
await batcher.submit(["C"]) # queue at capacity -> rejected synchronously
finally:
gate.set() # unblock the worker thread
await batcher.stop()
for task in (a, b):
task.cancel()
with contextlib.suppress(asyncio.CancelledError, Exception):
await task
async def test_request_timeout(make_model):
"""A request not served within request_timeout_s raises RequestTimeoutError, and the
worker survives (the cancelled future is skipped) to serve later requests."""
gate = threading.Event()
batcher = DynamicBatcher(
make_model(gate=gate), max_batch_size=1, max_wait_ms=50, request_timeout_s=0.1
)
batcher.start()
try:
with pytest.raises(RequestTimeoutError):
await batcher.submit(["slow"])
assert batcher._worker_task is not None and not batcher._worker_task.done()
gate.set() # let the stuck inference finish; the worker skips the cancelled future
result = await asyncio.wait_for(batcher.submit(["after"]), timeout=5)
assert result.embeddings == make_model().encode(["after"]).tolist()
assert result.ttfr_ms >= result.ttft_ms >= 0
finally:
gate.set()
await batcher.stop()
|