From 21056a2b098acf21fb863331eac9e97eb2991497 Mon Sep 17 00:00:00 2001 From: DiadkoShmek <207617272+DiadkoShmek@users.noreply.github.com> Date: Sun, 1 Nov 2026 00:08:51 +0200 Subject: [PATCH] fix(batching): preserve peers when a batch request is cancelled --- mlserver/batching/adaptive.py | 12 +++-- tests/batching/test_adaptive.py | 19 ++++++++ tests/batching/test_hooks.py | 83 +++++++++++++++++++++++++++++++++ 3 files changed, 111 insertions(+), 3 deletions(-) diff --git a/mlserver/batching/adaptive.py b/mlserver/batching/adaptive.py index be82fe03d..e00eb864c 100644 --- a/mlserver/batching/adaptive.py +++ b/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(): diff --git a/tests/batching/test_adaptive.py b/tests/batching/test_adaptive.py index 011981416..5a48c5221 100644 --- a/tests/batching/test_adaptive.py +++ b/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, diff --git a/tests/batching/test_hooks.py b/tests/batching/test_hooks.py index 4e6a43203..ee0d57a4e 100644 --- a/tests/batching/test_hooks.py +++ b/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