Motivation
I want to run some compute heavy tasks on a separate process, so that they don't hog the GIL and I can make effective use of a multi-core machine.
Where those tasks are pure functions, I'd just use the provided multiprocessing.Pool
. That, however, doesn't work as well for tasks that hold state. I'll assume an example of a process that's doing on-the-fly encryption of data and pumping it to a file. I'd like the keys, the block chaining parameters, and the open file handle (which can't be pickled and passed between processes) to reside as internal state of some EncryptedWriter
object. I'd like to be able to use the public interface of that object completley transparently. But I'd like that object to reside on the external process.
Overview
To that end, this code creates a decorator @process_wrap_object
which wraps a class. The new class will spawn an external process to instantiate an object of the wrapped class. The external process then calls methods on it in the required order, and passes back associated return values. The coordinating object that lives on the original process is responsible for forwarding these functions.
The function process_wrap_object
is the decorator itself, which takes a class and returns a class.
The function _process_wrap_event_loop
is the one which the worker process runs, which is tightly coupled to the process_wrap_object
.
Finally the function _process_disconnection_detector
just checks whether the process_wrap_object
coordinating object has been destroyed, whether by normal garbage collection or because the main process crashed. In either case, it should signal the worker process to close down cleanly.
Caveats
Please note that method calls are blocking, as normal method calls are. This means that on its own, this wrapper will not speed anything up: it just gets the work done elsewhere with more overhead. However, it cooperates effectively with the main process being divvied up with lighter intra-process threading.
Code
import inspect
from functools import partial
from multiprocessing import Process, Queue, Pipe
from threading import Thread
CLOSE_CODE = "_close"
def _process_disconnection_detector(pipe, instruction_queue):
"""Watcher thread function that triggers the process to close if its partner dies"""
try:
pipe.recv()
except EOFError:
instruction_queue.put((CLOSE_CODE, (), {}))
def _process_wrap_event_loop(new_cls, instruction_queue, output_queue, pipe, *args, **kwargs):
cls = new_cls.__wrapped__
obj = cls(*args, **kwargs)
routines = inspect.getmembers(obj, inspect.isroutine)
# Inform the partner class what instructions are valid
output_queue.put([r[0] for r in routines if not r[0].startswith("_")])
# and record them for the event loop
routine_lookup = dict(routines)
disconnect_monitor = Thread(target=_process_disconnection_detector, args=(pipe, instruction_queue))
disconnect_monitor.start()
while True:
instruction, inst_args, inst_kwargs = instruction_queue.get()
if instruction == CLOSE_CODE:
break
inst_op = routine_lookup[instruction]
res = inst_op(*inst_args, **inst_kwargs)
output_queue.put(res)
disconnect_monitor.join()
def process_wrap_object(cls):
"""
Class decorator which exposes the same public method interface as the original class,
but the object itself resides and runs on a separate process.
"""
class NewCls:
def __init__(self, *args, **kwargs):
self._instruction_queue = Queue() # Queue format is ({method_name}, {args}, {kwargs})
self._output_queue = Queue() # Totally generic queue, will carry the return type of the method
self._pipe1, pipe2 = Pipe() # Just a connection to indicate to the worker process when it can close
self._process = Process(
target=_process_wrap_event_loop,
args=([NewCls, self._instruction_queue, self._output_queue, pipe2] + list(args)),
kwargs=kwargs
)
self._process.start()
routine_names = self._output_queue.get()
assert CLOSE_CODE not in routine_names, "Cannot wrap class with reserved method name."
for r in routine_names:
self.__setattr__(
r,
partial(self.trigger_routine, routine_name=r)
)
def trigger_routine(self, *trigger_args, routine_name, **trigger_kwargs):
self._instruction_queue.put((routine_name, trigger_args, trigger_kwargs))
return self._output_queue.get()
def __del__(self):
# When the holding object gets destroyed,
# tell the process to shut down.
self._pipe1.close()
self._process.join()
for wa in ('__module__', '__name__', '__qualname__', '__doc__'):
setattr(NewCls, wa, getattr(cls, wa))
setattr(NewCls, "__wrapped__", cls)
return NewCls
Sample usage:
@process_wrap_object
class Example:
"""Sample class for demoing stuff"""
def __init__(self, a, b):
self._a = a
self._b = b
def inc_a(self):
self._a += 1
def inc_b(self, increment=1):
self._b += increment
def id(self):
return f"{self._a} - {self._b} = {self._a - self._b}"
proc_obj = Example(8, 6)
print(proc_obj.id())
proc_obj.inc_a()
proc_obj.inc_a()
print(proc_obj.id())
proc_obj.inc_b()
print(proc_obj.id())
proc_obj.inc_b(3)
print(proc_obj.id())
I'm after a general review of both high level approach and particular implementation, with particular interest in any subtleties around multiprocessing or correctly wrapping decorators for classes.
Any suggestions about additional functionality that would make this decorator markedly more useful are also welcome. One feature I'm considering, but have not yet implemented, is explicit support for __enter__
and __exit__
to work with with
blocks.
EncryptedWriter
doesn't seem to contain the mutating or printing. \$\endgroup\$