Skip to content

Commit 822f393

Browse files
isaacbmillerampagentIsaac Miller
authored
feat(gepa): parallelize candidate evaluation (#10210)
* feat(gepa): parallelize candidate evaluation Amp-Thread-ID: https://ampcode.com/threads/T-019ffe40-c240-70bb-8cb2-db33146ed8e6 Co-authored-by: Isaac Miller <isaac@cmpnd.ai> * fix(gepa): distribute candidate thread remainder Amp-Thread-ID: https://ampcode.com/threads/T-019ffe40-c240-70bb-8cb2-db33146ed8e6 Co-authored-by: Isaac Miller <isaac@cmpnd.ai> --------- Co-authored-by: Amp <amp@ampcode.com> Co-authored-by: Isaac Miller <isaac@cmpnd.ai>
1 parent aa9d2d0 commit 822f393

3 files changed

Lines changed: 75 additions & 8 deletions

File tree

‎dspy/teleprompt/gepa/gepa.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -279,7 +279,8 @@ def metric(
279279
add_format_failure_as_feedback: Whether to add format failures as feedback. Default is False.
280280
use_merge: Whether to use merge-based optimization. Default is True.
281281
max_merge_invocations: The maximum number of merge invocations to perform. Default is 5.
282-
num_threads: The number of threads to use for evaluation with `Evaluate`. Optional.
282+
num_threads: The total number of threads available for candidate and example evaluation. Multi-proposal
283+
candidate evaluations share this budget. Optional.
283284
failure_score: The score to assign to failed examples. Default is 0.0.
284285
perfect_score: The maximum score achievable by the metric. Default is 1.0. Used by GEPA
285286
to determine if all examples in a minibatch are perfect.
@@ -325,7 +326,7 @@ def metric(
325326
- mlflow_tracking_uri: The tracking URI to use for MLflow (when use_mlflow=True).
326327
- mlflow_experiment_name: The experiment name to use for MLflow (when use_mlflow=True).
327328
- sampling_strategy, selection_strategy, acceptance_criterion: GEPA 0.1.4 proposal controls.
328-
DSPy supports these using GEPA's sequential candidate-evaluation fallback.
329+
DSPy evaluates proposal candidates concurrently within the `num_threads` budget.
329330
- wandb_attach_existing, mlflow_attach_existing, tracking_key_prefix: GEPA 0.1.4 tracking controls.
330331
331332
`max_reflection_cost` is not supported yet because DSPy LMs do not expose the cumulative cost
@@ -346,8 +347,9 @@ def metric(
346347
Merge Configuration: GEPA can merge successful program variants using `use_merge=True`.
347348
The `max_merge_invocations` parameter controls how many merge attempts are made during optimization.
348349
349-
Evaluation Configuration: Use `num_threads` to parallelize evaluation. The `failure_score` and
350-
`perfect_score` parameters help GEPA understand your metric's range and optimize accordingly.
350+
Evaluation Configuration: `num_threads` controls total evaluation concurrency and is shared across
351+
candidates in multi-proposal batches. The `failure_score` and `perfect_score` parameters help GEPA
352+
understand your metric's range and optimize accordingly.
351353
352354
Logging Configuration: Set `log_dir` to save detailed logs and enable checkpoint resuming.
353355
Use `track_stats=True` to access detailed optimization results via the `detailed_results` attribute.

‎dspy/teleprompt/gepa/gepa_utils.py‎

Lines changed: 29 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
import logging
22
import random
3+
from concurrent.futures import ThreadPoolExecutor
4+
from contextvars import copy_context
35
from typing import Any, Callable, Protocol, TypedDict
46

57
from gepa import EvaluationBatch, GEPAAdapter
@@ -189,7 +191,8 @@ def build_program(self, candidate: dict[str, str]):
189191

190192
return new_prog
191193

192-
def evaluate(self, batch, candidate, capture_traces=False):
194+
def evaluate(self, batch, candidate, capture_traces=False, num_threads=None):
195+
num_threads = self.num_threads if num_threads is None else num_threads
193196
try:
194197
program = self.build_program(candidate)
195198
except (SyntaxError, CodeInterpreterError) as e:
@@ -217,7 +220,7 @@ def evaluate(self, batch, candidate, capture_traces=False):
217220
program,
218221
batch,
219222
metric_fn=self.metric_fn,
220-
num_threads=self.num_threads,
223+
num_threads=num_threads,
221224
failure_score=self.failure_score,
222225
callback_metadata=callback_metadata,
223226
capture_traces=capture_traces,
@@ -231,7 +234,7 @@ def evaluate(self, batch, candidate, capture_traces=False):
231234
program=program,
232235
dataset=batch,
233236
metric=self.metric_fn,
234-
num_threads=self.num_threads,
237+
num_threads=num_threads,
235238
raise_on_error=False,
236239
capture_failed_parses=True,
237240
failure_score=self.failure_score,
@@ -255,7 +258,7 @@ def evaluate(self, batch, candidate, capture_traces=False):
255258
evaluator = Evaluate(
256259
devset=batch,
257260
metric=self.metric_fn,
258-
num_threads=self.num_threads,
261+
num_threads=num_threads,
259262
return_all_scores=True,
260263
failure_score=self.failure_score,
261264
provide_traceback=True,
@@ -268,6 +271,28 @@ def evaluate(self, batch, candidate, capture_traces=False):
268271
scores = [s["score"] if hasattr(s, "score") else s for s in scores]
269272
return EvaluationBatch(outputs=outputs, scores=scores, trajectories=None)
270273

274+
def batch_evaluate(self, items, *, capture_traces=True):
275+
total_threads = self.num_threads or dspy.settings.num_threads
276+
candidate_workers = min(len(items), total_threads) or 1
277+
threads_per_candidate, extra_threads = divmod(total_threads, candidate_workers)
278+
279+
def evaluate(item):
280+
context, (candidate, batch), num_threads = item
281+
return context.run(
282+
self.evaluate,
283+
batch,
284+
candidate,
285+
capture_traces=capture_traces,
286+
num_threads=num_threads,
287+
)
288+
289+
work = [
290+
(copy_context(), item, threads_per_candidate + (index < extra_threads))
291+
for index, item in enumerate(items)
292+
]
293+
with ThreadPoolExecutor(max_workers=candidate_workers) as executor:
294+
return list(executor.map(evaluate, work))
295+
271296
def get_adapter_state(self) -> dict[str, Any]:
272297
return {"rng_state": self.rng.getstate()}
273298

‎tests/teleprompt/test_gepa.py‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -661,3 +661,43 @@ def test_adapter_state_round_trips_rng():
661661
adapter.rng.random()
662662
adapter.set_adapter_state(state)
663663
assert adapter.rng.random() == expected
664+
665+
666+
def test_batch_evaluate_is_parallel_and_preserves_context():
667+
from gepa import EvaluationBatch
668+
669+
from dspy.teleprompt.gepa.gepa_utils import DspyAdapter
670+
from dspy.utils.callback_context import ACTIVE_CALL_ID
671+
672+
adapter = DspyAdapter(SimpleModule("input -> output"), simple_metric, {}, num_threads=4)
673+
barrier = threading.Barrier(2)
674+
trace_modes = []
675+
thread_budgets = []
676+
677+
def evaluate(batch, candidate, capture_traces=False, num_threads=None):
678+
if capture_traces:
679+
barrier.wait(timeout=1)
680+
trace_modes.append(capture_traces)
681+
thread_budgets.append(num_threads)
682+
return EvaluationBatch(outputs=[(candidate, dspy.settings.lm, ACTIVE_CALL_ID.get())], scores=[1.0])
683+
684+
adapter.evaluate = evaluate
685+
lm = DummyLM([])
686+
token = ACTIVE_CALL_ID.set("parent")
687+
try:
688+
with dspy.context(lm=lm):
689+
results = adapter.batch_evaluate([("first", []), ("second", [])])
690+
finally:
691+
ACTIVE_CALL_ID.reset(token)
692+
693+
assert [result.outputs[0] for result in results] == [("first", lm, "parent"), ("second", lm, "parent")]
694+
695+
adapter.batch_evaluate([("first", [])], capture_traces=False)
696+
assert trace_modes == [True, True, False]
697+
assert thread_budgets == [2, 2, 4]
698+
699+
adapter.num_threads = 8
700+
adapter.batch_evaluate([("first", []), ("second", []), ("third", [])], capture_traces=False)
701+
assert trace_modes == [True, True, False, False, False, False]
702+
assert sorted(thread_budgets[-3:]) == [2, 3, 3]
703+
assert sum(thread_budgets[-3:]) == 8

0 commit comments

Comments
 (0)