mantisdk 0.1.0__py3-none-any.whl

mantisdk/algorithm/gepa/lib/core/state.py

@@ -0,0 +1,636 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+import hashlib
+import json
+import os
+from collections import defaultdict
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, field
+from typing import Any, ClassVar, Generic, Literal, TypeAlias
+
+from mantisdk.algorithm.gepa.lib.core.adapter import RolloutOutput
+from mantisdk.algorithm.gepa.lib.core.data_loader import DataId
+from mantisdk.algorithm.gepa.lib.gepa_utils import json_default
+from mantisdk.algorithm.gepa.lib.logging.logger import LoggerProtocol
+
+# Types for GEPAState
+ProgramIdx = int
+
+# Type aliases
+ObjectiveScores: TypeAlias = dict[str, float]
+FrontierType: TypeAlias = Literal["instance", "objective", "hybrid", "cartesian"]
+"""Strategy for tracking Pareto frontiers: 'instance' (per validation example), 'objective' (per objective metric), 'hybrid' (both), or 'cartesian' (per example × objective)."""
+FrontierKey: TypeAlias = DataId | str | tuple[str, DataId] | tuple[str, DataId, str]
+"""Key type for frontier mappings depending on frontier_type."""
+
+CandidateHash: TypeAlias = str
+CacheKey: TypeAlias = tuple[CandidateHash, DataId]
+
+
+def _candidate_hash(candidate: dict[str, str]) -> CandidateHash:
+    """Compute a deterministic hash of a candidate dictionary."""
+    return hashlib.sha256(json.dumps(sorted(candidate.items())).encode()).hexdigest()
+
+
+@dataclass
+class CachedEvaluation(Generic[RolloutOutput]):
+    """Cached evaluation result for a (candidate, example) pair."""
+
+    output: RolloutOutput
+    score: float
+    objective_scores: ObjectiveScores | None
+
+
+@dataclass
+class EvaluationCache(Generic[RolloutOutput, DataId]):
+    """Cache for storing evaluation results of (candidate, example) pairs."""
+
+    _cache: dict[CacheKey, CachedEvaluation[RolloutOutput]] = field(default_factory=dict)
+
+    def get(self, candidate: dict[str, str], example_id: DataId) -> CachedEvaluation[RolloutOutput] | None:
+        """Retrieve cached evaluation result if it exists."""
+        return self._cache.get((_candidate_hash(candidate), example_id))
+
+    def put(
+        self,
+        candidate: dict[str, str],
+        example_id: DataId,
+        output: RolloutOutput,
+        score: float,
+        objective_scores: ObjectiveScores | None = None,
+    ) -> None:
+        """Store an evaluation result in the cache."""
+        self._cache[(_candidate_hash(candidate), example_id)] = CachedEvaluation(output, score, objective_scores)
+
+    def get_batch(
+        self, candidate: dict[str, str], example_ids: list[DataId]
+    ) -> tuple[dict[DataId, CachedEvaluation[RolloutOutput]], list[DataId]]:
+        """Look up cached results for a batch. Returns (cached_results, uncached_ids)."""
+        h = _candidate_hash(candidate)
+        cached, uncached = {}, []
+        for eid in example_ids:
+            if entry := self._cache.get((h, eid)):
+                cached[eid] = entry
+            else:
+                uncached.append(eid)
+        return cached, uncached
+
+    def put_batch(
+        self,
+        candidate: dict[str, str],
+        example_ids: list[DataId],
+        outputs: list[RolloutOutput],
+        scores: list[float],
+        objective_scores_list: Sequence[ObjectiveScores] | None = None,
+    ) -> None:
+        """Store evaluation results for a batch of examples."""
+        h = _candidate_hash(candidate)
+        for i, eid in enumerate(example_ids):
+            self._cache[(h, eid)] = CachedEvaluation(
+                outputs[i], scores[i], objective_scores_list[i] if objective_scores_list else None
+            )
+
+    def evaluate_with_cache_full(
+        self,
+        candidate: dict[str, str],
+        example_ids: list[DataId],
+        fetcher: Callable[[list[DataId]], Any],
+        evaluator: Callable[[Any, dict[str, str]], tuple[Any, list[float], Sequence[ObjectiveScores] | None]],
+    ) -> tuple[dict[DataId, RolloutOutput], dict[DataId, float], dict[DataId, ObjectiveScores] | None, int]:
+        """
+        Evaluate using cache, returning full results.
+
+        Returns (outputs_by_id, scores_by_id, objective_scores_by_id, num_actual_evals).
+        """
+        cached, uncached_ids = self.get_batch(candidate, example_ids)
+
+        outputs_by_id: dict[DataId, RolloutOutput] = {eid: c.output for eid, c in cached.items()}
+        scores_by_id: dict[DataId, float] = {eid: c.score for eid, c in cached.items()}
+        objective_by_id: dict[DataId, ObjectiveScores] | None = None
+
+        # Populate objective scores from cache
+        for eid, c in cached.items():
+            if c.objective_scores is not None:
+                objective_by_id = objective_by_id or {}
+                objective_by_id[eid] = c.objective_scores
+
+        # Evaluate uncached examples
+        if uncached_ids:
+            batch = fetcher(uncached_ids)
+            outputs, scores, obj_scores = evaluator(batch, candidate)
+            for idx, eid in enumerate(uncached_ids):
+                outputs_by_id[eid] = outputs[idx]
+                scores_by_id[eid] = scores[idx]
+                if obj_scores is not None:
+                    objective_by_id = objective_by_id or {}
+                    objective_by_id[eid] = obj_scores[idx]
+            self.put_batch(candidate, uncached_ids, outputs, scores, obj_scores)
+
+        return outputs_by_id, scores_by_id, objective_by_id, len(uncached_ids)
+
+
+@dataclass(slots=True)
+class ValsetEvaluation(Generic[RolloutOutput, DataId]):
+    """Container for evaluation results on a validation set batch."""
+
+    outputs_by_val_id: dict[DataId, RolloutOutput]
+    scores_by_val_id: dict[DataId, float]
+    objective_scores_by_val_id: dict[DataId, ObjectiveScores] | None = None
+
+
+class GEPAState(Generic[RolloutOutput, DataId]):
+    """Persistent optimizer state tracking candidates, sparse validation coverage, and objective frontiers."""
+
+    _VALIDATION_SCHEMA_VERSION: ClassVar[int] = 4
+
+    program_candidates: list[dict[str, str]]
+    parent_program_for_candidate: list[list[ProgramIdx | None]]
+    prog_candidate_val_subscores: list[dict[DataId, float]]
+    prog_candidate_objective_scores: list[ObjectiveScores]
+
+    pareto_front_valset: dict[DataId, float]
+    program_at_pareto_front_valset: dict[DataId, set[ProgramIdx]]
+    objective_pareto_front: ObjectiveScores
+    program_at_pareto_front_objectives: dict[str, set[ProgramIdx]]
+    pareto_front_cartesian: dict[tuple[DataId, str], float]
+    program_at_pareto_front_cartesian: dict[tuple[DataId, str], set[ProgramIdx]]
+
+    list_of_named_predictors: list[str]
+    named_predictor_id_to_update_next_for_program_candidate: list[int]
+
+    i: int
+    num_full_ds_evals: int
+
+    total_num_evals: int
+
+    num_metric_calls_by_discovery: list[int]
+
+    full_program_trace: list[dict[str, Any]]
+    best_outputs_valset: dict[DataId, list[tuple[ProgramIdx, RolloutOutput]]] | None
+
+    validation_schema_version: int
+
+    # Optional evaluation cache for (candidate, example) pairs
+    evaluation_cache: "EvaluationCache[RolloutOutput, DataId] | None"
+
+    def __init__(
+        self,
+        seed_candidate: dict[str, str],
+        base_evaluation: ValsetEvaluation[RolloutOutput, DataId],
+        track_best_outputs: bool = False,
+        frontier_type: FrontierType = "instance",
+        evaluation_cache: "EvaluationCache[RolloutOutput, DataId] | None" = None,
+    ):
+        self.program_candidates = [dict(seed_candidate)]
+        self.prog_candidate_val_subscores = [dict(base_evaluation.scores_by_val_id)]
+
+        base_objective_aggregates = self._aggregate_objective_scores(base_evaluation.objective_scores_by_val_id)
+        self.prog_candidate_objective_scores = [base_objective_aggregates]
+
+        self.parent_program_for_candidate = [[None]]
+
+        self.frontier_type: FrontierType = frontier_type
+        self.pareto_front_valset = {val_id: score for val_id, score in base_evaluation.scores_by_val_id.items()}
+        self.program_at_pareto_front_valset = {val_id: {0} for val_id in base_evaluation.scores_by_val_id.keys()}
+        self.objective_pareto_front = dict(base_objective_aggregates)
+        self.program_at_pareto_front_objectives = {objective: {0} for objective in base_objective_aggregates.keys()}
+
+        # Validate that objective scores are provided for frontier types that require them
+        if frontier_type in ("objective", "hybrid", "cartesian"):
+            if not base_evaluation.objective_scores_by_val_id:
+                raise ValueError(
+                    f"frontier_type='{frontier_type}' requires objective_scores to be provided by the evaluator, "
+                    f"but none were found. Use an evaluator that returns objective_scores or use frontier_type='instance'."
+                )
+
+        # Cartesian frontier will be base_evaluation.objective_scores_by_val_id
+        if frontier_type == "cartesian":
+            assert base_evaluation.objective_scores_by_val_id is not None  # Already validated above
+            self.pareto_front_cartesian = {
+                (val_id, objective): objective_score
+                for val_id, objective_scores in base_evaluation.objective_scores_by_val_id.items()
+                for objective, objective_score in objective_scores.items()
+            }
+            self.program_at_pareto_front_cartesian = {
+                (val_id, objective): {0}
+                for val_id, objective_scores in base_evaluation.objective_scores_by_val_id.items()
+                for objective in objective_scores.keys()
+            }
+        else:
+            self.pareto_front_cartesian = {}
+            self.program_at_pareto_front_cartesian = {}
+
+        self.list_of_named_predictors = list(seed_candidate.keys())
+        self.named_predictor_id_to_update_next_for_program_candidate = [0]
+        self.i = -1
+
+        self.num_metric_calls_by_discovery = [0]
+
+        if track_best_outputs:
+            self.best_outputs_valset = {
+                val_id: [(0, output)] for val_id, output in base_evaluation.outputs_by_val_id.items()
+            }
+        else:
+            self.best_outputs_valset = None
+
+        self.full_program_trace = []
+        self.validation_schema_version = self._VALIDATION_SCHEMA_VERSION
+        self.evaluation_cache = evaluation_cache
+
+    def is_consistent(self) -> bool:
+        assert len(self.program_candidates) == len(self.parent_program_for_candidate)
+        assert len(self.program_candidates) == len(self.named_predictor_id_to_update_next_for_program_candidate)
+        assert len(self.program_candidates) == len(self.prog_candidate_val_subscores)
+        assert len(self.program_candidates) == len(self.prog_candidate_objective_scores)
+        assert len(self.program_candidates) == len(self.num_metric_calls_by_discovery)
+
+        assert len(self.pareto_front_valset) == len(self.program_at_pareto_front_valset)
+        assert set(self.pareto_front_valset.keys()) == set(self.program_at_pareto_front_valset.keys())
+        assert set(self.objective_pareto_front.keys()) == set(self.program_at_pareto_front_objectives.keys())
+
+        for front in self.program_at_pareto_front_valset.values():
+            for prog_idx in front:
+                assert prog_idx < len(self.program_candidates), (
+                    "Program index in valset pareto front exceeds number of program candidates"
+                )
+
+        return True
+
+    def save(self, run_dir: str | None, *, use_cloudpickle: bool = False) -> None:
+        if run_dir is None:
+            return
+        with open(os.path.join(run_dir, "gepa_state.bin"), "wb") as f:
+            if use_cloudpickle:
+                import cloudpickle as pickle  # type: ignore[import-not-found]
+            else:
+                import pickle
+            serialized = dict(self.__dict__.items())
+            serialized["validation_schema_version"] = GEPAState._VALIDATION_SCHEMA_VERSION
+            pickle.dump(serialized, f)
+
+    @staticmethod
+    def load(run_dir: str) -> "GEPAState[RolloutOutput, DataId]":
+        with open(os.path.join(run_dir, "gepa_state.bin"), "rb") as f:
+            import pickle
+
+            data = pickle.load(f)
+
+        # handle schema migration
+        version = data.get("validation_schema_version")
+        if version is None or version < 2:
+            GEPAState._migrate_from_legacy_state_v0(data)
+            version = data.get("validation_schema_version")
+        if version is None or version < GEPAState._VALIDATION_SCHEMA_VERSION:
+            GEPAState._upgrade_state_dict(data)
+
+        state = GEPAState.__new__(GEPAState)
+        state.__dict__.update(data)
+
+        state.validation_schema_version = GEPAState._VALIDATION_SCHEMA_VERSION
+        assert len(state.program_candidates) == len(state.prog_candidate_val_subscores)
+        assert len(state.program_candidates) == len(state.prog_candidate_objective_scores)
+        assert len(state.program_candidates) == len(state.num_metric_calls_by_discovery)
+        assert len(state.program_candidates) == len(state.parent_program_for_candidate)
+        assert len(state.program_candidates) == len(state.named_predictor_id_to_update_next_for_program_candidate)
+        assert len(state.pareto_front_valset) == len(state.program_at_pareto_front_valset)
+        assert set(state.pareto_front_valset.keys()) == set(state.program_at_pareto_front_valset.keys())
+        assert set(state.objective_pareto_front.keys()) == set(state.program_at_pareto_front_objectives.keys())
+        return state
+
+    @staticmethod
+    def _migrate_from_legacy_state_v0(d: dict[str, Any]) -> None:
+        assert isinstance(d, dict)
+        assert "prog_candidate_val_subscores" in d
+        assert isinstance(d["prog_candidate_val_subscores"], list)
+        assert all(isinstance(scores, list) for scores in d["prog_candidate_val_subscores"])
+        legacy_scores: list[list[float]] = d.pop("prog_candidate_val_subscores", [])
+        d["prog_candidate_val_subscores"] = [
+            {idx: score for idx, score in enumerate(scores)} for scores in legacy_scores
+        ]
+
+        pareto_front = d.get("pareto_front_valset")
+        if isinstance(pareto_front, list):
+            d["pareto_front_valset"] = {idx: score for idx, score in enumerate(pareto_front)}
+
+        program_at_front = d.get("program_at_pareto_front_valset")
+        if isinstance(program_at_front, list):
+            d["program_at_pareto_front_valset"] = {idx: set(front) for idx, front in enumerate(program_at_front)}
+
+        best_outputs = d.get("best_outputs_valset")
+        if isinstance(best_outputs, list):
+            d["best_outputs_valset"] = {idx: list(outputs) for idx, outputs in enumerate(best_outputs)}
+
+        d["validation_schema_version"] = 2
+
+    @staticmethod
+    def _upgrade_state_dict(d: dict[str, Any]) -> None:
+        num_candidates = len(d.get("program_candidates", []))
+        if "prog_candidate_objective_scores" not in d:
+            d["prog_candidate_objective_scores"] = [{} for _ in range(num_candidates)]
+        if "objective_pareto_front" not in d:
+            d["objective_pareto_front"] = {}
+        if "program_at_pareto_front_objectives" not in d:
+            d["program_at_pareto_front_objectives"] = {}
+        if "frontier_type" not in d:
+            d["frontier_type"] = "instance"
+            # Since frontier_type instance does not require "pareto_front_cartesian" and "program_at_pareto_front_cartesian", we can safely set them to empty dicts.
+            d["pareto_front_cartesian"] = {}
+            d["program_at_pareto_front_cartesian"] = {}
+        # evaluation_cache is not persisted across runs by default; initialize to None if missing
+        if "evaluation_cache" not in d:
+            d["evaluation_cache"] = None
+        d["validation_schema_version"] = GEPAState._VALIDATION_SCHEMA_VERSION
+
+    @staticmethod
+    def _aggregate_objective_scores(
+        val_objective_scores: dict[DataId, ObjectiveScores] | None,
+    ) -> ObjectiveScores:
+        if not val_objective_scores:
+            return {}
+        totals: dict[str, float] = {}
+        counts: dict[str, int] = {}
+        for objective_dict in val_objective_scores.values():
+            for objective, score in objective_dict.items():
+                totals[objective] = totals.get(objective, 0.0) + score
+                counts[objective] = counts.get(objective, 0) + 1
+        return {
+            objective: totals[objective] / counts[objective] for objective in totals.keys() if counts[objective] > 0
+        }
+
+    def get_program_average_val_subset(self, program_idx: int) -> tuple[float, int]:
+        # TODO: This should be only used/handled by the val_evaluation_policy, and never used directly.
+        scores = self.prog_candidate_val_subscores[program_idx]
+        if not scores:
+            return float("-inf"), 0
+        num_samples = len(scores)
+        avg = sum(scores.values()) / num_samples
+        return avg, num_samples
+
+    @property
+    def valset_evaluations(self) -> dict[DataId, list[ProgramIdx]]:
+        """
+        Valset examples by id and programs that have evaluated them. Keys include only validation
+        ids that have been scored at least once.
+        """
+        result: dict[DataId, list[ProgramIdx]] = defaultdict(list)
+        for program_idx, val_scores in enumerate(self.prog_candidate_val_subscores):
+            for val_id in val_scores.keys():
+                result[val_id].append(program_idx)
+        return result
+
+    @property
+    def program_full_scores_val_set(self) -> list[float]:
+        # TODO: This should be using the val_evaluation_policy instead of the get_program_average_val_subset method to calculate the scores.
+        return [
+            self.get_program_average_val_subset(program_idx)[0]
+            for program_idx in range(len(self.prog_candidate_val_subscores))
+        ]
+
+    @property
+    def per_program_tracked_scores(self) -> list[float]:
+        return [
+            self.get_program_average_val_subset(program_idx)[0]
+            for program_idx in range(len(self.prog_candidate_val_subscores))
+        ]
+
+    def _update_objective_pareto_front(self, objective_scores: ObjectiveScores, program_idx: ProgramIdx) -> None:
+        if not objective_scores:
+            return
+        for objective, score in objective_scores.items():
+            prev_score = self.objective_pareto_front.get(objective, float("-inf"))
+            if score > prev_score:
+                self.objective_pareto_front[objective] = score
+                self.program_at_pareto_front_objectives[objective] = {program_idx}
+            elif score == prev_score:
+                front = self.program_at_pareto_front_objectives.setdefault(objective, set())
+                front.add(program_idx)
+
+    def _update_pareto_front_for_val_id(
+        self,
+        val_id: DataId,
+        score: float,
+        program_idx: ProgramIdx,
+        output: RolloutOutput | None,
+        run_dir: str | None,
+        iteration: int,
+    ) -> None:
+        prev_score = self.pareto_front_valset.get(val_id, float("-inf"))
+        if score > prev_score:
+            self.pareto_front_valset[val_id] = score
+            self.program_at_pareto_front_valset[val_id] = {program_idx}
+            if self.best_outputs_valset is not None and output is not None:
+                self.best_outputs_valset[val_id] = [(program_idx, output)]
+                if run_dir is not None:
+                    task_dir = os.path.join(run_dir, "generated_best_outputs_valset", f"task_{val_id}")
+                    os.makedirs(task_dir, exist_ok=True)
+                    with open(os.path.join(task_dir, f"iter_{iteration}_prog_{program_idx}.json"), "w") as fout:
+                        json.dump(output, fout, indent=4, default=json_default)
+        elif score == prev_score:
+            pareto_front = self.program_at_pareto_front_valset.setdefault(val_id, set())
+            pareto_front.add(program_idx)
+            if self.best_outputs_valset is not None and output is not None:
+                self.best_outputs_valset[val_id].append((program_idx, output))
+
+    def _update_pareto_front_for_cartesian(
+        self,
+        val_id: DataId,
+        objective: str,
+        objective_score: float,
+        program_idx: ProgramIdx,
+    ) -> None:
+        prev_score = self.pareto_front_cartesian.get((val_id, objective), float("-inf"))
+        if objective_score > prev_score:
+            self.pareto_front_cartesian[(val_id, objective)] = objective_score
+            self.program_at_pareto_front_cartesian[(val_id, objective)] = {program_idx}
+        elif objective_score == prev_score:
+            front = self.program_at_pareto_front_cartesian.setdefault((val_id, objective), set())
+            front.add(program_idx)
+
+    def update_state_with_new_program(
+        self,
+        parent_program_idx: list[ProgramIdx],
+        new_program: dict[str, str],
+        valset_evaluation: ValsetEvaluation,
+        run_dir: str | None,
+        num_metric_calls_by_discovery_of_new_program: int,
+    ) -> ProgramIdx:
+        new_program_idx = len(self.program_candidates)
+        self.program_candidates.append(dict(new_program))
+        self.num_metric_calls_by_discovery.append(num_metric_calls_by_discovery_of_new_program)
+
+        max_predictor_id = max(
+            [self.named_predictor_id_to_update_next_for_program_candidate[p] for p in parent_program_idx],
+            default=0,
+        )
+        self.named_predictor_id_to_update_next_for_program_candidate.append(max_predictor_id)
+        self.parent_program_for_candidate.append(list(parent_program_idx))
+
+        valset_scores = dict(valset_evaluation.scores_by_val_id)
+        self.prog_candidate_val_subscores.append(valset_scores)
+        objective_scores = self._aggregate_objective_scores(valset_evaluation.objective_scores_by_val_id)
+        self.prog_candidate_objective_scores.append(objective_scores)
+
+        for val_id, score in valset_scores.items():
+            output = valset_evaluation.outputs_by_val_id.get(val_id) if valset_evaluation.outputs_by_val_id else None
+            self._update_pareto_front_for_val_id(
+                val_id,
+                score,
+                new_program_idx,
+                output,
+                run_dir,
+                self.i + 1,
+            )
+
+        self._update_objective_pareto_front(objective_scores, new_program_idx)
+
+        if self.frontier_type in ("objective", "hybrid", "cartesian"):
+            if not valset_evaluation.objective_scores_by_val_id:
+                raise ValueError(
+                    f"frontier_type='{self.frontier_type}' requires objective_scores to be provided by the evaluator, "
+                    f"but none were found in the evaluation result."
+                )
+
+        if self.frontier_type == "cartesian":
+            assert valset_evaluation.objective_scores_by_val_id is not None  # Validated above
+            for val_id, objective_scores in valset_evaluation.objective_scores_by_val_id.items():
+                for objective, objective_score in objective_scores.items():
+                    self._update_pareto_front_for_cartesian(
+                        val_id,
+                        objective,
+                        objective_score,
+                        new_program_idx,
+                    )
+
+        return new_program_idx
+
+    def _get_pareto_front_mapping(self, frontier_type: FrontierType) -> dict[FrontierKey, set[ProgramIdx]]:
+        if frontier_type == "instance":
+            return {val_id: set(front) for val_id, front in self.program_at_pareto_front_valset.items()}
+        if frontier_type == "objective":
+            return {objective: set(front) for objective, front in self.program_at_pareto_front_objectives.items()}
+        if frontier_type == "hybrid":
+            combined: dict[FrontierKey, set[ProgramIdx]] = {
+                ("val_id", val_id): set(front) for val_id, front in self.program_at_pareto_front_valset.items()
+            }
+            for objective, front in self.program_at_pareto_front_objectives.items():
+                combined[("objective", objective)] = set(front)
+            return combined
+        if frontier_type == "cartesian":
+            return {
+                ("cartesian", val_id, objective): set(front)
+                for (val_id, objective), front in self.program_at_pareto_front_cartesian.items()
+            }
+        raise ValueError(f"Unknown frontier_type: {frontier_type}")
+
+    def get_pareto_front_mapping(self) -> dict[FrontierKey, set[ProgramIdx]]:
+        """Return frontier key to best-program-indices mapping based on configured frontier_type."""
+        return self._get_pareto_front_mapping(self.frontier_type)
+
+    def cached_evaluate(
+        self,
+        candidate: dict[str, str],
+        example_ids: list[DataId],
+        fetcher: Callable[[list[DataId]], Any],
+        evaluator: Callable[[Any, dict[str, str]], tuple[Any, list[float], Sequence[ObjectiveScores] | None]],
+    ) -> tuple[list[float], int]:
+        """Evaluate with optional caching. Returns (scores, num_actual_evals)."""
+        _, scores_by_id, _, num_actual_evals = self.cached_evaluate_full(candidate, example_ids, fetcher, evaluator)
+        return [scores_by_id[eid] for eid in example_ids], num_actual_evals
+
+    def cached_evaluate_full(
+        self,
+        candidate: dict[str, str],
+        example_ids: list[DataId],
+        fetcher: Callable[[list[DataId]], Any],
+        evaluator: Callable[[Any, dict[str, str]], tuple[Any, list[float], Sequence[ObjectiveScores] | None]],
+    ) -> tuple[dict[DataId, RolloutOutput], dict[DataId, float], dict[DataId, ObjectiveScores] | None, int]:
+        """Evaluate with optional caching, returning full results."""
+        if self.evaluation_cache is not None:
+            return self.evaluation_cache.evaluate_with_cache_full(candidate, example_ids, fetcher, evaluator)
+        batch = fetcher(example_ids)
+        outputs, scores, objective_scores = evaluator(batch, candidate)
+        outputs_by_id = dict(zip(example_ids, outputs, strict=False))
+        scores_by_id = dict(zip(example_ids, scores, strict=False))
+        objective_by_id = dict(zip(example_ids, objective_scores, strict=False)) if objective_scores else None
+        return outputs_by_id, scores_by_id, objective_by_id, len(example_ids)
+
+
+def write_eval_scores_to_directory(scores: dict[DataId, float], output_dir: str) -> None:
+    for val_id, score in scores.items():
+        task_dir = os.path.join(output_dir, f"task_{val_id}")
+        os.makedirs(task_dir, exist_ok=True)
+        with open(os.path.join(task_dir, f"iter_{0}_prog_0.json"), "w") as f:
+            json.dump(score, f, indent=4, default=json_default)
+
+
+def write_eval_outputs_to_directory(outputs, output_dir: str) -> None:
+    """
+    Write generated rollout outputs (not scalar scores) to disk.
+
+    Structure:
+      {output_dir}/task_{val_id}/iter_0_prog_0.json
+
+    This directory is used to store best outputs for inspection/reuse.
+    """
+    for val_id, output in outputs.items():
+        task_dir = os.path.join(output_dir, f"task_{val_id}")
+        os.makedirs(task_dir, exist_ok=True)
+        with open(os.path.join(task_dir, "iter_0_prog_0.json"), "w") as f:
+            json.dump(output, f, indent=4, default=json_default)
+
+
+def initialize_gepa_state(
+    run_dir: str | None,
+    logger: LoggerProtocol,
+    seed_candidate: dict[str, str],
+    valset_evaluator: Callable[
+        [dict[str, str]],
+        ValsetEvaluation[RolloutOutput, DataId],
+    ],
+    track_best_outputs: bool = False,
+    frontier_type: FrontierType = "instance",
+    evaluation_cache: "EvaluationCache[RolloutOutput, DataId] | None" = None,
+) -> GEPAState[RolloutOutput, DataId]:
+    if run_dir is not None and os.path.exists(os.path.join(run_dir, "gepa_state.bin")):
+        logger.log("Loading gepa state from run dir")
+        gepa_state = GEPAState.load(run_dir)
+        if gepa_state.frontier_type != frontier_type:
+            raise ValueError(
+                f"Frontier type mismatch: requested '{frontier_type}' but loaded state has '{gepa_state.frontier_type}'. "
+                f"Use a different run_dir or match the frontier_type parameter."
+            )
+        # Sync cache with current run's cache_evaluation setting:
+        # - If caching is disabled (evaluation_cache is None), clear any loaded cache
+        #   to respect the current run's cache_evaluation=False setting
+        # - If caching is enabled and the loaded state has a cache, preserve it
+        #   (allows resuming with cached results from previous run)
+        # - If caching is enabled but no cache exists in loaded state, use the new empty cache
+        if evaluation_cache is None:
+            gepa_state.evaluation_cache = None
+        elif gepa_state.evaluation_cache is None:
+            gepa_state.evaluation_cache = evaluation_cache
+        # else: keep the loaded cache (gepa_state.evaluation_cache is already set)
+    else:
+        num_evals_run = 0
+
+        eval_result = valset_evaluator(seed_candidate)
+        if run_dir is not None:
+            write_eval_outputs_to_directory(
+                eval_result.outputs_by_val_id, os.path.join(run_dir, "generated_best_outputs_valset")
+            )
+
+        num_evals_run += len(eval_result.scores_by_val_id)
+
+        gepa_state = GEPAState(
+            seed_candidate,
+            eval_result,
+            track_best_outputs=track_best_outputs,
+            frontier_type=frontier_type,
+            evaluation_cache=evaluation_cache,
+        )
+
+        gepa_state.num_full_ds_evals = 1
+        gepa_state.total_num_evals = num_evals_run
+
+    return gepa_state

mantisdk/algorithm/gepa/lib/examples/aime.py

@@ -0,0 +1,24 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+
+def init_dataset():
+    import random
+
+    from datasets import load_dataset
+
+    train_split = [
+        {"input": x["problem"], "additional_context": {"solution": x["solution"]}, "answer": "### " + str(x["answer"])}
+        for x in load_dataset("AI-MO/aimo-validation-aime")["train"]
+    ]
+    random.Random(0).shuffle(train_split)
+    test_split = [
+        {"input": x["problem"], "answer": "### " + str(x["answer"])}
+        for x in load_dataset("MathArena/aime_2025")["train"]
+    ]
+
+    trainset = train_split[: len(train_split) // 2]
+    valset = train_split[len(train_split) // 2 :]
+    testset = test_split * 5
+
+    return trainset, valset, testset