diff --git a/src/akkudoktoreos/optimization/genetic/genetic.py b/src/akkudoktoreos/optimization/genetic/genetic.py index b85f1870..f76a5a6f 100644 --- a/src/akkudoktoreos/optimization/genetic/genetic.py +++ b/src/akkudoktoreos/optimization/genetic/genetic.py @@ -5,8 +5,10 @@ import ctypes.util import gc import math import multiprocessing +import multiprocessing.pool import os import random +import signal import sys import time from array import array @@ -277,7 +279,17 @@ def worker_cpu_set(workers: int, allowed: Optional[set[int]] = None) -> Optional def _evaluation_worker_init(cpus: Optional[set[int]] = None) -> None: - """Evaluate without the run cache in a worker; the parent owns the cache.""" + """Prepare a forked evaluation worker; the parent owns the run cache. + + The worker ignores SIGTERM and SIGINT. Sent to the whole process group + (Ctrl-C, a service manager stopping the server) they are meant for the + server process, which finishes or aborts the run and stops the workers + itself (see ``_EvaluationPool``). Without this, a worker would run the + signal handler it inherited from the server - uvicorn's only asks the + server to shut down - or die in the middle of a batch. + """ + signal.signal(signal.SIGTERM, signal.SIG_IGN) + signal.signal(signal.SIGINT, signal.SIG_IGN) if _WORKER_OPTIMIZER is not None: _WORKER_OPTIMIZER._fitness_cache_enabled = False if cpus: @@ -287,6 +299,24 @@ def _evaluation_worker_init(cpus: Optional[set[int]] = None) -> None: pass +class _EvaluationPool(multiprocessing.pool.Pool): + """Process pool whose ``terminate()`` stops the workers with SIGKILL. + + ``Pool.terminate()`` keeps the task queue's read lock and then stops the + workers with SIGTERM. A worker that survives it - because it ignores the + signal, still runs the server's inherited handler (uvicorn's only asks the + server to shut down) or got it before its initializer ran - waits for that + lock forever, and ``Pool.join()`` with it. SIGKILL cannot be caught; the + pool sends it at the same point it would send SIGTERM. + """ + + @staticmethod + def Process(ctx: Any, *args: Any, **kwds: Any) -> Any: # noqa: N802 - Pool API + process = ctx.Process(*args, **kwds) + setattr(process, "terminate", process.kill) + return process + + def _evaluation_worker(genome: list[int]) -> tuple[list[int], tuple[float], Any]: """Evaluate one genome in a worker and return the (repaired) genome and results.""" individual = creator.Individual(genome) @@ -2954,10 +2984,11 @@ class GeneticOptimization(OptimizationBase): cpus = worker_cpu_set(self._evaluation_workers) if self._pin_workers() else None _WORKER_OPTIMIZER = self try: - self._evaluation_pool = context.Pool( + self._evaluation_pool = _EvaluationPool( self._evaluation_workers, initializer=_evaluation_worker_init, initargs=(cpus,), + context=context, ) except Exception as exc: _WORKER_OPTIMIZER = None diff --git a/tests/test_genetic_parallel_evaluation.py b/tests/test_genetic_parallel_evaluation.py index c3db9176..96ff66a6 100644 --- a/tests/test_genetic_parallel_evaluation.py +++ b/tests/test_genetic_parallel_evaluation.py @@ -1,9 +1,15 @@ """Parallel fitness evaluation gives the same result as evaluating in one process.""" import multiprocessing +import os +import signal +import sys +import threading +import time from types import SimpleNamespace import pytest +import uvicorn from akkudoktoreos.config.config import ConfigEOS from akkudoktoreos.core.cache import CacheEnergyManagementStore @@ -146,6 +152,274 @@ def test_failed_batch_falls_back_to_serial(config_eos: ConfigEOS, monkeypatch): assert fallback.ac_charge == serial.ac_charge +class _PoolOnly(GeneticOptimization): + """Just the evaluation pool of an optimization: two unpinned workers.""" + + def _configured_evaluation_workers(self) -> int: + return 2 + + def _pin_workers(self) -> bool: + return False + + +def _square(value: int) -> int: + return value * value + + +_original_evaluation_worker = genetic_module._evaluation_worker + + +def _slow_evaluation_worker(genome): + time.sleep(5) # far beyond the (patched) batch timeout + return _original_evaluation_worker(genome) + + +@pytest.fixture +def server_signal_handlers(): + """Run the test under the signal handlers uvicorn.run() installs for the server. + + uvicorn's handler only asks the server to shut down; the evaluation workers + are forked from that process and inherit it. + """ + server = uvicorn.Server(uvicorn.Config("akkudoktoreos.server.eos:app")) + with server.capture_signals(): + assert signal.getsignal(signal.SIGTERM) == server.handle_exit + yield + + +def _in_thread(function, timeout: float, rescue) -> bool: + """Run ``function`` in a thread; on timeout call ``rescue`` and report a hang.""" + thread = threading.Thread(target=function, daemon=True) + thread.start() + thread.join(timeout) + if not thread.is_alive(): + return True + rescue() # kill the stuck workers so the test run itself does not hang + thread.join(timeout) + return False + + +def _kill(processes) -> None: + for process in processes: + if process.is_alive(): + os.kill(process.pid, signal.SIGKILL) + + +def _resume(pid: int) -> None: + """SIGCONT a paused worker; the stop may already have killed and reaped it.""" + try: + os.kill(pid, signal.SIGCONT) + except ProcessLookupError: + pass + + +def _worker_signal_dispositions(_task: int) -> tuple[bool, bool]: + """In a worker: (SIGTERM is ignored, SIGINT is ignored).""" + return ( + signal.getsignal(signal.SIGTERM) is signal.SIG_IGN, + signal.getsignal(signal.SIGINT) is signal.SIG_IGN, + ) + + +def _slow_square(value: int) -> int: + time.sleep(0.5) + return value * value + + +_original_evaluation_worker_init = genetic_module._evaluation_worker_init + + +def _late_evaluation_worker_init(*args): + time.sleep(2) # a worker that starts late on a busy machine + _original_evaluation_worker_init(*args) + + +@pytest.mark.skipif("fork" not in multiprocessing.get_all_start_methods(), reason="requires fork") +def test_workers_do_not_run_server_signal_handlers(server_signal_handlers): + """Workers ignore SIGTERM/SIGINT instead of running the server's handlers.""" + optimization = _PoolOnly.__new__(_PoolOnly) + starter = threading.Thread(target=optimization._start_evaluation_pool) + starter.start() # like the server: the run is not on the main thread + starter.join() + pool = optimization._evaluation_pool + assert pool is not None + try: + workers = list(pool._pool) + dispositions = pool.map(_worker_signal_dispositions, range(8), chunksize=1) + finally: + stopped = _in_thread( + optimization._stop_evaluation_pool, 30, lambda: _kill(list(pool._pool)) + ) + assert set(dispositions) == {(True, True)} + assert stopped, "stopping the pool hung" + assert not any(worker.is_alive() for worker in workers) + # The server process itself keeps its handlers. + assert signal.getsignal(signal.SIGTERM) not in (signal.SIG_DFL, signal.SIG_IGN) + + +@pytest.mark.skipif("fork" not in multiprocessing.get_all_start_methods(), reason="requires fork") +def test_sigterm_to_the_process_group_does_not_break_a_batch(server_signal_handlers): + """A service manager stops the server with SIGTERM to all its processes. + + The workers finish the running batch; the server process decides when the + run ends and stops them. + """ + optimization = _PoolOnly.__new__(_PoolOnly) + optimization._start_evaluation_pool() + pool = optimization._evaluation_pool + assert pool is not None + try: + pool.map(_square, range(4)) + workers = list(pool._pool) + batch = pool.map_async(_slow_square, range(8), chunksize=1) + time.sleep(0.2) + for worker in workers: + os.kill(worker.pid, signal.SIGTERM) + results = batch.get(timeout=30) + alive = [worker.is_alive() for worker in workers] + finally: + stopped = _in_thread( + optimization._stop_evaluation_pool, 30, lambda: _kill(list(pool._pool)) + ) + assert results == [value * value for value in range(8)] + assert alive == [True, True] + assert stopped, "stopping the pool hung" + + +@pytest.mark.skipif( + not sys.platform.startswith("linux"), reason="needs Linux semaphores and /proc" +) +def test_pool_stop_with_a_late_idle_worker_does_not_hang(server_signal_handlers): + """Stopping the pool must not depend on which process wins the task-queue lock. + + Pool.terminate() takes the read lock of the task queue for good. A worker + that has not got its stop sentinel by then (on a busy machine it may simply + not have been scheduled) can only end through SIGTERM. Freezing that worker + for a moment makes the parent win the lock every time. + """ + optimization = _PoolOnly.__new__(_PoolOnly) + optimization._start_evaluation_pool() + pool = optimization._evaluation_pool + assert pool is not None + pool.map(_square, range(20)) + workers = list(pool._pool) + time.sleep(0.5) # both idle: one reads the task pipe, one waits for its lock + + def waits_for_lock(worker) -> bool: + try: + with open(f"/proc/{worker.pid}/syscall") as syscall: + return syscall.read().split()[0] in ("202", "98") # futex x86_64/aarch64 + except (OSError, IndexError): + return False + + late = next((worker for worker in workers if waits_for_lock(worker)), None) + if late is None: + _in_thread(optimization._stop_evaluation_pool, 30, lambda: _kill(list(pool._pool))) + pytest.skip("cannot tell which worker waits for the lock") + os.kill(late.pid, signal.SIGSTOP) + resume = threading.Timer(1.0, _resume, (late.pid,)) + resume.start() + try: + stopped = _in_thread( + optimization._stop_evaluation_pool, 30, lambda: _kill(list(pool._pool)) + ) + finally: + resume.join() + assert stopped, "stopping the pool hung" + + +@pytest.mark.skipif( + not sys.platform.startswith("linux"), reason="needs Linux semaphores and /proc" +) +def test_pool_stop_never_kills_the_worker_holding_the_task_queue(server_signal_handlers): + """The idle worker that reads the task pipe holds the queue's read lock. + + Killing it would leave the lock taken for good and terminate() waiting for + it forever. While it is paused, the stop has to wait until it resumes, + takes its sentinel and releases the lock. + """ + optimization = _PoolOnly.__new__(_PoolOnly) + optimization._start_evaluation_pool() + pool = optimization._evaluation_pool + assert pool is not None + pool.map(_square, range(20)) + workers = list(pool._pool) + time.sleep(0.5) # both idle: one reads the task pipe, one waits for its lock + + def reads_task_pipe(worker) -> bool: + try: + with open(f"/proc/{worker.pid}/syscall") as syscall: + return syscall.read().split()[0] in ("0", "63") # read x86_64/aarch64 + except (OSError, IndexError): + return False + + reader = next((worker for worker in workers if reads_task_pipe(worker)), None) + if reader is None: + _in_thread(optimization._stop_evaluation_pool, 30, lambda: _kill(list(pool._pool))) + pytest.skip("cannot tell which worker reads the task pipe") + os.kill(reader.pid, signal.SIGSTOP) + resume = threading.Timer(3.0, _resume, (reader.pid,)) + started = time.monotonic() + resume.start() + try: + stopped = _in_thread( + optimization._stop_evaluation_pool, 30, lambda: _kill(list(pool._pool)) + ) + finally: + resume.join() + # Killing the paused reader would leave the lock taken and the stop hung + # (caught above). Once it resumed and released the lock, terminate() may + # still kill it before it exits on its own - that is fine. + assert stopped, "stopping the pool hung" + assert time.monotonic() - started >= 2.5, "the stop did not wait for the paused reader" + assert reader.exitcode is not None + + +@pytest.mark.skipif("fork" not in multiprocessing.get_all_start_methods(), reason="requires fork") +def test_batch_timeout_under_server_signal_handlers_falls_back_to_serial( + config_eos: ConfigEOS, monkeypatch, server_signal_handlers +): + """A batch that times out stops the busy workers and the run continues serially.""" + serial, _ = _run(config_eos, 1, None) + + monkeypatch.setattr(genetic_module, "EVALUATION_BATCH_TIMEOUT_SECONDS", 0.5) + monkeypatch.setattr(genetic_module, "_evaluation_worker", _slow_evaluation_worker) + outcome: dict = {} + + def run_parallel() -> None: + outcome["result"] = _run(config_eos, 2, None) + + finished = _in_thread(run_parallel, 60, lambda: _kill(multiprocessing.active_children())) + + assert finished, "stopping the busy workers after the batch timeout hung" + fallback, used = outcome["result"] + assert used == [2] + assert fallback.start_solution == serial.start_solution + assert fallback.ac_charge == serial.ac_charge + + +@pytest.mark.skipif("fork" not in multiprocessing.get_all_start_methods(), reason="requires fork") +def test_pool_stop_kills_a_worker_that_got_sigterm_before_its_init( + server_signal_handlers, monkeypatch +): + """A worker still starting up when the pool stops runs the server's handler. + + It would survive SIGTERM, finish its start and then wait for the task queue + that terminate() keeps locked. The pool stops its workers with SIGKILL. + """ + monkeypatch.setattr(genetic_module, "_evaluation_worker_init", _late_evaluation_worker_init) + optimization = _PoolOnly.__new__(_PoolOnly) + optimization._start_evaluation_pool() + pool = optimization._evaluation_pool + assert pool is not None + workers = list(pool._pool) + started = time.monotonic() + stopped = _in_thread(optimization._stop_evaluation_pool, 30, lambda: _kill(list(pool._pool))) + assert stopped, "stopping the pool hung" + assert time.monotonic() - started < 10 + assert not any(worker.is_alive() for worker in workers) + + def test_unsupported_fork_falls_back_to_serial(config_eos: ConfigEOS, monkeypatch): serial, _ = _run(config_eos, 1, None)