58 lines
2.0 KiB
Python
58 lines
2.0 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
Author: J.A. de Jong
|
|
|
|
Description: MonkeyPatch required to let the Multiprocessing library work properly.
|
|
Should be applied prior to running any other multiprocessing code. Comes from
|
|
Stackoverflow and is mainly used for managing a list of queues that can be
|
|
shared between processes.
|
|
|
|
For more information, see:
|
|
https://stackoverflow.com/questions/46779860/multiprocessing-managers-and-custom-classes
|
|
"""
|
|
from multiprocessing import managers
|
|
import logging
|
|
from functools import wraps
|
|
from inspect import signature
|
|
|
|
orig_AutoProxy = managers.AutoProxy
|
|
|
|
__all__ = ['apply_patch']
|
|
|
|
|
|
@wraps(managers.AutoProxy)
|
|
def AutoProxy(*args, incref=True, manager_owned=False, **kwargs):
|
|
# Create the autoproxy without the manager_owned flag, then
|
|
# update the flag on the generated instance. If the manager_owned flag
|
|
# is set, `incref` is disabled, so set it to False here for the same
|
|
# result.
|
|
autoproxy_incref = False if manager_owned else incref
|
|
proxy = orig_AutoProxy(*args, incref=autoproxy_incref, **kwargs)
|
|
proxy._owned_by_manager = manager_owned
|
|
return proxy
|
|
|
|
|
|
def apply_patch():
|
|
if "manager_owned" in signature(managers.AutoProxy).parameters:
|
|
return
|
|
|
|
logging.debug("Patching multiprocessing.managers.AutoProxy to add manager_owned")
|
|
managers.AutoProxy = AutoProxy
|
|
|
|
# re-register any types already registered to SyncManager without a custom
|
|
# proxy type, as otherwise these would all be using the old unpatched AutoProxy
|
|
SyncManager = managers.SyncManager
|
|
registry = managers.SyncManager._registry
|
|
for typeid, (callable, exposed, method_to_typeid, proxytype) in registry.items():
|
|
if proxytype is not orig_AutoProxy:
|
|
continue
|
|
create_method = hasattr(managers.SyncManager, typeid)
|
|
SyncManager.register(
|
|
typeid,
|
|
callable=callable,
|
|
exposed=exposed,
|
|
method_to_typeid=method_to_typeid,
|
|
create_method=create_method,
|
|
)
|
|
|