test: test batch queue flush all

Author: Joan Martinez

tests/unit/serve/dynamic_batching/test_batch_queue.py

@@ -9,7 +9,8 @@
 
 
 @pytest.mark.asyncio
-async def test_batch_queue_timeout():
+@pytest.mark.parametrize('flush_all', [False, True])
+async def test_batch_queue_timeout(flush_all):
     async def foo(docs, **kwargs):
         await asyncio.sleep(0.1)
         return DocumentArray([Document(text='Done') for _ in docs])
@@ -20,6 +21,7 @@ async def foo(docs, **kwargs):
         response_docarray_cls=DocumentArray,
         preferred_batch_size=4,
         timeout=2000,
+        flush_all=flush_all,
     )
 
     three_data_requests = [DataRequest() for _ in range(3)]
@@ -59,7 +61,8 @@ async def process_request(req):
 
 
 @pytest.mark.asyncio
-async def test_batch_queue_timeout_does_not_wait_previous_batch():
+@pytest.mark.parametrize('flush_all', [False, True])
+async def test_batch_queue_timeout_does_not_wait_previous_batch(flush_all):
     batches_lengths_computed = []
 
     async def foo(docs, **kwargs):
@@ -73,6 +76,7 @@ async def foo(docs, **kwargs):
         response_docarray_cls=DocumentArray,
         preferred_batch_size=5,
         timeout=3000,
+        flush_all=flush_all
     )
 
     data_requests = [DataRequest() for _ in range(3)]
@@ -105,7 +109,8 @@ async def process_request(req, sleep=0):
 
 
 @pytest.mark.asyncio
-async def test_batch_queue_req_length_larger_than_preferred():
+@pytest.mark.parametrize('flush_all', [False, True])
+async def test_batch_queue_req_length_larger_than_preferred(flush_all):
     async def foo(docs, **kwargs):
         await asyncio.sleep(0.1)
         return DocumentArray([Document(text='Done') for _ in docs])
@@ -116,6 +121,7 @@ async def foo(docs, **kwargs):
         response_docarray_cls=DocumentArray,
         preferred_batch_size=4,
         timeout=2000,
+        flush_all=flush_all,
     )
 
     data_requests = [DataRequest() for _ in range(3)]
@@ -142,7 +148,8 @@ async def process_request(req):
 
 
 @pytest.mark.asyncio
-async def test_exception():
+@pytest.mark.parametrize('flush_all', [False, True])
+async def test_exception(flush_all):
     BAD_REQUEST_IDX = [2, 6]
 
     async def foo(docs, **kwargs):
@@ -159,6 +166,7 @@ async def foo(docs, **kwargs):
         response_docarray_cls=DocumentArray,
         preferred_batch_size=1,
         timeout=500,
+        flush_all=flush_all,
     )
 
     data_requests = [DataRequest() for _ in range(35)]
@@ -188,7 +196,8 @@ async def process_request(req):
 
 
 @pytest.mark.asyncio
-async def test_exception_more_complex():
+@pytest.mark.parametrize('flush_all', [False, True])
+async def test_exception_more_complex(flush_all):
     TRIGGER_BAD_REQUEST_IDX = [2, 6]
     EXPECTED_BAD_REQUESTS = [2, 3, 6, 7]
 
@@ -208,6 +217,7 @@ async def foo(docs, **kwargs):
         request_docarray_cls=DocumentArray,
         response_docarray_cls=DocumentArray,
         preferred_batch_size=2,
+        flush_all=flush_all,
         timeout=500,
     )
 
@@ -240,7 +250,8 @@ async def process_request(req):
 
 
 @pytest.mark.asyncio
-async def test_exception_all():
+@pytest.mark.parametrize('flush_all', [False, True])
+async def test_exception_all(flush_all):
     async def foo(docs, **kwargs):
         raise AssertionError
 
@@ -249,6 +260,7 @@ async def foo(docs, **kwargs):
         request_docarray_cls=DocumentArray,
         response_docarray_cls=DocumentArray,
         preferred_batch_size=2,
+        flush_all=flush_all,
         timeout=500,
     )
 
@@ -287,8 +299,9 @@ async def foo(docs, **kwargs):
 @pytest.mark.parametrize('num_requests', [61, 127, 100])
 @pytest.mark.parametrize('preferred_batch_size', [7, 27, 61, 73, 100])
 @pytest.mark.parametrize('timeout', [0.3, 500])
+@pytest.mark.parametrize('flush_all', [False, True])
 @pytest.mark.asyncio
-async def test_return_proper_assignment(num_requests, preferred_batch_size, timeout):
+async def test_return_proper_assignment(num_requests, preferred_batch_size, timeout, flush_all):
     import random
 
     async def foo(docs, **kwargs):
@@ -301,6 +314,7 @@ async def foo(docs, **kwargs):
         request_docarray_cls=DocumentArray,
         response_docarray_cls=DocumentArray,
         preferred_batch_size=preferred_batch_size,
+        flush_all=flush_all,
         timeout=timeout,
     )
 
@@ -331,3 +345,58 @@ async def process_request(req):
         assert len(resp.docs) == length
         for j, d in enumerate(resp.docs):
             assert d.text == f'Text {j} from request {i} with len {length} Processed'
+
+
+@pytest.mark.parametrize('num_requests', [61, 127, 100])
+@pytest.mark.parametrize('preferred_batch_size', [7, 27, 61, 73, 100])
+@pytest.mark.parametrize('timeout', [0.3, 500])
+@pytest.mark.parametrize('flush_all', [False, True])
+@pytest.mark.asyncio
+async def test_length_processed_in_func(num_requests, preferred_batch_size, timeout, flush_all):
+    import random
+
+    async def foo(docs, **kwargs):
+        if not flush_all:
+            assert len(docs) <= preferred_batch_size
+        else:
+            assert len(docs) >= preferred_batch_size
+        await asyncio.sleep(0.1)
+        for doc in docs:
+            doc.text += ' Processed'
+
+    bq: BatchQueue = BatchQueue(
+        foo,
+        request_docarray_cls=DocumentArray,
+        response_docarray_cls=DocumentArray,
+        preferred_batch_size=preferred_batch_size,
+        flush_all=flush_all,
+        timeout=timeout,
+    )
+
+    data_requests = [DataRequest() for _ in range(num_requests)]
+    len_requests = []
+    for i, req in enumerate(data_requests):
+        len_request = random.randint(preferred_batch_size, preferred_batch_size * 10)
+        len_requests.append(len_request)
+        req.data.docs = DocumentArray(
+            [
+                Document(text=f'Text {j} from request {i} with len {len_request}')
+                for j in range(len_request)
+            ]
+        )
+
+    async def process_request(req):
+        q = await bq.push(req)
+        item = await q.get()
+        q.task_done()
+        return item
+
+    tasks = [asyncio.create_task(process_request(req)) for req in data_requests]
+    items = await asyncio.gather(*tasks)
+    for i, item in enumerate(items):
+        assert item is None
+
+    for i, (resp, length) in enumerate(zip(data_requests, len_requests)):
+        assert len(resp.docs) == length
+        for j, d in enumerate(resp.docs):
+            assert d.text == f'Text {j} from request {i} with len {length} Processed'