"""asyncio-based task runner for the infinite-swapping scheduler.
The runner drives the worker pool that executes the per-cycle
``run_md`` calls produced by
:py:mod:`pyretis.simulation.scheduler`.
"""
import asyncio
import concurrent.futures
import functools
import logging
import multiprocessing
import os
import threading
from collections.abc import Callable
from typing import Any, Dict, List, Optional, Set
from pyretis.setup.common import create_orderparameters
from pyretis.engines.factory import create_engines
from pyretis.inout.formats.formatter import get_log_formatter
logger = logging.getLogger("")
logger.setLevel(logging.DEBUG)
[docs]
class _ShutdownNoiseFilter(logging.Filter):
"""Drop the expected ``BrokenProcessPool`` teardown noise.
When the worker pool is force-released on an interrupt / SIGTERM
(a tutorial smoke run killed by its timeout, an HPC walltime
``kill``, a user Ctrl-C), an in-flight ``run_in_executor`` future can
finish with a :class:`concurrent.futures.BrokenExecutor` *after* its
wrapper task was already cancelled. asyncio then reports it -- usually
during interpreter shutdown -- as "Future exception was never
retrieved" at ERROR level with a full traceback. That is expected
teardown noise, not a real failure (a genuine pool failure mid-run is
retrieved onto the work future via ``future.set_exception`` and never
reaches this "never retrieved" path), so it is dropped. Everything
else passes through untouched.
"""
[docs]
def filter(self, record: logging.LogRecord) -> bool:
"""Return False only for the unretrieved-BrokenExecutor record."""
if "Future exception was never retrieved" not in record.getMessage():
return True
exc = record.exc_info[1] if record.exc_info else None
return not isinstance(exc, concurrent.futures.BrokenExecutor)
[docs]
def _install_shutdown_noise_filter() -> None:
"""Attach :class:`_ShutdownNoiseFilter` to the asyncio logger once.
asyncio's default exception handler logs through the ``asyncio``
logger, so the filter is installed there (a logger-level filter drops
the record before it propagates to the root handlers). Idempotent: a
second runner does not stack a duplicate filter.
"""
aio_logger = logging.getLogger("asyncio")
if not any(isinstance(filt, _ShutdownNoiseFilter)
for filt in aio_logger.filters):
aio_logger.addFilter(_ShutdownNoiseFilter())
# Process-local engine pool for the current worker. Populated by
# worker_initializer when the ProcessPool spins up each worker process and
# read back by the run_md bridge task. It is deliberately process-local
# (each worker imports this module fresh) so concurrent workers never share
# engine instances -- the pool is no longer a global of the core sampling
# module pyretis.core.moves.
_WORKER_ENGINES: Dict[str, Any] = {}
[docs]
def get_worker_engines() -> Dict[str, Any]:
"""Return the engine pool created for the current worker process."""
return _WORKER_ENGINES
[docs]
def set_worker_engines(engines: Dict[str, Any]) -> None:
"""Install the engine pool for the current worker process."""
_WORKER_ENGINES.clear()
_WORKER_ENGINES.update(engines)
[docs]
class RunnerError(Exception):
"""Exception class for the runner."""
[docs]
class aiorunner:
"""A light asynchronuous runner based on asyncio.
The runner manage an asyncio.queue with a pool of workers.
Upon instanciation, a dedicated event loop
is launched in a separate thread. The user can then
attach a worker function to the runner and start multiple
instances of that function in the background.
As work is submitted to the runner, it is picked up by
workers on-the-fly.
"""
[docs]
def __init__(self, config: Dict, n_workers: int = 1) -> None:
"""Init function of runner.
Parameters
----------
config : dict
The simulation configuration dictionary. It is
forwarded **unchanged** to every worker process via the
pool initializer (:func:`worker_initializer`) and must
therefore be picklable (the pool uses the ``spawn`` start
method). When ``config`` contains a ``"simulation"``
section the initializer builds that worker's engine pool
and order parameters from it (via
:func:`pyretis.engines.factory.create_engines` and
:func:`create_orderparameters`) and installs the engines
as process-local state with :func:`set_worker_engines`;
the engines are intentionally *not* part of the per-task
work units, so they never cross the process boundary.
n_workers : int
Number of worker processes in the pool.
"""
self._n_workers: int = n_workers
self._counter = multiprocessing.get_context("spawn").Value("i", 0)
self._executor: concurrent.futures.Executor = (
concurrent.futures.ProcessPoolExecutor(
max_workers=n_workers,
initializer=worker_initializer,
initargs=(self._counter, config),
mp_context=multiprocessing.get_context("spawn"),
)
)
self._loop = asyncio.new_event_loop()
self._queue: asyncio.Queue[Any] = asyncio.Queue()
self._state = "created"
self._state_lock = threading.Lock()
self._pending_lock = threading.Lock()
self._pending_futures: Set[concurrent.futures.Future] = set()
_install_shutdown_noise_filter()
self._thread = threading.Thread(
target=self._start_event_loop, daemon=True
)
self._thread.start()
self._task_f: Optional[Callable] = None
self._tasks: Optional[List[asyncio.Task[Any]]] = None
[docs]
def start(self) -> None:
"""Launch background tasks."""
with self._state_lock:
if self._state != "created":
raise RunnerError(
f"Unable to start runner in state '{self._state}'"
)
self._state = "starting"
future = asyncio.run_coroutine_threadsafe(
self._start_tasks(), self._loop
)
try:
# Task startup should be fast
future.result(5.0)
except TimeoutError as exc:
with self._state_lock:
self._state = "created"
raise RunnerError(
"Launching background tasks took too long") from exc
except Exception:
with self._state_lock:
self._state = "created"
raise
with self._state_lock:
self._state = "running"
[docs]
def _start_event_loop(self) -> None:
"""Start the event loop in a separate thread."""
asyncio.set_event_loop(self._loop)
self._loop.run_forever()
[docs]
def set_task(self, task_f: Callable) -> None:
"""Attach the task function to the runner.
Parameters
----------
task_f : callable
a callable function
"""
with self._state_lock:
if self._state in {"stopping", "stopped", "closing", "closed"}:
raise RunnerError(
f"Unable to set task in state '{self._state}'"
)
self._task_f = task_f
[docs]
async def _task_wrapper(
self,
queue: asyncio.Queue,
executor: concurrent.futures.Executor,
) -> None:
"""Wrap the sync task.
To enable running the sync task_f
from a dynamic list of tasks.
Parameters
----------
queue : asyncio.Queue
an asyncio queue to get work from
executor : concurrent.futures.Executor
an executor
"""
while True:
item = await queue.get()
try:
if item is None:
return
md_item, future = item
# Run the task in the event loop
if self._task_f is None:
if not future.done():
future.set_exception(
RuntimeError("worker has no task function set")
)
continue
loop = asyncio.get_running_loop()
try:
md_item = await loop.run_in_executor(
executor, functools.partial(self._task_f, md_item)
)
except asyncio.CancelledError:
future.cancel()
raise
except Exception as e:
# Pass the exception up in the future
if not future.done():
future.set_exception(e)
else:
if not future.done():
future.set_result(md_item)
finally:
queue.task_done()
[docs]
async def _add_work_to_queue(
self, work_unit: Dict[str, Any]
) -> concurrent.futures.Future:
"""Async function adding work to queue, returns a future.
Parameters
----------
work_unit : dict
a unit of work encapsulated in a dict
Returns
-------
concurrent.futures.Future
A future wih the results of the work
"""
future: concurrent.futures.Future = concurrent.futures.Future()
with self._pending_lock:
self._pending_futures.add(future)
future.add_done_callback(self._discard_pending_future)
await self._queue.put((work_unit, future))
return future
[docs]
def _discard_pending_future(
self, future: concurrent.futures.Future
) -> None:
"""Forget a completed or cancelled caller-visible future."""
with self._pending_lock:
self._pending_futures.discard(future)
[docs]
def submit_work(
self, work_unit: Dict[str, Any]
) -> concurrent.futures.Future:
"""Submit work to the runner.
Parameters
----------
task
a unit of work encapsulated in a dict
Returns
-------
concurrent.futures.Future
A future wih the results of the work
"""
with self._state_lock:
if self._state != "running":
raise RunnerError(
f"Unable to submit work in state '{self._state}'"
)
queued = asyncio.run_coroutine_threadsafe(
self._add_work_to_queue(work_unit), self._loop
)
try:
return queued.result(timeout=5.0)
except TimeoutError as exc:
queued.cancel()
raise RunnerError("Submitting work took too long") from exc
[docs]
async def _start_tasks(self) -> None:
"""Launch the background tasks."""
if not self._task_f:
raise RunnerError("Can't start task(s) without a task function.")
try:
self._tasks = [
asyncio.create_task(
self._task_wrapper(self._queue, self._executor)
)
for _ in range(self._n_workers)
]
except Exception as e:
raise e
[docs]
async def _drain_and_stop_tasks(self) -> None:
"""Drain submitted work and stop every queue consumer."""
await self._queue.join()
tasks = self._tasks or []
for _ in tasks:
await self._queue.put(None)
if tasks:
await asyncio.gather(*tasks)
self._tasks = []
[docs]
def n_workers(self) -> int:
"""Return runner number of workers."""
return self._n_workers
[docs]
def stop(self) -> None:
"""Terminate the runner and release the worker pool."""
with self._state_lock:
if self._state in {"stopped", "closed"}:
return
if self._state in {"stopping", "closing"}:
raise RunnerError(f"Runner is already {self._state}")
self._state = "stopping"
try:
if self._loop.is_running():
future = asyncio.run_coroutine_threadsafe(
self._drain_and_stop_tasks(), self._loop
)
future.result()
finally:
if self._loop.is_running():
self._loop.call_soon_threadsafe(self._loop.stop)
self._thread.join()
# ProcessPoolExecutor otherwise only reaps its workers via its
# interpreter-exit handler, which does not run on SIGTERM.
self._executor.shutdown(wait=True)
if not self._loop.is_closed():
self._loop.close()
with self._state_lock:
self._state = "stopped"
[docs]
async def _cancel_pending_tasks(self) -> None:
"""Cancel the worker-wrapper tasks and await their completion."""
tasks = [t for t in (self._tasks or []) if not t.done()]
for task in tasks:
task.cancel()
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
while True:
try:
item = self._queue.get_nowait()
except asyncio.QueueEmpty:
break
if item is not None:
_, future = item
future.cancel()
self._queue.task_done()
self._tasks = []
[docs]
def close(self) -> None:
"""Force-release runner resources without draining the queue.
Safe to call from an error/interrupt path (unlike :meth:`stop`,
which waits for the work queue to empty): it cancels the pending
worker tasks, stops the event loop and shuts the pool down so it
is never orphaned.
"""
with self._state_lock:
if self._state in {"stopped", "closed"}:
return
if self._state in {"stopping", "closing"}:
return
self._state = "closing"
if self._loop.is_running():
future = asyncio.run_coroutine_threadsafe(
self._cancel_pending_tasks(), self._loop
)
try:
future.result(timeout=5)
except Exception: # pylint: disable=broad-exception-caught
# Best-effort cleanup: never let shutdown hang on a stuck
# task; the pool is released unconditionally below.
logger.debug("shutdown: pending tasks did not finish in time")
with self._pending_lock:
pending = list(self._pending_futures)
for pending_future in pending:
pending_future.cancel()
if self._loop.is_running():
self._loop.call_soon_threadsafe(self._loop.stop)
self._thread.join(timeout=5)
self._executor.shutdown(wait=False, cancel_futures=True)
if not self._thread.is_alive() and not self._loop.is_closed():
self._loop.close()
with self._state_lock:
self._state = "closed"
[docs]
def prepare_streaming_engines(engines, config):
"""Give the internal integrators their file-backed streaming setup.
The coordinator hands every engine file-backed phase points (a
snapshot ``System`` whose ``config`` points at a trajectory file).
The external engines read that file directly; the internal
integrators (``langevin`` / ``velocityverlet`` / ``verlet`` /
``randomwalk``) integrate from in-memory particles instead, so they
need a streaming template -- their own box / particles / masses /
force field -- built from the carried-through ``[system]`` /
``[box]`` / ``[particles]`` / ``[potential]`` / ``[forcefield]``
sections. This mirrors how :py:class:`.TurtleMDEngine` builds its own
box / particles / potential in ``__init__``. Engines without a
``setup_streaming`` method (the external engines) are skipped.
Parameters
----------
engines : dict of lists
The per-worker engine pool, keyed by engine name.
config : dict
The coordinator configuration dictionary.
"""
for engine_key in engines:
for engine in engines[engine_key]:
setup = getattr(engine, "setup_streaming", None)
if callable(setup):
setup(config)
[docs]
def worker_initializer(counter, config):
"""Initialize function for each worker process."""
# load engines for a scheduler run
if "simulation" in config:
engines, _ = create_engines(config)
create_orderparameters(engines, config)
prepare_streaming_engines(engines, config)
# Install this worker's engine pool as process-local state, read
# back by the run_md bridge task (no module global in moves).
set_worker_engines(engines)
with counter.get_lock(): # Ensure that counter increment is thread-safe
worker_id = counter.value
counter.value += 1
# Unified logging: a spawn-started worker inherits no log handler, so it
# attaches one to the SAME run-level ``pyretis.log`` the main process
# writes -- not a separate ``worker<N>.log``. Append mode gives O_APPEND
# (each sub-PIPE_BUF log line is written atomically, so concurrent
# workers interleave by whole lines, never corrupt one). The worker is
# spawned with the run directory as CWD (before any per-move chdir), so
# the default relative name resolves to the main log; a custom ``-f``
# name is carried through ``[output] log_file`` when present.
log_name = "pyretis.log"
output_cfg = config.get("output") if isinstance(config, dict) else None
if isinstance(output_cfg, dict) and output_cfg.get("log_file"):
log_name = output_cfg["log_file"]
fileh = logging.FileHandler(os.path.abspath(log_name), mode="a")
# Workers contribute only WARNING+ to the shared log. Their INFO stream
# is per-move MD chatter (e.g. "Performing a shooting move ..."); it
# would flood pyretis.log and is not the run-level narrative -- the
# scheduler logs that from the main process. Errors/warnings from a
# worker, though, must not vanish now there is no per-worker file.
fileh.setLevel(logging.WARNING)
fileh.setFormatter(get_log_formatter(logging.WARNING))
logger.addHandler(fileh)
_ = worker_id # assigned for a unique, deterministic worker identity
[docs]
class future_list:
"""A managed list of future."""
[docs]
def __init__(self) -> None:
"""Initialize future list."""
self._futures: List[concurrent.futures.Future] = []
[docs]
def add(self, future: concurrent.futures.Future) -> None:
"""Add a future to list."""
self._futures.append(future)
[docs]
def as_completed(self) -> Optional[concurrent.futures.Future]:
"""Get future as they are done.
Returns
-------
concurrent.futures.Future, optional
return a future from the list, whenever it is done
or return None when the list is empty.
"""
if not self._futures:
return None
done, _ = concurrent.futures.wait(
self._futures, return_when=concurrent.futures.FIRST_COMPLETED
)
future_out = next(fut for fut in self._futures if fut in done)
self._futures.remove(future_out)
return future_out