From e2ca1495b5e35da7ae33a3c2293f237bbd9ce1b8 Mon Sep 17 00:00:00 2001 From: Tim Wang Date: Wed, 2 Sep 2026 00:50:46 +0000 Subject: [PATCH 1/2] fix(mp): release stage thread pools so evaluate() stops leaking workers MPRunner had no way to release its ThreadPoolExecutor. CPython pool workers have no idle timeout, so once started they park on the pool's work queue until the interpreter exits. Every evaluator builds its pools inside evaluate(), and every orchestrator builds a fresh evaluator per call, so each call added threads to a process that never restarts between requests in the server. StreamingOrchestrator is the worst case: eval_service hands each arriving input to evaluate_item, which constructs a new Evaluator and therefore four new pools for a dataset of one item. Measured on the streaming shape, 500 eval cases leaked 2000 threads and 47 MB of RSS; with the pools released it is 0 threads and 96 KB. A long-lived pod reaches RLIMIT_NPROC after roughly 33k cases and then fails to start new threads. Add MPRunner.shutdown() plus context-manager support, and release the pools from a finally block in Evaluator, AgentEvaluator, CortadoEvaluator, and DataEngineeringAgentEvaluator. The three agent evaluators now build their runner in evaluate() rather than __init__, since a runner is spent once shut down. Evaluator's sqlexec pool keeps its queued work instead of cancelling it. The dispatch loop takes a connection off db_queue before submitting each SQLExecWork and only SQLExecWork.run returns it, so a cancelled item would strand one; StreamingOrchestrator shares a single db_queue across inputs, where that would shrink the pool permanently. --- evalbench/evaluator/agentevaluator.py | 41 +- evalbench/evaluator/cortadoevaluator.py | 57 +-- .../dataengineeringagentevaluator.py | 10 +- evalbench/evaluator/evaluator.py | 387 ++++++++++-------- evalbench/mp/mprunner.py | 43 +- evalbench/test/mprunner_test.py | 210 ++++++++++ 6 files changed, 514 insertions(+), 234 deletions(-) create mode 100644 evalbench/test/mprunner_test.py diff --git a/evalbench/evaluator/agentevaluator.py b/evalbench/evaluator/agentevaluator.py index 0272badc..2edaed97 100644 --- a/evalbench/evaluator/agentevaluator.py +++ b/evalbench/evaluator/agentevaluator.py @@ -48,7 +48,6 @@ def __init__( runner_config = self.config.get("runners", {}) self.agent_runners = runner_config.get("agent_runners", 10) - self.agentrunner = mprunner.MPRunner(self.agent_runners) def evaluate( self, @@ -73,8 +72,6 @@ def _evaluate_agent_cli( generator_name = type(self.generator).__name__ logging.info(f"Running {generator_name} evaluation") - self.agentrunner.futures.clear() - # Extract generic metadata metadata = { "dialects": self.config.get("dialects", []), @@ -82,26 +79,30 @@ def _evaluate_agent_cli( "scorers": self.config.get("scorers", {}), } - for item in dataset: - simulated_user = SimulatedUser(self.config) - work = AgentGenWork( - processor=self.process_scenario, - eval_result=item, - job_id=job_id, - metadata=metadata, - simulated_user=simulated_user - ) - self.agentrunner.execute_work(work) + self.agentrunner = mprunner.MPRunner(self.agent_runners) + try: + for item in dataset: + simulated_user = SimulatedUser(self.config) + work = AgentGenWork( + processor=self.process_scenario, + eval_result=item, + job_id=job_id, + metadata=metadata, + simulated_user=simulated_user + ) + self.agentrunner.execute_work(work) - for future in concurrent.futures.as_completed(self.agentrunner.futures): - item = future.result() + for future in concurrent.futures.as_completed(self.agentrunner.futures): + item = future.result() - if hasattr(item, "agent_results"): - eval_outputs.extend(item.agent_results) - if hasattr(item, "scoring_results"): - scoring_results.extend(item.scoring_results) + if hasattr(item, "agent_results"): + eval_outputs.extend(item.agent_results) + if hasattr(item, "scoring_results"): + scoring_results.extend(item.scoring_results) - return eval_outputs, scoring_results + return eval_outputs, scoring_results + finally: + self.agentrunner.shutdown() def process_scenario( self, diff --git a/evalbench/evaluator/cortadoevaluator.py b/evalbench/evaluator/cortadoevaluator.py index 64ceb803..b58bf5e1 100644 --- a/evalbench/evaluator/cortadoevaluator.py +++ b/evalbench/evaluator/cortadoevaluator.py @@ -35,46 +35,47 @@ def __init__(self, config): runner_config = self.config.get("runners", {}) self.agent_runners = runner_config.get("agent_runners", 10) - self.agentrunner = mprunner.MPRunner(self.agent_runners) def evaluate(self, dataset: List[EvalCortadoRequest], job_id: str, run_time: datetime.datetime): eval_outputs: List[Any] = [] scoring_results: List[Any] = [] logging.info("Running Cortado gRPC evaluation") - self.agentrunner.futures.clear() - metadata = { "dialects": self.config.get("dialects", []), "database": self.config.get("database", "unknown"), "scorers": self.config.get("scorers", {}), } - # Spin up threads for concurrent conversation processing - for item in dataset: - simulated_user = SimulatedUser(self.config) - work = AgentGenWork( - processor=self.process_scenario, - eval_result=item, - job_id=job_id, - metadata=metadata, - simulated_user=simulated_user - ) - self.agentrunner.execute_work(work) - - for future in concurrent.futures.as_completed(self.agentrunner.futures): - try: - # This now contains the returned object from process_scenario - modified_item = future.result() - if hasattr(modified_item, "agent_results"): - eval_outputs.extend(modified_item.agent_results) - if hasattr(modified_item, "scoring_results"): - scoring_results.extend(modified_item.scoring_results) - except Exception as e: - logging.error( - f"Error getting result from future: {e}", exc_info=True) - - return eval_outputs, scoring_results + self.agentrunner = mprunner.MPRunner(self.agent_runners) + try: + # Spin up threads for concurrent conversation processing + for item in dataset: + simulated_user = SimulatedUser(self.config) + work = AgentGenWork( + processor=self.process_scenario, + eval_result=item, + job_id=job_id, + metadata=metadata, + simulated_user=simulated_user + ) + self.agentrunner.execute_work(work) + + for future in concurrent.futures.as_completed(self.agentrunner.futures): + try: + # This now contains the returned object from process_scenario + modified_item = future.result() + if hasattr(modified_item, "agent_results"): + eval_outputs.extend(modified_item.agent_results) + if hasattr(modified_item, "scoring_results"): + scoring_results.extend(modified_item.scoring_results) + except Exception as e: + logging.error( + f"Error getting result from future: {e}", exc_info=True) + + return eval_outputs, scoring_results + finally: + self.agentrunner.shutdown() def process_scenario( self, scenario: Dict[str, Any], eval_result: Any, job_id: str, diff --git a/evalbench/evaluator/dataengineeringagentevaluator.py b/evalbench/evaluator/dataengineeringagentevaluator.py index f2b97371..6a095ce9 100644 --- a/evalbench/evaluator/dataengineeringagentevaluator.py +++ b/evalbench/evaluator/dataengineeringagentevaluator.py @@ -55,7 +55,6 @@ def __init__(self, config: dict[str, Any]) -> None: runner_config = self.config.get("runners", {}) self.agent_runners = runner_config.get("agent_runners", 10) - self.agentrunner = mprunner.MPRunner(self.agent_runners) def _get_session_dir(self, job_id: str) -> str: """Resolves the session directory path for a given job ID.""" @@ -119,7 +118,7 @@ def evaluate( "'dataform_workspace', 'gcp_project_id', and 'gcp_region' in your run config." ) - self.agentrunner.futures.clear() + self.agentrunner = mprunner.MPRunner(self.agent_runners) metadata = { "dialects": self.config.get("dialects", []), @@ -150,6 +149,13 @@ def evaluate( except Exception as e: logger.exception(f"Error getting result from future: {e}") finally: + # Shut down before archiving. On the normal path as_completed has + # already drained every scenario, so this cancels nothing. On the + # error path it stops scenarios that are still queued from mutating + # the Dataform workspace while _archive_workspace_to_gcs zips it. A + # scenario already running is not interrupted, so the archive can + # still catch one mid-write. + self.agentrunner.shutdown() self._archive_workspace_to_gcs(workspace_uri, job_id, dataset) return eval_outputs, scoring_results diff --git a/evalbench/evaluator/evaluator.py b/evalbench/evaluator/evaluator.py index 0b89ee35..c2f9bb31 100644 --- a/evalbench/evaluator/evaluator.py +++ b/evalbench/evaluator/evaluator.py @@ -96,209 +96,232 @@ def evaluate( self.genrunner = mprunner.MPRunner(self.sqlgen_runners) self.sqlrunner = mprunner.MPRunner(self.sqlexec_runners) self.scoringrunner = mprunner.MPRunner(self.scoring_runners) - prompt_generator.setup() - - self.promptrunner.futures.clear() - self.genrunner.futures.clear() - self.sqlrunner.futures.clear() - self.scoringrunner.futures.clear() - - prompt_future_to_eval = {} - prompt_future_to_input = {} - for eval_input in dataset: - eval_output = EvalOutput(eval_input) - eval_output["job_id"] = job_id - eval_output["run_time"] = run_time - work = promptgenwork.SQLPromptGenWork( - prompt_generator, eval_output) - self.promptrunner.execute_work(work) - prompt_future_to_eval[self.promptrunner.futures[-1]] = eval_output - prompt_future_to_input[self.promptrunner.futures[-1]] = eval_input - - gen_future_to_eval = {} - for future, eval_output, timed_out in _process_futures_with_timeout( - self.promptrunner.futures, - prompt_future_to_eval, - timeout=self.task_timeout_seconds, - ): - if timed_out: - eval_output["prompt_generator_error"] = ( - "TimeoutError: Task hung for too long." - ) - else: - try: - future.result() - except Exception as e: + try: + prompt_generator.setup() + + self.promptrunner.futures.clear() + self.genrunner.futures.clear() + self.sqlrunner.futures.clear() + self.scoringrunner.futures.clear() + + prompt_future_to_eval = {} + prompt_future_to_input = {} + for eval_input in dataset: + eval_output = EvalOutput(eval_input) + eval_output["job_id"] = job_id + eval_output["run_time"] = run_time + work = promptgenwork.SQLPromptGenWork( + prompt_generator, eval_output) + self.promptrunner.execute_work(work) + prompt_future_to_eval[self.promptrunner.futures[-1]] = eval_output + prompt_future_to_input[self.promptrunner.futures[-1]] = eval_input + + gen_future_to_eval = {} + for future, eval_output, timed_out in _process_futures_with_timeout( + self.promptrunner.futures, + prompt_future_to_eval, + timeout=self.task_timeout_seconds, + ): + if timed_out: + eval_output["prompt_generator_error"] = ( + "TimeoutError: Task hung for too long." + ) + else: + try: + future.result() + except Exception as e: - logging.error(f"Promptgen future error: {e}") - eval_output["prompt_generator_error"] = str(e) - - record_successful_prompt_gen(progress_reporting) - - eval_input = prompt_future_to_input[future] - - query_type = eval_output.get("query_type", "dql").lower() - trials_to_run = self.num_trials if query_type == "dql" else 1 - for trial_idx in range(trials_to_run): - trial_output = EvalOutput(eval_input) - trial_output["job_id"] = job_id - trial_output["run_time"] = run_time - trial_output.update(eval_output) - - trial_output["prompt_id"] = eval_output["id"] - trial_output["trial_index"] = trial_idx - trial_output["id"] = f"{eval_output['id']}_trial_{trial_idx}" - - work = sqlgenwork.SQLGenWork(model_generator, trial_output) - self.genrunner.execute_work(work) - gen_future_to_eval[self.genrunner.futures[-1]] = trial_output - - exec_future_to_eval = {} - score_future_to_eval = {} - for future, eval_output, timed_out in _process_futures_with_timeout( - self.genrunner.futures, - gen_future_to_eval, - timeout=self.task_timeout_seconds, - ): - if timed_out: - eval_output["sql_generator_error"] = ( - "TimeoutError: Task hung for too long." - ) - else: - try: - future.result() - except Exception as e: + logging.error(f"Promptgen future error: {e}") + eval_output["prompt_generator_error"] = str(e) - logging.error(f"SQLgen future error: {e}") - eval_output["sql_generator_error"] = str(e) + record_successful_prompt_gen(progress_reporting) - record_successful_sql_gen(progress_reporting) + eval_input = prompt_future_to_input[future] - try: - db_conn = db_queue.get(timeout=180) - work = sqlexecwork.SQLExecWork( - db_conn, self.config, eval_output, db_queue - ) - self.sqlrunner.execute_work(work) - exec_future_to_eval[self.sqlrunner.futures[-1]] = eval_output - except queue.Empty: - error_msg = f"Timeout Error: Waited too long (queue.Empty) for database '{eval_output.get('database', 'unknown')}'" - logging.error(error_msg) - eval_output["generated_error"] = error_msg + query_type = eval_output.get("query_type", "dql").lower() + trials_to_run = self.num_trials if query_type == "dql" else 1 + for trial_idx in range(trials_to_run): + trial_output = EvalOutput(eval_input) + trial_output["job_id"] = job_id + trial_output["run_time"] = run_time + trial_output.update(eval_output) - record_successful_sql_exec(progress_reporting) - work = scorework.ScorerWork( - self.config, eval_output, scoring_results, global_models - ) - self.scoringrunner.execute_work(work) - score_future_to_eval[self.scoringrunner.futures[-1] - ] = eval_output - except Exception as e: - exc_msg = str(e) or type(e).__name__ - logging.error( - "Failed to acquire DB connection from queue for database" - f" '{eval_output.get('database')}': {exc_msg}" - ) - eval_output["generated_error"] = f"Failed to acquire DB connection: {exc_msg}" - record_successful_sql_exec(progress_reporting) - work = scorework.ScorerWork( - self.config, eval_output, scoring_results, global_models - ) - self.scoringrunner.execute_work(work) - score_future_to_eval[self.scoringrunner.futures[-1] - ] = eval_output - - for future, eval_output, timed_out in _process_futures_with_timeout( - self.sqlrunner.futures, - exec_future_to_eval, - timeout=self.task_timeout_seconds, - ): - if timed_out: - eval_output["generated_error"] = "TimeoutError: Task hung for too long." - else: - try: - future.result() - except Exception as e: + trial_output["prompt_id"] = eval_output["id"] + trial_output["trial_index"] = trial_idx + trial_output["id"] = f"{eval_output['id']}_trial_{trial_idx}" - logging.error(f"SQLExec future error: {e}") - eval_output["generated_error"] = str(e) + work = sqlgenwork.SQLGenWork(model_generator, trial_output) + self.genrunner.execute_work(work) + gen_future_to_eval[self.genrunner.futures[-1]] = trial_output + + exec_future_to_eval = {} + score_future_to_eval = {} + for future, eval_output, timed_out in _process_futures_with_timeout( + self.genrunner.futures, + gen_future_to_eval, + timeout=self.task_timeout_seconds, + ): + if timed_out: + eval_output["sql_generator_error"] = ( + "TimeoutError: Task hung for too long." + ) + else: + try: + future.result() + except Exception as e: + + logging.error(f"SQLgen future error: {e}") + eval_output["sql_generator_error"] = str(e) + + record_successful_sql_gen(progress_reporting) - record_successful_sql_exec(progress_reporting) - work = scorework.ScorerWork( - self.config, eval_output, scoring_results, global_models - ) - self.scoringrunner.execute_work(work) - score_future_to_eval[self.scoringrunner.futures[-1]] = eval_output - - for future, eval_output, timed_out in _process_futures_with_timeout( - self.scoringrunner.futures, - score_future_to_eval, - timeout=self.task_timeout_seconds, - ): - if timed_out: - eval_output["scoring_error"] = "TimeoutError: Task hung for too long." - else: try: - future.result() + db_conn = db_queue.get(timeout=180) + work = sqlexecwork.SQLExecWork( + db_conn, self.config, eval_output, db_queue + ) + self.sqlrunner.execute_work(work) + exec_future_to_eval[self.sqlrunner.futures[-1]] = eval_output + except queue.Empty: + error_msg = f"Timeout Error: Waited too long (queue.Empty) for database '{eval_output.get('database', 'unknown')}'" + logging.error(error_msg) + eval_output["generated_error"] = error_msg + + record_successful_sql_exec(progress_reporting) + work = scorework.ScorerWork( + self.config, eval_output, scoring_results, global_models + ) + self.scoringrunner.execute_work(work) + score_future_to_eval[self.scoringrunner.futures[-1] + ] = eval_output except Exception as e: + exc_msg = str(e) or type(e).__name__ + logging.error( + "Failed to acquire DB connection from queue for database" + f" '{eval_output.get('database')}': {exc_msg}" + ) + eval_output["generated_error"] = f"Failed to acquire DB connection: {exc_msg}" + record_successful_sql_exec(progress_reporting) + work = scorework.ScorerWork( + self.config, eval_output, scoring_results, global_models + ) + self.scoringrunner.execute_work(work) + score_future_to_eval[self.scoringrunner.futures[-1] + ] = eval_output + + for future, eval_output, timed_out in _process_futures_with_timeout( + self.sqlrunner.futures, + exec_future_to_eval, + timeout=self.task_timeout_seconds, + ): + if timed_out: + eval_output["generated_error"] = "TimeoutError: Task hung for too long." + else: + try: + future.result() + except Exception as e: - logging.error(f"Scoring future error: {e}") - eval_output["scoring_error"] = str(e) + logging.error(f"SQLExec future error: {e}") + eval_output["generated_error"] = str(e) - record_successful_scoring(progress_reporting) - try: - truncateExecutionOutputs( - eval_output, - self.config, - ) - except Exception as e: - - logging.error(f"Truncation error: {e}") - eval_outputs.append(eval_output) - - if self.num_trials > 1: - grouped_trials = collections.defaultdict(list) - for eo in eval_outputs: - prompt_id = eo.get("prompt_id") - if prompt_id: - grouped_trials[prompt_id].append(eo) - - multi_trial_futures = {} - for prompt_id, trials in grouped_trials.items(): - nl_prompt = trials[0].get("nl_prompt", "") - work = multi_trial_scorework.MultiTrialScorerWork( - prompt_id, - nl_prompt, - trials, - self.config, - multi_trial_scoring_results, - global_models, - progress_reporting, + record_successful_sql_exec(progress_reporting) + work = scorework.ScorerWork( + self.config, eval_output, scoring_results, global_models ) self.scoringrunner.execute_work(work) - multi_trial_futures[self.scoringrunner.futures[-1]] = trials[0] + score_future_to_eval[self.scoringrunner.futures[-1]] = eval_output - for future, _, timed_out in _process_futures_with_timeout( - list(multi_trial_futures.keys()), - multi_trial_futures, + for future, eval_output, timed_out in _process_futures_with_timeout( + self.scoringrunner.futures, + score_future_to_eval, timeout=self.task_timeout_seconds, ): if timed_out: - logging.error("Multi-trial scoring timed out.") + eval_output["scoring_error"] = "TimeoutError: Task hung for too long." else: try: future.result() except Exception as e: - logging.error(f"Multi-trial scoring future error: {e}") - if close_connections and db_queue: - while True: + logging.error(f"Scoring future error: {e}") + eval_output["scoring_error"] = str(e) + + record_successful_scoring(progress_reporting) try: - db = db_queue.get(block=False) - db.close_connections() - except queue.Empty: - break - except Exception: - break + truncateExecutionOutputs( + eval_output, + self.config, + ) + except Exception as e: - return eval_outputs, scoring_results, multi_trial_scoring_results + logging.error(f"Truncation error: {e}") + eval_outputs.append(eval_output) + + if self.num_trials > 1: + grouped_trials = collections.defaultdict(list) + for eo in eval_outputs: + prompt_id = eo.get("prompt_id") + if prompt_id: + grouped_trials[prompt_id].append(eo) + + multi_trial_futures = {} + for prompt_id, trials in grouped_trials.items(): + nl_prompt = trials[0].get("nl_prompt", "") + work = multi_trial_scorework.MultiTrialScorerWork( + prompt_id, + nl_prompt, + trials, + self.config, + multi_trial_scoring_results, + global_models, + progress_reporting, + ) + self.scoringrunner.execute_work(work) + multi_trial_futures[self.scoringrunner.futures[-1]] = trials[0] + + for future, _, timed_out in _process_futures_with_timeout( + list(multi_trial_futures.keys()), + multi_trial_futures, + timeout=self.task_timeout_seconds, + ): + if timed_out: + logging.error("Multi-trial scoring timed out.") + else: + try: + future.result() + except Exception as e: + logging.error(f"Multi-trial scoring future error: {e}") + + if close_connections and db_queue: + while True: + try: + db = db_queue.get(block=False) + db.close_connections() + except queue.Empty: + break + except Exception: + break + + return eval_outputs, scoring_results, multi_trial_scoring_results + finally: + self._shutdown_runners() + + def _shutdown_runners(self) -> None: + """Releases the worker threads owned by every stage runner. + + `evaluate` is called once per (dialect, database, query type) + sub-dataset and builds a fresh pool per stage, so leaving these pools + open would leak `promptgen + sqlgen + sqlexec + scoring` worker threads + per call for the remaining lifetime of the process. + + The sqlexec pool keeps its queued work rather than cancelling it. This + loop takes a connection off `db_queue` before submitting each + `SQLExecWork`, and only `SQLExecWork.run` returns that connection, so a + cancelled item would strand one. `StreamingOrchestrator` reuses a + single `db_queue` across eval inputs, where a stranded connection + shrinks the pool for every later input. Letting the queued items run + still releases the worker threads, just after the queue drains. + """ + for runner in (self.promptrunner, self.genrunner, self.scoringrunner): + runner.shutdown() + self.sqlrunner.shutdown(cancel_futures=False) diff --git a/evalbench/mp/mprunner.py b/evalbench/mp/mprunner.py index 91168fc4..3ad11e3f 100644 --- a/evalbench/mp/mprunner.py +++ b/evalbench/mp/mprunner.py @@ -23,10 +23,23 @@ def do_work(work_obj: work.Work, item_config: Any = None) -> Any: class MPRunner: """Multi-processing class that implements threadpool execution of work. + The runner owns a `ThreadPoolExecutor`, whose worker threads have no idle + timeout: once started they stay alive, blocked on the pool's internal work + queue, until the pool is shut down or the interpreter exits. Callers must + therefore release the runner when they are done with it, either explicitly + via `shutdown()` or by using it as a context manager:: + + with MPRunner(10) as runner: + runner.execute_work(work_obj) + ... + + Runners that are created per sub-dataset and never released leak their + worker threads for the remaining lifetime of the process. + Attributes: - executor: - futures: + executor: The thread pool backing this runner. + futures: The futures of every work item submitted so far. """ def __init__(self, concurrent_tests: int = 10) -> None: @@ -46,3 +59,29 @@ def execute_work(self, work_obj: work.Work) -> None: """ ctx = contextvars.copy_context() self.futures.append(self.executor.submit(ctx.run, do_work, work_obj)) + + def shutdown(self, wait: bool = False, cancel_futures: bool = True) -> None: + """Release the pool's worker threads. + + Idle workers exit as soon as they pick up the shutdown sentinel. A + worker that is still executing a work item exits once that item + returns. The runner is spent afterwards: `execute_work` raises + `RuntimeError`, so a caller that runs more than once needs a fresh + runner per run. + + Args: + wait: Whether to block until every running work item has finished. + Defaults to False so a work item that has hung (and that the caller + has already abandoned, e.g. after a stage timeout) cannot stall the + rest of the evaluation. + cancel_futures: Whether to cancel work items that are queued but have + not started running. Defaults to True, since the caller is done + with this runner and any remaining queued work is dead work. + """ + self.executor.shutdown(wait=wait, cancel_futures=cancel_futures) + + def __enter__(self) -> "MPRunner": + return self + + def __exit__(self, exc_type, exc_value, exc_traceback) -> None: + self.shutdown() diff --git a/evalbench/test/mprunner_test.py b/evalbench/test/mprunner_test.py new file mode 100644 index 00000000..9c0e0497 --- /dev/null +++ b/evalbench/test/mprunner_test.py @@ -0,0 +1,210 @@ +import concurrent.futures +import threading +import time +import unittest + +from mp import mprunner +from work.work import Work + + +class _RecordingWork(Work): + """Work item that records the thread it ran on.""" + + def __init__(self, threads: list, barrier: threading.Event | None = None): + self.threads = threads + self.barrier = barrier + self.started = threading.Event() + + def run(self, work_config=None): + self.started.set() + if self.barrier is not None: + self.barrier.wait() + self.threads.append(threading.current_thread()) + return "done" + + +def _live(threads) -> int: + return sum(1 for t in threads if t.is_alive()) + + +class TestMPRunnerShutdown(unittest.TestCase): + + def test_workers_stay_alive_until_shutdown(self): + """Idle workers persist after their work finishes, and exit on shutdown.""" + threads = [] + runner = mprunner.MPRunner(3) + for _ in range(3): + runner.execute_work(_RecordingWork(threads)) + concurrent.futures.wait(runner.futures, timeout=30) + + self.assertEqual(len(threads), 3) + worker_threads = set(threads) + # All work is complete, but the pool keeps its workers parked. + self.assertEqual(_live(worker_threads), len(worker_threads)) + + runner.shutdown() + for t in worker_threads: + t.join(timeout=30) + self.assertEqual(_live(worker_threads), 0) + + def test_context_manager_shuts_down_on_exit(self): + threads = [] + with mprunner.MPRunner(2) as runner: + for _ in range(2): + runner.execute_work(_RecordingWork(threads)) + concurrent.futures.wait(runner.futures, timeout=30) + + for t in set(threads): + t.join(timeout=30) + self.assertEqual(_live(set(threads)), 0) + + def test_context_manager_shuts_down_on_exception(self): + threads = [] + runner = mprunner.MPRunner(2) + with self.assertRaises(ValueError): + with runner: + runner.execute_work(_RecordingWork(threads)) + concurrent.futures.wait(runner.futures, timeout=30) + raise ValueError("boom") + + for t in set(threads): + t.join(timeout=30) + self.assertEqual(_live(set(threads)), 0) + + def test_shutdown_does_not_block_on_hung_work(self): + """A hung work item must not stall shutdown of the rest of the pool.""" + release = threading.Event() + threads = [] + runner = mprunner.MPRunner(2) + hung = _RecordingWork(threads, barrier=release) + runner.execute_work(hung) + self.assertTrue(hung.started.wait(timeout=30)) + + try: + start = time.monotonic() + runner.shutdown() + elapsed = time.monotonic() - start + self.assertLess(elapsed, 5.0) + self.assertFalse(runner.futures[0].done()) + finally: + release.set() + concurrent.futures.wait(runner.futures, timeout=30) + + def test_shutdown_cancels_queued_work(self): + """Work still queued behind a busy worker is cancelled, not run.""" + release = threading.Event() + threads = [] + runner = mprunner.MPRunner(1) + blocker = _RecordingWork(threads, barrier=release) + runner.execute_work(blocker) + self.assertTrue(blocker.started.wait(timeout=30)) + # Second item cannot start: the single worker is blocked on the first. + runner.execute_work(_RecordingWork(threads)) + + try: + runner.shutdown() + self.assertTrue(runner.futures[1].cancelled()) + finally: + release.set() + concurrent.futures.wait(runner.futures, timeout=30) + self.assertEqual(len(threads), 1) + + def test_shutdown_keeps_queued_work_when_cancel_disabled(self): + """`cancel_futures=False` lets queued work run so it can free resources.""" + release = threading.Event() + threads = [] + runner = mprunner.MPRunner(1) + blocker = _RecordingWork(threads, barrier=release) + runner.execute_work(blocker) + self.assertTrue(blocker.started.wait(timeout=30)) + runner.execute_work(_RecordingWork(threads)) + + runner.shutdown(cancel_futures=False) + self.assertFalse(runner.futures[1].cancelled()) + release.set() + concurrent.futures.wait(runner.futures, timeout=30) + self.assertEqual(len(threads), 2) + + +class TestEvaluatorReleasesRunners(unittest.TestCase): + """The stage pools are built per sub-dataset, so they must be released.""" + + def _runner_names(self): + return ["promptrunner", "genrunner", "sqlrunner", "scoringrunner"] + + def test_evaluate_shuts_down_every_stage_runner(self): + from unittest.mock import MagicMock, patch + + from evaluator.evaluator import Evaluator + + created = [] + + def make_runner(*args, **kwargs): + runner = MagicMock() + runner.futures = [] + created.append(runner) + return runner + + with patch("evaluator.evaluator.mprunner.MPRunner", side_effect=make_runner): + evaluator = Evaluator({"runners": {}}) + prompt_generator = MagicMock() + evaluator.evaluate( + dataset=[], + db_queue=None, + prompt_generator=prompt_generator, + model_generator=MagicMock(), + job_id="job", + run_time=None, + progress_reporting=None, + global_models={}, + ) + + self.assertEqual(len(created), 4) + for runner in created: + runner.shutdown.assert_called_once() + # Pools are built in stage order, so created[2] is the sqlexec pool. + # Its queued work holds DB connections that only SQLExecWork.run + # returns, so it must not be cancelled. + created[2].shutdown.assert_called_once_with(cancel_futures=False) + + def test_evaluate_shuts_down_runners_when_pipeline_raises(self): + from unittest.mock import MagicMock, patch + + from evaluator.evaluator import Evaluator + + created = [] + + def make_runner(*args, **kwargs): + runner = MagicMock() + runner.futures = [] + created.append(runner) + return runner + + prompt_generator = MagicMock() + prompt_generator.setup.side_effect = RuntimeError("setup failed") + + with patch("evaluator.evaluator.mprunner.MPRunner", side_effect=make_runner): + evaluator = Evaluator({"runners": {}}) + with self.assertRaises(RuntimeError): + evaluator.evaluate( + dataset=[], + db_queue=None, + prompt_generator=prompt_generator, + model_generator=MagicMock(), + job_id="job", + run_time=None, + progress_reporting=None, + global_models={}, + ) + + self.assertEqual(len(created), 4) + for runner in created: + runner.shutdown.assert_called_once() + # Pools are built in stage order, so created[2] is the sqlexec pool. + # Its queued work holds DB connections that only SQLExecWork.run + # returns, so it must not be cancelled. + created[2].shutdown.assert_called_once_with(cancel_futures=False) + + +if __name__ == "__main__": + unittest.main() From 9e3078ea3514d3da4e26fc1475607e462f137ef9 Mon Sep 17 00:00:00 2001 From: Tim Wang Date: Wed, 2 Sep 2026 01:27:59 +0000 Subject: [PATCH 2/2] fix(mp): give MPRunner.__exit__ the standard Executor semantics Exiting a `with` block called shutdown() with its own defaults, which cancel queued work and return without waiting. That inverts Executor.__exit__ and breaks the very pattern the class docstring advertises: a caller who submits more items than the pool has workers would lose the queued ones silently. Exit now drains and waits; a caller that needs to abandon work in progress calls shutdown() directly. Also stop test_shutdown_cancels_queued_work from waiting on a cancelled future. Cancelling through shutdown() leaves the future CANCELLED rather than CANCELLED_AND_NOTIFIED, so concurrent.futures.wait never sees it finish and burned the full 30s timeout. The suite drops from 45s to 14s. --- evalbench/mp/mprunner.py | 17 ++++++++++++++--- evalbench/test/mprunner_test.py | 19 +++++++++++++------ 2 files changed, 27 insertions(+), 9 deletions(-) diff --git a/evalbench/mp/mprunner.py b/evalbench/mp/mprunner.py index 3ad11e3f..aaa27c04 100644 --- a/evalbench/mp/mprunner.py +++ b/evalbench/mp/mprunner.py @@ -26,13 +26,17 @@ class MPRunner: The runner owns a `ThreadPoolExecutor`, whose worker threads have no idle timeout: once started they stay alive, blocked on the pool's internal work queue, until the pool is shut down or the interpreter exits. Callers must - therefore release the runner when they are done with it, either explicitly - via `shutdown()` or by using it as a context manager:: + therefore release the runner when they are done with it. Exiting a `with` + block runs every submitted work item to completion and then releases the + threads:: with MPRunner(10) as runner: runner.execute_work(work_obj) ... + A caller that must abandon work in progress, such as one that has already + timed a stage out, calls `shutdown()` directly instead. + Runners that are created per sub-dataset and never released leak their worker threads for the remaining lifetime of the process. @@ -69,6 +73,10 @@ def shutdown(self, wait: bool = False, cancel_futures: bool = True) -> None: `RuntimeError`, so a caller that runs more than once needs a fresh runner per run. + Both defaults are the opposite of `Executor.shutdown`, which waits and + keeps queued work. They suit a caller that has already collected the + results it wants and now only needs the threads back. + Args: wait: Whether to block until every running work item has finished. Defaults to False so a work item that has hung (and that the caller @@ -84,4 +92,7 @@ def __enter__(self) -> "MPRunner": return self def __exit__(self, exc_type, exc_value, exc_traceback) -> None: - self.shutdown() + # Matches `Executor.__exit__`: run everything that was submitted and + # block until it finishes. A caller that wants the abandoning + # behaviour asks for it by calling `shutdown()` directly. + self.shutdown(wait=True, cancel_futures=False) diff --git a/evalbench/test/mprunner_test.py b/evalbench/test/mprunner_test.py index 9c0e0497..ee52765d 100644 --- a/evalbench/test/mprunner_test.py +++ b/evalbench/test/mprunner_test.py @@ -47,13 +47,18 @@ def test_workers_stay_alive_until_shutdown(self): t.join(timeout=30) self.assertEqual(_live(worker_threads), 0) - def test_context_manager_shuts_down_on_exit(self): + def test_context_manager_drains_before_shutting_down(self): + """`with` runs queued work to completion, like Executor.__exit__.""" threads = [] - with mprunner.MPRunner(2) as runner: - for _ in range(2): - runner.execute_work(_RecordingWork(threads)) - concurrent.futures.wait(runner.futures, timeout=30) + with mprunner.MPRunner(1) as runner: + # The second item cannot start until the first returns, so it is + # still queued when the block exits. + runner.execute_work(_RecordingWork(threads)) + runner.execute_work(_RecordingWork(threads)) + self.assertEqual(len(threads), 2) + self.assertTrue(all(f.done() and not f.cancelled() + for f in runner.futures)) for t in set(threads): t.join(timeout=30) self.assertEqual(_live(set(threads)), 0) @@ -106,7 +111,9 @@ def test_shutdown_cancels_queued_work(self): self.assertTrue(runner.futures[1].cancelled()) finally: release.set() - concurrent.futures.wait(runner.futures, timeout=30) + # Only the running item can finish. A cancelled future never reaches + # CANCELLED_AND_NOTIFIED, so waiting on it would burn the full timeout. + concurrent.futures.wait([runner.futures[0]], timeout=30) self.assertEqual(len(threads), 1) def test_shutdown_keeps_queued_work_when_cancel_disabled(self):