Skip to content

Commit b28621a

Browse files
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
1 parent e2a97b9 commit b28621a

File tree

1 file changed

+32
-5
lines changed

1 file changed

+32
-5
lines changed

memory_profiler.py

Lines changed: 32 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,40 @@
1818
import logging
1919
import traceback
2020
from signal import SIGKILL
21-
21+
import importlib
22+
import distutils
23+
from distutils.util import strtobool
2224

2325
# TODO: provide alternative when multiprocessing is not available
24-
try:
25-
from multiprocessing import Process, Pipe
26-
except ImportError:
27-
from multiprocessing.dummy import Process, Pipe
26+
27+
# import Process and Pipe from a multiprocessing library in the following order:
28+
#
29+
# billiard - if and only if MEMPROF_PREFERS_BILLIARD envvar is set and true
30+
# multiprocessing
31+
# multiprocessing.dummy
32+
#
33+
MULTILIBS = (
34+
("billiard", (lambda: True if "MEMPROF_PREFERS_BILLIARD" in os.environ and
35+
strtobool(os.environ["MEMPROF_PREFERS_BILLIARD"])
36+
else False)),
37+
("multiprocessing", None,),
38+
("multiprocessing.dummy", None,)
39+
)
40+
for module_tuple in MULTILIBS:
41+
multi_module_name, use_if_test = module_tuple
42+
try:
43+
multi_module = importlib.import_module(multi_module_name)
44+
except ImportError:
45+
if multi_module_name == "multiprocessing.dummy":
46+
raise
47+
continue
48+
if use_if_test is None or use_if_test():
49+
try:
50+
Process = multi_module.Process
51+
Pipe = multi_module.Pipe
52+
break
53+
except AttributeError:
54+
continue
2855

2956
try:
3057
from IPython.core.magic import Magics, line_cell_magic, magics_class

0 commit comments

Comments
 (0)