5
\$\begingroup\$

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.

\$\endgroup\$
2
  • \$\begingroup\$ Interesting question, could you fill out the "sample usage" so it's like a convoluted "hello world" using your class and passing between the threads? It seems like you started to, but EncryptedWriter doesn't seem to contain the mutating or printing. \$\endgroup\$ – Peilonrayz Feb 15 '20 at 0:47
  • \$\begingroup\$ I've changed the example to be runnable (and a bit more hello world like). There shouldn't be anything more to do than run it as is: the decorator takes care of working on another process, and using the new object uses the same interface as the old one. \$\endgroup\$ – Josiah Feb 15 '20 at 8:24
4
\$\begingroup\$

There are a couple of issues with destructors, firstly is the issue with closing pipes on Linux/Unix as discussed here (though the pipe is actually no longer necessary once we fix the second issue). Secondly the functools.partial method appears to capture a reference to self which causes the wrapper object to not be destructed when expected, I have fixed this using __getattr__ however this has the drawback of allowing the class to respond to any method, perhaps a check in that function that the routine_name is valid would be wise, alternatively if this solution is acceptable the inspect code could all be removed simply calling getattr in _process_wrap_event_loop.

This means we can remove the pipe

import inspect
from multiprocessing import Process, Queue, Pipe
from threading import Thread

CLOSE_CODE = "_close"

def _process_wrap_event_loop(new_cls, instruction_queue, output_queue, *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)

    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)

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._process = Process(
                target=_process_wrap_event_loop,
                args=([NewCls, self._instruction_queue, self._output_queue] + 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."

        def __getattr__(self, routine_name):
            def f(*trigger_args, **trigger_kwargs):
                self._instruction_queue.put((routine_name, trigger_args, trigger_kwargs))
                return self._output_queue.get()
            return f

        def __del__(self):
            # When the holding object gets destroyed,
            # tell the process to shut down.
            self._instruction_queue.put((CLOSE_CODE, (), {}))
            self._process.join()

    for wa in ('__module__', '__name__', '__qualname__', '__doc__'):
        setattr(NewCls, wa, getattr(cls, wa))
    setattr(NewCls, "__wrapped__", cls)

    return NewCls

This can be demonstrated with the following example

@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}"

    def __del__(self):
        print("Deleting example")

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())

del proc_obj
\$\endgroup\$
1
  • \$\begingroup\$ Thank you. I'll spin up a Linux VM for further work, since I've hitherto been testing exclusively on Windows. The problem with partial capturing self is also a very good spot. I've been looking at using partialmethod and have most of a solution. I just want to nail down some of the details around metadata, and in particular making sure that help is as helpful as it can be. \$\endgroup\$ – Josiah Feb 15 '20 at 23:45

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.