FIX: kernprof to always run code in sys.modules['__main__']'s namespace

CHANGELOG.rst

@@ -1,8 +1,14 @@
 Changes
 =======
 
+5.0.3
+-----
+* FIX: Make sure that the profiled code is run in the
+  ``sys.modules['__main__']`` namespace to avoid issues w/e.g. pickling
+  (#423)
 
-5.0.1
+
+5.0.2
 ~~~~~
 * ENH: improved type annotations and moved them inline
 

kernprof.py

@@ -207,7 +207,7 @@ def main():
 from pprint import pformat
 from shlex import quote
 from textwrap import indent, dedent
-from types import MethodType, SimpleNamespace
+from types import MethodType, ModuleType, SimpleNamespace
 
 # NOTE: This version needs to be manually maintained in
 # line_profiler/line_profiler.py and line_profiler/__init__.py as well
@@ -1331,6 +1331,7 @@ def _main_profile(options, module=False, exit_on_error=True):
     after initial parsing of options; not to be invoked on its own.
     """
     script_file, prof = _pre_profile(options, module, exit_on_error)
+    call = functools.partial(_call_with_diagnostics, options)
     try:
         rmod = functools.partial(
             run_module, run_name='__main__', alter_sys=True
@@ -1345,8 +1346,7 @@ def _main_profile(options, module=False, exit_on_error=True):
         if options.prof_mod and options.line_by_line:
             from line_profiler.autoprofile import autoprofile
 
-            _call_with_diagnostics(
-                options,
+            call(
                 autoprofile.run,
                 script_file,
                 ns,
@@ -1355,21 +1355,34 @@ def _main_profile(options, module=False, exit_on_error=True):
                 as_module=module is not None,
             )
         else:
+            # Note: to reduce complications (e.g. whenever something
+            # needs to be pickled), regardless of whether the code is to
+            # be run as a module, we always create a mock module object
+            # for `sys.modules['__main__']` and execute the code in its
+            # context;
+            # similar handling is already used for
+            # `~.autoprofile.autoprofile.run()`, and we do the same for
+            # the other execution modes here.
+            module_obj = ModuleType('__main__')
+            module_ns = vars(module_obj)
+            module_ns.update(ns)
             if module:
                 runner, target = 'rmod', options.script
             else:
                 runner, target = 'execfile', script_file
-            assert runner in ns
-            if options.builtin:
-                _call_with_diagnostics(options, ns[runner], target, ns)
-            else:
-                _call_with_diagnostics(
-                    options,
-                    prof.runctx,
-                    f'{runner}({target!r}, globals())',
-                    ns,
-                    ns,
-                )
+            assert runner in module_ns
+
+            with _restore.mapping(sys.modules):
+                sys.modules['__main__'] = module_obj
+                if options.builtin:
+                    call(module_ns[runner], target, module_ns)
+                else:
+                    call(
+                        prof.runctx,
+                        f'{runner}({target!r}, globals())',
+                        module_ns,
+                        module_ns,
+                    )
     finally:
         _post_profile(options, prof)
 

line_profiler/autoprofile/autoprofile.py

@@ -46,10 +46,7 @@ def main():
 """
 
 from __future__ import annotations
-import contextlib
-import functools
 import importlib.util
-import operator
 import sys
 import types
 from collections.abc import MutableMapping
@@ -112,15 +109,13 @@ def run(
     """
 
     class restore_dict:
-        def __init__(self, d: MutableMapping[str, Any], target=None):
+        def __init__(self, d: MutableMapping[str, Any]):
             self.d = d
-            self.target = target
             self.copy: Mapping[str, Any] | None = None
 
         def __enter__(self):
             assert self.copy is None
             self.copy = dict(self.d)
-            return self.target
 
         def __exit__(self, *_, **__):
             self.d.clear()
@@ -129,8 +124,6 @@ def __exit__(self, *_, **__):
             self.copy = None
 
     Profiler: type[AstTreeModuleProfiler] | type[AstTreeProfiler]
-    namespace: MutableMapping[str, Any]
-    ctx: ContextManager
 
     if as_module:
         Profiler = AstTreeModuleProfiler
@@ -141,28 +134,23 @@ def __exit__(self, *_, **__):
             )
 
         module_obj = types.ModuleType(module_name)
-        namespace = vars(module_obj)
-        namespace.update(ns)
-
         # Set the `__spec__` correctly
         module_obj.__spec__ = importlib.util.find_spec(module_name)
-
-        # Set the module object to `sys.modules` via a callback, and
-        # then restore it via the context manager
-        callback = functools.partial(
-            operator.setitem, sys.modules, '__main__', module_obj
-        )
-        ctx = restore_dict(sys.modules, callback)
     else:
         Profiler = AstTreeProfiler
-        namespace = ns
-        ctx = contextlib.nullcontext(lambda: None)
+        module_obj = types.ModuleType('__main__')
+
+    namespace: MutableMapping[str, Any] = vars(module_obj)
+    namespace.update(ns)
 
     profiler = Profiler(script_file, prof_mod, profile_imports)
     tree_profiled = profiler.profile()
 
     _extend_line_profiler_for_profiling_imports(ns[PROFILER_LOCALS_NAME])
     code_obj = compile(tree_profiled, script_file, 'exec')
-    with ctx as callback:
-        callback()
+    with restore_dict(sys.modules):
+        # Always set the module object to `sys.modules['__main__']` and
+        # then restore it via the context manager, so that the executed
+        # code is run as `__main__`
+        sys.modules['__main__'] = module_obj
         exec(code_obj, cast(Dict[str, Any], namespace), namespace)

tests/test_child_procs.py

@@ -0,0 +1,233 @@
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from collections.abc import Callable, Generator, Mapping
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from textwrap import dedent, indent
+
+import pytest
+import ubelt as ub
+
+
+NUM_NUMBERS = 100
+NUM_PROCS = 4
+TEST_MODULE_BODY = dedent(f"""
+from __future__ import annotations
+from argparse import ArgumentParser
+from multiprocessing import Pool
+
+
+def my_sum(x: list[int]) -> int:
+    result: int = 0
+    for item in x:
+        result += item
+    return result
+
+
+def sum_in_child_procs(length: int, n: int) -> int:
+    my_list: list[int] = list(range(1, length + 1))
+    sublists: list[list[int]] = []
+    subsums: list[int]
+    sublength = length // n
+    if sublength * n < length:
+        sublength += 1
+    while my_list:
+        sublist, my_list = my_list[:sublength], my_list[sublength:]
+        sublists.append(sublist)
+    with Pool(n) as pool:
+        subsums = pool.map(my_sum, sublists)
+        pool.close()
+        pool.join()
+    return my_sum(subsums)
+
+
+def main(args: list[str] | None = None) -> None:
+    parser = ArgumentParser()
+    parser.add_argument('-l', '--length', type=int, default={NUM_NUMBERS})
+    parser.add_argument('-n', type=int, default={NUM_PROCS})
+    options = parser.parse_args(args)
+    print(sum_in_child_procs(options.length, options.n))
+
+
+if __name__ == '__main__':
+    main()
+""").strip('\n')
+
+
+@pytest.fixture(scope='session')
+def test_module() -> Generator[Path, None, None]:
+    with TemporaryDirectory() as mydir_str:
+        my_dir = Path(mydir_str)
+        my_dir.mkdir(exist_ok=True)
+        my_module = my_dir / 'my_test_module.py'
+        with my_module.open('w') as fobj:
+            fobj.write(TEST_MODULE_BODY + '\n')
+        yield my_module
+
+
+@pytest.mark.parametrize('as_module', [True, False])
+@pytest.mark.parametrize(
+    ('nnums', 'nprocs'), [(None, None), (None, 3), (200, None)],
+)
+def test_multiproc_script_sanity_check(
+    test_module: Path,
+    tmp_path_factory: pytest.TempPathFactory,
+    nnums: int,
+    nprocs: int,
+    as_module: bool,
+) -> None:
+    """
+    Sanity check that the test module functions as expected when run
+    with vanilla Python.
+    """
+    _run_test_module(
+        _run_as_module if as_module else _run_as_script,
+        test_module, tmp_path_factory, [sys.executable], None, False,
+        nnums=nnums, nprocs=nprocs,
+    )
+
+
+# Note:
+# Currently code execution in child processes is not properly profiled;
+# these tests are just for checking that `kernprof` doesn't impair the
+# proper execution of `multiprocessing` code
+
+
+fuzz_invocations = pytest.mark.parametrize(
+    ('runner', 'outfile', 'profile',
+     'label'),  # Dummy argument to make `pytest` output more legible
+    [
+        (['kernprof', '-q'], 'out.prof', False, 'cProfile'),
+        # Run with `line_profiler` with and w/o profiling targets
+        (['kernprof', '-q', '-l'], 'out.lprof', False,
+         'line_profiler-inactive'),
+        (['kernprof', '-q', '-l'], 'out.lprof', True,
+         'line_profiler-active'),
+    ],
+)
+
+
+@fuzz_invocations
+def test_running_multiproc_script(
+    test_module: Path,
+    tmp_path_factory: pytest.TempPathFactory,
+    runner: str | list[str],
+    outfile: str | None,
+    profile: bool,
+    label: str,
+) -> None:
+    """
+    Check that `kernprof` can run the test module as a script
+    (`kernprof [...] <path>`).
+    """
+    _run_test_module(
+        _run_as_script,
+        test_module, tmp_path_factory, runner, outfile, profile,
+    )
+
+
+@fuzz_invocations
+def test_running_multiproc_module(
+    test_module: Path,
+    tmp_path_factory: pytest.TempPathFactory,
+    runner: str | list[str],
+    outfile: str | None,
+    profile: bool,
+    label: str,
+) -> None:
+    """
+    Check that `kernprof` can run the test module as a module
+    (`kernprof [...] -m <module>`).
+    """
+    _run_test_module(
+        _run_as_module,
+        test_module, tmp_path_factory, runner, outfile, profile,
+    )
+
+
+def _run_as_script(
+    runner_args: list[str], test_args: list[str], test_module: Path, **kwargs
+) -> subprocess.CompletedProcess:
+    cmd = runner_args + [str(test_module)] + test_args
+    return subprocess.run(cmd, **kwargs)
+
+
+def _run_as_module(
+    runner_args: list[str],
+    test_args: list[str],
+    test_module: Path,
+    *,
+    env: Mapping[str, str] | None = None,
+    **kwargs
+) -> subprocess.CompletedProcess:
+    cmd = runner_args + ['-m', test_module.stem] + test_args
+    env_dict = {**os.environ, **(env or {})}
+    python_path = env_dict.pop('PYTHONPATH', '')
+    if python_path:
+        env_dict['PYTHONPATH'] = '{}:{}'.format(
+            test_module.parent, python_path,
+        )
+    else:
+        env_dict['PYTHONPATH'] = str(test_module.parent)
+    return subprocess.run(cmd, env=env_dict, **kwargs)
+
+
+def _run_test_module(
+    run_helper: Callable[..., subprocess.CompletedProcess],
+    test_module: Path,
+    tmp_path_factory: pytest.TempPathFactory,
+    runner: str | list[str],
+    outfile: str | None,
+    profile: bool,
+    *,
+    nnums: int | None = None,
+    nprocs: int | None = None,
+    check: bool = True,
+) -> tuple[subprocess.CompletedProcess, Path | None]:
+    """
+    Return
+    ------
+    `(process_running_the_test_module, path_to_profiling_output | None)`
+    """
+    if isinstance(runner, str):
+        runner_args: list[str] = [runner]
+    else:
+        runner_args = list(runner)
+    if profile:
+        runner_args.extend(['--prof-mod', str(test_module)])
+
+    test_args: list[str] = []
+    if nnums is None:
+        nnums = NUM_NUMBERS
+    else:
+        test_args.extend(['-l', str(nnums)])
+    if nprocs is not None:
+        test_args.extend(['-n', str(nprocs)])
+
+    with ub.ChDir(tmp_path_factory.mktemp('mytemp')):
+        if outfile is not None:
+            runner_args.extend(['--outfile', outfile])
+        proc = run_helper(
+            runner_args, test_args, test_module,
+            text=True, capture_output=True,
+        )
+        try:
+            if check:
+                proc.check_returncode()
+        finally:
+            print(f'stdout:\n{indent(proc.stdout, "  ")}')
+            print(f'stderr:\n{indent(proc.stderr, "  ")}', file=sys.stderr)
+
+        assert proc.stdout == f'{nnums * (nnums + 1) // 2}\n'
+
+        prof_result: Path | None = None
+        if outfile is None:
+            assert not list(Path.cwd().iterdir())
+        else:
+            prof_result = Path(outfile).resolve()
+            assert prof_result.exists()
+            assert prof_result.stat().st_size
+    return proc, prof_result