fix(dspy): correct Refine/BestOfN fail_count budget accounting

dspy/predict/best_of_n.py

@@ -52,6 +52,7 @@ def forward(self, **kwargs):
         start = lm.kwargs.get("rollout_id", 0)
         rollout_ids = [start + i for i in range(self.N)]
         best_pred, best_trace, best_reward = None, None, -float("inf")
+        fail_count = self.fail_count
 
         for idx, rid in enumerate(rollout_ids):
             lm_ = lm.copy(rollout_id=rid, temperature=1.0)
@@ -74,9 +75,9 @@ def forward(self, **kwargs):
 
             except Exception as e:
                 print(f"BestOfN: Attempt {idx + 1} failed with rollout id {rid}: {e}")
-                if idx > self.fail_count:
+                fail_count -= 1
+                if fail_count < 0:
                     raise e
-                self.fail_count -= 1
 
         if best_trace:
             dspy.settings.trace.extend(best_trace)

dspy/predict/refine.py

@@ -102,6 +102,7 @@ def forward(self, **kwargs):
         best_pred, best_trace, best_reward = None, None, -float("inf")
         advice = None
         adapter = dspy.settings.adapter or dspy.ChatAdapter()
+        fail_count = self.fail_count
 
         for idx, rid in enumerate(rollout_ids):
             lm_ = lm.copy(rollout_id=rid, temperature=1.0)
@@ -169,9 +170,9 @@ def __call__(self, lm, lm_kwargs, signature, demos, inputs):
 
             except Exception as e:
                 print(f"Refine: Attempt failed with rollout id {rid}: {e}")
-                if idx > self.fail_count:
+                fail_count -= 1
+                if fail_count < 0:
                     raise e
-                self.fail_count -= 1
         if best_trace:
             dspy.settings.trace.extend(best_trace)
         return best_pred

tests/predict/test_best_of_n.py

@@ -55,8 +55,12 @@ def always_raise(self, **kwargs):
     predict = DummyModule("question -> answer", always_raise)
 
     best_of_n = BestOfN(module=predict, N=3, reward_fn=lambda _, __: 1.0, threshold=0.0)
-    with pytest.raises(ValueError):
-        best_of_n(question="What is the capital of Belgium?")
+    # fail_count defaults to N: up to N failures are tolerated, so all N failing
+    # rollouts are absorbed and the best prediction seen so far (None) is returned.
+    result = best_of_n(question="What is the capital of Belgium?")
+    assert result is None
+    # the instance budget must not leak across forward() calls
+    assert best_of_n.fail_count == 3
 
 
 def test_refine_module_custom_fail_count():
@@ -78,3 +82,73 @@ def raise_on_second_call(self, **kwargs):
     assert module_call_count[0] == 2, (
         "Module should have been called exactly 2 times, but was called %d times" % module_call_count[0]
     )
+
+
+def test_best_of_n_fail_count_tolerates_single_late_failure():
+    # threshold=2.0 is unreachable (reward capped at 1.0), so all N rollouts run and
+    # the single failure on the final rollout is actually reached. With
+    # fail_count=2 a single late failure must be tolerated (not re-raised).
+    dspy.configure(lm=DummyLM([{"answer": "Brussels"}] * 4))
+
+    rollout_count = [0]
+
+    def succeed_four_then_fail(self, **kwargs):
+        rollout_count[0] += 1
+        if rollout_count[0] <= 4:
+            return self.predictor(**kwargs)
+        raise ValueError("transient failure on rollout %d" % rollout_count[0])
+
+    predict = DummyModule("question -> answer", succeed_four_then_fail)
+    best_of_n = BestOfN(module=predict, N=5, reward_fn=lambda _, __: 1.0, threshold=2.0, fail_count=2)
+
+    result = best_of_n(question="What is the capital of Belgium?")
+
+    assert result.answer == "Brussels"
+    assert best_of_n.fail_count == 2
+
+
+def test_best_of_n_fail_count_still_raises_when_exceeded():
+    # A single late failure is tolerated with fail_count=1, but the *second*
+    # late failure exceeds the budget and must be re-raised.
+    dspy.configure(lm=DummyLM([{"answer": "Brussels"}] * 4))
+
+    rollout_count = [0]
+
+    def succeed_four_then_fail_twice(self, **kwargs):
+        rollout_count[0] += 1
+        if rollout_count[0] <= 4:
+            return self.predictor(**kwargs)
+        raise ValueError("failure on rollout %d" % rollout_count[0])
+
+    predict = DummyModule("question -> answer", succeed_four_then_fail_twice)
+    best_of_n = BestOfN(module=predict, N=6, reward_fn=lambda _, __: 1.0, threshold=2.0, fail_count=1)
+
+    with pytest.raises(ValueError, match="failure on rollout 6"):
+        best_of_n(question="What is the capital of Belgium?")
+    assert best_of_n.fail_count == 1
+
+
+def test_best_of_n_fail_count_does_not_leak_across_calls():
+    # Each forward() call fails twice then succeeds; with fail_count=2 both
+    # leading failures must be tolerated on *every* call (no cross-call leak).
+    dspy.configure(lm=DummyLM([{"answer": "Brussels"}] * 20))
+
+    state = {"rollout_in_call": 0}
+
+    def fail_twice_per_call(self, **kwargs):
+        if state["rollout_in_call"] < 2:
+            state["rollout_in_call"] += 1
+            raise ValueError("leading failure #%d" % state["rollout_in_call"])
+        return self.predictor(**kwargs)
+
+    predict = DummyModule("question -> answer", fail_twice_per_call)
+    best_of_n = BestOfN(module=predict, N=5, reward_fn=lambda _, __: 1.0, threshold=2.0, fail_count=2)
+
+    r1 = best_of_n(question="Q?")
+    assert r1.answer == "Brussels"
+    assert best_of_n.fail_count == 2
+
+    state["rollout_in_call"] = 0
+    r2 = best_of_n(question="Q?")
+    assert r2.answer == "Brussels"
+    assert best_of_n.fail_count == 2

tests/predict/test_refine.py

@@ -55,8 +55,12 @@ def always_raise(self, **kwargs):
     predict = DummyModule("question -> answer", always_raise)
 
     refine = Refine(module=predict, N=3, reward_fn=lambda _, __: 1.0, threshold=0.0)
-    with pytest.raises(ValueError):
-        refine(question="What is the capital of Belgium?")
+    # fail_count defaults to N: up to N failures are tolerated, so all N failing
+    # rollouts are absorbed and the best prediction seen so far (None) is returned.
+    result = refine(question="What is the capital of Belgium?")
+    assert result is None
+    # the instance budget must not leak across forward() calls
+    assert refine.fail_count == 3
 
 
 def test_refine_module_custom_fail_count():
@@ -78,3 +82,82 @@ def raise_on_second_call(self, **kwargs):
     assert module_call_count[0] == 2, (
         "Module should have been called exactly 2 times, but was called %d times" % module_call_count[0]
     )
+
+
+# A single LM response that satisfies both the `question -> answer` predictor
+# (which extracts `answer`) and the `OfferFeedback` predictor (which extracts
+# `discussion` and `advice`). Using one combined response makes the LM answer
+# order irrelevant, so consumption stays correct across multiple forward() calls
+# (where each call consumes an odd number of responses: predictor + advice pairs
+# plus a final predictor without advice).
+ANSWER_AND_ADVICE = {"answer": "Brussels", "discussion": "no blame", "advice": {"predictor": "N/A"}}
+
+
+def test_refine_fail_count_tolerates_single_late_failure():
+    # threshold=2.0 is unreachable (reward capped at 1.0), so all N rollouts run and
+    # the single failure on the final rollout is actually reached. With
+    # fail_count=2 a single late failure must be tolerated (not re-raised).
+    dspy.configure(lm=DummyLM([ANSWER_AND_ADVICE] * 8))
+
+    rollout_count = [0]
+
+    def succeed_four_then_fail(self, **kwargs):
+        rollout_count[0] += 1
+        if rollout_count[0] <= 4:
+            return self.predictor(**kwargs)
+        raise ValueError("transient failure on rollout %d" % rollout_count[0])
+
+    predict = DummyModule("question -> answer", succeed_four_then_fail)
+    refine = Refine(module=predict, N=5, reward_fn=lambda _, __: 1.0, threshold=2.0, fail_count=2)
+
+    result = refine(question="What is the capital of Belgium?")
+
+    assert result.answer == "Brussels"
+    assert refine.fail_count == 2
+
+
+def test_refine_fail_count_still_raises_when_exceeded():
+    # A single late failure is tolerated with fail_count=1, but the *second*
+    # late failure exceeds the budget and must be re-raised.
+    dspy.configure(lm=DummyLM([ANSWER_AND_ADVICE] * 8))
+
+    rollout_count = [0]
+
+    def succeed_four_then_fail_twice(self, **kwargs):
+        rollout_count[0] += 1
+        if rollout_count[0] <= 4:
+            return self.predictor(**kwargs)
+        raise ValueError("failure on rollout %d" % rollout_count[0])
+
+    predict = DummyModule("question -> answer", succeed_four_then_fail_twice)
+    refine = Refine(module=predict, N=6, reward_fn=lambda _, __: 1.0, threshold=2.0, fail_count=1)
+
+    with pytest.raises(ValueError, match="failure on rollout 6"):
+        refine(question="What is the capital of Belgium?")
+    assert refine.fail_count == 1
+
+
+def test_refine_fail_count_does_not_leak_across_calls():
+    # Each forward() call fails twice then succeeds; with fail_count=2 both
+    # leading failures must be tolerated on *every* call (no cross-call leak).
+    dspy.configure(lm=DummyLM([ANSWER_AND_ADVICE] * 20))
+
+    state = {"rollout_in_call": 0}
+
+    def fail_twice_per_call(self, **kwargs):
+        if state["rollout_in_call"] < 2:
+            state["rollout_in_call"] += 1
+            raise ValueError("leading failure #%d" % state["rollout_in_call"])
+        return self.predictor(**kwargs)
+
+    predict = DummyModule("question -> answer", fail_twice_per_call)
+    refine = Refine(module=predict, N=5, reward_fn=lambda _, __: 1.0, threshold=2.0, fail_count=2)
+
+    r1 = refine(question="Q?")
+    assert r1.answer == "Brussels"
+    assert refine.fail_count == 2
+
+    state["rollout_in_call"] = 0
+    r2 = refine(question="Q?")
+    assert r2.answer == "Brussels"
+    assert refine.fail_count == 2