from __future__ import annotations
import heapq
from collections.abc import Iterator, Mapping
from enum import Enum, auto
from multiprocessing import Event as ProcessEvent, Pipe, get_context
from multiprocessing.process import BaseProcess
from pathlib import Path
from threading import Event as ThreadEvent, Thread
from time import monotonic
from typing import TYPE_CHECKING, Any, cast
from edri.abstract.manager.worker import (
Worker, WorkerExit, WorkerFactoryProtocol, WorkerNotFoundError, WorkerOutcome,
WorkerPolicy, WorkerRuntimeProtocol, WorkerState, WorkerStopReason,
WorkerTimeout,
)
from edri.abstract.worker import WorkerProcess
from edri.dataclass.event import Event
from edri.events.edri.manager import WorkerQuit
from edri.utility.storage import Storage
if TYPE_CHECKING:
from edri.abstract.manager.manager_base import ManagerBase
class _DeadlineType(Enum):
QUEUE = auto()
START = auto()
RUN = auto()
TOTAL = auto()
CANCEL_GRACE = auto()
TERMINATE_GRACE = auto()
_TIMEOUT_KIND = {
_DeadlineType.QUEUE: WorkerTimeout.QUEUE,
_DeadlineType.START: WorkerTimeout.START,
_DeadlineType.RUN: WorkerTimeout.RUN,
_DeadlineType.TOTAL: WorkerTimeout.TOTAL,
}
[docs]
class WorkerView(Mapping[str, Worker]):
"""A stable, live, read-only view over selected controller states."""
def __init__(self, controller: WorkerController, states: frozenset[WorkerState]) -> None:
self._controller = controller
self._states = states
def __getitem__(self, worker_id: str) -> Worker:
worker = self._controller[worker_id]
if worker.state not in self._states:
raise KeyError(worker_id)
return worker
def __iter__(self) -> Iterator[str]:
return (worker_id for worker_id, worker in self._controller.items() if worker.state in self._states)
def __len__(self) -> int:
return sum(worker.state in self._states for worker in self._controller.values())
[docs]
class WorkerController(Mapping[str, Worker]):
"""Owns a Manager's worker registry, scheduling, lifecycle and deadlines."""
def __init__(self, manager: ManagerBase, policy: WorkerPolicy | None = None) -> None:
self._manager = manager
self._workers: Storage[Worker] = Storage()
self.policy = policy or WorkerPolicy()
self._accepting = True
self._deadline_heap: list[tuple[float, int, str, _DeadlineType]] = []
self._deadlines: dict[tuple[str, _DeadlineType], tuple[float, int]] = {}
self._deadline_sequence = 0
self._pending = WorkerView(self, frozenset({WorkerState.PENDING}))
self._running = WorkerView(self, frozenset({WorkerState.RUNNING}))
self._stopping = WorkerView(
self, frozenset({WorkerState.STOPPING, WorkerState.TERMINATING})
)
def __getitem__(self, worker_id: str) -> Worker:
worker: Worker = self._workers[worker_id]
return worker
def __iter__(self) -> Iterator[str]:
return iter(self._workers)
def __len__(self) -> int:
return len(self._workers)
@property
def pending(self) -> WorkerView:
return self._pending
@property
def running(self) -> WorkerView:
return self._running
@property
def stopping(self) -> WorkerView:
return self._stopping
[docs]
def start(
self,
event: Event | None,
resolver: WorkerFactoryProtocol,
/,
*args: Any,
environment: Path | None = None,
policy: WorkerPolicy | None = None,
**kwargs: Any,
) -> str:
if not self._accepting:
raise RuntimeError("worker controller is closing")
now = monotonic()
worker = Worker(
event=event,
resolver=resolver,
args=args,
kwargs=kwargs,
environment=environment,
policy=policy or self.policy,
created=now,
)
worker_id: str = self._workers.append(worker)
if worker.policy.total_timeout is not None:
self._schedule_deadline(worker_id, _DeadlineType.TOTAL, now + worker.policy.total_timeout)
if worker.policy.queue_timeout is not None:
self._schedule_deadline(worker_id, _DeadlineType.QUEUE, now + worker.policy.queue_timeout)
self._start_pending()
return worker_id
def _active_count(self) -> int:
return sum(worker.state is not WorkerState.PENDING for worker in self.values())
def _has_capacity(self) -> bool:
limit: int = self._manager.max_workers
return limit == -1 or self._active_count() < limit
def _start_pending(self) -> None:
if not self._accepting:
return
while self._has_capacity():
pending = next(iter(self.pending), None)
if pending is None:
return
self._start(pending)
def _start(self, worker_id: str) -> None:
worker = self._workers.get(worker_id)
if worker is None or worker.state is not WorkerState.PENDING:
return
self._clear_deadline(worker_id, _DeadlineType.QUEUE)
pipe_local, pipe_remote = Pipe()
name = f"{self._manager.name}-Worker_{worker_id}"
resolver = worker.resolver
try:
# Public worker factories are classes. Keep construction typed via
# the callable protocol while using the concrete class only for the
# custom multiprocessing environment bootstrap.
resolver_class = cast(type[Any], resolver)
runtime: WorkerRuntimeProtocol
if issubclass(resolver_class, WorkerProcess) and worker.environment:
context = get_context("spawn")
context.set_executable(str(worker.environment / "bin" / "python"))
process_runtime: WorkerProcess[Any] = object.__new__(resolver_class)
resolver_class.__init__(
process_runtime, pipe_remote, worker.event, name,
*worker.args, **worker.kwargs,
)
setattr(process_runtime, "_Popen", context.Process._Popen)
runtime = process_runtime
cancel_event: Any = context.Event()
else:
runtime = resolver(pipe_remote, worker.event, name, *worker.args, **worker.kwargs)
cancel_event = ProcessEvent() if isinstance(runtime, BaseProcess) else ThreadEvent()
runtime.name = name
runtime._bind_lifecycle(cancel_event)
# Thread runtimes share this process's address space, so they can
# publish to the manager's router queue directly. A process runtime
# cannot be handed the queue object this way -- binding one does not
# merely go unused, it breaks the spawn bootstrap -- so it is left
# unbound and falls back to the worker-manager pipe.
#
# getattr, because ManagerBase.__init__ only annotates router_queue
# when it is constructed without one; starting a worker must not be
# what discovers that. Such a worker sees no binding, like a process
# worker does.
if not isinstance(runtime, BaseProcess):
router_queue = getattr(self._manager, "router_queue", None)
if router_queue is not None:
runtime._bind_router_queue(router_queue)
worker.pipe = pipe_local
worker.runtime = runtime
worker.state = WorkerState.STARTING
worker.spawned = monotonic()
if worker.policy.start_timeout is not None:
self._schedule_deadline(
worker_id, _DeadlineType.START, worker.spawned + worker.policy.start_timeout
)
runtime.start()
if isinstance(runtime, BaseProcess):
pipe_remote.close()
self._manager._worker_runtime_started(worker_id, worker)
self._manager.logger.debug("Worker %s is starting", name)
except BaseException:
pipe_local.close()
pipe_remote.close()
self._manager.logger.exception("Worker %s could not be started", name)
self._finish(worker_id, outcome=WorkerOutcome.FAILED, exit=WorkerExit.LOST)
raise
[docs]
def started(self, worker_id: str, at: float | None = None) -> None:
worker = self._workers.get(worker_id)
if worker is None or worker.state not in (WorkerState.STARTING, WorkerState.STOPPING):
return
started = monotonic() if at is None else at
worker.started = started
self._clear_deadline(worker_id, _DeadlineType.START)
if worker.state is WorkerState.STARTING:
worker.state = WorkerState.RUNNING
if worker.policy.run_timeout is not None:
self._schedule_deadline(
worker_id, _DeadlineType.RUN, started + worker.policy.run_timeout
)
[docs]
def cancel(
self, worker_id: str, reason: WorkerStopReason = WorkerStopReason.REQUESTED
) -> None:
worker = self._workers.get(worker_id)
if worker is None:
raise WorkerNotFoundError(worker_id)
if worker.state in (WorkerState.STOPPING, WorkerState.TERMINATING):
return
worker.stop_reason = reason
worker.stop_requested = monotonic()
self._clear_execution_deadlines(worker_id)
if worker.state is WorkerState.PENDING:
outcome = WorkerOutcome.TIMED_OUT if reason is WorkerStopReason.TIMEOUT else WorkerOutcome.CANCELLED
self._finish(worker_id, outcome=outcome, exit=WorkerExit.CANCELLED)
return
worker.state = WorkerState.STOPPING
runtime = worker.runtime
if runtime is not None:
runtime._request_cancel()
if isinstance(runtime, BaseProcess):
self._schedule_deadline(
worker_id, _DeadlineType.CANCEL_GRACE,
monotonic() + worker.policy.cancel_grace,
)
elif worker.policy.cancel_grace is not None:
self._schedule_deadline(
worker_id, _DeadlineType.CANCEL_GRACE,
monotonic() + worker.policy.cancel_grace,
)
[docs]
def cancel_all(self) -> int:
targets = [
worker_id for worker_id, worker in self.items()
if worker.state in (WorkerState.PENDING, WorkerState.STARTING, WorkerState.RUNNING)
]
for worker_id in targets:
self.cancel(worker_id)
return len(targets)
[docs]
def finish_from_quit(self, worker_id: str, event: WorkerQuit) -> None:
worker = self._workers.get(worker_id)
if worker is None:
return
if worker.stop_reason is WorkerStopReason.TIMEOUT:
outcome = WorkerOutcome.TIMED_OUT
elif worker.stop_reason is not None or event.exit is WorkerExit.CANCELLED:
outcome = WorkerOutcome.CANCELLED
elif event.exit is WorkerExit.RAISED:
outcome = WorkerOutcome.FAILED
else:
outcome = WorkerOutcome.SUCCEEDED
self._finish(worker_id, outcome=outcome, exit=event.exit)
[docs]
def finish_lost(self, worker_id: str) -> None:
worker = self._workers.get(worker_id)
if worker is None:
return
if worker.stop_reason is WorkerStopReason.TIMEOUT:
outcome = WorkerOutcome.TIMED_OUT
elif worker.stop_reason is not None:
outcome = WorkerOutcome.CANCELLED
else:
outcome = WorkerOutcome.FAILED
exit = WorkerExit.TERMINATED if worker.state is WorkerState.TERMINATING else WorkerExit.LOST
self._finish(worker_id, outcome=outcome, exit=exit)
def _finish(self, worker_id: str, *, outcome: WorkerOutcome, exit: WorkerExit) -> bool:
worker = self._workers.get(worker_id)
if worker is None:
return False
worker.outcome = outcome
worker.exit = exit
worker.finished = monotonic()
self._clear_worker_deadlines(worker_id)
if worker.pending_responses:
self._manager.logger.warning(
"Worker %s exited with %d unresolved response(s) pending",
worker_id, len(worker.pending_responses),
)
for pending_id in tuple(worker.pending_responses):
entry = self._manager._pending_responses.pop(pending_id, None)
if entry is not None:
self._manager._release_pending_response_type(entry[0])
if worker.pipe is not None:
# Async managers must remove add_reader() while fileno() is still
# valid; their hook is idempotent for WorkerQuit's eager removal.
self._manager._worker_runtime_finishing(worker_id, worker)
try:
worker.pipe.close()
except OSError:
pass
if isinstance(worker.runtime, BaseProcess) and worker.runtime.pid is not None:
worker.runtime.join(timeout=0.05)
if not worker.runtime.is_alive():
try:
worker.runtime.close()
except ValueError:
pass
self._workers.pop(worker_id, None)
self._manager.logger.debug(
"Worker %s finished: outcome=%s exit=%s", worker_id, outcome.name, exit.name
)
self._start_pending()
return True
[docs]
def pipes(self) -> set[Any]:
return {worker.pipe for worker in self.values() if worker.pipe is not None}
[docs]
def find_by_pipe(self, pipe: Any) -> tuple[str, Worker]:
for worker_id, worker in self.items():
if worker.pipe == pipe:
return worker_id, worker
raise ValueError
def _schedule_deadline(self, worker_id: str, kind: _DeadlineType, at: float) -> None:
self._discard_stale_deadlines()
previous = self._deadline_heap[0][0] if self._deadline_heap else None
sequence = self._deadline_sequence
self._deadline_sequence += 1
self._deadlines[(worker_id, kind)] = (at, sequence)
heapq.heappush(self._deadline_heap, (at, sequence, worker_id, kind))
if previous is None or at < previous:
self._manager._worker_deadline_changed()
def _clear_deadline(self, worker_id: str, kind: _DeadlineType) -> None:
self._deadlines.pop((worker_id, kind), None)
def _clear_worker_deadlines(self, worker_id: str) -> None:
for key in [key for key in self._deadlines if key[0] == worker_id]:
self._deadlines.pop(key, None)
def _clear_execution_deadlines(self, worker_id: str) -> None:
for kind in (_DeadlineType.QUEUE, _DeadlineType.START, _DeadlineType.RUN, _DeadlineType.TOTAL):
self._clear_deadline(worker_id, kind)
def _discard_stale_deadlines(self) -> None:
while self._deadline_heap:
at, sequence, worker_id, kind = self._deadline_heap[0]
if self._deadlines.get((worker_id, kind)) == (at, sequence):
return
heapq.heappop(self._deadline_heap)
def _next_timeout(self, maximum: float | None = None) -> float | None:
self._discard_stale_deadlines()
deadline = None if not self._deadline_heap else max(0.0, self._deadline_heap[0][0] - monotonic())
if maximum is None:
return deadline
return maximum if deadline is None else min(maximum, deadline)
def _process_timeouts(self) -> None:
now = monotonic()
self._discard_stale_deadlines()
while self._deadline_heap and self._deadline_heap[0][0] <= now:
at, sequence, worker_id, kind = heapq.heappop(self._deadline_heap)
if self._deadlines.get((worker_id, kind)) != (at, sequence):
continue
self._deadlines.pop((worker_id, kind), None)
worker = self._workers.get(worker_id)
if worker is None:
continue
if kind in _TIMEOUT_KIND:
expected = {
_DeadlineType.QUEUE: WorkerState.PENDING,
_DeadlineType.START: WorkerState.STARTING,
_DeadlineType.RUN: WorkerState.RUNNING,
}.get(kind)
if expected is not None and worker.state is not expected:
continue
worker.timeout = _TIMEOUT_KIND[kind]
self.cancel(worker_id, WorkerStopReason.TIMEOUT)
elif kind is _DeadlineType.CANCEL_GRACE:
self._cancel_grace_expired(worker_id, worker)
elif kind is _DeadlineType.TERMINATE_GRACE:
self._kill(worker_id, worker)
self._discard_stale_deadlines()
def _cancel_grace_expired(self, worker_id: str, worker: Worker) -> None:
runtime = worker.runtime
if runtime is None:
self.finish_lost(worker_id)
elif not runtime.is_alive():
# A process/thread can exit immediately after successfully queuing
# WorkerQuit. Let the already-readable control frame win over the
# fallback death detector so its precise exit metadata is retained.
try:
quit_is_pending = worker.pipe is not None and worker.pipe.poll()
except (OSError, EOFError):
quit_is_pending = False
if not quit_is_pending:
self.finish_lost(worker_id)
elif isinstance(runtime, BaseProcess):
self._terminate(worker_id, worker)
else:
self._manager.logger.warning("Worker thread %s is stuck while stopping", worker_id)
def _terminate(self, worker_id: str, worker: Worker) -> None:
runtime = worker.runtime
if not isinstance(runtime, BaseProcess):
return
worker.state = WorkerState.TERMINATING
runtime.terminate()
self._schedule_deadline(
worker_id, _DeadlineType.TERMINATE_GRACE,
monotonic() + worker.policy.terminate_grace,
)
def _kill(self, worker_id: str, worker: Worker) -> None:
runtime = worker.runtime
if not isinstance(runtime, BaseProcess):
return
if runtime.is_alive():
runtime.kill()
runtime.join()
exit = WorkerExit.KILLED
else:
exit = WorkerExit.TERMINATED
outcome = WorkerOutcome.TIMED_OUT if worker.stop_reason is WorkerStopReason.TIMEOUT else WorkerOutcome.CANCELLED
self._finish(worker_id, outcome=outcome, exit=exit)
def _close(self) -> None:
if not self._accepting:
return
self._accepting = False
ids = list(self)
for worker_id in ids:
worker = self._workers.get(worker_id)
if worker is None:
continue
if worker.state is WorkerState.PENDING:
worker.stop_reason = WorkerStopReason.MANAGER_SHUTDOWN
self._finish(worker_id, outcome=WorkerOutcome.CANCELLED, exit=WorkerExit.CANCELLED)
elif worker.state not in (WorkerState.STOPPING, WorkerState.TERMINATING):
self.cancel(worker_id, WorkerStopReason.MANAGER_SHUTDOWN)
# Shutdown is already outside the event loop, so reap processes here.
for worker_id in list(self):
worker = self._workers.get(worker_id)
if worker is None or not isinstance(worker.runtime, BaseProcess):
continue
terminated = False
worker.runtime.join(worker.policy.cancel_grace)
if worker.runtime.is_alive():
worker.state = WorkerState.TERMINATING
worker.runtime.terminate()
terminated = True
worker.runtime.join(worker.policy.terminate_grace)
if worker.runtime.is_alive():
worker.runtime.kill()
worker.runtime.join()
exit = WorkerExit.KILLED
else:
try:
pending = worker.pipe.recv() if worker.pipe is not None and worker.pipe.poll() else None
except (EOFError, OSError):
pending = None
if isinstance(pending, WorkerQuit):
self.finish_from_quit(worker_id, pending)
continue
exit = WorkerExit.TERMINATED if terminated else WorkerExit.LOST
outcome = WorkerOutcome.TIMED_OUT if worker.stop_reason is WorkerStopReason.TIMEOUT else WorkerOutcome.CANCELLED
self._finish(worker_id, outcome=outcome, exit=exit)
for worker_id, worker in list(self.items()):
if isinstance(worker.runtime, Thread):
if worker.runtime.ident is not None:
worker.runtime.join(worker.policy.cancel_grace)
try:
pending = worker.pipe.recv() if worker.pipe is not None and worker.pipe.poll() else None
except (EOFError, OSError):
pending = None
if isinstance(pending, WorkerQuit):
self.finish_from_quit(worker_id, pending)
elif worker.runtime.is_alive():
self._manager.logger.warning("Worker thread %s remains alive at shutdown", worker_id)
outcome = WorkerOutcome.TIMED_OUT if worker.stop_reason is WorkerStopReason.TIMEOUT else WorkerOutcome.CANCELLED
self._finish(worker_id, outcome=outcome, exit=WorkerExit.LOST)
else:
self.finish_lost(worker_id)
else:
self.finish_lost(worker_id)
__all__ = ["WorkerController", "WorkerView"]