diff --git a/openevolve/config.py b/openevolve/config.py index 560dfe058..aea6e7dc7 100644 --- a/openevolve/config.py +++ b/openevolve/config.py @@ -457,6 +457,7 @@ class Config: # General settings max_iterations: int = 10000 checkpoint_interval: int = 100 + checkpoint_on_improvement: bool = False log_level: str = "INFO" log_dir: Optional[str] = None # Optional prefix applied to all OpenEvolve logger names (e.g. "run1" logs diff --git a/openevolve/process_parallel.py b/openevolve/process_parallel.py index 33c19303e..814dcc03b 100644 --- a/openevolve/process_parallel.py +++ b/openevolve/process_parallel.py @@ -642,12 +642,21 @@ def _select_island(self, iteration: int, island_pending: Dict[int, List[int]]) - ) return int(island_id) - def _checkpoint_if_due(self, iteration: int, callback=None) -> None: - if iteration > 0 and iteration % self.config.checkpoint_interval == 0: + def _checkpoint_if_due(self, iteration: int, callback=None, is_new_best: bool = False) -> None: + # Don't checkpoint at iteration 0 (that's just the initial program) + if iteration <= 0: + return + interval_hit = iteration % self.config.checkpoint_interval == 0 + improvement_hit = self.config.checkpoint_on_improvement and is_new_best + if not (interval_hit or improvement_hit): + return + if interval_hit: logger.info("Checkpoint interval reached at iteration %d", iteration) - self.database.log_island_status() - if callback: - callback(iteration) + else: + logger.info("Checkpointing new best solution at iteration %d", iteration) + self.database.log_island_status() + if callback: + callback(iteration) async def run_evolution( self, @@ -877,13 +886,16 @@ async def run_evolution( self._warned_about_combined_score = True # Check for new best - if self.database.best_program_id == child_program.id: + is_new_best = self.database.best_program_id == child_program.id + if is_new_best: logger.info( f"🌟 New best solution found at iteration {completed_iteration}: " f"{child_program.id}" ) - self._checkpoint_if_due(completed_iteration, checkpoint_callback) + self._checkpoint_if_due( + completed_iteration, checkpoint_callback, is_new_best=is_new_best + ) # Check target score if target_score is not None and child_program.metrics: diff --git a/tests/test_process_parallel.py b/tests/test_process_parallel.py index 23c6f9219..75ff16be4 100644 --- a/tests/test_process_parallel.py +++ b/tests/test_process_parallel.py @@ -213,6 +213,73 @@ async def run_test(): # Run the async test asyncio.run(run_test()) + def test_checkpoint_on_improvement_only_fires_for_new_best(self): + """Test checkpoint callback fires on improvement when interval is not reached""" + + async def run_test(): + self.config.checkpoint_on_improvement = True + self.config.checkpoint_interval = 10000 + controller = ProcessParallelController(self.config, self.eval_file, self.database) + checkpoint_calls = [] + + with patch.object(controller, "_submit_iteration") as mock_submit: + mock_future1 = MagicMock() + mock_result1 = SerializableResult( + child_program_dict={ + "id": "child_best", + "code": "def evolved(): return 1", + "language": "python", + "parent_id": "test_0", + "generation": 1, + "metrics": {"score": 1.0, "performance": 1.0}, + "iteration_found": 1, + "metadata": {"changes": "improved", "island": 0}, + }, + parent_id="test_0", + iteration_time=0.1, + iteration=1, + target_island=0, + ) + mock_future1.done.return_value = True + mock_future1.result.return_value = mock_result1 + mock_future1.cancel.return_value = True + + mock_future2 = MagicMock() + mock_result2 = SerializableResult( + child_program_dict={ + "id": "child_not_best", + "code": "def evolved(): return 0", + "language": "python", + "parent_id": "test_0", + "generation": 1, + "metrics": {"score": 0.1, "performance": 0.1}, + "iteration_found": 2, + "metadata": {"changes": "not improved", "island": 1}, + }, + parent_id="test_0", + iteration_time=0.1, + iteration=2, + target_island=1, + ) + mock_future2.done.return_value = True + mock_future2.result.return_value = mock_result2 + mock_future2.cancel.return_value = True + + mock_submit.side_effect = [mock_future1, mock_future2] + + controller.start() + await controller.run_evolution( + start_iteration=1, + max_iterations=2, + target_score=None, + checkpoint_callback=checkpoint_calls.append, + ) + controller.stop() + + self.assertEqual(checkpoint_calls, [1]) + + asyncio.run(run_test()) + def test_request_shutdown(self): """Test graceful shutdown request""" controller = ProcessParallelController(self.config, self.eval_file, self.database) diff --git a/tests/test_worker_config_serialization.py b/tests/test_worker_config_serialization.py index 32ae5930e..f7d5abd61 100644 --- a/tests/test_worker_config_serialization.py +++ b/tests/test_worker_config_serialization.py @@ -40,6 +40,7 @@ def setUp(self): self.config = Config() self.config.max_iterations = 7 self.config.checkpoint_interval = 3 + self.config.checkpoint_on_improvement = True self.config.log_level = "DEBUG" self.config.log_dir = "/tmp/openevolve-test-logs" self.config.logger_prefix = "run1"