Reverting sys.modules causes pickling issues with concurrent.futures.ProcessPoolWorker #436

Summary

When working on #431 I ran into a weird situation where the in-process integration tests using concurrent.futures would occasionally fail when one picks a random subset thereof to run, and pass if the entire test module is run, indicating some global-state pollution. After some spelunking it seems that the cause is:

• kernprof imports concurrent.futures, which (< Python 3.15) uses a module-level __getattr__() to lazy-load the various Executor subclasses (e.g. ProcessPoolExecutor) when requested.
• In kernprof.py::_main_profile() and likewise line_profiler/autoprofile/autoprofile.py::run(), we essentially force-revert the state of sys.modules to before these functions are called.
• If concurrent.futures.process hasn't already been loaded before kernprof.py::main() is called, but the profiled code imports ProcessPoolExecutor from concurrent.futures:
  • concurrent.futures.process will be loaded inside kernprof.py::main(), and its copy of ProcessPoolExecutor will be cached in the concurrent.futures namespace.
  • When kernprof.py::main() exits, concurrent.futures.process will be dropped from sys.modules, but its ProcessPoolExecutor remains in concurrent.futures.
• If one is to go on and make a second identical call to kernprof.py::main():
  • The code will load the copy of ProcessPoolExecutor belonging to the now-dead copy of concurrent.futures.process which (1) its methods all refer to and (2) the import system doesn't know about.
  • ProcessPoolExecutor will then create processes with Process(target=concurrent.futures.process._process_worker, ...), which then promptly goes on to fail with a PicklingError because its copy of _process_worker() is stale and would not be the same object as sys.modules['concurrent.futures.process']._process_worker.

The freezing of sys.modules was originally (#339, #423) intended to mirror the functionalities of runpy.run_module(alter_sys=True), but it seems that we've been overzealous: should've stuck to just reverting sys.modules['__main__'], and tempering with the rest may after all cause more problems than it solves. While this bug may seem specific, a package lazy-loading from its submodules via __getattr__() is a common pattern, and it's not hard to imagine other such cases where import orders causes unexpected behaviors.

Will write a small PR fix.

Example

import os
from tempfile import TemporaryDirectory

from kernprof import main as kp_main


script = """
from __future__ import annotations

from multiprocessing import get_context
from concurrent.futures import ProcessPoolExecutor


def my_sum(x: list[int]) -> int:
    result = 0
    for n in x:
        result += n
    return result


def main() -> None:
    with ProcessPoolExecutor(mp_context=get_context('spawn')) as ex:
        print(list(ex.map(my_sum, [[1, 2], [3, 4], [5, 6]])))


if __name__ == '__main__':
    main()
"""


def main(n: int = 2) -> None:
    for _ in range(n):
        with TemporaryDirectory() as tmpdir:
            kp_main([
                '-lzv',
                f'--outfile={os.path.join(tmpdir, "out.lprof")}',
                '-c',
                script,
            ])


if __name__ == '__main__':
    main()

$ python demo.py
[3, 7, 11]
Wrote profile results to '[...]/out.lprof'
Timer unit: ...
...
Function: main at line 15
...
Wrote profile results to '[...]/out.lprof'
Timer unit: ...
...
Function: main at line 15
...
Traceback (most recent call last):
  ...
_pickle.PicklingError: Can't pickle <function _process_worker at ...>: it's not the same object as concurrent.futures.process._process_worker
when serializing dict item '_target'
when serializing multiprocessing.context.SpawnProcess state
when serializing multiprocessing.context.SpawnProcess object
...