fix(batching): preserve peers when a batch request is cancelled

mlserver/batching/adaptive.py

@@ -95,7 +95,8 @@ def _batching_task_callback(self, batching_task: Task):
     def _clear_queue(self, err: BaseException):
         # Cancel all pending async responses
         for async_response in self._async_responses.values():
-            async_response.set_exception(err)
+            if not async_response.done():
+                async_response.set_exception(err)
 
         # Empty queue
         for _ in range(self._requests.qsize()):
@@ -116,10 +117,15 @@ def _predict_callback(self, batched: BatchedRequests, predict_task: Task):
             batched_response = predict_task.result()
             responses = batched.split_response(batched_response)
             for internal_id, response in responses.items():
-                self._async_responses[internal_id].set_result(response)
+                # A cancelled request may have already removed its response.
+                async_response = self._async_responses.get(internal_id)
+                if async_response is not None and not async_response.done():
+                    async_response.set_result(response)
         except Exception as err:
             for internal_id in batched.inference_requests.keys():
-                self._async_responses[internal_id].set_exception(err)
+                async_response = self._async_responses.get(internal_id)
+                if async_response is not None and not async_response.done():
+                    async_response.set_exception(err)
 
     async def _batch_requests(self) -> AsyncIterator[BatchedRequests]:
         while not self._requests.empty():

tests/batching/test_adaptive.py

@@ -12,6 +12,25 @@
 from .conftest import TestRequestSender
 
 
+async def test_clear_queue_with_cancelled_response(
+    adaptive_batcher: AdaptiveBatcher, send_request: TestRequestSender
+):
+    cancelled_id, _ = await send_request()
+    pending_id, _ = await send_request()
+    cancelled_response = adaptive_batcher._async_responses[cancelled_id]
+    pending_response = adaptive_batcher._async_responses[pending_id]
+    cancelled_response.cancel()
+    error = RuntimeError("batching failed")
+
+    adaptive_batcher._clear_queue(error)
+
+    assert cancelled_response.cancelled()
+    assert adaptive_batcher._requests.empty()
+    with pytest.raises(RuntimeError) as raised:
+        await pending_response
+    assert raised.value is error
+
+
 async def test_batch_requests(
     adaptive_batcher: AdaptiveBatcher,
     send_request: TestRequestSender,

tests/batching/test_hooks.py

@@ -1,9 +1,92 @@
+import asyncio
 import pytest
 
 from mlserver.model import MLModel
 from mlserver.types import InferenceRequest, InferenceResponse
 from mlserver.batching.hooks import load_batching
 
+from ..fixtures import SumModel
+
+
+@pytest.mark.parametrize("cancelled_index", [None, 0, 1])
+@pytest.mark.parametrize("fail_prediction", [False, True])
+async def test_batching_cancelled_request(
+    sum_model_settings, cancelled_index, fail_prediction
+):
+    started = asyncio.Event()
+    release = asyncio.Event()
+    model_error = RuntimeError("prediction failed")
+    batch_shapes = []
+
+    class GatedSumModel(SumModel):
+        async def predict(self, payload):
+            batch_shapes.append(payload.inputs[0].shape)
+            started.set()
+            await release.wait()
+            if fail_prediction:
+                raise model_error
+            return await super().predict(payload)
+
+    sum_model_settings.max_batch_size = 2
+    sum_model_settings.max_batch_time = 0.01
+    model = await load_batching(GatedSumModel(sum_model_settings))
+    requests = [
+        InferenceRequest(
+            id=f"request-{idx}",
+            inputs=[
+                {
+                    "name": "input-0",
+                    "shape": [1, 3],
+                    "datatype": "INT32",
+                    "data": [idx, idx + 1, idx + 2],
+                }
+            ],
+        )
+        for idx in range(2)
+    ]
+    tasks = [asyncio.create_task(model.predict(request)) for request in requests]
+    loop = asyncio.get_running_loop()
+    callback_errors = []
+    previous_handler = loop.get_exception_handler()
+    loop.set_exception_handler(lambda loop, context: callback_errors.append(context))
+    try:
+        await asyncio.wait_for(started.wait(), timeout=1)
+        assert batch_shapes == [[2, 3]]
+        if cancelled_index is not None:
+            tasks[cancelled_index].cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await tasks[cancelled_index]
+
+        release.set()
+        active_indices = [idx for idx in range(2) if idx != cancelled_index]
+        responses = await asyncio.wait_for(
+            asyncio.gather(
+                *(tasks[idx] for idx in active_indices), return_exceptions=True
+            ),
+            timeout=1,
+        )
+        for idx, response in zip(active_indices, responses):
+            if fail_prediction:
+                assert response is model_error
+            else:
+                assert response.id == requests[idx].id
+                assert response.outputs[0].shape == [1, 1]
+                assert response.outputs[0].data.root == [3 * idx + 3]
+        assert not callback_errors
+
+        # The same batcher must remain usable after cancellation or a model error.
+        fail_prediction = False
+        following = await asyncio.wait_for(model.predict(requests[0]), timeout=1)
+        assert following.id == requests[0].id
+        assert following.outputs[0].data.root == [3]
+        assert not callback_errors
+    finally:
+        release.set()
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+        loop.set_exception_handler(previous_handler)
+
 
 async def test_batching_predict(
     sum_model: MLModel, inference_request: InferenceRequest