From b28621acae0972bb6138d1b3334d6dc829863469 Mon Sep 17 00:00:00 2001 From: Fred McDavid Date: Mon, 16 Oct 2017 14:00:17 -0400 Subject: [PATCH] adds code to enable profiling of Celery tasks * as a workaround for multiprocessing refusal to fork daemon processes, incorporates changes to allow importing of Process and Pipe from billiard, the Celery project's fork of the multiprocessing module if MEMPROF_PREFERS_BILLIARD envvar is set to true --- memory_profiler.py | 37 ++++++++++++++++++++++++++++++++----- 1 file changed, 32 insertions(+), 5 deletions(-) diff --git a/memory_profiler.py b/memory_profiler.py index 9e74846..544f411 100644 --- a/memory_profiler.py +++ b/memory_profiler.py @@ -18,13 +18,40 @@ import logging import traceback from signal import SIGKILL - +import importlib +import distutils +from distutils.util import strtobool # TODO: provide alternative when multiprocessing is not available -try: - from multiprocessing import Process, Pipe -except ImportError: - from multiprocessing.dummy import Process, Pipe + +# import Process and Pipe from a multiprocessing library in the following order: +# +# billiard - if and only if MEMPROF_PREFERS_BILLIARD envvar is set and true +# multiprocessing +# multiprocessing.dummy +# +MULTILIBS = ( + ("billiard", (lambda: True if "MEMPROF_PREFERS_BILLIARD" in os.environ and + strtobool(os.environ["MEMPROF_PREFERS_BILLIARD"]) + else False)), + ("multiprocessing", None,), + ("multiprocessing.dummy", None,) +) +for module_tuple in MULTILIBS: + multi_module_name, use_if_test = module_tuple + try: + multi_module = importlib.import_module(multi_module_name) + except ImportError: + if multi_module_name == "multiprocessing.dummy": + raise + continue + if use_if_test is None or use_if_test(): + try: + Process = multi_module.Process + Pipe = multi_module.Pipe + break + except AttributeError: + continue try: from IPython.core.magic import Magics, line_cell_magic, magics_class