Support context argument in batch and unbatch methods (#711)

* Support context argument in batch and unbatch methods

The batched loops called lit_api.batch and lit_api.unbatch directly, so user implementations declaring a context parameter failed with a TypeError. Route both calls through _inject_context, the same way decode_request, predict and encode_response are already handled, in both the regular and streaming batched loops. Document the optional context argument in the LitAPI.batch/unbatch docstrings and add regression tests covering both loops.

Fixes #617

* test: exercise a real batch in the batch/unbatch context test

The test set max_batch_size=2 but sent a single request, so every batch
had size 1 and the ordering between contexts and requests was never
checked. Send two concurrent requests with different inputs instead and
assert each one gets its own context value back.

Each context also records the batch size, so the test fails loudly if the
two requests ever stop sharing a batch rather than silently degrading.

* docs: show how to use context in batch/unbatch docstrings

Say which signature to write and note that the argument has to be named
context, since _inject_context matches on the parameter name and `**kwargs`
never receives it. Drops the RST markup for the plain prose and indented
examples used elsewhere in the file.

---------

Co-authored-by: Bhimraj Yadav <bhimrajyadav977@gmail.com>

Author: discobot

src/litserve/api.py

@@ -252,7 +252,12 @@ def decode_request(self, request, **kwargs):
         return request
 
     def batch(self, inputs):
-        """Convert a list of inputs to a batched input."""
+        """Convert a list of inputs to a batched input.
+
+        Add a context argument, batch(self, inputs, context), to also get the request context: one dict per input, in
+        the same order. The argument must be named context.
+
+        """
         # consider assigning an implementation when starting server
         # to avoid the runtime cost of checking (should be negligible)
         if hasattr(inputs[0], "__torch_function__"):
@@ -325,6 +330,14 @@ def unbatch(self, output):
         If you need to return dictionaries, return a list of dicts:
             [{"key1": val1, "key2": val3}, {"key1": val2, "key2": val4}]  # Correct
 
+        Add a context argument to also get the request context, one dict per request in the same
+        order as the outputs:
+
+            def unbatch(self, output, context):
+                return [ctx["input"] for ctx in context]
+
+        The argument must be named context.
+
         """
         if self._default_unbatch is None:
             raise ValueError(

src/litserve/loops/simple_loops.py

@@ -366,13 +366,13 @@ def run_batched_loop(
                 ]
                 callback_runner.trigger_event(EventTypes.AFTER_DECODE_REQUEST.value, lit_api=lit_api)
 
-                x = lit_api.batch(x)
+                x = _inject_context(contexts, lit_api.batch, x)
 
                 callback_runner.trigger_event(EventTypes.BEFORE_PREDICT.value, lit_api=lit_api)
                 y = _inject_context(contexts, lit_api.predict, x)
                 callback_runner.trigger_event(EventTypes.AFTER_PREDICT.value, lit_api=lit_api)
 
-                outputs = lit_api.unbatch(y)
+                outputs = _inject_context(contexts, lit_api.unbatch, y)
 
                 if len(outputs) != num_inputs:
                     actual = len(outputs)

src/litserve/loops/streaming_loops.py

@@ -347,13 +347,13 @@ def run_batched_streaming_loop(
                 ]
                 callback_runner.trigger_event(EventTypes.AFTER_DECODE_REQUEST.value, lit_api=lit_api)
 
-                x = lit_api.batch(x)
+                x = _inject_context(contexts, lit_api.batch, x)
 
                 callback_runner.trigger_event(EventTypes.BEFORE_PREDICT.value, lit_api=lit_api)
                 y_iter = _inject_context(contexts, lit_api.predict, x)
                 callback_runner.trigger_event(EventTypes.AFTER_PREDICT.value, lit_api=lit_api)
 
-                unbatched_iter = lit_api.unbatch(y_iter)
+                unbatched_iter = _inject_context(contexts, lit_api.unbatch, y_iter)
 
                 callback_runner.trigger_event(EventTypes.BEFORE_ENCODE_RESPONSE.value, lit_api=lit_api)
                 y_enc_iter = _inject_context(contexts, lit_api.encode_response, unbatched_iter)

tests/unit/test_lit_server.py

@@ -395,6 +395,39 @@ def encode_response(self, output_stream, context):
             yield [{"output": ctx["input"]} for ctx in context]
 
 
+class BatchUnbatchContextAPI(ls.test_examples.SimpleBatchedAPI):
+    def batch(self, inputs, context):
+        for c, x in zip(context, inputs):
+            c["input"] = float(x)
+            c["batch_size"] = len(inputs)
+        return super().batch(inputs)
+
+    def unbatch(self, output, context):
+        return [{"input": c["input"], "batch_size": c["batch_size"]} for c in context]
+
+    def encode_response(self, output):
+        return {"output": output["input"], "batch_size": output["batch_size"]}
+
+
+class BatchUnbatchContextStreamingAPI(ls.test_examples.SimpleBatchedAPI):
+    def batch(self, inputs, context):
+        for c, x in zip(context, inputs):
+            c["input"] = float(x)
+            c["batch_size"] = len(inputs)
+        return super().batch(inputs)
+
+    def predict(self, x_batch):
+        yield self.model(x_batch)
+
+    def unbatch(self, output_stream, context):
+        for _ in output_stream:
+            yield [{"input": c["input"], "batch_size": c["batch_size"]} for c in context]
+
+    def encode_response(self, output_stream):
+        for outputs in output_stream:
+            yield [{"output": output["input"], "batch_size": output["batch_size"]} for output in outputs]
+
+
 class PredictErrorAPI(ls.test_examples.SimpleLitAPI):
     def predict(self, x, y, context):
         context["input"] = x
@@ -444,6 +477,36 @@ async def test_inject_context():
         assert resp.status_code == 500, "predict() missed 1 required positional argument: 'y'"
 
 
+@pytest.mark.asyncio
+async def test_inject_context_in_batch_and_unbatch():
+    # Two requests at once, so each one must get its own context back. batched loop:
+    server = LitServer(BatchUnbatchContextAPI(max_batch_size=2, batch_timeout=4), 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 ac,
+        ):
+            resp1 = ac.post("/predict", json={"input": 5.0})
+            resp2 = ac.post("/predict", json={"input": 7.0})
+            resp1, resp2 = await asyncio.gather(resp1, resp2)
+    # batch_size == 2 proves they really shared a batch
+    assert resp1.json() == {"output": 5.0, "batch_size": 2}
+    assert resp2.json() == {"output": 7.0, "batch_size": 2}
+
+    # ...and the same for the batched streaming loop:
+    server = LitServer(BatchUnbatchContextStreamingAPI(max_batch_size=2, batch_timeout=4, stream=True), 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 ac,
+        ):
+            resp1 = ac.post("/predict", json={"input": 5.0})
+            resp2 = ac.post("/predict", json={"input": 7.0})
+            resp1, resp2 = await asyncio.gather(resp1, resp2)
+    assert resp1.json() == {"output": 5.0, "batch_size": 2}
+    assert resp2.json() == {"output": 7.0, "batch_size": 2}
+
+
 def test_custom_api_path():
     with pytest.raises(ValueError, match="api_path must start with '/'"):
         LitServer(ls.test_examples.SimpleLitAPI(api_path="predict"))