fix(dspy): Unbatchify rejects batch_fn output length mismatch

dspy/utils/unbatchify.py

@@ -7,10 +7,7 @@
 
 class Unbatchify:
     def __init__(
-        self,
-        batch_fn: Callable[[list[Any]], list[Any]],
-        max_batch_size: int = 32,
-        max_wait_time: float = 0.1
+        self, batch_fn: Callable[[list[Any]], list[Any]], max_batch_size: int = 32, max_wait_time: float = 0.1
     ):
         """
         Initializes the Unbatchify.
@@ -68,6 +65,8 @@ def _worker(self):
             if batch:
                 try:
                     outputs = self.batch_fn(batch)
+                    if len(outputs) != len(futures):
+                        raise ValueError(f"batch_fn returned {len(outputs)} outputs for {len(futures)} inputs")
                     for output, future in zip(outputs, futures, strict=False):
                         future.set_result(output)
                 except Exception as e:

tests/utils/test_unbatchify.py

@@ -1,7 +1,10 @@
+import threading
 import time
 from concurrent.futures import Future
 from unittest.mock import MagicMock
 
+import pytest
+
 from dspy.utils.unbatchify import Unbatchify
 
 
@@ -95,3 +98,112 @@ def test_unbatchify_honors_max_wait_time_under_trickling_input():
     assert elapsed < wait_time * 1.5
 
     unbatcher.close()
+
+
+def test_unbatchify_raises_on_short_batch_fn_output():
+    """When batch_fn returns fewer outputs than inputs, every future must receive a
+    ValueError rather than a (mis-paired) result, and no future may be left pending."""
+
+    def short_batch_fn(batch):
+        # drop the second input -> len(batch) - 1 outputs
+        return [item + 1 for i, item in enumerate(batch) if i != 1]
+
+    batch_fn_mock = MagicMock(wraps=short_batch_fn)
+    # max_batch_size=4 fills immediately on the 4 submits below regardless of
+    # max_wait_time; a small max_wait_time just keeps close() teardown fast.
+    unbatcher = Unbatchify(batch_fn=batch_fn_mock, max_batch_size=4, max_wait_time=0.2)
+
+    futures = [unbatcher.submit(10), unbatcher.submit(20), unbatcher.submit(30), unbatcher.submit(40)]
+
+    for f in futures:
+        with pytest.raises(ValueError, match="batch_fn returned 3 outputs for 4 inputs"):
+            f.result(timeout=2)
+    # No future received a result: mis-pairing would have resolved some futures with a value
+    # instead of an exception.
+    for f in futures:
+        assert isinstance(f.exception(), ValueError)
+    assert batch_fn_mock.call_count == 1
+    unbatcher.close()
+
+
+def test_unbatchify_raises_on_long_batch_fn_output():
+    """When batch_fn returns more outputs than inputs, every future must receive a
+    ValueError rather than receiving one of the extra results out of position."""
+
+    def long_batch_fn(batch):
+        # add a spurious trailing output -> len(batch) + 1 outputs
+        return [item + 1 for item in batch] + [999]
+
+    batch_fn_mock = MagicMock(wraps=long_batch_fn)
+    unbatcher = Unbatchify(batch_fn=batch_fn_mock, max_batch_size=3, max_wait_time=0.2)
+
+    futures = [unbatcher.submit(10), unbatcher.submit(20), unbatcher.submit(30)]
+
+    for f in futures:
+        with pytest.raises(ValueError, match="batch_fn returned 4 outputs for 3 inputs"):
+            f.result(timeout=2)
+    for f in futures:
+        assert isinstance(f.exception(), ValueError)
+    assert batch_fn_mock.call_count == 1
+    unbatcher.close()
+
+
+def test_unbatchify_call_raises_and_does_not_hang_on_short_output():
+    """End-to-end via __call__: when batch_fn returns short, every caller must raise
+    ValueError. No caller may receive another caller's result, and no caller may hang."""
+
+    def short_batch_fn(batch):
+        return [f"out-{x}" for i, x in enumerate(batch) if i != 1]
+
+    ub = Unbatchify(short_batch_fn, max_batch_size=4, max_wait_time=1.0)
+
+    results = {}
+    errors = {}
+
+    def call(idx):
+        try:
+            results[idx] = ub(f"in-{idx}")
+        except Exception as e:
+            errors[idx] = e
+
+    threads = [threading.Thread(target=call, args=(i,), daemon=True) for i in range(4)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join(timeout=2.0)
+
+    assert not any(t.is_alive() for t in threads), "a caller is still hanging"
+    assert results == {}, f"callers received (mis-paired) results: {results}"
+    assert set(errors) == {0, 1, 2, 3}
+    for i in range(4):
+        assert isinstance(errors[i], ValueError)
+        assert "batch_fn returned 3 outputs for 4 inputs" in str(errors[i])
+
+    ub.close()
+
+
+def test_unbatchify_survives_after_mismatched_batch():
+    """A length mismatch fails only the affected batch; the worker thread must survive
+    and the same Unbatchify instance must keep servicing subsequent calls."""
+
+    def short_batch_fn(batch):
+        return [item + 1 for i, item in enumerate(batch) if i != 1]
+
+    def good_batch_fn(batch):
+        return [item + 1 for item in batch]
+
+    ub = Unbatchify(short_batch_fn, max_batch_size=4, max_wait_time=0.2)
+
+    futures = [ub.submit(1), ub.submit(2), ub.submit(3), ub.submit(4)]
+    for f in futures:
+        with pytest.raises(ValueError):
+            f.result(timeout=2)
+
+    # The worker thread must still be alive after rejecting the mismatched batch.
+    assert ub.worker_thread.is_alive()
+
+    # Swap to a well-behaved batch_fn; the same instance must service the next call.
+    ub.batch_fn = good_batch_fn
+    assert ub(7) == 8
+
+    ub.close()