feat(config): add checkpoint_on_improvement option to save only on new best

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

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:

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)

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"