multiprocessing_cpython_jupyter_threading.py

"""
Import this first, before *anything* else, if you are using `multiprocessing`
directly or indirectly and you encounter "freezing" when any threaded code uses
things like NumPy, OpenCV, etc.
This should only be be necessary when using the "fork" start method; other methods
should (hopefully) work fine without it:
https://docs.python.org/3.6/library/multiprocessing.html#contexts-and-start-methods

For more information, see:
https://stackoverflow.com/questions/17053671/python-how-do-you-stop-numpy-from-multithreading  # noqa

For clarity, this should *only* be used by the *main* module! (see noisier but more robust
example below)
"""

import os
os.environ.update(
    OMP_NUM_THREADS = '1',
    OPENBLAS_NUM_THREADS = '1',
    NUMEXPR_NUM_THREADS = '1',
    MKL_NUM_THREADS = '1',
)

import numpy as np
import cv2
cv2.setNumThreads(0)

multiprocessing_threading_workaround_more_verbose.py

"""
Same as above, but has assertions to try and ensure that you only call this
in a `__main__` file.

Otherwise, you may "infect" programs that don't need it! (e.g. you have a program
that *doesn't* use `mp.get_context(method="fork")`, but now NumPy, OpenCV, etc.
will go super slow because you've constrained the number of threads)

Only tested on CPython 3.6.x on Ubuntu 18.04.
"""

import importlib
import inspect
import os
import sys

_IGNORE_CALLING_FILENAMES = {
    # Skip importlib modules, as they're used to parse and execute the code.
    '<frozen importlib._bootstrap_external>',
    '<frozen importlib._bootstrap>',
    importlib.__file__,
}


def _get_calling_module_name():
    calling_stack = inspect.stack()[3:]
    for frame_info in calling_stack:
        if frame_info.filename not in _IGNORE_CALLING_FILENAMES:
            break
    else:
        assert False, "Stack too shallow?"
    m = inspect.getmodule(frame_info.frame)
    return m.__name__


def _assert_main_and_import_order():
    calling_name = _get_calling_module_name()
    assert calling_name == "__main__", (
        f"This should only be imported when the calling module is '__main__', "
        f"but it is '{calling_name}'"
    )

    this_should_be_imported_before = {"cv2", "numpy"}

    toplevel_modules = set([name for name in sys.modules if "." not in name])
    already_imported = this_should_be_imported_before & toplevel_modules
    assert len(already_imported) == 0, (
        f"Modules should be imported *after* '{__name__}': {already_imported}"
    )


def _set_env():
    # This implements the actual workaround.
    os.environ.update(
        OMP_NUM_THREADS='1',
        OPENBLAS_NUM_THREADS='1',
        NUMEXPR_NUM_THREADS='1',
        MKL_NUM_THREADS='1',
    )

    import numpy as np  # noqa
    import cv2  # noqa
    cv2.setNumThreads(0)


_assert_main_and_import_order()
_set_env()

Thanks for documenting this! Saved my day!

Welcome! FWIW I've updated this with some extra docs and a more robust "import check".

Just FYI, `threadpoolctl` seems to be the preferred way to handle this according to the resolution of this issue: numpy/numpy#11826 (comment).

…

On Fri, Jun 4, 2021 at 6:14 AM Eric Cousineau ***@***.***> wrote: ***@***.**** commented on this gist. ------------------------------ Welcome! FWIW I've updated this with some extra docs and a more robust "import check". — You are receiving this because you commented. Reply to this email directly, view it on GitHub <https://gist.github.com/8a2d1550f5fa4be4fed87d55a522dbf2#gistcomment-3768843>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ACANLPPIFKVLXNBBIYMKUR3TRDGUDANCNFSM46CE6DUA> .