fix(loops): isolate request errors in batches

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)

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,18 +392,26 @@ 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)
 
                 for response_queue_id, uid, y_enc in y_enc_list:
                     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__(

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)}

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"),
     [