from __future__ import annotations
import os
import uuid
import enum
import time
import signal
import weakref
import asyncio
import logging
import inspect
import functools
import multiprocessing
from asyncio import Task
import aiomultiprocess
from collections import defaultdict
from collections.abc import Iterable
from typing import (
Dict,
Set,
Tuple,
Any,
Union,
Callable,
Optional,
DefaultDict,
)
from palaestrai.core.protocol import (
DelayedResultResponse,
DelayedResultRequest,
ErrorIndicator,
AttachDescriptor,
AttachServiceRequest,
AttachServiceResponse,
)
from palaestrai.core import RuntimeConfig, MajorDomoClient, MajorDomoWorker
from palaestrai.util.async_inspection import dump_all_tasks
LOG = logging.getLogger(__name__)
def _safe_attr(obj: Any, name: str, default: Any = None) -> Any:
"""Best-effort attribute read that tolerates custom Mock subclasses."""
if obj is None:
return default
obj_dict = getattr(obj, "__dict__", None)
if isinstance(obj_dict, dict) and name in obj_dict:
return obj_dict.get(name, default)
# unittest.mock dynamically fabricates attributes via __getattr__, which
# can produce non-serializable placeholders or even raise for subclasses.
if type(obj).__module__.startswith("unittest.mock"):
return default
try:
return getattr(obj, name)
except Exception:
return default
class Flags(enum.Enum):
"""Flags to force custom behavior of the EventStateMachine
Usually, the EventStateMachine tries to get the behavior right.
But sometimes, there are edge cases in which the wanted behavior derives
from the defaults of the ESM.
For this, these flags exist. They can be returned additionally from event
handlers, e.g., a method decorated with ``@spawns`` would then not only
return a process, but a tuple of process and flag.
Attributes
----------
KEEPALIVE :
Usually, the ESM stops the MDP transceiving loop when a response
object ends with ``ShutdownResponse``. If this behavior is not desired,
supply this keepalive flag.
"""
KEEPALIVE = enum.auto()
[docs]
class EventStateMachine:
"""An event-triggered state machine
The EventStateMachine (ESM) can be used to transparently handle events
within palaestrAI. An ESM wraps another class and callbacks can be defined
with method decorators for events. Events are:
* A message received,
* a signal received (SIGCHLD, SIGTERM, etc.)
* setup
* enter (the initial event)
* teardown
The initial event *enter* is issued immediately after the main event/state
loop commences in order to provide an entrypoint for operation.
The *enter* event can be used to, e.g., send out the first request.
For example::
@ESM.monitor()
class Foo:
@ESM.enter
async def _enter(self):
_ = await self._request_initialization()
@ESM.requests
async def _request_initialization(self):
# ...
return InitRequest(
# ...
)
It is not strictly necessary to provide an *enter* event.
If the monitored class is exclusively an MDP worker, then there is no need
for the *enter* event, because the worker reacts on the first request it
receives and not on its own volition.
In order to make a class use the ESM, you must decorate it with
::`~.monitor`. The ``monitor`` decorator can also inject all necessary
code to handle ZMQ MDP workers.
If the monitored class does not have a ``run`` method, the ESM will also
inject it. The ``run`` method then serves as an event/state loop that
continues until it is stopped. At the start of the ``run`` method, the
target objects ``setup`` method is called if it exists. Likewise,
a ``teardown`` method will be called immediately after the loop ends.
The ESM also adds a ``stop`` method to the target object. It serves to
terminate the event/state loop.
In order to react to a specific event, users of the ESM can decorate their
methods with ``on(event)``. The ::`~.on` decorator takes as parameter
the class of what is handled. E.g., the class of a particular message, or
``signal.SIGCHLD`` to react to a process that has ended. For example::
from palaestrai.core import EventStateMachine as ESM
import signal
@ESM.monitor()
class Foo:
@ESM.on(SomeRequest)
async def handle_some_request(self, request):
# ...
pass
@ESM.on(signal.SIGCHLD)
async def handle_process_termination(self, process):
# ...
pass
Spawning processes is also handled through a decorator: ``spawns``. If a
method decorated with ``spawns`` returns a ::`Process` object, this
process will automatically be monitored. E.g.,::
# ...
@ESM.spawns
def start_some_fancy_process(self):
p = multiprocessing.Process(target=somefunc)
p.start()
return p
The ESM also handles the sending of requests. ESM-monitored classes do not
need to instantiate and monitor MDP client objects themselves. Instead,
they simply need methods to be decorated with ``requests``. The so
decorated method must return a message object that has the ``receiver``
property, so that ::`~.requests` can handle sending. E.g.,::
# ...
@ESM.requests
def get_something_from_a_worker(self):
req = SomeRequest()
req.receiver = "Foo"
return req
@ESM.on(SomeResponse) # also handle the response!
def handle_response_from_worker(self, response):
# ...
pass
The ESM also supports classes that act as workers. For this, the ESM's
``monitor`` decorator needs the flag ``is_mdp_worker=True``. Then, the
ESM injects the property ``mdp_service``. Setting this property connects
the MDP worker, and ``ESM.on`` can be used to handle requests from clients.
For example::
@ESM.monitor(is_mdp_worker=True)
class Foo:
async def setup(self):
self.mdp_service = "Foo"
@ESM.on(SomeRequest)
def handle_request_from_client(self, req):
do_something_with(request)
rsp = SomeResponse()
rsp.receiver = req.sender
return rsp
"""
_decorated_methods: Dict[Callable, Any] = dict()
_monitored_objects: Dict[Tuple[int, weakref.ref], EventStateMachine] = (
dict()
)
# Tracks attached local services in this process.
# Local dispatch may bypass transport while preserving broker visibility
# via asynchronous mirroring.
_local_registry: Dict[str, Tuple["EventStateMachine", Task]] = dict()
@staticmethod
def _cleanup_monitored_objects(ref: weakref.ref):
pid = os.getpid()
LOG.debug("EventStateMachine cleaning %s", (pid, ref))
del EventStateMachine._monitored_objects[(pid, ref)]
@staticmethod
def esm_for(monitored: Any) -> EventStateMachine:
"""Returns the ESM instance for any monitored object.
This method retrieves the ESM instance responsible for a monitored
object. It does not check whether the object has been decoreted with
::`~.monitored`, though.
Parameters
----------
monitored : Any
A monitored object
Returns
-------
EventStateMachine
The ESM instance responsible for the monitored object. A new
instance will be created if it does not already exist.
"""
pid = os.getpid()
ref = weakref.ref(
monitored, EventStateMachine._cleanup_monitored_objects
)
try:
esm = EventStateMachine._monitored_objects[(pid, ref)]
except KeyError:
esm = EventStateMachine(monitored)
EventStateMachine._monitored_objects[(pid, ref)] = esm
return esm
@staticmethod
def _make_mdp_service_property():
@property
def mdp_service(self) -> str:
return self.__esm__._mdp_worker_service
@mdp_service.setter
def mdp_service(self, value: str):
self.__esm__._mdp_worker_service = value
self.__esm__._connect_worker_and_listen()
return mdp_service
@staticmethod
def monitor(is_mdp_worker=False):
"""Decorates a class to monitor instances of it with the ESM.
This decorator is the minimal required decoration of any class that
makes use of the ESM. It injects the ESM instance into new objects of
that class, and also adds relevant methods. The usage is::
from palaestrai.core import EventStateMachine as ESM
@ESM.monitor()
class Foo:
pass
@ESM.monitor(is_mdp_worker=True)
class Bar:
pass
The ``@monitor`` decorator injects methods to the target class, namely:
* ``run()``: The default run method that kicks off the event/state
loop of the target class.
* ``stop()``: Stops the event/state loop of the class and can be
called from any handler.
If ``is_mdp_worker=True`` was given, then the ESM also takes care of
handling MDP requests for the target class. Then, another property
is injected: ``mdp_service``. This is then the service name the worker
will listen on. Setting the property instanciates a ::`MajorDomoWorker`
and connects it to the broker.
Parameters
----------
is_mdp_worker : bool
If ``True``, the monitored class will act as MDP worker. The ESM
will inject a property ``mdp_service``. Setting the property will
create a ::`MajorDomoWorker` instance and connect it to the broker.
"""
def _wraps(clazz):
attrs = dir(clazz)
setattr(
clazz,
"__esm__",
property(lambda self: EventStateMachine.esm_for(self)),
)
if is_mdp_worker:
setattr(
clazz,
"mdp_service",
EventStateMachine._make_mdp_service_property(),
)
if "run" not in attrs:
setattr(clazz, "run", EventStateMachine.run)
setattr(clazz, "stop", EventStateMachine.stop)
return clazz
return _wraps
@staticmethod
def on(sig_or_msg_or_str):
"""Register an event/state transition handler.
``on`` is a decorator used to register handlers for any kind of event.
Typical usage is::
from palaestrai.core import EventStateMachine as ESM
@ESM.monitor()
class Foo:
@ESM.on(SomeRequest)
def bar(self, req):
pass # ...
Typical arguments to ``on`` are:
* A class: When the ESM receives an MDP request or response, it will
check whether the message's class has a handler registered. The
registered method is then called and the message passed.
* An exception class: The handler is triggered when the exception is
thrown.
* A signal: Handles signals such as ``SIGCHLD`` (when a child process
terminates), ``SIGINT``, or ``SIGTERM``.
"""
try:
sig_or_msg_or_str = sig_or_msg_or_str.__name__
except AttributeError:
pass
def _register_func(func):
EventStateMachine._decorated_methods[func] = sig_or_msg_or_str
return func
return _register_func
@staticmethod
def enter(func):
"""Decorates a method to be the very first state
Each state machine needs an initial state; the ESM is no exception. A
method decorated with ``enter`` is called immediately at the beginning
of the event/state loop.
Usage example::
@ESM.monitor()
class Foo:
@ESM.enter
def _enter(self):
pass # Do something, like launching a process.
``enter`` is not used for setup purposes: If the target class has a
``setup`` method, this one is called immediately *before* the
event/state loop commences. Thus, the enter method is optional. E.g.,
a class that simply acts as MDP worker does not need it; it is
sufficient to set the MDP service name in the ``setup`` method.
"""
EventStateMachine._decorated_methods[func] = "ENTER"
return func
def _handle_enter(self):
pass # Intentional a noop to suppress a warning for the ENTER event.
def _handle_terminated_child(
self, process: Union[aiomultiprocess.Process, multiprocessing.Process]
):
LOG.debug(
"%s saw termination of process: %s. No other handler is "
"installed.",
self,
process,
)
if process.exitcode != 0:
LOG.error(
"Process %s died with exit code %d",
process,
process.exitcode,
)
def _handle_sigint(self):
LOG.debug("%s handles SIGINT for %s", self, self._monitored)
self.stop(self._monitored)
def _handle_sigterm(self):
LOG.debug("%s handles SIGTERM for %s", self, self._monitored)
self.stop(self._monitored)
def _handle_sigtrap(self):
LOG.debug("%s handles SIGTRAP for %s", self, self._monitored)
dump_all_tasks()
self._trap_pycharm_debugger()
self._trap_debugger()
@staticmethod
def spawns(func):
"""Signify that a method creates (spawns) new sub-processes
Child processes are also monitored by the ESM. In order to find out
which processes to monitor, the ESM checks the return values of all
methods that are decorated with ``@spawns``. For example::
from palaestrai.core import EventStateMachine as ESM
@ESM.monitor()
class Foo:
@ESM.spawns
async def some_method(self):
p = multiprocessing.Process(target=foofunc)
p.start()
return p
@ESM.enter
def _enter(self):
_= await self.spawns()
**Note:** Processes that are returned from a spawning function are not
automatically started, just monitored.
"""
def _monitor_spawned_processes(self, ret):
for process in [
x
for x in (ret if isinstance(ret, Iterable) else [ret])
if isinstance(x, multiprocessing.Process)
or isinstance(x, aiomultiprocess.Process)
]:
self.__esm__.monitor_process(process)
def _wraps(self, *args, **kwargs):
from palaestrai.util.otel import get_tracer
service = self.__esm__._mdp_worker_service or ""
func_name = getattr(func, "__name__", str(func))
span_name = (
f"{service}.spawns.{func_name}"
if service
else f"spawns.{func_name}"
)
span_attrs = {
"palaestrai.service": service,
"palaestrai.handler": func_name,
"palaestrai.esm.decorator": "spawns",
}
with get_tracer().start_as_current_span(
span_name, attributes=span_attrs
):
ret = func(self, *args, **kwargs)
_monitor_spawned_processes(self, ret)
return ret
async def _awraps(self, *args, **kwargs):
from palaestrai.util.otel import get_tracer
service = self.__esm__._mdp_worker_service or ""
func_name = getattr(func, "__name__", str(func))
span_name = (
f"{service}.spawns.{func_name}"
if service
else f"spawns.{func_name}"
)
span_attrs = {
"palaestrai.service": service,
"palaestrai.handler": func_name,
"palaestrai.esm.decorator": "spawns",
}
with get_tracer().start_as_current_span(
span_name, attributes=span_attrs
):
ret = await func(self, *args, **kwargs)
_monitor_spawned_processes(self, ret)
return ret
return _awraps if asyncio.iscoroutinefunction(func) else _wraps
@staticmethod
def requests(func):
"""Signify that the returned request message object awaits an answer.
When the object monitored by the ESM sends out requests, it will
also want to react to responses. In order to manage tracking of
requests and responses, the ESM uses the ``requests`` method decorator.
For example,::
from palaestrai.core import EventStateMachine as ESM
@ESM.monitor()
class Foo:
@ESM.requests
def send_some_request(self):
# ...
return SomeRequest(receiver="SomeWorker")
@ESM.on(SomeResponse)
def handle_some_response(self, response):
# ...
pass
The message object's class needs to end with ``Request``. It can be
passed along with other objects as well. So if the method returns
a tuple or a list, the ESM will inspect each object to be a message
object, and track that.
"""
def _trace_and_schedule(self, ret):
for mdp_request in [
x
for x in (ret if isinstance(ret, Iterable) else [ret])
if type(x).__name__.endswith("Request")
]:
# Capture context at request creation time to avoid context-loss
# across async task boundaries and lock handovers.
try:
from opentelemetry import propagate
carrier = {}
propagate.inject(carrier)
if hasattr(mdp_request, "__dict__"):
mdp_request._otel_trace_context = dict(carrier)
LOG.debug(
"OTEL_CTX request-decorator fn=%s msg=%s sender=%s receiver=%s traceparent=%s",
getattr(func, "__name__", str(func)),
type(mdp_request).__name__,
_safe_attr(mdp_request, "sender"),
_safe_attr(mdp_request, "receiver"),
mdp_request._otel_trace_context.get("traceparent"),
)
except ImportError:
pass
self.__esm__._tasks.add(
asyncio.create_task(self.__esm__.send_request(mdp_request))
)
def _wraps(self, *args, **kwargs):
from palaestrai.util.otel import get_tracer
service = self.__esm__._mdp_worker_service or ""
func_name = getattr(func, "__name__", str(func))
span_name = (
f"{service}.requests.{func_name}"
if service
else f"requests.{func_name}"
)
span_attrs = {
"palaestrai.service": service,
"palaestrai.handler": func_name,
"palaestrai.esm.decorator": "requests",
}
with get_tracer().start_as_current_span(
span_name, attributes=span_attrs
):
ret = func(self, *args, **kwargs)
_trace_and_schedule(self, ret)
return ret
async def _awraps(self, *args, **kwargs):
from palaestrai.util.otel import get_tracer
service = self.__esm__._mdp_worker_service or ""
func_name = getattr(func, "__name__", str(func))
span_name = (
f"{service}.requests.{func_name}"
if service
else f"requests.{func_name}"
)
span_attrs = {
"palaestrai.service": service,
"palaestrai.handler": func_name,
"palaestrai.esm.decorator": "requests",
}
with get_tracer().start_as_current_span(
span_name, attributes=span_attrs
):
ret = await func(self, *args, **kwargs)
_trace_and_schedule(self, ret)
return ret
return _awraps if asyncio.iscoroutinefunction(func) else _wraps
@staticmethod
def attaches(func):
"""Signify that a method attaches a service to another process.
The decorated method returns AttachDescriptor(s). The ESM sends
an AttachServiceRequest via MDP to the target service, which
unpickles the instance and runs it in its process.
"""
def _wraps(self, *args, **kwargs):
ret = func(self, *args, **kwargs)
esm_mode = RuntimeConfig().esm_mode
for descriptor in [
x
for x in (ret if isinstance(ret, Iterable) else [ret])
if isinstance(x, AttachDescriptor)
]:
if esm_mode == "spawn":
# Reachable only via esm_mode = "spawn". The default
# (esm_mode = "attach") takes the else branch below and
# sends an AttachServiceRequest instead of forking a
# process.
from palaestrai.util.spawn import spawn_wrapper
instance = descriptor.instance
uid = (
_safe_attr(instance, "uid", None)
or type(instance).__name__
)
process_name = f"{self.__esm__._mdp_worker_service}.{uid}"
process = aiomultiprocess.Process(
name=process_name,
target=spawn_wrapper,
args=(
f"palaestrAI[{uid}]",
RuntimeConfig().to_dict(),
instance.run,
),
)
process.start()
self.__esm__.monitor_process(process)
LOG.debug(
"%s spawned attached instance %s in process %s "
"(descriptor target service=%s)",
self,
instance,
process_name,
descriptor.target_service_uid,
)
else:
request = AttachServiceRequest(
sender=self.__esm__._mdp_worker_service,
receiver=descriptor.target_service_uid,
instance=descriptor.instance,
)
self.__esm__._tasks.add(
asyncio.create_task(self.__esm__.send_request(request))
)
return ret
return _wraps
@staticmethod
async def run(
monitored,
_skip_signal_setup=False,
service_queue: Optional[asyncio.Queue] = None,
):
"""Main event/state loop of the ESM
This ``run`` method is injected into monitored classes if they do not
have one already. The structure of ``run`` is as follows:
1. It resets the handlers for SIGCHLD, SIGINT, and SIGTERM to the OS'
default.
2. It calls ``monitored.setup()``, if it exists.
3. It creates an ESM instance for the monitored object and adds signal
handlers for SIGCHLD, SIGINT, and SIGTERM according to what the
monitored class defines (via ``@ESM.on(signal.SIGINT)``, etc.)
4. It transides to the first state, defined by ``@ESM.enter``. It then
waits for state changes/events until ``monitored.stop()`` is called.
5. Finally, once the main event/state loop concludes,
``monitored.teardown()`` is called (if present).
"""
from palaestrai.util.otel import get_tracer
with get_tracer().start_as_current_span(
"esm.run",
attributes={
"palaestrai.monitored_type": type(monitored).__name__,
"palaestrai.skip_signal_setup": bool(_skip_signal_setup),
},
):
if not _skip_signal_setup:
signal.signal(signal.SIGCHLD, signal.SIG_DFL)
signal.signal(signal.SIGINT, signal.SIG_DFL)
signal.signal(signal.SIGTERM, signal.SIG_DFL)
signal.signal(signal.SIGTRAP, signal.SIG_DFL)
if "setup" in dir(monitored):
LOG.debug("Running %s.setup()...", monitored)
try:
with get_tracer().start_as_current_span(
"esm.run.setup",
attributes={
"palaestrai.monitored_type": type(
monitored
).__name__,
},
):
if asyncio.iscoroutinefunction(monitored.setup):
await monitored.setup()
else:
monitored.setup()
if service_queue is not None:
service_queue.put_nowait(monitored.mdp_service)
except Exception as e:
LOG.exception("%s.setup() failed with %s", monitored, e)
return
esm = EventStateMachine.esm_for(monitored)
loop = asyncio.get_running_loop()
if not _skip_signal_setup:
loop.add_signal_handler(
signal.SIGINT,
lambda: asyncio.create_task(
esm._handle_event(signal.SIGINT)
),
)
loop.add_signal_handler(
signal.SIGTERM,
lambda: asyncio.create_task(
esm._handle_event(signal.SIGTERM)
),
)
loop.add_signal_handler(
signal.SIGTRAP,
lambda: asyncio.create_task(
esm._handle_event(signal.SIGTRAP)
),
)
esm._future = asyncio.get_running_loop().create_future()
LOG.debug("%s commencing loop: Waiting for my own future…", esm)
asyncio.create_task(esm._handle_event("ENTER"))
with get_tracer().start_as_current_span("esm.run.wait_for_stop"):
await esm._future
LOG.debug("%s: The future is now!", esm)
if "teardown" in dir(monitored):
try:
with get_tracer().start_as_current_span(
"esm.run.teardown",
attributes={
"palaestrai.monitored_type": type(
monitored
).__name__,
},
):
if asyncio.iscoroutinefunction(monitored.teardown):
await monitored.teardown()
else:
monitored.teardown()
except Exception as e:
LOG.exception("%s.teardown() failed with %s", monitored, e)
await esm._cleanup()
exc = esm._future.exception()
if exc is not None:
raise exc
@staticmethod
def stop(monitored, error=None):
"""Stops the ESM.
Stopping the ESM also means shutting down all running processes and
cancelling all outstanding tasks (e.g., request monitors).
Paramters
---------
error : Exception
If given, the ESM will raise this after cleaning up.
"""
esm = EventStateMachine.esm_for(monitored)
esm._stop(error)
def __init__(self, monitored: Any):
self._future: Optional[asyncio.Future] = None
self._monitored = monitored
self._monitored_processes: Dict[
Union[aiomultiprocess.Process, multiprocessing.Process],
asyncio.Task,
] = dict()
self._tasks: Set[asyncio.Task] = {
asyncio.create_task(self._watch_tasks(), name="Tasks Watcher")
}
self._pending_tasks: Dict[str, asyncio.Task] = {} # If worker pending
self._delayed_result_retry_attempts: Dict[str, int] = {}
self._delayed_result_retry_initial_interval_s = 0.2
self._delayed_result_retry_max_interval_s = 0.5
self._request_queues: Dict[str, asyncio.Queue] = {}
self._request_dispatchers: Dict[str, asyncio.Task] = {}
self._local_request_queue: asyncio.Queue = asyncio.Queue()
self._local_request_dispatcher: Optional[asyncio.Task] = None
self._local_store_forward_queue: asyncio.Queue = asyncio.Queue()
self._local_store_forward_dispatcher: Optional[asyncio.Task] = None
self._local_store_forward_service = "mmi.service"
# Keep request batching lightweight and low-latency. We derive the
# batch size from existing store buffering defaults (256 -> 32).
try:
store_buffer_size = max(
0, int(getattr(RuntimeConfig(), "store_buffer_size", 256))
)
except Exception:
store_buffer_size = 256
self._request_batch_size = (
max(8, min(64, store_buffer_size // 8))
if store_buffer_size > 0
else 8
)
self._request_batch_flush_interval = 0.005
self._local_store_forward_batch_size = max(
16, min(128, self._request_batch_size * 2)
)
self._local_store_forward_flush_interval = max(
self._request_batch_flush_interval, 0.02
)
# Ack window for AttachServiceRequest: how long we wait for the
# attached worker's setup() to either complete or fail before
# replying to the caller. A synchronous setup() that raises still
# takes a couple of event-loop hops (task scheduling + ESM.run
# scaffolding), roughly 100-180 ms on CI. We wait a bit longer than
# that so a fast failure surfaces as success=False in the ACK, while
# legitimately slow (>= ~200 ms) async setups still return
# immediately and finish registration in the background.
self._attach_service_ack_timeout = 0.2
try:
client_timeout = float(
getattr(RuntimeConfig(), "major_domo_client_timeout", 300_000)
)
except Exception:
client_timeout = 300_000.0
self._request_flush_timeout = max(
5.0,
min(30.0, client_timeout / 1000.0),
)
self._local_store_forward_flush_timeout = self._request_flush_timeout
self._local_request_flush_timeout = self._request_flush_timeout
self._handlers = {
signal.SIGCHLD: self._handle_terminated_child,
signal.SIGINT: self._handle_sigint,
signal.SIGTERM: self._handle_sigterm,
signal.SIGTRAP: self._handle_sigtrap,
"ENTER": self._handle_enter,
DelayedResultRequest.__name__: self._try_get_delayed_result,
DelayedResultResponse.__name__: self._retry_delayed_request,
AttachServiceRequest.__name__: self._handle_attach_service,
Any: None,
}
self._mdp_worker_service: Optional[str] = None
self._mdp_worker: Optional[MajorDomoWorker] = None
self._mdp_clients: DefaultDict[
str, Tuple[MajorDomoClient, asyncio.Lock]
] = defaultdict(
lambda: (
MajorDomoClient(RuntimeConfig().broker_uri),
asyncio.Lock(),
)
)
# Update handlers, match to methods of the _monitored object:
injected_methods = [ # These are injected by us, ignore them here:
"__esm__",
"mdp_service",
"run",
]
directory = dir(self._monitored)
all_attributes = [
getattr(self._monitored, x, None) # None instead of AttributeError
for x in directory
if not x in injected_methods
]
self._handlers.update(
{
EventStateMachine._decorated_methods[x.__func__]: x
for x in all_attributes
if x is not None
and inspect.ismethod(x)
and x.__func__ in EventStateMachine._decorated_methods
}
)
self._handlers.update(
{
EventStateMachine._decorated_methods[x]: x
for x in all_attributes
if inspect.isfunction(x)
and x in EventStateMachine._decorated_methods
}
)
async def _watch_tasks(self):
while self._tasks:
done, pending = await asyncio.wait(
self._tasks, return_when=asyncio.FIRST_COMPLETED
)
exceptionals = [t for t in done if t.exception() is not None]
for e in exceptionals:
LOG.error(
"%s saw task %s raise exception: %s",
self,
e,
e.exception(),
)
await self._handle_event(
type(e.exception()).__name__, e.exception()
)
self._tasks = pending
async def _handle_event(self, event: Any, *args, **kwargs) -> Any:
try:
handler: Any = self._handlers[event]
except KeyError:
handler = self._handlers[Any]
if handler is None: # Default handler
LOG.warning("%s has no handler for %s", self._monitored, event)
return
# --- OTel span wrapping (transparent to handlers) ---
from palaestrai.util.otel import get_tracer
tracer = get_tracer()
# Extract parent context from message if present
parent_ctx = None
incoming_traceparent = None
if args:
msg = args[0]
trace_ctx = _safe_attr(msg, "_otel_trace_context", None)
if isinstance(trace_ctx, dict):
incoming_traceparent = trace_ctx.get("traceparent")
try:
from opentelemetry import propagate
parent_ctx = propagate.extract(trace_ctx)
except ImportError:
pass
# Build span name: "ServiceName.handler_name"
service = self._mdp_worker_service or ""
handler_name = getattr(handler, "__name__", str(event))
span_name = f"{service}.{handler_name}" if service else handler_name
# Build span attributes
span_attrs = {
"palaestrai.service": service,
"palaestrai.handler": handler_name,
}
if args:
msg = args[0]
for attr in ("sender", "receiver", "experiment_run_id"):
val = _safe_attr(msg, attr, None)
if val is not None:
span_attrs[f"palaestrai.{attr}"] = str(val)
try:
with tracer.start_as_current_span(
span_name, context=parent_ctx, attributes=span_attrs
):
LOG.debug(
"OTEL_CTX handle-enter service=%s handler=%s event=%s sender=%s receiver=%s parent_traceparent=%s",
service,
handler_name,
event,
_safe_attr(args[0], "sender", None) if args else None,
_safe_attr(args[0], "receiver", None) if args else None,
incoming_traceparent,
)
if asyncio.iscoroutinefunction(handler):
result = await handler(*args, **kwargs)
else:
result = handler(*args, **kwargs)
# Capture active span context for all outgoing message objects
# while the handler span is still live.
try:
from opentelemetry import propagate
carrier: dict[str, str] = {}
propagate.inject(carrier)
for msg in (
result
if isinstance(result, (tuple, list))
else [result]
):
if (
msg is not None
and type(msg).__name__.endswith(
("Request", "Response")
)
and hasattr(msg, "__dict__")
):
msg._otel_trace_context = dict(carrier)
LOG.debug(
"OTEL_CTX handle-exit service=%s handler=%s out_msg=%s sender=%s receiver=%s traceparent=%s",
service,
handler_name,
type(msg).__name__,
_safe_attr(msg, "sender", None),
_safe_attr(msg, "receiver", None),
msg._otel_trace_context.get("traceparent"),
)
except ImportError:
pass
return result
except Exception as e:
LOG.exception(
"%s encountered exception from the handler for %s",
self,
event,
)
# Perhaps there is a handler for the exception...?
# Except, of course, we're already trying to handle the
# exception...
if (
not isinstance(event, Exception)
and type(e).__name__ in self._handlers
):
_ = await self._handle_event(type(e).__name__, e)
else:
if RuntimeConfig().wait_for_debugger_on_error:
self._trap_debugger()
assert self._future is not None
self._future.set_exception(e)
def _trap_debugger(self):
try:
import socket
import debugpy # type: ignore[import-untyped]
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.bind(("", 0))
port = s.getsockname()[1]
s.close()
debugpy.listen(port)
LOG.critical(
"%s started debugpy, waiting for debugger client "
"to connect on port %s.",
self,
port,
)
debugpy.wait_for_client()
debugpy.breakpoint()
except ImportError:
LOG.critical(
"debugpy package not available, cannot wait for debugger to attach."
)
def _trap_pycharm_debugger(self, host="127.0.0.1", port=5690):
LOG.critical(">>> _trap_pycharm_debugger() ENTERED")
try:
import pydevd_pycharm # type: ignore[import-not-found]
except ImportError:
LOG.critical(
"pydevd_pycharm not available, skipping debugger attach"
)
return
LOG.critical(
"%s connecting to Pycharm debugger at %s:%s",
self,
host,
port,
)
try:
pydevd_pycharm.settrace(
host=host,
port=port,
stdoutToServer=True,
stderrToServer=True,
suspend=True,
)
# pydevd_pycharm.breakpoint()
LOG.critical(">>> pydevd_pycharm.settrace() RETURNED")
except Exception as e:
import traceback
LOG.critical("pydevd_pycharm.settrace failed: %s", e)
LOG.critical("traceback: %s", traceback.format_exc())
def monitor_process(
self, process: Union[aiomultiprocess.Process, multiprocessing.Process]
):
task = asyncio.create_task(
self._watch_process(process),
name=f"Process watcher for child {process.pid}",
)
self._monitored_processes[process] = task
LOG.debug("%s now monitors process %s", self, process)
async def _watch_process(
self, process: Union[aiomultiprocess.Process, multiprocessing.Process]
):
from palaestrai.util.otel import get_tracer
with get_tracer().start_as_current_span(
"esm.watch_process",
attributes={
"palaestrai.process_name": process.name,
"palaestrai.service": self._mdp_worker_service or "",
},
):
LOG.debug("%s starts to watch process: %s", self, process)
if isinstance(process, aiomultiprocess.Process):
await process.join()
else:
process.join()
LOG.debug(
"%s saw a process end: %s, calling handler...",
self,
process.name,
)
await self._handle_event(signal.SIGCHLD, process)
del self._monitored_processes[process] # Cleanup.
async def _wait_for_response(self, service: str, request: Any):
from palaestrai.util.otel import get_tracer
parent_ctx = None
trace_ctx = _safe_attr(request, "_otel_trace_context", None)
if isinstance(trace_ctx, dict):
try:
from opentelemetry import propagate
parent_ctx = propagate.extract(trace_ctx)
except ImportError:
pass
with get_tracer().start_as_current_span(
"esm.wait_for_response",
context=parent_ctx,
attributes={
"palaestrai.service": str(service),
"palaestrai.request_type": type(request).__name__,
"palaestrai.sender": str(
_safe_attr(request, "sender", "") or ""
),
"palaestrai.receiver": str(
_safe_attr(request, "receiver", "") or ""
),
},
) as wait_span:
local_target = self._resolve_local_service(service)
dispatch_mode = "broker"
if local_target is not None:
dispatch_mode = self._local_dispatch_mode_for_request(
type(request).__name__
)
wait_span.set_attribute(
"palaestrai.local_dispatch_mode", dispatch_mode
)
resp = None
if local_target is not None and dispatch_mode != "broker":
wait_span.set_attribute("palaestrai.dispatch_kind", "local")
target_esm, _ = local_target
if dispatch_mode == "queued":
resp = await self._dispatch_local_request_queued(
service, request, target_esm
)
else:
resp = await self._dispatch_local_request_inline(
service, request, target_esm
)
else:
# Broker (MDP) round-trip. Taken when the target is not a
# local ESM (spawn topology) or when a request type is routed
# off the inline default via esm_local_dispatch_overrides =
# "broker". With the attach + inline defaults and co-located
# services this branch is not exercised.
wait_span.set_attribute("palaestrai.dispatch_kind", "broker")
mdp_client, mdp_client_lock = self._mdp_clients[service]
try:
async with mdp_client_lock:
resp = await mdp_client.send(service, request)
except Exception as e:
LOG.exception("Sending request failed: %s", e)
if resp is None:
return
if isinstance(request, DelayedResultRequest) and not isinstance(
resp, DelayedResultResponse
):
self._delayed_result_retry_attempts.pop(
request.task_uuid, None
)
await self._handle_event(type(resp).__name__, resp)
def _resolve_local_service(
self, service: str
) -> Optional[Tuple["EventStateMachine", Task]]:
local_target = EventStateMachine._local_registry.get(service)
if local_target is None:
return None
target_esm, target_task = local_target
if target_task.done():
EventStateMachine._local_registry.pop(service, None)
return None
if target_esm is self:
return None
return target_esm, target_task
def _local_dispatch_mode_for_request(self, request_type: str) -> str:
rc = RuntimeConfig()
override = rc.esm_local_dispatch_overrides.get(request_type)
if override is not None:
return override
return rc.esm_local_dispatch_mode
async def _dispatch_local_request_inline(
self,
service: str,
request: Any,
target_esm: "EventStateMachine",
):
from palaestrai.util.otel import get_tracer
with get_tracer().start_as_current_span(
"esm.local_dispatch_inline",
attributes={
"palaestrai.service": str(service),
"palaestrai.request_type": type(request).__name__,
"palaestrai.sender": str(
_safe_attr(request, "sender", "") or ""
),
"palaestrai.receiver": str(
_safe_attr(request, "receiver", "") or ""
),
},
):
resp = await target_esm._execute_request_handler(request)
if resp is None:
self._enqueue_local_store_messages([request])
else:
self._enqueue_local_store_messages([request, resp])
return resp
async def _dispatch_local_request_queued(
self,
service: str,
request: Any,
target_esm: "EventStateMachine",
):
from palaestrai.util.otel import get_tracer
with get_tracer().start_as_current_span(
"esm.local_dispatch_queued",
attributes={
"palaestrai.service": str(service),
"palaestrai.request_type": type(request).__name__,
"palaestrai.sender": str(
_safe_attr(request, "sender", "") or ""
),
"palaestrai.receiver": str(
_safe_attr(request, "receiver", "") or ""
),
},
):
loop = asyncio.get_running_loop()
response_future: asyncio.Future = loop.create_future()
target_esm._enqueue_local_request(request, response_future)
with get_tracer().start_as_current_span(
"esm.local_dispatch_wait_response",
attributes={
"palaestrai.service": str(service),
"palaestrai.request_type": type(request).__name__,
},
):
try:
resp = await response_future
except asyncio.CancelledError:
if not response_future.done():
response_future.cancel()
raise
if resp is None:
self._enqueue_local_store_messages([request])
else:
self._enqueue_local_store_messages([request, resp])
return resp
def _enqueue_local_request(
self, request: Any, response_future: asyncio.Future
) -> None:
self._local_request_queue.put_nowait(
(request, response_future, time.perf_counter_ns())
)
self._ensure_local_request_dispatcher()
def _ensure_local_request_dispatcher(self):
task = self._local_request_dispatcher
if task is not None and not task.done():
return
task = asyncio.create_task(
self._dispatch_local_request_queue(),
name="Local request dispatcher",
)
self._local_request_dispatcher = task
self._tasks.add(task)
async def _dispatch_local_request_queue(self):
from palaestrai.util.otel import get_tracer
queue = self._local_request_queue
try:
with get_tracer().start_as_current_span(
"esm.dispatch_local_request_queue",
attributes={
"palaestrai.service": str(self._mdp_worker_service or ""),
},
):
while True:
request, response_future, enqueued_at_ns = (
await queue.get()
)
queue_wait_ms = max(
0.0,
(time.perf_counter_ns() - enqueued_at_ns)
/ 1_000_000.0,
)
with get_tracer().start_as_current_span(
"esm.dispatch_local_request",
attributes={
"palaestrai.service": str(
self._mdp_worker_service or ""
),
"palaestrai.request_type": type(request).__name__,
"palaestrai.local_queue_wait_ms": queue_wait_ms,
},
):
try:
resp = await self._execute_request_handler(request)
if not response_future.done():
response_future.set_result(resp)
except Exception as e:
if not response_future.done():
response_future.set_exception(e)
else:
LOG.exception(
"%s local queued request raised after caller completion: %s",
self,
e,
)
queue.task_done()
if queue.empty():
await asyncio.sleep(0)
if queue.empty():
break
finally:
current = self._local_request_dispatcher
if current == asyncio.current_task():
self._local_request_dispatcher = None
def _enqueue_local_store_messages(self, messages: list[Any]) -> None:
from palaestrai.util.otel import get_tracer
any_message = False
with get_tracer().start_as_current_span(
"esm.local_mirror_enqueue",
attributes={
"palaestrai.message_count": len(messages),
"palaestrai.message_types": ",".join(
sorted(
{type(m).__name__ for m in messages if m is not None}
)
),
},
):
for message in messages:
if message is None:
continue
any_message = True
self._local_store_forward_queue.put_nowait(message)
if any_message:
self._ensure_local_store_forwarder()
def _ensure_local_store_forwarder(self):
task = self._local_store_forward_dispatcher
if task is not None and not task.done():
return
task = asyncio.create_task(
self._dispatch_local_store_forward_queue(),
name="Local store forwarder",
)
self._local_store_forward_dispatcher = task
self._tasks.add(task)
async def _dispatch_local_store_forward_queue(self):
from palaestrai.util.otel import get_tracer
queue = self._local_store_forward_queue
try:
with get_tracer().start_as_current_span(
"esm.dispatch_local_store_queue",
attributes={
"palaestrai.batch_size_limit": self._local_store_forward_batch_size,
"palaestrai.batch_flush_interval_s": self._local_store_forward_flush_interval,
},
):
while True:
message = await queue.get()
batch = [message]
collect_start_ns = time.perf_counter_ns()
while len(batch) < self._local_store_forward_batch_size:
try:
batch.append(queue.get_nowait())
continue
except asyncio.QueueEmpty:
pass
if self._local_store_forward_flush_interval <= 0:
break
try:
batch.append(
await asyncio.wait_for(
queue.get(),
timeout=self._local_store_forward_flush_interval,
)
)
except asyncio.TimeoutError:
break
collect_wait_ms = (
time.perf_counter_ns() - collect_start_ns
) / 1_000_000.0
while not await self._forward_local_store_batch(
batch,
collect_wait_ms=collect_wait_ms,
):
LOG.warning(
"%s failed to forward local store batch, retrying",
self,
)
await asyncio.sleep(
max(0.01, self._local_store_forward_flush_interval)
)
for _ in batch:
queue.task_done()
if queue.empty():
await asyncio.sleep(
self._local_store_forward_flush_interval
)
if queue.empty():
break
finally:
current = self._local_store_forward_dispatcher
if current == asyncio.current_task():
self._local_store_forward_dispatcher = None
async def _forward_local_store_batch(
self,
batch: list[Any],
collect_wait_ms: float = 0.0,
) -> bool:
from palaestrai.util.otel import get_tracer
from palaestrai.core import MDP
from palaestrai.core.serialisation import serialize
with get_tracer().start_as_current_span(
"esm.forward_local_store_batch",
attributes={
"palaestrai.service": self._local_store_forward_service,
"palaestrai.batch_size": len(batch),
"palaestrai.batch_collect_wait_ms": collect_wait_ms,
},
):
mdp_client, mdp_client_lock = self._mdp_clients[
self._local_store_forward_service
]
try:
async with mdp_client_lock:
request = serialize(batch)
service = bytes(
str(self._local_store_forward_service), "ascii"
)
if not mdp_client._socket:
mdp_client.reconnect_to_broker()
retries = mdp_client.retries
while retries > 0:
assert mdp_client._socket is not None
await mdp_client._socket.send_multipart(
[MDP.C_CLIENT, service, request]
)
items = await mdp_client._poller.poll(
mdp_client.timeout
)
if items:
msg = await mdp_client._socket.recv_multipart()
return len(msg) >= 1 and msg[0] == MDP.C_CLIENT
retries -= 1
if retries > 0:
mdp_client.reconnect_to_broker()
return False
except Exception as e:
LOG.exception(
"%s failed to forward local store batch: %s", self, e
)
return False
def _ensure_request_dispatcher(self, service: str):
task = self._request_dispatchers.get(service, None)
if task is not None and not task.done():
return
task = asyncio.create_task(
self._dispatch_request_queue(service),
name=f"Request dispatcher for {service}",
)
self._request_dispatchers[service] = task
self._tasks.add(task)
async def _dispatch_request_queue(self, service: str):
from palaestrai.util.otel import get_tracer
queue = self._request_queues[service]
try:
with get_tracer().start_as_current_span(
"esm.dispatch_request_queue",
attributes={
"palaestrai.service": str(service),
"palaestrai.batch_size_limit": self._request_batch_size,
"palaestrai.batch_flush_interval_s": self._request_batch_flush_interval,
},
):
while True:
request = await queue.get()
batch = [request]
collect_start_ns = time.perf_counter_ns()
while len(batch) < self._request_batch_size:
try:
batch.append(queue.get_nowait())
continue
except asyncio.QueueEmpty:
pass
if self._request_batch_flush_interval <= 0:
break
try:
batch.append(
await asyncio.wait_for(
queue.get(),
timeout=self._request_batch_flush_interval,
)
)
except asyncio.TimeoutError:
break
with get_tracer().start_as_current_span(
"esm.dispatch_request_batch",
attributes={
"palaestrai.service": str(service),
"palaestrai.batch_size": len(batch),
"palaestrai.batch_collect_wait_ms": (
time.perf_counter_ns() - collect_start_ns
)
/ 1_000_000.0,
},
) as batch_span:
queue_ages_ms = []
for req in batch:
enqueued_at_ns = _safe_attr(
req, "_esm_enqueued_at_ns", None
)
if enqueued_at_ns is None:
continue
queue_ages_ms.append(
max(
0.0,
(time.perf_counter_ns() - enqueued_at_ns)
/ 1_000_000.0,
)
)
if queue_ages_ms:
batch_span.set_attribute(
"palaestrai.queue_age_first_ms",
queue_ages_ms[0],
)
batch_span.set_attribute(
"palaestrai.queue_age_mean_ms",
sum(queue_ages_ms) / len(queue_ages_ms),
)
batch_span.set_attribute(
"palaestrai.queue_age_max_ms",
max(queue_ages_ms),
)
batch_span.set_attribute(
"palaestrai.queue_size_after_dequeue",
queue.qsize(),
)
for req in batch:
try:
await self._wait_for_response(service, req)
finally:
queue.task_done()
if queue.empty():
await asyncio.sleep(0)
if queue.empty():
break
finally:
current = self._request_dispatchers.get(service, None)
if current == asyncio.current_task():
del self._request_dispatchers[service]
async def send_request(self, request: Any):
from palaestrai.util.otel import get_tracer
parent_ctx = None
trace_ctx = _safe_attr(request, "_otel_trace_context", None)
if isinstance(trace_ctx, dict):
try:
from opentelemetry import propagate
parent_ctx = propagate.extract(trace_ctx)
except ImportError:
pass
with get_tracer().start_as_current_span(
"esm.send_request",
context=parent_ctx,
attributes={
"palaestrai.request_type": type(request).__name__,
"palaestrai.sender": str(
_safe_attr(request, "sender", "") or ""
),
"palaestrai.receiver": str(
_safe_attr(request, "receiver", "") or ""
),
},
):
try:
service = request.receiver
if service is None:
raise ValueError
if service not in self._request_queues:
self._request_queues[service] = asyncio.Queue()
if hasattr(request, "__dict__"):
request._esm_enqueued_at_ns = time.perf_counter_ns()
self._request_queues[service].put_nowait(request)
self._ensure_request_dispatcher(service)
except (AttributeError, ValueError):
LOG.error(
"%s cannot determine target service for %s. Please "
"extend %s to provide the 'receiver' property.",
self,
request,
type(request),
)
def _remove_child_service(self, remove_task: Task):
remove_services = [
mdp_service
for mdp_service, (esm, task) in self._local_registry.items()
if task == remove_task
]
if len(remove_services) > 1:
LOG.error("Internal error, multiple services held same task!")
for service in remove_services:
del self._local_registry[service]
# An attached (co-located) service that ends with an exception must
# fail its host process too. In the spawn topology a crashing child is
# a dead subprocess that the parent detects via SIGCHLD and reacts to;
# in the attach topology the child is merely a task in the host's event
# loop, so its exception would otherwise be swallowed and the host
# would hang. Propagating it here reproduces the spawn behaviour: the
# host process exits with an error, which the RunGovernor sees as a
# prematurely-ended subprocess and turns into a run failure.
if remove_task.cancelled():
return
child_exc = remove_task.exception()
if child_exc is not None:
LOG.error(
"%s: attached service task %s ended with an exception; "
"propagating to fail this process.",
self,
remove_task.get_name(),
)
self._stop(child_exc)
async def _wait_for_attached_service(
self, task: Task, service_queue: asyncio.Queue
) -> str:
service_wait_task = asyncio.create_task(
service_queue.get(), name="Wait for attached service registration"
)
done, _ = await asyncio.wait(
{task, service_wait_task},
return_when=asyncio.FIRST_COMPLETED,
)
if service_wait_task in done:
return service_wait_task.result()
service_wait_task.cancel()
try:
await service_wait_task
except asyncio.CancelledError:
pass
if task.cancelled():
raise RuntimeError(
"Attached service task was cancelled before registration"
)
task_exc = task.exception()
if task_exc is not None:
raise task_exc
raise RuntimeError(
"Attached service ended before setting an MDP service id"
)
async def _handle_attach_service(self, request: AttachServiceRequest):
from palaestrai.util.otel import get_tracer
instance = request.instance
LOG.info("%s attaching %s to this process", self, instance)
with get_tracer().start_as_current_span(
"esm.handle_attach_service",
attributes={
"palaestrai.sender": str(
_safe_attr(request, "sender", "") or ""
),
"palaestrai.receiver": str(
_safe_attr(request, "receiver", "") or ""
),
},
):
try:
service_queue: asyncio.Queue = asyncio.Queue()
task = asyncio.create_task(
EventStateMachine.run(
instance,
_skip_signal_setup=True,
service_queue=service_queue,
)
)
self._tasks.add(task)
task.add_done_callback(self._remove_child_service)
register_task = asyncio.create_task(
self._wait_for_attached_service(task, service_queue),
name=f"Attach registration wait for {type(instance).__name__}",
)
done, _ = await asyncio.wait(
{register_task},
timeout=self._attach_service_ack_timeout,
return_when=asyncio.FIRST_COMPLETED,
)
if register_task in done:
mdp_service = register_task.result()
if not task.done():
self._local_registry[mdp_service] = (
EventStateMachine.esm_for(instance),
task,
)
else:
async def _finalize_registration():
from palaestrai.util.otel import get_tracer as _tracer
with _tracer().start_as_current_span(
"esm.finalize_attach_registration",
attributes={
"palaestrai.sender": str(
_safe_attr(request, "sender", "") or ""
),
"palaestrai.receiver": str(
_safe_attr(request, "receiver", "") or ""
),
},
):
try:
mdp_service = await register_task
if task.done():
return
self._local_registry[mdp_service] = (
EventStateMachine.esm_for(instance),
task,
)
except asyncio.CancelledError:
if not register_task.done():
register_task.cancel()
raise
except Exception as e:
LOG.exception(
"Failed to finalize attachment registration "
"for %s: %s",
instance,
e,
)
finalize_task = asyncio.create_task(
_finalize_registration(),
name=f"Finalize attach registration for {type(instance).__name__}",
)
self._tasks.add(finalize_task)
return AttachServiceResponse(
sender=self._mdp_worker_service,
receiver=request.sender,
success=True,
)
except Exception as e:
LOG.exception("Failed to attach %s", instance)
return AttachServiceResponse(
sender=self._mdp_worker_service,
receiver=request.sender,
success=False,
error=str(e),
)
@staticmethod
def _handler_poll_timeout_for_request(req: Any) -> float:
timeout_s = min(
1.0,
RuntimeConfig().major_domo_client_timeout / 2000.0,
)
# Simulation startup can run for several seconds while additional
# setup/control messages are queued for the same service. Poll the
# handler task more frequently so these queued requests are not
# delayed by up to one full second each.
if type(req).__name__ == "SimulationStartRequest":
return min(timeout_s, 0.1)
return timeout_s
async def _mdp_worker_transceive(self):
from palaestrai.util.otel import get_tracer
with get_tracer().start_as_current_span(
"esm.mdp_worker_transceive",
attributes={
"palaestrai.service": str(self._mdp_worker_service or ""),
},
):
reply = None
flags = None
while True:
if isinstance(reply, tuple):
reply, flags = reply
else:
flags = None
shutdown_resp_but_no_keepalive = type(reply).__name__.endswith(
"ShutdownResponse"
) and (
flags is None
or (flags is not None and flags != Flags.KEEPALIVE)
)
if shutdown_resp_but_no_keepalive:
LOG.debug("%s giving final %s", self, reply)
req = await self._mdp_worker.transceive(
reply,
skip_recv=shutdown_resp_but_no_keepalive,
)
if req is None:
LOG.info(
"Going to stop ESM for %s based on %s",
self._monitored,
reply,
)
if (
shutdown_resp_but_no_keepalive
and not self._future.done()
):
LOG.error(
"ESM received None request for %s based on %s"
"trying to recover by implicit stopping",
self._monitored,
reply,
)
self._stop(
RuntimeError(
f"Received {reply} without proper ESM shutdown"
)
)
break
if type(reply).__name__.endswith("ShutdownRequest"):
LOG.debug("%s got %s", self, reply)
reply = await self._execute_request_handler(req)
async def _execute_request_handler(self, req: Any):
worker_task = asyncio.create_task(
self._handle_event(type(req).__name__, req)
)
poll_timeout_s = self._handler_poll_timeout_for_request(req)
done, pending = await asyncio.wait(
[worker_task],
timeout=poll_timeout_s,
)
if done:
return worker_task.result()
LOG.debug("%s is still waiting for response to %s", self, req)
assert len(pending) == 1
task_uuid = str(uuid.uuid4())
t = pending.pop()
self._pending_tasks[task_uuid] = t
self._tasks.add(t)
return DelayedResultResponse(
sender=req.receiver,
receiver=req.sender,
task_uuid=task_uuid,
)
def _connect_worker_and_listen(self):
from palaestrai.util.otel import get_tracer
if self._mdp_worker_service is None:
raise ValueError(f"{self}._mdp_worker_service string unset")
if self._mdp_worker is not None:
LOG.warning("%s already has an MDP Worker", self)
return
with get_tracer().start_as_current_span(
"esm.connect_worker",
attributes={
"palaestrai.service": str(self._mdp_worker_service or ""),
},
):
self._mdp_worker = MajorDomoWorker(
RuntimeConfig().broker_uri,
self._mdp_worker_service,
)
self._tasks.add(
asyncio.create_task(
self._mdp_worker_transceive(), name="Transceiver"
)
)
def _try_get_delayed_result(self, req: DelayedResultRequest):
try:
task = self._pending_tasks[req.task_uuid]
except KeyError as e:
self._delayed_result_retry_attempts.pop(req.task_uuid, None)
return ErrorIndicator(
sender=req.receiver,
receiver=req.sender,
error_message=f"Task UUID {req.task_uuid} unknown.",
exception=e,
)
if task.done():
del self._pending_tasks[req.task_uuid]
self._delayed_result_retry_attempts.pop(req.task_uuid, None)
return task.result()
return DelayedResultResponse(
sender=req.receiver,
receiver=req.sender,
task_uuid=req.task_uuid,
)
async def _retry_delayed_request(self, rsp: DelayedResultResponse):
LOG.debug(
"%s was notified that a worker needs some more time: "
"sleeping and retrying.",
self,
)
retry_attempt = self._delayed_result_retry_attempts.get(
rsp.task_uuid, 0
)
retry_sleep_s = min(
self._delayed_result_retry_initial_interval_s * (2**retry_attempt),
self._delayed_result_retry_max_interval_s,
)
self._delayed_result_retry_attempts[rsp.task_uuid] = retry_attempt + 1
await asyncio.sleep(retry_sleep_s)
await self.send_request( # Creates a task to wait for the response
DelayedResultRequest(
sender=rsp.receiver,
receiver=rsp.sender,
task_uuid=rsp.task_uuid,
)
)
def _stop(self, error: Optional[BaseException] = None):
if self._future is None:
return
try:
if error is not None:
self._future.set_exception(error)
else:
self._future.set_result(True)
except asyncio.exceptions.InvalidStateError:
# Doubly-stop.
pass
async def _cleanup(self):
from palaestrai.util.otel import get_tracer
LOG.debug(
"%s cleaning up: tasks %s; processes %s",
self,
self._tasks,
self._monitored_processes,
)
with get_tracer().start_as_current_span(
"esm.cleanup",
attributes={
"palaestrai.service": str(self._mdp_worker_service or ""),
"palaestrai.tasks_count": len(self._tasks),
"palaestrai.monitored_process_count": len(
self._monitored_processes
),
},
):
await self._stop_all_processes()
await self._flush_request_queues()
await self._flush_local_request_queue()
await self._flush_local_store_forward_queue()
self._delayed_result_retry_attempts.clear()
# Now give all tasks a reasonable amount of time to finish.
# All should terminate by themselves, except for the task watcher,
# which terminates when there's nothing left in self._tasks.
watcher = next(
t for t in self._tasks if t.get_name() == "Tasks Watcher"
)
self._tasks -= {watcher}
pending = set()
if self._tasks:
_, pending = await asyncio.wait(self._tasks, timeout=5)
for task in pending: # Now cancel what is left
if task.get_name() != "Transceiver":
# The transceiver might linger (e.g., during error state, or
# when the RunGovernor shuts down), so don't print a warning
# for the Transceiver task, just cancel it silently.
LOG.warning(
"%s: Task %s did not end, terminating", self, task
)
task.cancel()
for mdp_client, _ in self._mdp_clients.values():
mdp_client.close()
if self._mdp_worker is not None:
await self._mdp_worker.disconnect()
for sig in (signal.SIGINT, signal.SIGTERM, signal.SIGTRAP):
asyncio.get_running_loop().remove_signal_handler(sig)
LOG.debug("%s: done cleaning up.", self)
async def _flush_request_queues(self):
from palaestrai.util.otel import get_tracer
# Ensure queued outbound requests are drained before cleanup cancels
# remaining tasks.
with get_tracer().start_as_current_span(
"esm.flush_request_queues",
attributes={
"palaestrai.queue_count": len(self._request_queues),
"palaestrai.dispatcher_count": len(self._request_dispatchers),
"palaestrai.flush_timeout_s": self._request_flush_timeout,
},
):
join_tasks = [
asyncio.create_task(
queue.join(), name=f"Flush request queue for {service}"
)
for service, queue in self._request_queues.items()
if (not queue.empty() or service in self._request_dispatchers)
]
if join_tasks:
_, pending_joins = await asyncio.wait(
join_tasks, timeout=self._request_flush_timeout
)
for task in pending_joins:
LOG.warning(
"%s: Request queue flush timed out for task %s",
self,
task,
)
task.cancel()
active_dispatchers = [
task
for task in self._request_dispatchers.values()
if not task.done()
]
if active_dispatchers:
_, pending_dispatchers = await asyncio.wait(
active_dispatchers, timeout=self._request_flush_timeout
)
for task in pending_dispatchers:
LOG.warning(
"%s: Dispatcher %s did not end, terminating",
self,
task,
)
task.cancel()
async def _flush_local_request_queue(self):
from palaestrai.util.otel import get_tracer
with get_tracer().start_as_current_span(
"esm.flush_local_request_queue",
attributes={
"palaestrai.queue_size": self._local_request_queue.qsize(),
"palaestrai.flush_timeout_s": self._local_request_flush_timeout,
},
):
if (
self._local_request_queue.empty()
and self._local_request_dispatcher is None
):
return
join_task = asyncio.create_task(
self._local_request_queue.join(),
name="Flush local request queue",
)
done, _ = await asyncio.wait(
{join_task},
timeout=self._local_request_flush_timeout,
return_when=asyncio.ALL_COMPLETED,
)
if join_task not in done:
LOG.warning("%s: Local request queue flush timed out", self)
join_task.cancel()
dispatcher = self._local_request_dispatcher
if dispatcher is not None and not dispatcher.done():
_, pending = await asyncio.wait(
{dispatcher},
timeout=self._local_request_flush_timeout,
return_when=asyncio.ALL_COMPLETED,
)
for task in pending:
LOG.warning(
"%s: Local request dispatcher %s did not end, terminating",
self,
task,
)
task.cancel()
async def _flush_local_store_forward_queue(self):
from palaestrai.util.otel import get_tracer
with get_tracer().start_as_current_span(
"esm.flush_local_store_queue",
attributes={
"palaestrai.queue_size": self._local_store_forward_queue.qsize(),
"palaestrai.flush_timeout_s": self._local_store_forward_flush_timeout,
},
):
if (
self._local_store_forward_queue.empty()
and self._local_store_forward_dispatcher is None
):
return
join_task = asyncio.create_task(
self._local_store_forward_queue.join(),
name="Flush local store forward queue",
)
done, _ = await asyncio.wait(
{join_task},
timeout=self._local_store_forward_flush_timeout,
return_when=asyncio.ALL_COMPLETED,
)
if join_task not in done:
LOG.warning("%s: Local store queue flush timed out", self)
join_task.cancel()
dispatcher = self._local_store_forward_dispatcher
if dispatcher is not None and not dispatcher.done():
_, pending = await asyncio.wait(
{dispatcher},
timeout=self._local_store_forward_flush_timeout,
return_when=asyncio.ALL_COMPLETED,
)
for task in pending:
LOG.warning(
"%s: Local store forwarder %s did not end, terminating",
self,
task,
)
task.cancel()
async def _stop_all_processes(self):
from palaestrai.util.otel import get_tracer
with get_tracer().start_as_current_span(
"esm.stop_all_processes",
attributes={
"palaestrai.monitored_process_count": len(
self._monitored_processes
),
},
):
all_processes = list(
self._monitored_processes.keys()
) # Dict changes
for process in all_processes:
# First see whether the process exits all by itself:
if process.is_alive():
if asyncio.iscoroutinefunction(process.join):
try:
await process.join(3.0)
except asyncio.CancelledError:
pass # This is okay
except asyncio.TimeoutError:
# The process seems unwilling to terminate, but no
# worries, we're not done yet...
pass
else:
process.join(3.0)
if not process.is_alive():
continue # Yay, it ended as we wished.
# Then, send SIGTERM and wait for it to finish:
LOG.warning(
"Process %s did not exit by itself, sending SIGTERM.",
process.name,
)
process.terminate()
if asyncio.iscoroutinefunction(process.join):
try:
await process.join(3.0)
except asyncio.CancelledError:
pass
except asyncio.TimeoutError:
pass # Same as above, but now we kill
else:
process.join(3.0)
if not process.is_alive():
continue # Okay, not as nice as it could be, but still...
if all(
not process.is_alive() for process in self._monitored_processes
):
return # Don't wait
# Still someone here? Let's draw the big friggin' gun:
for process in all_processes:
if process.is_alive():
LOG.error(
"Process %s is still there, killing it.", process.name
)
process.kill()
if asyncio.iscoroutinefunction(process.join):
try:
await process.join() # This has to terminate.
except asyncio.CancelledError:
pass
except asyncio.TimeoutError:
pass # Yeah, well, we tried. Hand it to the reaper.
else:
process.join()
for task in self._monitored_processes.values():
if not task.done():
task.cancel()
def __str__(self):
return (
f"EventStateMachine(pid={os.getpid()}, "
f"monitored={self._monitored})"
)
def __del__(self):
if not hasattr(self, "_tasks"):
return
for t in self._tasks:
t.cancel()