Combining a multiprocessing Process with decorators in python

renatodamas 0
Viewed 41

Desperately in need for help here ...

I am trying to use multiprocessing in combination with decorators. I know there are some similar questions out there, however nothing seems to solve my problem.

I have the following class:

import functools
from multiprocessing import Pipe, get_context


def target(conn, fn, *args, **kwargs):
    conn.send(fn(*args, **kwargs))
    conn.close()


def decor(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        ctx = get_context("fork")

        p_conn, c_conn = Pipe()
        p = ctx.Process(target=target, args=(c_conn, fn, *args), kwargs=kwargs)
        p.start()
        ret = p_conn.recv()
        p.join()
        return ret

    return wrapper


class MyClass:
    @decor
    def printer(self, string=""):
        print(string)


MyClass().printer("foo")

In this scenario, I am gonna get this error:

_pickle.PicklingError: Can't pickle <function MyClass.printer at ...>: it's not the same object as __main__.MyClass.printer

I understand that the pickler complains due to the wrapper "transforming" my function so that it's not my original function anymore. To workaround it, I provide the function name to target:

def target(conn, name, *args, **kwargs):
    # :name: is the name of the function
    fn = ...  # translation from name to fn (suggestions on how it's done are appreciated)
    conn.send(fn(*args, **kwargs))
    conn.close()

Now I run the program and it enters in a never ending loop of process creation. The reason is that when I call fn(*args, **kwargs) in the target, the decorator gets triggered which creates a new process to run target and so on and so forth.

How can I break the loop and have my function processed one single time?

decorator multiprocessing python
2 Answers

The reason this happens is because even though you are using functools.wraps, it only acts as a sort of band aid which changes the __name__ and __qualname__ of the decorated function, but the decorated function (the inner function of your decorator) itself will always be a different object from the actual decorated function (method printer in this case). Pickle internally checks this before pickling, and hence the error.

However, this is actually good because if the check wasn't performed then you likely would end up getting a recursion error since the call to the inner function would never be reached as the inner function would be pickled as pointing towards the decorator in the child process.

So, as a workaround, you would need to explicitly import the decorated function in your local scope (so it passes the checks pickle makes), and then call it from the child process using the functions __wrapped__ attribute (so it does not result in infinite recursion). This attribute is set when you use functools.wraps and points to the original function (so make sure to use functools.wraps on your decorator). Here is a solution:

from functools import wraps
from multiprocessing import Pipe, get_context
import pickle, sys
from pickle import _getattribute, whichmodule


def target(conn, fn, *args, **kwargs):
    conn.send(fn.__wrapped__(*args, **kwargs))  # Use __wrapped__ !
    conn.close()

def decor(fn):

    @wraps(fn)  # Don't forget to use functools.wraps
    def wrapper(*args, **kwargs):

        # Import the original function, works even if it's imported from another file
        module = whichmodule(fn, fn.__qualname__)
        __import__(module, level=0)
        module = sys.modules[module]
        original_func, _ = _getattribute(module, fn.__qualname__)

        ctx = get_context()
        p_conn, c_conn = Pipe()
        p = ctx.Process(target=target, args=(c_conn, original_func, *args), kwargs=kwargs)
        p.start()
        ret = p_conn.recv()
        p.join()
        return ret

    return wrapper


class MyClass:
    @decor
    def printer(self, string=""):
        print(string)

if __name__ == '__main__':

    MyClass().printer("foo")

Tested it on Windows (so using spawn rather than fork), but I don't think it should make a difference in this case.

Charchit Agarwal0

The answer @charchit does the job, just wanna provide a twist to it with minor improvements

import functools
import sys
from multiprocessing import Pipe, get_context


def target(conn, fn, *args, **kwargs):
    conn.send(fn.__wrapped__(*args, **kwargs))
    conn.close()


def decor(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):

        cls_name = fn.__qualname__.split('.')[0]
        cls = vars(sys.modules[fn.__module__])[cls_name]
        func = getattr(cls, fn.__name__)

        ctx = get_context("fork")

        p_conn, c_conn = Pipe()
        p = ctx.Process(target=target, args=(c_conn, func, *args), kwargs=kwargs)
        p.start()
        ret = p_conn.recv()
        p.join()
        return ret

    return wrapper


class MyClass:
    @decor
    def printer(self, string=""):
        print(string)
renatodamas0
Related