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..aaa27c04 100644 --- a/evalbench/mp/mprunner.py +++ b/evalbench/mp/mprunner.py @@ -23,10 +23,27 @@ 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. 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. + 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 +63,36 @@ 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. + + 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 + 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: + # 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 new file mode 100644 index 00000000..ee52765d --- /dev/null +++ b/evalbench/test/mprunner_test.py @@ -0,0 +1,217 @@ +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_drains_before_shutting_down(self): + """`with` runs queued work to completion, like Executor.__exit__.""" + threads = [] + 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) + + 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() + # 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): + """`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()