from __future__ import annotations
import asyncio
import asyncio.exceptions
import collections
import dataclasses
import enum
import logging
import multiprocessing
import os
import signal
import uuid
from collections import deque
from datetime import datetime
from typing import (
Dict,
List,
Optional,
Set,
Union,
Sequence,
TYPE_CHECKING,
Deque,
)
import aiomultiprocess
import setproctitle
import zmq.error
from palaestrai.version import __version__
from palaestrai.core.protocol import (
ExperimentRunCancelRequest,
ExperimentRunCancelResponse,
ExperimentRunScheduleRequest,
ExperimentRunScheduleResponse,
ExperimentRunShutdownRequest,
ExperimentRunShutdownResponse,
ExperimentRunStartRequest,
ExperimentRunStartResponse,
ShutdownRequest,
ShutdownResponse,
)
from palaestrai.util import LogServer, spawn_wrapper
from palaestrai.core import EventStateMachine as ESM
from palaestrai.core import MajorDomoBroker, MajorDomoClient, RuntimeConfig
from palaestrai.types import ExperimentRunInstanceStatus
if TYPE_CHECKING:
import palaestrai.experiment
from palaestrai.store.receiver import StoreReceiver
LOG = logging.getLogger(__name__)
#: Number of attempts to persist the ``RUNNING`` status of an experiment run
#: instance, retried because the instance row is created asynchronously by the
#: StoreReceiver once it sees the ``ExperimentRunStartRequest``.
_RUNNING_STATUS_RETRIES = 30
#: Delay (seconds) between attempts to persist the ``RUNNING`` status.
_RUNNING_STATUS_RETRY_DELAY = 1.0
[docs]
class ExecutorState(enum.Enum):
PRISTINE = 0
INITIALIZED = 1
RUNNING = 2
SHUTDOWN = 3
EXITED = 4
SIGINT = 5
SIGABRT = 6
SIGTERM = 7
@dataclasses.dataclass
class _RunGovernorPCB:
"""Simple Process Control Block for RunGovernor control."""
run_governor_id: str
started_at: datetime
run_governor_process: aiomultiprocess.Process
experiment_run: palaestrai.experiment.ExperimentRun
experiment_run_id: Optional[str] = None
# Set once we have successfully persisted the ``RUNNING`` status for this
# run's experiment run instance. The instance row is created
# asynchronously by the StoreReceiver, so the write may have to be retried
# a few times until the row exists.
instance_status_running_marked: bool = False
[docs]
class ExperimentRunStartError(RuntimeError):
def __init__(self, experiment_run_id, run_governor_id, message):
super().__init__(message)
self.experiment_run_id = experiment_run_id
self.run_governor_id = run_governor_id
self.message = message
[docs]
class DeadMajorDomoBrokerError(RuntimeError):
def __init__(self):
super().__init__()
[docs]
class InterruptSignal(RuntimeError):
def __init__(self):
super().__init__()
async def _execute_run_governor(uid: str):
"""Executes the ::`RunGovernor` main loop, catching errors
This is a wrapper function around that creates a ::`RunGovernor` and
handles ::`RunGovernor.run`. It takes care of clearing signal handlers,
setting the proctitle, and generally catching errors in a meaningful in
order to report it to the Executor without simply dying.
This method belongs to the ::`Executor` logically, but is not part of the
class in order to avoid serialization/deserialization of the whole Executor
each time a new ::`RunGovernor` process is spawned.
Parameters
----------
uid : str
UID of the new ::`RunGovernor`
Return
------
Nothing.
"""
pid = os.getpid()
os.setpgrp()
try:
from .run_governor import RunGovernor
run_gov = RunGovernor(
uid=uid,
)
await run_gov.run() # type: ignore
except Exception as e:
LOG.critical("Execution of RunGovernor(uid=%s) failed: %s.", uid, e)
LOG.debug("Execution of RunGovernor(uid=%s) failed:", exc_info=e)
try:
os.killpg(pid, signal.SIGKILL)
except ProcessLookupError:
pass
raise e
[docs]
@ESM.monitor(is_mdp_worker=True)
class Executor:
"""The executor is the entrypoint for every run execution.
The role of the executor is to receive new experiment runs and
distribute them to existing :class:`RunGovernor` instances. If
palaestrAI is used in local run mode, the executor will initialize
a :class:`RunGovernor`.
The executor is an :class:`EventStateMachine`-monitored MDP worker. It
listens on the well-known service name
:data:`Executor.EXECUTOR_SERVICE_NAME`
and reacts to:
* :class:`ExperimentRunScheduleRequest` -- enqueue a new experiment run;
* :class:`ShutdownRequest` -- shut the executor down;
* ``signal.SIGCHLD`` -- a :class:`RunGovernor` process has ended;
* ``signal.SIGINT``/``SIGTERM``/``SIGABRT`` -- shut down on interruption.
Unlike its sibling workers, the executor additionally owns two
long-lived subprocesses that are *not* ESM-monitored: the
:class:`MajorDomoBroker` (the central message broker for the whole
palaestrAI process tree) and the :class:`StoreReceiver`. These are brought
up in :meth:`setup` and torn down in :meth:`teardown`.
Parameters
----------
parallel_runs : int, default: 1
Number of experiment runs the executor executes in
parallel.
is_service : bool, default: False
Controls the termination behaviour. When ``False`` (the default), the
executor behaves as it traditionally has: as soon as there are no more
scheduled runs and no active run governors, it shuts down and
``execute()`` returns. When ``True``, the executor is run as a
long-lived *service* (as used by ``palaestrai serve``): it stays online
even when idle, so that newly scheduled runs (added at runtime through
an :class:`ExperimentRunScheduleRequest`) are picked up. It only shuts
down on an explicit :class:`ShutdownRequest` or on receipt of
SIGINT/SIGABRT/SIGTERM.
Notes
-----
At some point, when the core protocol has progressed far enough,
we will be able to run several experiments at once from an
executor. But, until we're sure we can, the public API accepts
only one experiment.
"""
EXECUTOR_SERVICE_NAME = "executor"
def __init__(
self,
parallel_runs: int = 1,
is_service: bool = False,
):
# (At some point, when the core protocol has progressed far enough,
# we will be able to run several experiments at once from an executor.
# But, until we're sure we can, the public API accepts only one
# experiment.)
self._state: ExecutorState = ExecutorState.PRISTINE
self.is_service: bool = is_service
self._store: Optional[StoreReceiver] = None
self._store_process: Optional[aiomultiprocess.Process] = None
self._messsage_storage_queue: Optional[multiprocessing.Queue] = None
self._client: Union[None, MajorDomoClient] = None
self._broker: Union[None, MajorDomoBroker] = None
self._broker_process: Union[None, aiomultiprocess.Process] = None
self._broker_ctrl = multiprocessing.Pipe()
self._num_parallel_runs = parallel_runs
self._run_governors: Dict[str, _RunGovernorPCB] = {}
self._runs_scheduled: Deque[palaestrai.experiment.ExperimentRun] = (
deque()
)
# Background tasks that persist the RUNNING instance status; tracked so
# they can be cancelled cleanly on teardown.
self._running_status_tasks: Set[asyncio.Task] = set()
self._log_server = LogServer("127.0.0.1", RuntimeConfig().logger_port)
# Set default start method:
try:
multiprocessing.set_start_method(RuntimeConfig().fork_method)
aiomultiprocess.set_start_method(RuntimeConfig().fork_method)
except RuntimeError as e:
# When running nested or in tests, the context has already been
# set, so we cannot do this anymore.
# This is okay and not an error, we just need to roll with it.
LOG.debug("Cannot set start method: %s", e)
[docs]
async def setup(self):
"""Brings up the executor and connects it as an MDP worker.
Starts the log server, the store subprocess, and the
:class:`MajorDomoBroker`, then connects this executor as an MDP worker
on the well-known :data:`EXECUTOR_SERVICE_NAME` service. The broker
must be up before the worker connects, which is guaranteed because
:meth:`_init_communication` blocks until the broker has reported its
bound URI back over the control pipe.
"""
self._state = ExecutorState.PRISTINE
setproctitle.setproctitle("palaestrAI[Executor]")
await self._init_logging()
LOG.info("This is palaestrAI, version %s", __version__)
self._init_store()
self._init_communication()
self._install_sigabrt_handler()
self._state = ExecutorState.INITIALIZED
self.mdp_service = Executor.EXECUTOR_SERVICE_NAME
self._state = ExecutorState.RUNNING
LOG.debug("Executor(id=0x%x) is set up and running.", id(self))
[docs]
async def execute(self):
"""Executes scheduled experiment runs.
This is the public entrypoint, kept for backwards compatibility with
callers that ``await executor.execute()`` (the local runner and the
``serve`` child). It drives the ESM event loop (via the injected
:meth:`run`) and returns the final :class:`ExecutorState` once the loop
concludes.
Returns
-------
ExecutorState
The state the executor is now in, either ``EXITED`` if everything
exited normally, or one of the SIG* states if a signal was
received.
"""
await self.run() # type: ignore[attr-defined]
return self._state
@ESM.enter
async def _enter(self):
"""Initial state: start any already-scheduled runs."""
LOG.debug("Executor(id=0x%x) entering main loop.", id(self))
self._start_scheduled_runs()
self._maybe_stop_when_idle()
[docs]
async def teardown(self):
"""Conducts an orderly shutdown of the executor-owned subprocesses.
ESM cleans up the ESM-monitored RunGovernor processes itself, but the
executor still owns the broker and store subprocesses (and the
RunGovernors' process *groups*, which ESM does not know about), so we
tear those down here.
"""
LOG.info("palaestrAI executor initiating shutdown procedure.")
for task in self._running_status_tasks:
task.cancel()
self._running_status_tasks.clear()
await self._shutdown_all_run_governors()
try:
if self._client is not None:
await self._client.destroy()
except zmq.error.ZMQError:
# This can happen on ^C. It's actually not that bad, so we just
# do a debug log entry here.
LOG.debug(
"%s could not send destroy message via MajorDomoClient to "
"MajorDomoBroker; ignoring that anyways and dragging on.",
self,
)
if self._broker_process is not None:
self._broker_process.terminate()
while self._broker_process.is_alive():
try:
await self._broker_process.join(1)
except asyncio.TimeoutError:
await asyncio.sleep(0.05)
LOG.info("Major domo broker has shut down.")
if self._store_process is not None:
while self._store_process.is_alive():
try:
await self._store_process.join(1)
except asyncio.TimeoutError:
LOG.info("Waiting for the storage backend to shut down...")
await asyncio.sleep(0.05)
if self._messsage_storage_queue is not None:
self._messsage_storage_queue.close()
LOG.info("Storage backend has shut down.")
self._remove_sigabrt_handler()
if self._state == ExecutorState.SHUTDOWN:
self._state = ExecutorState.EXITED
await self._log_server.stop()
def _install_sigabrt_handler(self):
"""Installs a SIGABRT handler in the loop.
The ESM installs handlers for SIGINT/SIGTERM/SIGTRAP itself, but not
for SIGABRT, which the executor traditionally reacts to. We therefore
wire it up manually so that :meth:`_handle_sigabrt` is invoked.
"""
loop = asyncio.get_running_loop()
loop.add_signal_handler(
signal.SIGABRT,
lambda: asyncio.create_task(
self.__esm__._handle_event(signal.SIGABRT) # type: ignore[attr-defined]
),
)
def _remove_sigabrt_handler(self):
try:
asyncio.get_running_loop().remove_signal_handler(signal.SIGABRT)
except (ValueError, RuntimeError):
pass
@ESM.on(signal.SIGINT)
def _handle_sigint(self):
LOG.info(
"palaestrAI executor has received signal %s, shutting down",
signal.SIGINT,
)
self._state = ExecutorState.SIGINT
self.stop() # type: ignore[attr-defined]
@ESM.on(signal.SIGTERM)
def _handle_sigterm(self):
LOG.info(
"palaestrAI executor has received signal %s, shutting down",
signal.SIGTERM,
)
self._state = ExecutorState.SIGTERM
self.stop() # type: ignore[attr-defined]
@ESM.on(signal.SIGABRT)
def _handle_sigabrt(self):
LOG.info(
"palaestrAI executor has received signal %s, shutting down",
signal.SIGABRT,
)
self._state = ExecutorState.SIGABRT
self.stop() # type: ignore[attr-defined]
async def _init_logging(self):
"""Starts the log server and configures log filters"""
await self._log_server.start()
RuntimeConfig().logger_port = self._log_server.listen_port
def filter_record_above_debug(record: logging.LogRecord):
if record.levelname != "DEBUG":
return False
return True
for h in [h for h in logging.root.handlers if "debug" in h.name]:
h.addFilter(filter_record_above_debug)
# Service mode (``palaestrai serve``) automatically persists log
# records to a separate SQLite database. Because every subprocess
# forwards its records to this LogServer, which re-injects them into
# this process' loggers, a single handler on the root logger here
# captures the whole process tree. The handler itself only stores
# records that carry an ``experiment_run_instance_uid`` (stamped by
# the RunGovernor's MetadataLogFilter); everything else is dropped by
# the handler and only reaches stdout via the console handlers. This
# is NOT enabled for the normal ``experiment-start`` path.
if self.is_service:
from palaestrai.util.sqlite_log_handler import SQLiteLogHandler
if not any(
isinstance(h, SQLiteLogHandler) for h in logging.root.handlers
):
logging.root.addHandler(
SQLiteLogHandler(
db_uri=RuntimeConfig().log_store_uri,
level="INFO",
)
)
def _init_store(self):
LOG.info("Starting storage backend...")
# We load & init the storage backend in any case, even if it is
# disabled. If we don't, we create a queue that eats messages but
# is never able to release them.
from palaestrai.store.receiver import StoreReceiver
self._messsage_storage_queue = multiprocessing.Queue()
self._store = StoreReceiver(message_queue=self._messsage_storage_queue)
if not self._store._is_enabled:
LOG.info("Storage backend disabled by runtime configuration.")
if self._messsage_storage_queue is not None:
self._messsage_storage_queue.close()
self._messsage_storage_queue = None
self._store = None
return
self._store_process = aiomultiprocess.Process(
target=spawn_wrapper,
args=(
"StoreReceiver",
RuntimeConfig().to_dict(),
self._store.run,
),
name="Executor-StoreReceiver",
)
self._store_process.start()
def _init_communication(self):
"""Initialization of all core components"""
LOG.info("Starting Major Domo Broker...")
self._broker = MajorDomoBroker(
uri=None,
ctrl=self._broker_ctrl[1],
store_queue=self._messsage_storage_queue,
)
self._broker_process = aiomultiprocess.Process(
daemon=True,
target=spawn_wrapper,
args=(
"MajorDomoBroker",
RuntimeConfig().to_dict(),
self._broker.mediate,
),
name="Executor-MajorDomoBroker",
)
self._broker_process.start()
broker_uri = self._broker_ctrl[0].recv()
RuntimeConfig().broker_uri = broker_uri
LOG.info(
"Major Domo Broker is bound to %s.",
RuntimeConfig().broker_uri,
)
self._client = MajorDomoClient(RuntimeConfig().broker_uri)
@staticmethod
def _set_instance_status(instance_uid: str, status) -> bool:
"""Persist the lifecycle status of an experiment run instance.
This writes the given
:class:`~palaestrai.types.ExperimentRunInstanceStatus` to the
``experiment_run_instances`` row identified by ``instance_uid``. It is
the authoritative status writer used while running as a service: the
REST API only ever *reads* this column.
The instance row itself is created asynchronously by the
:class:`~palaestrai.store.receiver.StoreReceiver` (in a separate
process) once it sees the ``ExperimentRunStartRequest``. There is
therefore an inherent race when transitioning to ``RUNNING``: the row
may not exist yet. This method tolerates that gracefully -- it returns
``False`` if the row is not (yet) present or if the store is disabled,
and ``True`` on a successful update. It never raises, so a status
bookkeeping problem can never take down the executor.
Parameters
----------
instance_uid : str
UID of the :class:`ExperimentRunInstance` to update.
status : palaestrai.types.ExperimentRunInstanceStatus
The new status to persist.
Returns
-------
bool
``True`` if a row was found and updated, ``False`` otherwise.
"""
if not RuntimeConfig().store_uri:
return False
try:
from palaestrai.store import Session
from palaestrai.store import database_model as dbm
session = Session()
try:
instance = (
session.query(dbm.ExperimentRunInstance)
.filter(dbm.ExperimentRunInstance.uid == instance_uid)
.one_or_none()
)
if instance is None:
return False
instance.status = status
session.commit()
LOG.debug(
"Executor(id=0x%x) set status of "
"ExperimentRunInstance(uid=%s) to %s.",
id(Executor),
instance_uid,
status,
)
return True
finally:
session.close()
except Exception as e: # noqa: BLE001 -- never break the executor
LOG.warning(
"Could not update status of ExperimentRunInstance(uid=%s) "
"to %s: %s. Continuing regardless.",
instance_uid,
status,
e,
)
return False
[docs]
def experiment_runs(self) -> List[ExperimentRunRuntimeInformation]:
runs = list()
for run in self._runs_scheduled:
runs.append(ExperimentRunRuntimeInformation(experiment_run=run))
for pcb in self._run_governors.values():
runs.append(
ExperimentRunRuntimeInformation(
experiment_run=pcb.experiment_run,
started_at=pcb.started_at,
experiment_run_id=pcb.experiment_run_id,
)
)
return runs
# ----------------------------------------------------------------- #
# Scheduling and process supervision
# ----------------------------------------------------------------- #
[docs]
def schedule(
self,
experiment_run: Union[
palaestrai.experiment.ExperimentRun,
Sequence[palaestrai.experiment.ExperimentRun],
],
):
"""Schedules an experiment run to be executed.
This method schedules experiment runs, i.e., puts them in the
waiting queue. The ESM event loop (started by ::`execute`) picks up
experiment run objects and executes them.
Parameters
----------
experiment_run : Union[palaestrai.experiment.ExperimentRun,
Sequence[palaestrai.experiment.ExperimentRun]]
One or many :class:`palaestrai.experiment.ExperimentRun` objects,
which are added to the queue.
"""
if isinstance(experiment_run, collections.abc.Sequence):
self._runs_scheduled.extend(experiment_run)
else:
self._runs_scheduled.append(experiment_run)
def _start_scheduled_runs(self):
"""Starts as many scheduled runs as the parallelism budget allows.
For each run that can be started, a :class:`RunGovernor` process is
spawned (and ESM-monitored) and an :class:`ExperimentRunStartRequest`
is sent to it. The reply is handled by
:meth:`_handle_run_start_response`.
"""
while (
len(self._run_governors) < self._num_parallel_runs
and self._runs_scheduled
):
experiment_run = self._runs_scheduled.popleft()
run_gov_uid = str(uuid.uuid4())
self._spawn_run_governor(run_gov_uid, experiment_run)
LOG.info(
'Starting experiment run "%s". Our business is life itself',
experiment_run.uid,
)
self._request_experiment_run_start(run_gov_uid)
@ESM.spawns
def _spawn_run_governor(
self,
run_gov_uid: str,
experiment_run: palaestrai.experiment.ExperimentRun,
) -> aiomultiprocess.Process:
"""Launches and ESM-monitors a new :class:`RunGovernor` process."""
run_gov_process = aiomultiprocess.Process(
target=spawn_wrapper,
args=(
f"palaestrAI[RunGovernor-{run_gov_uid[-6:]}",
RuntimeConfig().to_dict(),
_execute_run_governor,
[run_gov_uid],
),
name=f"RunGovernor-{run_gov_uid[-6:]}",
)
run_gov_process.start()
self._run_governors[run_gov_uid] = _RunGovernorPCB(
experiment_run=experiment_run,
started_at=datetime.now(),
run_governor_id=run_gov_uid,
run_governor_process=run_gov_process,
)
LOG.debug("Launched new RunGovernor %s", run_gov_uid)
return run_gov_process
@ESM.requests
def _request_experiment_run_start(self, run_governor_id: str):
"""Asks a spawned :class:`RunGovernor` to start its experiment run."""
pcb = self._run_governors[run_governor_id]
return ExperimentRunStartRequest(
sender_executor_id=Executor.EXECUTOR_SERVICE_NAME,
receiver_run_governor_id=run_governor_id,
experiment_run_id=pcb.experiment_run.uid,
experiment_run=pcb.experiment_run,
)
@ESM.on(ExperimentRunStartResponse)
async def _handle_run_start_response(
self, response: ExperimentRunStartResponse
):
"""Handles the reply to an :class:`ExperimentRunStartRequest`."""
run_governor_id = response.sender
pcb = self._run_governors.get(run_governor_id)
if pcb is None:
LOG.debug(
"Executor(id=0x%x) got ExperimentRunStartResponse from "
"unknown RunGovernor(uid=%s); ignoring.",
id(self),
run_governor_id,
)
return
if response.successful:
pcb.experiment_run_id = response.experiment_run_id
LOG.debug(
"Executor(id=0x%x) received successful ExperimentRunStart"
"Response from RunGovernor(uid=%s) for ExperimentRun(id=%s).",
id(self),
run_governor_id,
response.experiment_run_id,
)
# The StoreReceiver creates the instance row asynchronously; mark
# it RUNNING via a bounded retry task so the REST API reflects the
# live state.
task = asyncio.create_task(self._persist_running_status(pcb))
self._running_status_tasks.add(task)
task.add_done_callback(self._running_status_tasks.discard)
else:
LOG.error(
"Could not start experiment run on RunGovernor(uid=%s): %s",
run_governor_id,
response.error,
)
self._set_instance_status(
pcb.experiment_run.instance_uid,
ExperimentRunInstanceStatus.ERROR,
)
# Terminate the governor; the SIGCHLD handler cleans up the PCB.
pcb.run_governor_process.terminate()
async def _persist_running_status(self, pcb: _RunGovernorPCB):
"""Persists the RUNNING status, retrying until the row exists."""
for _ in range(_RUNNING_STATUS_RETRIES):
if pcb.instance_status_running_marked:
return
if not pcb.run_governor_process.is_alive():
return
if self._set_instance_status(
pcb.experiment_run.instance_uid,
ExperimentRunInstanceStatus.RUNNING,
):
pcb.instance_status_running_marked = True
return
try:
await asyncio.sleep(_RUNNING_STATUS_RETRY_DELAY)
except asyncio.CancelledError:
return
@ESM.on(signal.SIGCHLD)
async def _handle_child(
self,
process: Union[aiomultiprocess.Process, multiprocessing.Process],
):
"""Reaps a terminated :class:`RunGovernor` process."""
pcb = next(
(
p
for p in self._run_governors.values()
if p.run_governor_process is process
),
None,
)
if pcb is None:
LOG.debug(
"Executor(id=0x%x) saw a process end that is not a tracked "
"RunGovernor: %s",
id(self),
getattr(process, "name", process),
)
return
if process.exitcode != 0:
LOG.error("RunGovernor process %s died.", pcb.run_governor_id)
self._set_instance_status(
pcb.experiment_run.instance_uid,
ExperimentRunInstanceStatus.ERROR,
)
else:
self._set_instance_status(
pcb.experiment_run.instance_uid,
ExperimentRunInstanceStatus.FINISHED,
)
self._run_governors.pop(pcb.run_governor_id, None)
LOG.debug(
"Executor(id=0x%x) has %d active run governors and %d "
"experiment runs scheduled.",
id(self),
len(self._run_governors),
len(self._runs_scheduled),
)
self._start_scheduled_runs()
self._maybe_stop_when_idle()
def _maybe_stop_when_idle(self):
"""Shuts the executor down when idle, unless running as a service."""
if self.is_service:
return
if not self._run_governors and not self._runs_scheduled:
LOG.debug("Executor(id=0x%x) is idle, shutting down.", id(self))
self._state = ExecutorState.SHUTDOWN
self.stop() # type: ignore[attr-defined]
@ESM.on(ExperimentRunScheduleRequest)
def _handle_schedule_request(self, request: ExperimentRunScheduleRequest):
"""Enqueues a run requested over the broker and starts it if possible."""
LOG.info(
"Executor(id=0x%x) received ExperimentRunScheduleRequest for "
"ExperimentRun(uid=%s).",
id(self),
request.experiment_run.uid,
)
try:
self.schedule(request.experiment_run)
self._start_scheduled_runs()
response = ExperimentRunScheduleResponse(
sender=Executor.EXECUTOR_SERVICE_NAME,
receiver=request.sender,
experiment_run_id=request.experiment_run.uid,
successful=True,
)
except Exception as e: # noqa: BLE001 -- report failure to the client
LOG.error(
"Executor(id=0x%x) could not schedule ExperimentRun(uid=%s): "
"%s",
id(self),
request.experiment_run.uid,
e,
)
response = ExperimentRunScheduleResponse(
sender=Executor.EXECUTOR_SERVICE_NAME,
receiver=request.sender,
experiment_run_id=request.experiment_run.uid,
successful=False,
error=str(e),
)
return response
@ESM.on(ExperimentRunCancelRequest)
async def _handle_cancel_request(
self, request: ExperimentRunCancelRequest
):
"""Cancels a single run instance requested over the broker.
Two cases are handled:
* The instance is currently *running*: the associated
:class:`RunGovernor` is shut down gracefully via
:meth:`cancel`. The SIGCHLD handler then reaps the process and
writes the terminal status.
* The instance is still *scheduled* (queued but not yet started):
it is removed from the scheduling queue and marked ``FINISHED``
directly, since no process exists for it.
Note
----
The :class:`~palaestrai.types.ExperimentRunInstanceStatus` enum has
no dedicated ``CANCELLED`` state, so a cancelled instance ends up as
``FINISHED`` (queued case) or as whatever the RunGovernor's exit code
maps to (``FINISHED``/``ERROR``) for the running case.
"""
instance_id = request.experiment_run_instance_id
LOG.info(
"Executor(id=0x%x) received ExperimentRunCancelRequest for "
"ExperimentRunInstance(uid=%s).",
id(self),
instance_id,
)
try:
pcb = next(
(
p
for p in self._run_governors.values()
if p.experiment_run.instance_uid == instance_id
),
None,
)
if pcb is not None:
await self._cancel_run_governor(pcb)
else:
queued = next(
(
r
for r in self._runs_scheduled
if r.instance_uid == instance_id
),
None,
)
if queued is not None:
try:
self._runs_scheduled.remove(queued)
except ValueError:
pass
self._set_instance_status(
instance_id,
ExperimentRunInstanceStatus.FINISHED,
)
else:
LOG.warning(
"Executor(id=0x%x) cannot cancel "
"ExperimentRunInstance(uid=%s): not running or "
"scheduled here.",
id(self),
instance_id,
)
response = ExperimentRunCancelResponse(
sender=Executor.EXECUTOR_SERVICE_NAME,
receiver=request.sender,
experiment_run_instance_id=instance_id,
successful=True,
)
except Exception as e: # noqa: BLE001 -- report failure to the client
LOG.error(
"Executor(id=0x%x) could not cancel "
"ExperimentRunInstance(uid=%s): %s",
id(self),
instance_id,
e,
)
response = ExperimentRunCancelResponse(
sender=Executor.EXECUTOR_SERVICE_NAME,
receiver=request.sender,
experiment_run_instance_id=instance_id,
successful=False,
error=str(e),
)
return response
async def _cancel_run_governor(self, pcb: _RunGovernorPCB):
"""Gracefully shuts the RunGovernor behind ``pcb`` down.
Sends an :class:`ExperimentRunShutdownRequest` to the governor over
the executor's MDP client. The SIGCHLD handler reaps the process and
writes the terminal instance status once it exits.
"""
msg = ExperimentRunShutdownRequest(
sender_executor_id=Executor.EXECUTOR_SERVICE_NAME,
receiver_run_governor_id=pcb.run_governor_id,
experiment_run_id=pcb.experiment_run_id or pcb.experiment_run.uid,
)
LOG.debug(
"Executor(id=0x%x) sending ExperimentRunShutdownRequest to "
"RunGovernor(uid=%s) to cancel ExperimentRunInstance(uid=%s).",
id(self),
pcb.run_governor_id,
pcb.experiment_run.instance_uid,
)
assert self._client is not None
response = await self._client.send(pcb.run_governor_id, msg)
if (
isinstance(response, ExperimentRunShutdownResponse)
and not response.successful
):
raise RuntimeError(response.error)
@ESM.on(ShutdownRequest)
def _handle_shutdown_request(self, request: ShutdownRequest):
"""Shuts the executor down on request from a client."""
LOG.info(
"Executor(id=0x%x) received ShutdownRequest, shutting down.",
id(self),
)
self._runs_scheduled.clear()
self._state = ExecutorState.SHUTDOWN
self.stop() # type: ignore[attr-defined]
return ShutdownResponse(
sender=Executor.EXECUTOR_SERVICE_NAME,
receiver=request.sender,
experiment_run_id=request.experiment_run_id,
)
[docs]
async def cancel(self, experiment_run_id):
"""Shuts an experiment run down prematurely.
This method sends a :py:class:`ExperimentRunShutdownRequest` to
the :py:class:`RunGovernor` responsible for executing the
associated experiment run. This allows for a graceful, yet
premature shutdown of a running experiment run.
Parameters
----------
experiment_run_id: str
UID of the experiment run to shut down.
"""
rg_dict_item = next(
filter(
lambda i: i[1].experiment_run_id == experiment_run_id,
self._run_governors.items(),
),
None,
)
if rg_dict_item is None:
LOG.error(
"Executor (id=0x%x) cannot terminate ExperimentRun("
"run_id=%s): Cannot find RunGovernor",
id(self),
experiment_run_id,
)
return
run_governor_id = rg_dict_item[0]
LOG.debug(
"Executor (id=0x%x) sending ExperimentRunShutdownRequest for "
"ExperimentRun(run_id=%s).",
id(self),
experiment_run_id,
)
msg = ExperimentRunShutdownRequest(
sender_executor_id=Executor.EXECUTOR_SERVICE_NAME,
receiver_run_governor_id=run_governor_id,
experiment_run_id=experiment_run_id,
)
assert self._client is not None
response = await self._client.send(run_governor_id, msg)
LOG.debug("Executor(id=0x%x) received %s", id(self), response)
if not isinstance(response, ExperimentRunShutdownResponse):
return
if not response.successful:
raise RuntimeError(response.error)
async def _shutdown_all_run_governors(self):
"""Shuts all ::`RunGovernor` instances down, cleaning up
This method terminates all running ::`RunGovernor` instances. They
get a nice message first, but are forcefully terminated if they don't
react. Because RunGovernors call ``os.setpgrp()``, they are taken down
by process *group* (which the ESM's own process cleanup does not do).
"""
for run_gov_uid in list(self._run_governors):
pcb = self._run_governors[run_gov_uid]
if not pcb.run_governor_process.is_alive():
continue
LOG.info(
"Signalling RunGovernor for experiment run '%s' to shut down.",
pcb.experiment_run_id,
)
if pcb.experiment_run_id:
try:
await asyncio.wait_for(
self.cancel(pcb.experiment_run_id), timeout=5
)
except asyncio.TimeoutError:
LOG.debug(
"Executor(id=%s) has encountered a "
"RunGovernor(uid=%s, run_id=%s) "
"that seems to be still active, trying to abort.",
id(self),
pcb.run_governor_id,
pcb.experiment_run_id,
)
LOG.info(
"Sending SIGTERM to process group %d...",
pcb.run_governor_process.pid,
)
pgpid = pcb.run_governor_process.pid
try:
os.killpg(pgpid, signal.SIGTERM)
except ProcessLookupError:
pass # Already down - ok, don't race here.
try:
await pcb.run_governor_process.join(3)
except asyncio.TimeoutError:
pass # Okay, we send SIGKILL anyways...
try:
# Zombies will be shot:
os.killpg(pgpid, signal.SIGKILL)
await pcb.run_governor_process.join(3)
except (asyncio.TimeoutError, ProcessLookupError):
pass # Just a shot in the dark.
def __str__(self):
return "Executor(id=%s)" % id(self)