use process pool for testing mutations

Author: Otto-AA

mutmut/__main__.py

@@ -1,10 +1,13 @@
+import platform
+from mutmut.custom_process_pool import Task
+from mutmut.custom_process_pool import CustomProcessPool
 import ast
 import fnmatch
 import gc
 import inspect
 import itertools
 import json
-from multiprocessing import Pool, Process, set_start_method
+from multiprocessing import JoinableQueue, Pool, Process, Queue, set_start_method
 import multiprocessing
 import multiprocessing.connection
 import os
@@ -47,7 +50,9 @@
 )
 from typing import (
     Dict,
+    Generic,
     List,
+    TypeVar,
     Union,
 )
 
@@ -373,6 +378,8 @@ def new_tests(self):
         return self.ids - collected_test_names()
 
 
+_pytest_initialized = False
+
 class PytestRunner(TestRunner):
     # noinspection PyMethodMayBeStatic
     def execute_pytest(self, params: list[str], **kwargs):
@@ -386,6 +393,10 @@ def execute_pytest(self, params: list[str], **kwargs):
             print('    exit code', exit_code)
         if exit_code == 4:
             raise BadTestExecutionCommandsException(params)
+
+        global _pytest_initialized
+        _pytest_initialized = True
+
         return exit_code
 
     def run_stats(self, *, tests):
@@ -889,7 +900,10 @@ def inner_timout_checker():
 @click.option('--max-children', type=int)
 @click.argument('mutant_names', required=False, nargs=-1)
 def run(mutant_names, *, max_children):
-    set_start_method('spawn')
+    if platform.system() == 'Windows':
+        set_start_method('spawn')
+    else:
+        set_start_method('fork')
 
     assert isinstance(mutant_names, (tuple, list)), mutant_names
     _run(mutant_names, max_children)
@@ -956,23 +970,6 @@ def _run(mutant_names: Union[tuple, list], max_children: Union[None, int]):
 
     running_processes: set[Process] = set()
 
-    def handle_finished_processes() -> int:
-        nonlocal running_processes
-        sentinels = [p.sentinel for p in running_processes]
-        multiprocessing.connection.wait(sentinels)
-
-        finished_processes = {p for p in running_processes if not p.is_alive()}
-        running_processes -= finished_processes
-
-        for p in finished_processes:
-            if mutmut.config.debug:
-                print('    worker exit code', p.exitcode)
-            source_file_mutation_data_by_pid[p.pid].register_result(pid=p.pid, exit_code=p.exitcode)
-
-            p.close()
-
-        return len(finished_processes)
-
     source_file_mutation_data_by_pid: dict[int, SourceFileMutationData] = {}  # many pids map to one MutationData
     running_children = 0
     count_tried = 0
@@ -996,6 +993,8 @@ def handle_finished_processes() -> int:
         # TODO: implement timeout for windows + unix
         # Thread(target=timeout_checker(mutants), daemon=True).start()
 
+        args: list[tuple[TestRunner, SourceFileMutationData, str, list[str], Config]] = []
+
         # Now do mutation
         for m, mutant_name, result in mutants:
             print_stats(source_file_mutation_data_by_path)
@@ -1015,37 +1014,32 @@ def handle_finished_processes() -> int:
                 m.save()
                 continue
 
-            p = Process(target=_test_mutation, args=(runner, m, mutant_name, tests, mutmut.config))
-            running_processes.add(p)
-            p.start()
-            pid = p.pid
-            # in the parent
-            source_file_mutation_data_by_pid[pid] = m
-            m.register_pid(pid=pid, key=mutant_name, estimated_time_of_tests=estimated_time_of_tests)
-            running_children += 1
-
-            if running_children >= max_children:
-                count_finished = handle_finished_processes()
-                count_tried += count_finished
-                running_children -= count_finished
-
-        try:
-            while running_children:
-                print_stats(source_file_mutation_data_by_path)
-                count_finished = handle_finished_processes()
-                count_tried += count_finished
-                running_children -= count_finished
-        except ChildProcessError:
-            pass
+            args.append((runner, m, mutant_name, tests, mutmut.config))
+            source_file_mutation_data_by_pid[mutant_name] = m
+            m.register_pid(pid=mutant_name, key=mutant_name, estimated_time_of_tests=estimated_time_of_tests)
+
+        tasks: list[Task] = []
+        for arg in args:
+            tasks.append(Task(id=arg[2], args=arg, timeout_seconds=1000))
+        pool = CustomProcessPool(tasks, _test_mutation, max_children)
+        done = 0
+        for finished_task in pool.run():
+            done += 1
+            # print(f'Finished {done} tasks')
+            if finished_task.error:
+                print(finished_task)
+            # print(finished_task)
+            source_file_mutation_data_by_pid[finished_task.id].register_result(pid=finished_task.id, exit_code=finished_task.result)
+            print_stats(source_file_mutation_data_by_path)
     except KeyboardInterrupt:
+        pool.shutdown()
         print('Stopping...')
-        stop_all_children(mutants)
 
     t = datetime.now() - start
 
     print_stats(source_file_mutation_data_by_path, force_output=True)
     print()
-    print(f'{count_tried / t.total_seconds():.2f} mutations/second')
+    print(f'{len(tasks) / t.total_seconds():.2f} mutations/second')
 
     if mutant_names:
         print()
@@ -1062,12 +1056,18 @@ def handle_finished_processes() -> int:
         print()
 
 
-def _test_mutation(runner: TestRunner, m: SourceFileMutationData, mutant_name: str, tests, config):
+def _test_mutation(task: Task):
+    args: tuple[TestRunner, SourceFileMutationData, str, list[str], Config] = task.args
+    runner, m, mutant_name, tests, config = args
     try:
         mutmut.config = config
 
-        with CatchOutput():
-            runner.list_all_tests()
+        # ensure that we imported all files at least once per process
+        # before we set MUTANT_UNDER_TEST (so everything that runs at import
+        # time is not mutated)
+        if not _pytest_initialized:
+            with CatchOutput():
+                runner.list_all_tests()
 
         os.environ['MUTANT_UNDER_TEST'] = mutant_name
         setproctitle(f'mutmut: {mutant_name}')
@@ -1082,13 +1082,14 @@ def _test_mutation(runner: TestRunner, m: SourceFileMutationData, mutant_name: s
 
             with CatchOutput():
                 result = runner.run_tests(mutant_name=mutant_name, tests=tests)
-        os._exit(result)
+        
+        return result
+        # os._exit(result)
     except Exception as e:
         with open(f'error.{mutant_name}.log', 'w') as log:
             log.write(str(e))
-        os._exit(-1)
-
-
+            log.flush()
+        return -24
 
 def tests_for_mutant_names(mutant_names):
     tests = set()

mutmut/custom_process_pool.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+from typing import Generic, Union, Any, Callable, Iterable, TypeVar
+from typing_extensions import ParamSpec
+from dataclasses import dataclass
+
+import multiprocessing.connection
+from multiprocessing import Queue, Process
+import queue
+import os
+
+
+TaskArgs = ParamSpec('TaskArgs')
+TaskResult = TypeVar('TaskResult')
+
+@dataclass
+class Task:
+    id: str
+    args: tuple[Any, ...]
+    # this timeout is real time, not process cpu time
+    timeout_seconds: int
+
+@dataclass
+class TaskError:
+    id: str
+    error: Exception
+
+@dataclass
+class FinishedTask(Generic[TaskResult]):
+    id: str
+    result: Union[TaskResult, None]
+    error: Union[Exception, None]
+
+class JobTimeoutException(Exception):
+    pass
+
+class CustomProcessPool(Generic[TaskArgs, TaskResult]):
+    def __init__(self, tasks: list[Task], job: Callable[TaskArgs, TaskResult], max_workers: int):
+        self._tasks = tasks
+        self._job = job
+        self._remaining_tasks_queue: Queue[Task] = Queue()
+        self._remaining_tasks_count = len(tasks)
+        self._results: Queue[FinishedTask[TaskResult]] = Queue()
+        self._max_workers = max_workers
+        self._workers: set[Process] = set()
+        self._killed_workers = 0
+        self._shutdown = False
+
+    def run(self) -> Iterable[FinishedTask]:
+        for task in self._tasks:
+            self._remaining_tasks_queue.put(task)
+
+        self._start_missing_workers()
+
+        while not self.done() and not self._shutdown:
+            self._remove_stopped_workers()
+            self._start_missing_workers()
+
+            yield from self._get_new_results(timeout=1)
+
+        self.shutdown()
+
+    def shutdown(self):
+        # TODO: is this a good way to shutdown processes?
+        for p in self._workers:
+            if p.is_alive():
+                p.kill()
+        for p in self._workers:
+            p.join()
+        self._remaining_tasks_queue.close()
+        self._results.close()
+        self._shutdown = True
+
+    def _start_missing_workers(self):
+        self._workers = {p for p in self._workers if p.is_alive()}
+
+        desired_workers = min(self._max_workers, self._remaining_tasks_count)
+        missing_workers = desired_workers - len(self._workers)
+
+        for _ in range(missing_workers):
+            self._start_worker()
+
+    def _remove_stopped_workers(self):
+        """Start a new worker for all stopped workers. We kill workers for timeouts."""
+        killed_workers = {p for p in self._workers if not p.is_alive()}
+        self._workers -= killed_workers
+
+        for worker in killed_workers:
+            print(f'Worker {worker.pid} stopped with exitcode {worker.exitcode}')
+
+    def _get_new_results(self, timeout: int) -> Iterable[FinishedTask]:
+        try:
+            result = self._results.get(timeout=timeout)
+            self._remaining_tasks_count -= 1
+            yield result
+        except queue.Empty:
+            pass
+
+    def _start_worker(self):
+        p = Process(target=CustomProcessPool._pool_job_executor, args=(self._job, self._remaining_tasks_queue, self._results))
+        p.start()
+        self._workers.add(p)
+
+    def done(self) -> bool:
+        return self._remaining_tasks_count == 0
+
+    @staticmethod
+    def _pool_job_executor(job: Callable[..., TaskResult], task_queue: Queue[Task], results: Queue[FinishedTask[TaskResult]]):
+        while True:
+            try:
+                task = task_queue.get(timeout=1)
+                # f = open(f'logs/log-{task.id}.txt', 'w')
+                # pid = os.getpid()
+            except queue.Empty:
+                os._exit(0)
+
+            try:
+                result = job(task)
+                finished_task: FinishedTask[TaskResult] = FinishedTask(id=task.id, result=result, error=None)
+            except Exception as e:
+                finished_task = FinishedTask(id=task.id, result=None, error=e)
+            finally:
+                # f.write(f'Finished job: {finished_task}\n')
+                # f.flush()
+                results.put(finished_task)
+                # f.write(f'Added job to queue\n')
+                # f.write(f'Finished qsize: {results.qsize()}\n')
+                # f.flush()
+
+
+

mutmut/file_mutation.py

@@ -176,6 +176,8 @@ def combine_mutations_to_source(module: cst.Module, mutations: Sequence[Mutation
     :param mutations: Mutations that should be applied.
     :return: Mutated code and list of mutation names"""
 
+    # mutations = mutations[0:10]
+
     # copy start of the module (in particular __future__ imports)
     result: list[MODULE_STATEMENT] = get_statements_until_func_or_class(module.body)
     mutation_names: list[str] = []

tests/e2e/test_e2e_result_snapshots.py

@@ -78,3 +78,6 @@ def test_my_lib_result_snapshot():
 def test_config_result_snapshot():
     mutmut._reset_globals()
     asserts_results_did_not_change("config")
+
+if __name__ == '__main__':
+    test_my_lib_result_snapshot()

tests/test_custom_pool.py

@@ -0,0 +1,31 @@
+from mutmut.__main__ import CustomProcessPool, Task
+import pytest
+import time
+
+def test_custom_process_pool():
+    tasks = [
+        Task(id='a-small', args=(1, 2), timeout_seconds=1000),
+        Task(id='b-medium', args=(30, 20), timeout_seconds=1000),
+        Task(id='c-neg', args=(-2, -2), timeout_seconds=1000),
+        Task(id='d-div-by-zero', args=(-2, 0), timeout_seconds=1000),
+    ]
+    pool = CustomProcessPool(tasks, _divide, max_workers=2)
+
+    results = []
+    for result in pool.run():
+        print(result)
+        results.append(result)
+
+    assert len(results) == 4
+
+    results = sorted(results, key=lambda result: result.id)
+    assert results[0].result == pytest.approx(0.5)
+    assert results[1].result == pytest.approx(1.5)
+    assert results[2].result == pytest.approx(1)
+    assert results[3].result == None
+    assert isinstance(results[3].error, ZeroDivisionError)
+
+def _divide(task: Task):
+    a, b = task.args
+    # time.sleep(timeout)
+    return a / b