feat(dspy): support gepa 0.1.4 — parallel proposals, reflection-cost budget, batch evaluation, checkpoint-safe adapter state

dspy/teleprompt/gepa/gepa.py

@@ -3,12 +3,18 @@
 import math
 import random
 from dataclasses import dataclass
-from typing import Any, Literal, Optional, Protocol, Union
+from typing import TYPE_CHECKING, Any, Literal, Optional, Protocol, Union
 
 from gepa import GEPAResult
 from gepa.core.adapter import ProposalFn
 from gepa.proposer.reflective_mutation.base import ReflectionComponentSelector
 
+if TYPE_CHECKING:
+    from gepa.strategies.acceptance import AcceptanceCriterion
+    from gepa.strategies.proposal_sampling import SamplingStrategy
+    from gepa.strategies.proposal_selection import SelectionStrategy
+
+import dspy
 from dspy.clients.lm import LM
 from dspy.primitives import Example, Module, Prediction
 from dspy.teleprompt.gepa.gepa_utils import DspyAdapter, DSPyTrace, PredictorFeedbackFn, ScoreWithFeedback
@@ -276,6 +282,27 @@ def metric(
             using LLM-driven selection logic based on optimization state and trajectories.
             See [gepa component selectors](https://github.com/gepa-ai/gepa/blob/main/src/gepa/strategies/component_selector.py)
             for available built-in selectors and the ReflectionComponentSelector protocol for implementing custom selectors.
+        sampling_strategy: Controls how many (parent, minibatch) proposal tasks GEPA samples per iteration.
+            Accepts a gepa `SamplingStrategy` such as `SameParentSampling(n)`, `IndependentSampling(n)`, or
+            `PxNSampling(p, n)`. Defaults to classic single-mutation GEPA. Multi-proposal iterations
+            evaluate up to 8 candidates concurrently (see `DspyAdapter.batch_evaluate`), each using up to
+            `num_threads` threads, so effective LM concurrency is multiplied accordingly — size
+            `num_threads` and provider rate limits with this in mind.
+        selection_strategy: Decides which of an iteration's improving proposals enter the candidate pool.
+            Accepts a gepa `SelectionStrategy` such as `BestImprovement` or `TopKImprovements(k)`. Defaults
+            to accepting all improvements.
+        acceptance_criterion: 'strict_improvement' (default), 'improvement_or_equal', or a custom gepa
+            `AcceptanceCriterion` instance controlling when a mutated candidate is considered an improvement.
+        max_reflection_cost: Maximum cumulative reflection-LM spend in USD before GEPA stops. Requires
+            `reflection_lm`. Cost is read from the reflection LM's call history, so it reflects actual
+            litellm-reported costs. Requires LM history to be enabled (`dspy.settings.disable_history`
+            must be False and `dspy.settings.max_history_size` must be nonzero); `compile()` raises if
+            history is disabled. This complements, but does not replace, the metric-call budget: one of
+            `auto`, `max_full_evals`, or `max_metric_calls` is still required.
+        wandb_attach_existing: Log into the already-active wandb run instead of creating one. Use when GEPA
+            is embedded in a training loop that owns the run.
+        mlflow_attach_existing: Log into the already-active MLflow run instead of creating one.
+        tracking_key_prefix: Prefix applied to metric keys logged to wandb/MLflow, useful when sharing a run.
         add_format_failure_as_feedback: Whether to add format failures as feedback. Default is False.
         use_merge: Whether to use merge-based optimization. Default is True.
         max_merge_invocations: The maximum number of merge invocations to perform. Default is 5.
@@ -366,6 +393,12 @@ def __init__(
         add_format_failure_as_feedback: bool = False,
         instruction_proposer: "ProposalFn | None" = None,
         component_selector: "ReflectionComponentSelector | str" = "round_robin",
+        # Proposal strategy configuration (gepa >= 0.1.4)
+        sampling_strategy: "SamplingStrategy | None" = None,
+        selection_strategy: "SelectionStrategy | None" = None,
+        acceptance_criterion: "AcceptanceCriterion | Literal['strict_improvement', 'improvement_or_equal']" = "strict_improvement",
+        # Budget
+        max_reflection_cost: float | None = None,
         # Merge-based configuration
         use_merge: bool = True,
         max_merge_invocations: int | None = 5,
@@ -382,6 +415,9 @@ def __init__(
         track_best_outputs: bool = False,
         warn_on_score_mismatch: bool = True,
         use_mlflow: bool = False,
+        wandb_attach_existing: bool = False,
+        mlflow_attach_existing: bool = False,
+        tracking_key_prefix: str = "",
         # Reproducibility
         seed: int | None = 0,
         # GEPA passthrough kwargs
@@ -418,6 +454,18 @@ def __init__(
             "Reflection LM is used by GEPA to reflect on the behavior of the program and propose new instructions, and will benefit from a strong model. "
         )
 
+        if max_reflection_cost is not None:
+            assert reflection_lm is not None, (
+                "max_reflection_cost requires reflection_lm to be set, so GEPA can track reflection spend."
+            )
+        self.sampling_strategy = sampling_strategy
+        self.selection_strategy = selection_strategy
+        self.acceptance_criterion = acceptance_criterion
+        self.max_reflection_cost = max_reflection_cost
+        self.wandb_attach_existing = wandb_attach_existing
+        self.mlflow_attach_existing = mlflow_attach_existing
+        self.tracking_key_prefix = tracking_key_prefix
+
         self.reflection_lm = reflection_lm
         self.skip_perfect_score = skip_perfect_score
         self.add_format_failure_as_feedback = add_format_failure_as_feedback
@@ -600,23 +648,42 @@ def feedback_fn(
         for path, flex in flex_submodules.items():
             seed_candidate[path] = flex.module_src
 
+        # TrackedReflectionLM derives max_reflection_cost's spend totals from the reflection LM's
+        # call history. If history is disabled (or capped at 0 entries), those totals stay 0.0
+        # forever and the stopper never fires, silently allowing unbounded spend. Checked here
+        # (not __init__) because settings can change between construction and compile.
+        if self.max_reflection_cost is not None and (
+            dspy.settings.disable_history or not dspy.settings.max_history_size
+        ):
+            raise ValueError(
+                "max_reflection_cost requires LM call history to be enabled, since reflection spend is "
+                "computed from the reflection LM's history. Currently "
+                f"dspy.settings.disable_history={dspy.settings.disable_history} and "
+                f"dspy.settings.max_history_size={dspy.settings.max_history_size}. Either unset "
+                "max_reflection_cost, or ensure disable_history is False and max_history_size is nonzero."
+            )
+
         gepa_result: GEPAResult = optimize(
             seed_candidate=seed_candidate,
             trainset=trainset,
             valset=valset,
             adapter=adapter,
             # Reflection-based configuration
-            reflection_lm=(lambda x: adapter.stripped_lm_call(x)[0]) if self.reflection_lm is not None else None,
+            reflection_lm=adapter.tracked_reflection_lm,
             candidate_selection_strategy=self.candidate_selection_strategy,
             skip_perfect_score=self.skip_perfect_score,
             reflection_minibatch_size=self.reflection_minibatch_size,
             module_selector=self.component_selector,
             perfect_score=self.perfect_score,
+            sampling_strategy=self.sampling_strategy,
+            selection_strategy=self.selection_strategy,
+            acceptance_criterion=self.acceptance_criterion,
             # Merge-based configuration
             use_merge=self.use_merge,
             max_merge_invocations=self.max_merge_invocations,
             # Budget
             max_metric_calls=self.max_metric_calls,
+            max_reflection_cost=self.max_reflection_cost,
             # Logging
             logger=LoggerAdapter(logger),
             run_dir=self.log_dir,
@@ -625,6 +692,9 @@ def feedback_fn(
             wandb_init_kwargs=self.wandb_init_kwargs,
             use_mlflow=self.use_mlflow,
             track_best_outputs=self.track_best_outputs,
+            wandb_attach_existing=self.wandb_attach_existing,
+            mlflow_attach_existing=self.mlflow_attach_existing,
+            tracking_key_prefix=self.tracking_key_prefix,
             display_progress_bar=True,
             raise_on_exception=True,
             # Reproducibility

dspy/teleprompt/gepa/gepa_utils.py

@@ -1,5 +1,7 @@
 import logging
 import random
+from concurrent.futures import ThreadPoolExecutor
+from copy import deepcopy
 from typing import Any, Callable, Protocol, TypedDict
 
 from gepa import EvaluationBatch, GEPAAdapter
@@ -10,6 +12,7 @@
 from dspy.adapters.chat_adapter import ChatAdapter
 from dspy.adapters.types import History
 from dspy.adapters.types.base_type import Type
+from dspy.dsp.utils.settings import thread_local_overrides
 from dspy.evaluate import Evaluate
 from dspy.primitives import Example, Prediction
 from dspy.primitives.code_interpreter import CodeInterpreterError
@@ -83,6 +86,73 @@ def __call__(
         ...
 
 
+def _stripped_lm_outputs(lm, x: str) -> list[str]:
+    raw_outputs = lm(x)
+    outputs = []
+    for raw_output in raw_outputs:
+        if type(raw_output) == str:
+            outputs.append(raw_output)
+        elif type(raw_output) == dict:
+            if "text" not in raw_output:
+                raise KeyError("Missing 'text' field in the output from the base LM!")
+            outputs.append(raw_output["text"])
+        else:
+            raise TypeError("Unexpected output type from the base LM! Expected str or dict")
+    return outputs
+
+
+class TrackedReflectionLM:
+    """Reflection callable handed to `gepa.optimize` that exposes real cost totals.
+
+    gepa's `max_reflection_cost` stopper reads `total_cost` from the reflection LM;
+    plain callables get wrapped in a gepa `TrackingLM` that always reports 0.0, so
+    the stopper would never fire. Totals here are derived from the underlying
+    `dspy.LM.history`, which also captures reflection calls made by the instruction
+    and code proposers that never go through this callable. Requires LM history to
+    be enabled (it is unless `dspy.settings.disable_history` is set).
+
+    Totals are scoped to this instance's lifetime: the sums present in `lm.history`
+    at construction time are snapshotted as a baseline and subtracted back out, so
+    reusing a `reflection_lm` across multiple `compile()` calls doesn't carry over
+    spend from a previous run (`DspyAdapter.__init__` constructs a fresh
+    `TrackedReflectionLM` per compile). This does NOT protect against sharing the
+    reflection LM instance as both the task LM and the reflection LM within the
+    same run: concurrent task-LM calls append to the same `history` list, so their
+    cost/usage would still be counted toward the reflection budget.
+    """
+
+    def __init__(self, lm):
+        self.lm = lm
+        self._baseline_cost = self._sum_cost(lm.history)
+        self._baseline_tokens_in = self._sum_usage(lm.history, "prompt_tokens")
+        self._baseline_tokens_out = self._sum_usage(lm.history, "completion_tokens")
+
+    def __call__(self, x: str) -> str:
+        return _stripped_lm_outputs(self.lm, x)[0]
+
+    @staticmethod
+    def _sum_cost(history) -> float:
+        return sum(entry.get("cost") or 0.0 for entry in history if isinstance(entry, dict))
+
+    @staticmethod
+    def _sum_usage(history, key: str) -> int:
+        return sum((entry.get("usage") or {}).get(key) or 0 for entry in history if isinstance(entry, dict))
+
+    @property
+    def total_cost(self) -> float:
+        # max(0, ...) also tolerates history eviction at settings.max_history_size: if entries this
+        # instance counted toward the baseline get evicted, the raw sum can dip below the baseline.
+        return max(0.0, self._sum_cost(self.lm.history) - self._baseline_cost)
+
+    @property
+    def total_tokens_in(self) -> int:
+        return max(0, self._sum_usage(self.lm.history, "prompt_tokens") - self._baseline_tokens_in)
+
+    @property
+    def total_tokens_out(self) -> int:
+        return max(0, self._sum_usage(self.lm.history, "completion_tokens") - self._baseline_tokens_out)
+
+
 class DspyAdapter(GEPAAdapter[Example, TraceData, Prediction]):
     def __init__(
         self,
@@ -106,6 +176,7 @@ def __init__(
         self.add_format_failure_as_feedback = add_format_failure_as_feedback
         self.rng = rng or random.Random(0)
         self.reflection_lm = reflection_lm
+        self.tracked_reflection_lm = TrackedReflectionLM(reflection_lm) if reflection_lm is not None else None
         self.custom_instruction_proposer = custom_instruction_proposer
         self.warn_on_score_mismatch = warn_on_score_mismatch
         self.reflection_minibatch_size = reflection_minibatch_size
@@ -268,6 +339,46 @@ def evaluate(self, batch, candidate, capture_traces=False):
             scores = [s["score"] if hasattr(s, "score") else s for s in scores]
             return EvaluationBatch(outputs=outputs, scores=scores, trajectories=None)
 
+    def batch_evaluate(self, items):
+        """Evaluate multiple (candidate, batch) pairs concurrently.
+
+        gepa's multi-proposal iterations (`sampling_strategy`) evaluate several
+        candidates per step; its fallback runs them one at a time, and each
+        `evaluate` call only parallelizes across the (small) minibatch. Candidates
+        are independent, so run them on worker threads. dspy settings overrides are
+        thread-local and must be copied into each worker.
+        """
+        if len(items) <= 1:
+            return [self.evaluate(batch, candidate, capture_traces=True) for candidate, batch in items]
+
+        parent_overrides = thread_local_overrides.get().copy()
+
+        def _evaluate_pair(pair):
+            candidate, batch = pair
+            new_overrides = {**thread_local_overrides.get(), **parent_overrides}
+            if new_overrides.get("usage_tracker"):
+                # Mirror dspy's ParallelExecutor: deep-copy the usage tracker per worker so each
+                # thread tracks its own usage instead of contending over one shared tracker.
+                new_overrides["usage_tracker"] = deepcopy(new_overrides["usage_tracker"])
+            token = thread_local_overrides.set(new_overrides)
+            try:
+                return self.evaluate(batch, candidate, capture_traces=True)
+            finally:
+                thread_local_overrides.reset(token)
+
+        with ThreadPoolExecutor(max_workers=min(len(items), 8)) as executor:
+            return list(executor.map(_evaluate_pair, items))
+
+    def get_adapter_state(self) -> dict[str, Any]:
+        """Snapshot adapter state into gepa's checkpoint (gepa persists it via pickle)."""
+        return {"rng_state": self.rng.getstate()}
+
+    def set_adapter_state(self, state: dict[str, Any]) -> None:
+        """Restore adapter state on checkpoint resume."""
+        rng_state = state.get("rng_state")
+        if rng_state is not None:
+            self.rng.setstate(rng_state)
+
     def make_reflective_dataset(
         self, candidate, eval_batch, components_to_update
     ) -> dict[str, list[ReflectiveExample]]:
@@ -402,17 +513,5 @@ def make_reflective_dataset(
     # Always return strings from the LM outputs
     # Even when it returns a dict with e.g., "text" and "reasoning" fields
     def stripped_lm_call(self, x: str) -> list[str]:
-        raw_outputs = self.reflection_lm(x)
-        outputs = []
-        for raw_output in raw_outputs:
-            if type(raw_output) == str:
-                outputs.append(raw_output)
-            elif type(raw_output) == dict:
-                if "text" not in raw_output:
-                    raise KeyError("Missing 'text' field in the output from the base LM!")
-                outputs.append(raw_output["text"])
-            else:
-                raise TypeError("Unexpected output type from the base LM! Expected str or dict")
-
-        return outputs
+        return _stripped_lm_outputs(self.reflection_lm, x)
 

pyproject.toml

@@ -34,7 +34,7 @@ dependencies = [
     "anyio",
     "cachetools>=5.5.0",
     "cloudpickle>=3.1.2",
-    "gepa[dspy]==0.1.1",
+    "gepa[dspy]==0.1.4",
 ]
 
 [project.optional-dependencies]