diff --git a/dspy/utils/unbatchify.py b/dspy/utils/unbatchify.py index 0c0c1c818f..345b151cb8 100644 --- a/dspy/utils/unbatchify.py +++ b/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: diff --git a/tests/utils/test_unbatchify.py b/tests/utils/test_unbatchify.py index 63987fd0ca..43b51fcd60 100644 --- a/tests/utils/test_unbatchify.py +++ b/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()