Add warning for unexpected model output in batched prediction (#300)

* fix: add warning unexpected output from HF model (closes #294)

* add: warning if batched loop return string

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* move output check to unbatch_no_stream

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* restore format loops.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: lint E501 Line too long

* Update tests/test_batch.py

Co-authored-by: Aniket Maurya <theaniketmaurya@gmail.com>

* Update src/litserve/api.py

Co-authored-by: Aniket Maurya <theaniketmaurya@gmail.com>

* Update test to match new warning string

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Delete whitespace in warning string test_batch

* Update test_batch.py

* Update test_batch.py

* Update warning copy

* update test

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix test

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Aniket Maurya <theaniketmaurya@gmail.com>

src/litserve/api.py

@@ -63,6 +63,13 @@ def predict(self, x, **kwargs):
         pass
 
     def _unbatch_no_stream(self, output):
+        if isinstance(output, str):
+            warnings.warn(
+                "The 'predict' method returned a string instead of a list of predictions. "
+                "When batching is enabled, 'predict' must return a list to handle multiple inputs correctly. "
+                "Please update the 'predict' method to return a list of predictions to avoid unexpected behavior.",
+                UserWarning,
+            )
         return list(output)
 
     def _unbatch_stream(self, output_stream):

tests/test_batch.py

@@ -148,6 +148,22 @@ def test_max_batch_size_warning():
         LitServer(SimpleTorchAPI(), accelerator="cpu", devices=1, timeout=2)
 
 
+def test_batch_predict_string_warning():
+    api = ls.test_examples.SimpleBatchedAPI()
+    api._sanitize(2, None)
+    api.predict = MagicMock(return_value="This is a string")
+
+    mock_input = torch.tensor([[1.0], [2.0]])
+
+    with pytest.warns(
+        UserWarning,
+        match="When batching is enabled, 'predict' must return a list to handle multiple inputs correctly.",
+    ):
+        # Simulate the behavior in run_batched_loop
+        y = api.predict(mock_input)
+        api.unbatch(y)
+
+
 class FakeResponseQueue:
     def put(self, *args):
         raise Exception("Exit loop")