fix(dspy): Unbatchify rejects batch_fn output length mismatch

[detail-app]

Warning

GitHub issue creation failed

Detail attempted to publish this bug to GitHub, but the issue could not be created. This fix PR was created without that issue, and missing tracker references are shown as Unknown issue.

You can review and merge this PR normally. Please review your tracker integration settings before the next publish run.

Detail bug report: View on Detail

📝 Changes Description

This MR/PR contains the following changes:

• Bug: Unbatchify (dspy/utils/unbatchify.py) routes batch_fn's returned outputs back to pending caller futures with zip(outputs, futures, strict=False). When batch_fn returns a list whose length differs from the number of submitted inputs, the zip silently truncates to the shorter side: outputs shift onto the wrong callers (callers receive each other's results), and any trailing future with no corresponding output is left forever unset. Because __call__ waits on future.result() with no timeout, the orphaned caller hangs indefinitely with no error raised. close() cannot rescue it once the (input, future) pair has been dequeued into the worker.
• Fix: Validate len(outputs) == len(futures) before calling any set_result, raising ValueError so the existing except path safely propagates the exception to every still-pending future. Because no set_result runs before the check passes, mis-pairing is impossible, no future is orphaned, the worker thread survives, and all callers receive a clear error.
• Why this and not strict=True: A one-character zip(..., strict=True) change is strictly worse — strict=True raises only after yielding every short-side pair, so set_result has already run on the leading futures; the existing except block then calls set_exception on already-FINISHED futures, raising InvalidStateError inside the handler, which escapes _worker and crashes the worker thread (converting a per-batch failure into a permanent per-object failure). Validating the length before iterating avoids this entirely.

Closes Unknown issue

✅ Contributor Checklist

• Pre-Commit checks are passing (locally and remotely)
• Title of your PR / MR corresponds to the required format
• Commit message follows required format {label}(dspy): {message}

⚠️ Warnings

None. This is a latent hardening fix: the only in-tree batch_fn passed to Unbatchify (Embeddings._batch_forward → _rerank_and_predict) always returns exactly len(queries) tuples by construction, so no current code path is affected. The fix closes a dormant trap for future batch_fn implementations that could return a mismatched-length list.

Testing:

• New regression tests in tests/utils/test_unbatchify.py covering: short batch_fn output rejected, long batch_fn output rejected, end-to-end __call__ raises and does not hang on a short output, and worker thread survives a mismatched batch to service subsequent calls.
• Existing tests/utils/test_unbatchify.py tests (batch-size trigger, timeout trigger, max_wait_time re-arming) and tests/retrievers/test_embeddings.py (the only in-tree consumer) pass unchanged — no regression.
• Ruff lint and format checks pass; pre-commit hooks pass.
• No typecheck step exists in this repo (no mypy/pyright config, CI job, or hook), so none was run.

Authored by Detail.

---

Automatic Fixes PRs can be configured here.

[greptile-apps]

Greptile Summary

This PR prevents Unbatchify from partially resolving a batch when batch_fn returns the wrong number of outputs.

• Validates output cardinality before resolving any futures.
• Propagates a clear ValueError to every caller affected by a mismatched batch.
• Adds regression coverage for short and long outputs, non-hanging callers, and worker reuse after failure.

Confidence Score: 5/5

The PR appears safe to merge and correctly rejects mismatched batch output without partially resolving callers.

The length check runs before any future is resolved, and the existing exception path delivers the mismatch error to every pending future while allowing the worker to continue.

Important Files Changed

Filename	Overview
dspy/utils/unbatchify.py	Validates batch output cardinality before resolving futures, preventing partial delivery and orphaned callers.
tests/utils/test_unbatchify.py	Adds focused regression tests for both mismatch directions, caller completion, and worker survival.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Collect queued inputs and futures] --> B[Invoke batch_fn]
    B --> C{Output count matches future count?}
    C -->|Yes| D[Resolve each future with its output]
    C -->|No| E[Raise ValueError]
    E --> F[Set exception on all batch futures]
    F --> G[Worker continues processing later batches]

Loading

_{Reviews (1): Last reviewed commit: "fix(dspy): Unbatchify rejects batch_fn o..." | Re-trigger Greptile}