From 0967e59e6a13feaf432a6f07bd9437e35d92d1bb Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Thu, 5 Mar 2026 01:35:11 +0100 Subject: [PATCH 1/4] `~.autoprofile`: always running code as `__main__` line_profiler/autoprofile/autoprofile.py::run() - Now always creating a temporary module object for `sys.modules['__main__']` (not only when `as_module=True`), executing the code in its namespace and context - Simplified internal code and dropped some imports due to the reduced branching --- line_profiler/autoprofile/autoprofile.py | 32 ++++++++---------------- 1 file changed, 10 insertions(+), 22 deletions(-) diff --git a/line_profiler/autoprofile/autoprofile.py b/line_profiler/autoprofile/autoprofile.py index 558cd385..1e20fbb7 100644 --- a/line_profiler/autoprofile/autoprofile.py +++ b/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) From 97bf143ce384b45b86db8055aae5b564636177c4 Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Thu, 5 Mar 2026 04:09:46 +0100 Subject: [PATCH 2/4] kernprof.py: always run code as `__main__` kernprof.py::_main_profile() As with `line_profiler.autoprofile.autoprofile.run()`, the code execution now always happens in the namespace of a temporary module object at `sys.modules['__main__']`; this now covers cases without with `--prof-mod=...` and/or `--line-by-line` flags --- kernprof.py | 41 +++++++++++++++++++++++++++-------------- 1 file changed, 27 insertions(+), 14 deletions(-) diff --git a/kernprof.py b/kernprof.py index 222f05f0..8a7c4d6a 100755 --- a/kernprof.py +++ b/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) From b9f7942f76cd99dc80ed8742b49c139207777dd6 Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Thu, 5 Mar 2026 09:17:26 +0100 Subject: [PATCH 3/4] Added tests against issue #422 tests/test_child_procs.py New test module (to be expanded) checking that `kernprof` doesn't choke when running code using `multiprocessing` - test_module `pytest` fixture for a script/module, in which a function is called parallel-ly with `multiprocessing.Pool.map()` - test_multiproc_script_sanity_check() Test that the test code works as intended with vanilla Python - test_running_multiproc_{script,module}() Test that the test code works as intended with `kernprof`, both when invoked as a module and a script TODO: - XFailed tests for incomplete profiling results - Actual implementation of child-process profiling --- tests/test_child_procs.py | 233 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 233 insertions(+) create mode 100644 tests/test_child_procs.py diff --git a/tests/test_child_procs.py b/tests/test_child_procs.py new file mode 100644 index 00000000..10be3f78 --- /dev/null +++ b/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 [...] `). + """ + _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 `). + """ + _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 From 1120c7fe7315fd1cba5a85079b57a9006d9f0c76 Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Thu, 5 Mar 2026 09:32:37 +0100 Subject: [PATCH 4/4] CHANGELOG --- CHANGELOG.rst | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index aef183ba..d28eb7cf 100644 --- a/CHANGELOG.rst +++ b/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