diff --git a/openevolve/process_parallel.py b/openevolve/process_parallel.py index 33c19303e..0e89baf7b 100644 --- a/openevolve/process_parallel.py +++ b/openevolve/process_parallel.py @@ -22,6 +22,39 @@ logger = logging.getLogger(__name__) +_worker_event_loop = None + + +def _close_worker_event_loop(): + """Release the worker's loop on orderly process exit.""" + global _worker_event_loop + loop = _worker_event_loop + _worker_event_loop = None + if loop is not None and not loop.is_closed(): + try: + pending = asyncio.all_tasks(loop) + for task in pending: + task.cancel() + if pending: + loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.run_until_complete(loop.shutdown_default_executor()) + finally: + loop.close() + asyncio.set_event_loop(None) + + +def _run_worker_async(coroutine): + """Keep cached async clients on one loop for the lifetime of a worker.""" + global _worker_event_loop + if _worker_event_loop is None: + from multiprocessing.util import Finalize + + _worker_event_loop = asyncio.new_event_loop() + asyncio.set_event_loop(_worker_event_loop) + Finalize(None, _close_worker_event_loop, exitpriority=10) + return _worker_event_loop.run_until_complete(coroutine) + class _CandidateRejected(Exception): """Stop processing an iteration after admission rejects its child.""" @@ -56,6 +89,10 @@ def _worker_init(config_dict: dict, evaluation_file: str, parent_env: dict = Non global _worker_evaluator global _worker_llm_ensemble global _worker_prompt_sampler + global _worker_event_loop + + # A forked worker must not inherit a loop owned by the parent process. + _worker_event_loop = None # Store config for later use # Reconstruct Config object from nested dictionaries @@ -208,7 +245,7 @@ def _run_iteration_worker( # Generate code modification (sync wrapper for async) try: - llm_response = asyncio.run( + llm_response = _run_worker_async( _worker_llm_ensemble.generate_with_context( system_message=prompt["system"], messages=[{"role": "user", "content": prompt["user"]}], @@ -364,7 +401,7 @@ def _run_iteration_worker( import uuid child_id = str(uuid.uuid4()) - child_metrics = asyncio.run(_worker_evaluator.evaluate_program(child_code, child_id)) + child_metrics = _run_worker_async(_worker_evaluator.evaluate_program(child_code, child_id)) # Get artifacts artifacts = _worker_evaluator.get_pending_artifacts(child_id) diff --git a/tests/test_worker_loop_lifecycle.py b/tests/test_worker_loop_lifecycle.py new file mode 100644 index 000000000..8b6170cbb --- /dev/null +++ b/tests/test_worker_loop_lifecycle.py @@ -0,0 +1,107 @@ +"""Regression tests for persistent worker event-loop ownership.""" + +import asyncio +import unittest + +from openevolve import process_parallel as worker + + +class LoopBoundClient: + """Simulates an asynchronous client retaining a loop-bound resource.""" + + def __init__(self): + self.loop = None + self.calls = 0 + + async def request(self): + current = asyncio.get_running_loop() + + if self.loop is None: + self.loop = current + elif self.loop is not current: + raise RuntimeError("Event loop is closed") + + await asyncio.sleep(0) + self.calls += 1 + return self.calls + + +class WorkerLoopLifecycleTests(unittest.TestCase): + + def tearDown(self): + worker._close_worker_event_loop() + + def test_reuses_loop_across_twelve_calls(self): + client = LoopBoundClient() + + for count in range(1, 13): + self.assertEqual( + worker._run_worker_async(client.request()), + count, + ) + + self.assertIsNotNone(client.loop) + self.assertFalse(client.loop.is_closed()) + + def test_generation_and_evaluation_share_loop(self): + generation_client = LoopBoundClient() + evaluation_client = LoopBoundClient() + + for _ in range(8): + worker._run_worker_async( + generation_client.request() + ) + worker._run_worker_async( + evaluation_client.request() + ) + + self.assertIs( + generation_client.loop, + evaluation_client.loop, + ) + + def test_exception_propagates_without_closing_loop(self): + async def fail(): + raise ValueError("deliberate failure") + + with self.assertRaisesRegex( + ValueError, "deliberate failure" + ): + worker._run_worker_async(fail()) + + client = LoopBoundClient() + self.assertEqual( + worker._run_worker_async(client.request()), 1 + ) + + def test_cleanup_closes_loop_and_resets_state(self): + client = LoopBoundClient() + worker._run_worker_async(client.request()) + + loop = client.loop + worker._close_worker_event_loop() + + self.assertTrue(loop.is_closed()) + self.assertIsNone(worker._worker_event_loop) + + def test_cleanup_cancels_pending_tasks(self): + cancelled = [] + + async def pending(): + try: + await asyncio.Event().wait() + finally: + cancelled.append(True) + + async def schedule(): + asyncio.create_task(pending()) + await asyncio.sleep(0) + + worker._run_worker_async(schedule()) + worker._close_worker_event_loop() + + self.assertEqual(cancelled, [True]) + + +if __name__ == "__main__": + unittest.main()