From 9b0ded6e1af6e405acb4df22fa0c017c45b07051 Mon Sep 17 00:00:00 2001 From: bitkira <163083654+bitkira@users.noreply.github.com> Date: Mon, 1 Jun 2026 11:04:09 +0800 Subject: [PATCH] fix(evoprompt): isolate registry state during concurrent evaluation --- evoagentx/optimizers/engine/base.py | 11 +- evoagentx/optimizers/engine/decorators.py | 8 +- evoagentx/optimizers/evoprompt_optimizer.py | 104 +++++++--- .../test_evoprompt_state_isolation.py | 185 ++++++++++++++++++ 4 files changed, 272 insertions(+), 36 deletions(-) create mode 100644 tests/src/optimizers/test_evoprompt_state_isolation.py diff --git a/evoagentx/optimizers/engine/base.py b/evoagentx/optimizers/engine/base.py index 7cd032d0..c21210f8 100644 --- a/evoagentx/optimizers/engine/base.py +++ b/evoagentx/optimizers/engine/base.py @@ -1,7 +1,12 @@ -from typing import Any, Callable, Dict, List, Optional +from __future__ import annotations + import abc +from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING + from .decorators import EntryPoint -from .registry import ParamRegistry + +if TYPE_CHECKING: + from .registry import ParamRegistry class BaseOptimizer(abc.ABC): # def __init__( @@ -67,4 +72,4 @@ def optimize(self): if self.program is None: raise RuntimeError("No entry function provided or registered.") print(f"Starting optimization from entry: {self.program.__name__}") - raise NotImplementedError \ No newline at end of file + raise NotImplementedError diff --git a/evoagentx/optimizers/engine/decorators.py b/evoagentx/optimizers/engine/decorators.py index 44baaa1d..376bc5d5 100644 --- a/evoagentx/optimizers/engine/decorators.py +++ b/evoagentx/optimizers/engine/decorators.py @@ -1,5 +1,9 @@ -from typing import Any, Callable, List, Tuple, Optional -from .registry import ParamRegistry +from __future__ import annotations + +from typing import Any, Callable, List, Tuple, Optional, TYPE_CHECKING + +if TYPE_CHECKING: + from .registry import ParamRegistry # --------- EntryPoint decorator --------- class EntryPoint: diff --git a/evoagentx/optimizers/evoprompt_optimizer.py b/evoagentx/optimizers/evoprompt_optimizer.py index 7bdeceba..caad88c8 100644 --- a/evoagentx/optimizers/evoprompt_optimizer.py +++ b/evoagentx/optimizers/evoprompt_optimizer.py @@ -12,6 +12,8 @@ # https://opensource.microsoft.com/codeofconduct/ # ----------------------------------------------------------------------------- +from __future__ import annotations + import asyncio import json import random @@ -20,19 +22,19 @@ import csv import time import itertools -from typing import Callable, Dict, List +from typing import Callable, Dict, List, TYPE_CHECKING from datetime import datetime import numpy as np from tqdm.asyncio import tqdm as aio_tqdm -import matplotlib.pyplot as plt -from evoagentx.agents import CustomizeAgent -from evoagentx.benchmark.bigbenchhard import BIGBenchHard from evoagentx.core.logging import logger -from evoagentx.models import OpenAILLMConfig from evoagentx.optimizers.engine.base import BaseOptimizer -from evoagentx.optimizers.engine.registry import ParamRegistry + +if TYPE_CHECKING: + from evoagentx.benchmark.bigbenchhard import BIGBenchHard + from evoagentx.models import OpenAILLMConfig + from evoagentx.optimizers.engine.registry import ParamRegistry class EvopromptOptimizer(BaseOptimizer): @@ -77,6 +79,7 @@ def __init__(self, self.iterations = iterations self.llm_config = llm_config self.semaphore = asyncio.Semaphore(concurrency_limit) + self._program_config_lock = asyncio.Lock() self.combination_sample_size = combination_sample_size # Logging configuration @@ -100,6 +103,7 @@ def __init__(self, self.avg_combo_scores_per_gen: Dict[str, float] = {} # Initialize paraphrase agent for prompt generation + from evoagentx.agents import CustomizeAgent self.paraphrase_agent = CustomizeAgent( name="ParaphraseAgent", description="An agent that paraphrases a given instruction.", @@ -215,6 +219,8 @@ def _log_detailed_evaluation(self, generation: int, combinations: List[Dict[str, def _create_single_metric_plot(self, metric_name: str, generations: List[int], best_scores: List[float], avg_scores: List[float], algorithm_name: str, plot_dir: str): + import matplotlib.pyplot as plt + fig, ax = plt.subplots(figsize=(12, 7)) ax.plot(generations, best_scores, marker='o', linestyle='-', linewidth=2, markersize=8, label='Best Score') ax.plot(generations, avg_scores, marker='x', linestyle='--', linewidth=2, markersize=8, label='Average Score') @@ -242,9 +248,12 @@ def _create_single_metric_plot(self, metric_name: str, generations: List[int], plt.close(fig) def _plot_and_save_performance_graph(self, algorithm_name: str): - if not self.enable_logging or plt is None: - if plt is None: - logger.warning("Matplotlib not found, skipping plot generation.") + if not self.enable_logging: + return + try: + import matplotlib.pyplot as plt + except ImportError: + logger.warning("Matplotlib not found, skipping plot generation.") return if not self.best_scores_per_gen and not self.best_combo_scores_per_gen: logger.warning("No performance data to plot.") @@ -460,8 +469,7 @@ async def _evaluate_combination_list(self, combinations: List[Dict], benchmark: all_scores = [] pbar = aio_tqdm(total=len(combinations), desc="Evaluating batch", leave=False) for combo in combinations: - tasks = [self._evaluate_combination_on_example(combo, benchmark, ex) for ex in eval_dev_set] - example_scores = await asyncio.gather(*tasks) + example_scores = await self._evaluate_combination_on_examples(combo, benchmark, eval_dev_set) avg_score = sum(example_scores) / len(example_scores) if example_scores else 0.0 all_scores.append(avg_score) pbar.update(1) @@ -507,38 +515,71 @@ def _generate_combinations(self, node_populations: Dict[str, List[str]]) -> List logger.info(f"Generated {len(sampled_combinations)} unique combinations") return sampled_combinations - async def _evaluate_combination_on_example(self, combination: Dict[str, str], - benchmark: BIGBenchHard, example: Dict) -> float: + def _get_eval_cache_key(self, combination: Dict[str, str], example: Dict): combo_key = tuple(sorted(combination.items())) example_key = str(hash(str(example))) - cache_key = hash((combo_key, example_key)) - - if not hasattr(self, '_eval_cache'): - self._eval_cache = {} + return hash((combo_key, example_key)) - if cache_key in self._eval_cache: - return self._eval_cache[cache_key] + def _cache_eval_score(self, cache_key, score: float): + self._eval_cache[cache_key] = score + if len(self._eval_cache) > 5000: + keys_to_del = list(self._eval_cache.keys())[:1000] + for key in keys_to_del: + del self._eval_cache[key] + async def _evaluate_current_config_on_example(self, benchmark: BIGBenchHard, example: Dict) -> float: async with self.semaphore: try: - original_config = self.get_current_cfg() - self.apply_cfg(combination) inputs = {k: v for k, v in example.items() if k in benchmark.get_input_keys()} prediction, _ = await asyncio.to_thread(self.program, **inputs) label = benchmark.get_label(example) score_dict = benchmark.evaluate(prediction, label) - score = score_dict.get("em", 0.0) - self.apply_cfg(original_config) - self._eval_cache[cache_key] = score - if len(self._eval_cache) > 5000: - keys_to_del = list(self._eval_cache.keys())[:1000] - for key in keys_to_del: - del self._eval_cache[key] - return score + return score_dict.get("em", 0.0) except Exception as e: logger.error(f"Error evaluating combination: {e}") return 0.0 + async def _evaluate_combination_on_examples(self, combination: Dict[str, str], + benchmark: BIGBenchHard, examples: List[Dict]) -> List[float]: + if not examples: + return [] + + if not hasattr(self, '_eval_cache'): + self._eval_cache = {} + + scores = [0.0] * len(examples) + pending_examples = [] + for idx, example in enumerate(examples): + cache_key = self._get_eval_cache_key(combination, example) + if cache_key in self._eval_cache: + scores[idx] = self._eval_cache[cache_key] + else: + pending_examples.append((idx, cache_key, example)) + + if pending_examples: + async with self._program_config_lock: + original_config = self.get_current_cfg() + try: + self.apply_cfg(combination) + tasks = [ + self._evaluate_current_config_on_example(benchmark, example) + for _, _, example in pending_examples + ] + evaluated_scores = await asyncio.gather(*tasks) + finally: + self.apply_cfg(original_config) + + for (idx, cache_key, _), score in zip(pending_examples, evaluated_scores): + scores[idx] = score + self._cache_eval_score(cache_key, score) + + return scores + + async def _evaluate_combination_on_example(self, combination: Dict[str, str], + benchmark: BIGBenchHard, example: Dict) -> float: + scores = await self._evaluate_combination_on_examples(combination, benchmark, [example]) + return scores[0] if scores else 0.0 + async def _evaluate_combinations_and_update_node_scores(self, combinations: List[Dict[str, str]], benchmark: BIGBenchHard, dev_set: list) -> List[float]: eval_dev_set = dev_set[:50] if len(dev_set) > 50 else dev_set @@ -546,8 +587,7 @@ async def _evaluate_combinations_and_update_node_scores(self, combinations: List print(f"Evaluating {len(combinations)} combinations on {len(eval_dev_set)} examples...") combo_pbar = aio_tqdm(total=len(combinations), desc="Evaluating Combinations") for combination in combinations: - tasks = [self._evaluate_combination_on_example(combination, benchmark, ex) for ex in eval_dev_set] - example_scores = await asyncio.gather(*tasks) + example_scores = await self._evaluate_combination_on_examples(combination, benchmark, eval_dev_set) avg_score = sum(example_scores) / len(example_scores) if example_scores else 0.0 combination_scores.append(avg_score) combo_pbar.update(1) @@ -687,6 +727,7 @@ def __init__(self, *args, full_evaluation: bool = False, **kwargs): logger.info(f"GAOptimizer initialized with '{mode_str}' mode.") # Initialize genetic algorithm agent for prompt evolution + from evoagentx.agents import CustomizeAgent self.ga_agent = CustomizeAgent( name="ga_agent", description="An agent that evolves a new prompt from two parent prompts.", @@ -961,6 +1002,7 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # Initialize differential evolution agent for prompt mutation + from evoagentx.agents import CustomizeAgent self.de_agent = CustomizeAgent( name="DE_Agent", description="Generates a new trial prompt using the Differential Evolution strategy.", diff --git a/tests/src/optimizers/test_evoprompt_state_isolation.py b/tests/src/optimizers/test_evoprompt_state_isolation.py new file mode 100644 index 00000000..8658b98c --- /dev/null +++ b/tests/src/optimizers/test_evoprompt_state_isolation.py @@ -0,0 +1,185 @@ +import asyncio +import threading + +import pytest + +from evoagentx.optimizers.engine.base import BaseOptimizer +from evoagentx.optimizers.evoprompt_optimizer import EvopromptOptimizer + + +class FakeBenchmark: + def get_input_keys(self): + return ["case"] + + def get_label(self, example): + return example["target"] + + def evaluate(self, prediction, label): + return {"em": float(prediction == label)} + + +class RaisingBenchmark(FakeBenchmark): + def evaluate(self, prediction, label): + raise RuntimeError("evaluation failed") + + +class SimpleProgram: + def __init__(self): + self.prompt = "base" + + def __call__(self, case): + return self.prompt, {"case": case} + + +class PromptRegistry: + def __init__(self, program): + self.program = program + self.fields = {"prompt": object()} + + def get(self, name): + assert name == "prompt" + return self.program.prompt + + def set(self, name, value): + assert name == "prompt" + self.program.prompt = value + + def names(self): + return ["prompt"] + + +class CoordinatedExampleProgram: + def __init__(self): + self._prompt = "base" + self.slow_started = threading.Event() + self.fast_returned = threading.Event() + self.base_restored_after_fast = threading.Event() + self.observed = {} + + @property + def prompt(self): + return self._prompt + + @prompt.setter + def prompt(self, value): + self._prompt = value + if value == "base" and self.fast_returned.is_set(): + self.base_restored_after_fast.set() + + def __call__(self, case): + if case == "fast": + self.slow_started.wait(timeout=1.0) + observed = self.prompt + self.observed[case] = observed + self.fast_returned.set() + return observed, {"case": case} + + if case == "slow": + self.slow_started.set() + self.fast_returned.wait(timeout=1.0) + self.base_restored_after_fast.wait(timeout=0.05) + observed = self.prompt + self.observed[case] = observed + return observed, {"case": case} + + observed = self.prompt + self.observed[case] = observed + return observed, {"case": case} + + +class CrossCombinationProgram: + def __init__(self): + self.prompt = "base" + self.combo_a_entered = threading.Event() + self.combo_b_entered = threading.Event() + self.observed = [] + + def __call__(self, case): + initial_prompt = self.prompt + if initial_prompt == "combo-a": + self.combo_a_entered.set() + self.combo_b_entered.wait(timeout=0.05) + elif initial_prompt == "combo-b": + self.combo_b_entered.set() + self.combo_a_entered.wait(timeout=0.05) + + observed = self.prompt + self.observed.append((case, initial_prompt, observed)) + return observed, {"case": case} + + +def make_optimizer(program, concurrency_limit=2): + registry = PromptRegistry(program) + + class TestOptimizer(EvopromptOptimizer): + def __init__(self): + BaseOptimizer.__init__(self, registry=registry, program=program) + self.semaphore = asyncio.Semaphore(concurrency_limit) + self._program_config_lock = asyncio.Lock() + self._eval_cache = {} + + async def optimize(self): + return None + + return TestOptimizer() + + +@pytest.mark.asyncio +async def test_combination_config_is_kept_until_all_examples_finish(): + program = CoordinatedExampleProgram() + optimizer = make_optimizer(program, concurrency_limit=2) + benchmark = FakeBenchmark() + + scores = await optimizer._evaluate_combination_on_examples( + {"prompt": "optimized"}, + benchmark, + [ + {"case": "fast", "target": "optimized"}, + {"case": "slow", "target": "optimized"}, + ], + ) + + assert scores == [1.0, 1.0] + assert program.observed == {"fast": "optimized", "slow": "optimized"} + assert program.prompt == "base" + + +@pytest.mark.asyncio +async def test_combination_config_is_restored_after_evaluation_error(): + program = SimpleProgram() + optimizer = make_optimizer(program) + + scores = await optimizer._evaluate_combination_on_examples( + {"prompt": "optimized"}, + RaisingBenchmark(), + [{"case": "single", "target": "optimized"}], + ) + + assert scores == [0.0] + assert program.prompt == "base" + + +@pytest.mark.asyncio +async def test_concurrent_combinations_do_not_share_temporary_config(): + program = CrossCombinationProgram() + optimizer = make_optimizer(program, concurrency_limit=2) + benchmark = FakeBenchmark() + + scores_a, scores_b = await asyncio.gather( + optimizer._evaluate_combination_list( + [{"prompt": "combo-a"}], + benchmark, + [{"case": "a", "target": "combo-a"}], + ), + optimizer._evaluate_combination_list( + [{"prompt": "combo-b"}], + benchmark, + [{"case": "b", "target": "combo-b"}], + ), + ) + + assert scores_a == [1.0] + assert scores_b == [1.0] + assert ("a", "combo-a", "combo-a") in program.observed + assert ("b", "combo-b", "combo-b") in program.observed + assert program.prompt == "base"