diff --git a/src/litserve/api.py b/src/litserve/api.py index 5d6371d9f..93a86eaf6 100644 --- a/src/litserve/api.py +++ b/src/litserve/api.py @@ -251,6 +251,9 @@ def decode_request(self, request, **kwargs): Add a context argument, decode_request(self, request, context), to also get the request context. The HTTP request headers are available as context["headers"], with lower-cased names. The argument must be named context. + With non-streaming batching, a decoding error fails only this request. The remaining decoded inputs are passed + to batch and predict. If every request fails to decode, prediction is skipped. + """ if self._spec: return self._spec.decode_request(request, **kwargs) @@ -356,6 +359,9 @@ def encode_response(self, output, **kwargs): To enable streaming, it should yield the output. + With non-streaming batching, an encoding error fails only this request. Other responses in the batch are still + returned. Errors in batch, predict, or unbatch affect all requests that decoded successfully. + """ if self._spec: return self._spec.encode_response(output, **kwargs) diff --git a/src/litserve/loops/simple_loops.py b/src/litserve/loops/simple_loops.py index c8ccb26be..b8c3fefdc 100644 --- a/src/litserve/loops/simple_loops.py +++ b/src/litserve/loops/simple_loops.py @@ -348,25 +348,30 @@ def run_batched_loop( if not batches: continue logger.debug(f"{len(batches)} batched requests received") - response_queue_ids, uids, inputs, headers_list = zip(*batches) - num_inputs = len(inputs) + pending_requests = {uid: response_queue_id for response_queue_id, uid, _, _ in batches} try: - contexts = [{"headers": headers} for headers in headers_list] - if hasattr(lit_spec, "populate_context"): - for input, context in zip(inputs, contexts): - lit_spec.populate_context(context, input) - callback_runner.trigger_event(EventTypes.BEFORE_DECODE_REQUEST.value, lit_api=lit_api) - x = [ - _inject_context( - context, - lit_api.decode_request, - input, - ) - for input, context in zip(inputs, contexts) - ] + x, contexts, decoded_requests = [], [], [] + for response_queue_id, uid, input, headers in batches: + try: + context = {"headers": headers} + if hasattr(lit_spec, "populate_context"): + lit_spec.populate_context(context, input) + decoded = _inject_context(context, lit_api.decode_request, input) + except Exception as e: + if not isinstance(e, HTTPException): + logger.exception("LitAPI failed to decode request uid=%s.", uid) + self.put_error_response(transport, response_queue_id, uid, e) + del pending_requests[uid] + continue + x.append(decoded) + contexts.append(context) + decoded_requests.append((response_queue_id, uid)) callback_runner.trigger_event(EventTypes.AFTER_DECODE_REQUEST.value, lit_api=lit_api) + if not x: + continue + num_inputs = len(x) x = _inject_context(contexts, lit_api.batch, x) callback_runner.trigger_event(EventTypes.BEFORE_PREDICT.value, lit_api=lit_api) @@ -387,8 +392,15 @@ def run_batched_loop( callback_runner.trigger_event(EventTypes.BEFORE_ENCODE_RESPONSE.value, lit_api=lit_api) y_enc_list = [] - for response_queue_id, y, uid, context in zip(response_queue_ids, outputs, uids, contexts): - y_enc = _inject_context(context, lit_api.encode_response, y) + for (response_queue_id, uid), y, context in zip(decoded_requests, outputs, contexts): + try: + y_enc = _inject_context(context, lit_api.encode_response, y) + except Exception as e: + if not isinstance(e, HTTPException): + logger.exception("LitAPI failed to encode response uid=%s.", uid) + self.put_error_response(transport, response_queue_id, uid, e) + del pending_requests[uid] + continue y_enc_list.append((response_queue_id, uid, y_enc)) callback_runner.trigger_event(EventTypes.AFTER_ENCODE_RESPONSE.value, lit_api=lit_api) @@ -396,9 +408,10 @@ def run_batched_loop( self.put_response( transport, response_queue_id, uid, y_enc, LitAPIStatus.OK, LoopResponseType.REGULAR ) + del pending_requests[uid] except HTTPException as e: - for response_queue_id, uid in zip(response_queue_ids, uids): + for uid, response_queue_id in pending_requests.items(): self.put_response( transport, response_queue_id, @@ -415,7 +428,7 @@ def run_batched_loop( "LitAPI ran into an error while processing the batched request.\n" "Please check the error trace for more details." ) - for response_queue_id, uid in zip(response_queue_ids, uids): + for uid, response_queue_id in pending_requests.items(): self.put_error_response(transport, response_queue_id, uid, e, LoopResponseType.REGULAR) def __call__( diff --git a/tests/integration/test_batch_errors.py b/tests/integration/test_batch_errors.py new file mode 100644 index 000000000..e7bacdbd5 --- /dev/null +++ b/tests/integration/test_batch_errors.py @@ -0,0 +1,73 @@ +# Copyright The Lightning AI team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import asyncio + +import pytest +from asgi_lifespan import LifespanManager +from httpx import ASGITransport, AsyncClient + +from litserve import LitAPI, LitServer +from litserve.utils import PickleableHTTPException, wrap_litserve_start + + +class BatchErrorAPI(LitAPI): + def __init__(self, error_stage): + super().__init__(max_batch_size=3, batch_timeout=1) + self.error_stage = error_stage + + def setup(self, device): + pass + + def decode_request(self, request): + value = request["input"] + if self.error_stage == "decode" and value < 0: + raise PickleableHTTPException(422, "invalid input") + return value + + def batch(self, inputs): + return inputs + + def predict(self, inputs): + return [{"output": value, "batch_size": len(inputs)} for value in inputs] + + def unbatch(self, output): + return output + + def encode_response(self, output, context): + if self.error_stage == "encode" and output["output"] < 0: + raise PickleableHTTPException(422, "invalid input") + return {**output, "header": context["headers"]["x-request-id"]} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error_stage", ["decode", "encode"]) +async def test_batch_request_error_isolation(error_stage): + server = LitServer(BatchErrorAPI(error_stage), accelerator="cpu", devices=1, timeout=10) + with wrap_litserve_start(server) as server: + async with ( + LifespanManager(server.app) as manager, + AsyncClient(transport=ASGITransport(app=manager.app), base_url="http://test") as client, + ): + responses = await asyncio.gather( + *[ + client.post("/predict", json={"input": value}, headers={"x-request-id": str(value)}) + for value in [3, -1, 5] + ] + ) + + assert [response.status_code for response in responses] == [200, 422, 200] + assert responses[1].json() == {"detail": "invalid input"} + batch_size = 2 if error_stage == "decode" else 3 + for response, value in [(responses[0], 3), (responses[2], 5)]: + assert response.json() == {"output": value, "batch_size": batch_size, "header": str(value)} diff --git a/tests/unit/test_batch.py b/tests/unit/test_batch.py index 34dbd31b1..733a87cc5 100644 --- a/tests/unit/test_batch.py +++ b/tests/unit/test_batch.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import asyncio +import pickle import time from queue import Queue from unittest.mock import ANY, MagicMock @@ -20,7 +21,7 @@ import torch import torch.nn as nn from asgi_lifespan import LifespanManager -from fastapi import Request, Response +from fastapi import HTTPException, Request, Response from httpx import ASGITransport, AsyncClient import litserve as ls @@ -30,7 +31,7 @@ from litserve.loops.simple_loops import BatchedLoop from litserve.test_examples import SimpleBatchedAPI from litserve.transport.base import MessageTransport -from litserve.utils import LoopResponseType, wrap_litserve_start +from litserve.utils import LitAPIStatus, LoopResponseType, PickleableHTTPException, wrap_litserve_start NOOP_CB_RUNNER = CallbackRunner() @@ -241,6 +242,136 @@ def test_batched_loop(): lit_api_mock.unbatch.assert_called_once() +class BatchErrorAPI(LitAPI): + def __init__(self, error_stage, error, shared_error_stage=None, **kwargs): + super().__init__(**kwargs) + self.error_stage = error_stage + self.error = error + self.shared_error_stage = shared_error_stage + self.calls = [] + self.decode_inputs = [] + self.batch_inputs = None + self.batch_headers = None + + def setup(self, device): + pass + + def decode_request(self, request, context): + self.decode_inputs.append(request) + if self.error_stage == "decode" and request < 0: + raise self.error + return request + + def batch(self, inputs, context): + self.calls.append("batch") + self.batch_inputs = inputs + self.batch_headers = [item["headers"]["x-id"] for item in context] + if self.shared_error_stage == "batch": + raise RuntimeError("batch failed") + return inputs + + def predict(self, inputs, context): + self.calls.append("predict") + assert [item["headers"]["x-id"] for item in context] == self.batch_headers + if self.shared_error_stage == "predict": + raise RuntimeError("predict failed") + return inputs + + def unbatch(self, outputs, context): + self.calls.append("unbatch") + assert [item["headers"]["x-id"] for item in context] == self.batch_headers + if self.shared_error_stage == "unbatch": + raise RuntimeError("unbatch failed") + return outputs + + def encode_response(self, output, context): + if self.error_stage == "encode" and output < 0: + raise self.error + return {"output": output, "header": context["headers"]["x-id"]} + + +def run_error_batch(api, inputs): + api.request_timeout = 30 + request_queue = Queue() + for i, value in enumerate(inputs): + request_queue.put((0, str(i), time.monotonic(), value, {"x-id": str(value)})) + request_queue.put(_SENTINEL_VALUE) + transport = FakeTransport() + BatchedLoop().run_batched_loop(api, request_queue, transport, NOOP_CB_RUNNER) + assert len(transport.responses) == len(inputs) + responses = dict(transport.responses) + assert set(responses) == {str(i) for i in range(len(inputs))} + return responses + + +@pytest.mark.parametrize("error_stage", ["decode", "encode"]) +@pytest.mark.parametrize("error_type", [HTTPException, PickleableHTTPException, ValueError]) +def test_batched_loop_isolates_request_errors(error_stage, error_type): + error = error_type(422, "invalid input") if issubclass(error_type, HTTPException) else error_type("invalid input") + api = BatchErrorAPI(error_stage, error, max_batch_size=3) + + responses = run_error_batch(api, [3, -1, 5]) + + for uid, value in [("0", 3), ("2", 5)]: + assert responses[uid] == ( + {"output": value, "header": str(value)}, + LitAPIStatus.OK, + LoopResponseType.REGULAR, + ANY, + ) + data, status, response_type, _ = responses["1"] + assert status == LitAPIStatus.ERROR + assert response_type == LoopResponseType.REGULAR + if issubclass(error_type, HTTPException): + exception = pickle.loads(data) if isinstance(data, bytes) else data + assert isinstance(exception, HTTPException) + assert exception.status_code == 422 + assert exception.detail == "invalid input" + else: + exception = pickle.loads(data) + assert isinstance(exception, ValueError) + assert str(exception) == "invalid input" + expected_inputs = [3, 5] if error_stage == "decode" else [3, -1, 5] + assert api.batch_inputs == expected_inputs + assert api.batch_headers == [str(value) for value in expected_inputs] + assert api.calls == ["batch", "predict", "unbatch"] + + +def test_batched_loop_skips_batch_when_all_decodes_fail(): + api = BatchErrorAPI("decode", HTTPException(422, "invalid input"), max_batch_size=2) + + responses = run_error_batch(api, [-1, -2]) + + assert api.calls == [] + assert api.decode_inputs == [-1, -2] + for data, status, _, _ in responses.values(): + assert status == LitAPIStatus.ERROR + exception = pickle.loads(data) if isinstance(data, bytes) else data + assert exception.status_code == 422 + + +@pytest.mark.parametrize("shared_error_stage", ["batch", "predict", "unbatch"]) +def test_batched_loop_shared_failure_only_affects_decoded_requests(shared_error_stage): + api = BatchErrorAPI( + "decode", HTTPException(422, "invalid input"), shared_error_stage=shared_error_stage, max_batch_size=3 + ) + + responses = run_error_batch(api, [3, -1, 5]) + + data = responses["1"][0] + exception = pickle.loads(data) if isinstance(data, bytes) else data + assert exception.status_code == 422 + for uid in ["0", "2"]: + data, status, _, _ = responses[uid] + assert status == LitAPIStatus.ERROR + assert isinstance(data, bytes) + exception = pickle.loads(data) + assert isinstance(exception, RuntimeError) + assert str(exception) == f"{shared_error_stage} failed" + assert api.batch_inputs == [3, 5] + assert api.batch_headers == ["3", "5"] + + @pytest.mark.parametrize( ("batch_timeout", "batch_size"), [