diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index f318f5ba..49506c9b 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -91,6 +91,7 @@ jobs: - name: Test full loose sdist env: COVERAGE_CORE: ctrace + timeout-minutes: 10 run: |- pwd ls -al @@ -154,6 +155,7 @@ jobs: CIBW_ARCHS_LINUX: ${{ matrix.arch }} PYTHONUTF8: '1' VSCMD_ARG_TGT_ARCH: '' + timeout-minutes: 60 - name: Show built files shell: bash run: ls -la wheelhouse @@ -426,6 +428,7 @@ jobs: env: CI_PYTHON_VERSION: py${{ matrix.python-version }} COVERAGE_CORE: ctrace + timeout-minutes: 10 run: |- echo "Creating test sandbox directory" export WORKSPACE_DNAME="testdir_${CI_PYTHON_VERSION}_${GITHUB_RUN_ID}_${RUNNER_OS}" diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 30156f63..118e2fca 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -12,6 +12,8 @@ Changes is available (#427) * FIX: Bytecodes of profiled functions now always labeled to prevent confusion with non-profiled "twins" (#425) +* FEAT: Experimental support for profiling child processes with + ``kernprof --prof-child-procs`` (#431) 5.0.2 diff --git a/_line_profiler_hooks.py b/_line_profiler_hooks.py new file mode 100644 index 00000000..8418e4a4 --- /dev/null +++ b/_line_profiler_hooks.py @@ -0,0 +1,79 @@ +""" +Additional hooks installed by :py:mod:`line_profiler`. + +Notes: + - This file and its content should be considered an implementation + detail of :py:mod:`line_profiler`; currently we just use this to + set up shop in a child Python process, and extend profiling to + therein. + + - This current implementation writes temporary .pth files to the + site-packages directory, which are executed for all Python + processes referring to the same ``lib/``. However, only processes + originating from a parent which set the requisite environment + variables will execute to the profiling code. + + - Said .pth file always import this module; hence, this file is kept + intentionally lean and separate from the main + :py:mod:`line_profiler` package to reduce overhead; e.g. + imports in this file are deferred to being as late as possible. + + - Inspired by similar code in :py:mod:`coverage.control` and + :py:mod:`pytest_autoprofile.startup_hook`. +""" +import os + + +__all__ = ('load_pth_hook',) + +INHERITED_PID_ENV_VARNAME = ( + 'LINE_PROFILER_PROFILE_CHILD_PROCESSES_CACHE_PID' +) + + +def load_pth_hook(ppid: int) -> None: + """ + Function imported and called by the written .pth file; to reduce + overhead, we immediately return if ``ppid`` doesn't match + :env:`LINE_PROFILER_PROFILE_CHILD_PROCESSES_CACHE_PID`. + """ + try: + env_ppid = int(os.environ[INHERITED_PID_ENV_VARNAME]) + except (KeyError, ValueError): + return + if env_ppid != ppid: + return + # Note: .pth files may be double-loaded in a virtual environment + # (see https://stackoverflow.com/questions/58807569), so work around + # that; + # also see similar check in `coverage.control.process_startup()` + if getattr(load_pth_hook, 'called', False): + return + + # If we're here, we're most probably in a descendent process of a + # profiled Python process, so we can be more liberal with the + # imports without worrying about overhead + import warnings + from line_profiler._diagnostics import log + from line_profiler._child_process_profiling.cache import LineProfilingCache + + try: + cache = LineProfilingCache.load() + cache._setup_in_child_process(True, 'pth') + except Exception as e: # nocover + # A child which has the profiling environment but cannot set + # up profiling is an abnormal condition the user asked to know + # about, so always warn (but still let the process run + # unprofiled) + msg = ( + f'line_profiler child-process profiling setup failed in ' + f'PID {os.getpid()} ({type(e).__name__}: {e}); ' + f'this process runs unprofiled' + ) + # Write log before issuing the warning, in case the warning is + # promoted to an exception + log.warning(msg) + warnings.warn(msg) + load_pth_hook.called = True # type: ignore + else: + cache.patch(load_pth_hook, 'called', True) diff --git a/kernprof.py b/kernprof.py index 8a7c4d6a..82a4f031 100755 --- a/kernprof.py +++ b/kernprof.py @@ -79,6 +79,7 @@ def main(): [-s SETUP] [-p {path/to/script | object.dotted.path}[,...]] [--preimports [Y[es] | N[o] | T[rue] | F[alse] | on | off | 1 | 0]] [--prof-imports [Y[es] | N[o] | T[rue] | F[alse] | on | off | 1 | 0]] + [--prof-child-procs [Y[es] | N[o] | T[rue] | F[alse] | on | off | 1 | 0]] [-o OUTFILE] [-v] [-q] [--rich [Y[es] | N[o] | T[rue] | F[alse] | on | off | 1 | 0]] [-u UNIT] @@ -137,6 +138,10 @@ def main(): If the script/module profiled is in `--prof-mod`, autoprofile all its imports. Only works with line profiling (`-l`/`--line- by-line`). (Default: False) + --prof-child-procs [Y[es] | N[o] | T[rue] | F[alse] | on | off | 1 | 0] + Extend profiling into child Python processes. Only works with + line profiling (`-l`/`--line-by-line`). (EXPERIMENTAL; + default: False) output options: -o, --outfile OUTFILE @@ -187,7 +192,6 @@ def main(): """ # noqa: E501 import atexit -import builtins import functools import os import sys @@ -198,9 +202,7 @@ def main(): import shutil import tempfile import time -import warnings -from argparse import ArgumentParser -from io import StringIO +from argparse import ArgumentParser, SUPPRESS from operator import methodcaller from runpy import run_module from pathlib import Path @@ -228,12 +230,19 @@ def main(): positive_float, short_string_path, ) +from line_profiler.line_profiler_utils import ( + make_tempfile as _touch_tempfile, # Compatibility +) from line_profiler.profiler_mixin import ByCountProfilerMixin +from line_profiler._child_process_profiling.cache import LineProfilingCache from line_profiler._logger import Logger from line_profiler import _diagnostics as diagnostics DIAGNOSITICS_VERBOSITY = 2 +CLEANUP_PRIORITIES = { # More negative number -> more delayed + 'gather_logs': -1, +} def execfile(filename, globals=None, locals=None): @@ -330,6 +339,7 @@ def resolve_module_path(mod_name): # type: (str) -> str | None fname = mod_spec.origin # type: str | None if fname and os.path.exists(fname): return fname + return None get_module_path = modname_to_modpath if static else resolve_module_path @@ -681,6 +691,14 @@ def _add_core_parser_arguments(parser): 'Only works with line profiling (`-l`/`--line-by-line`). ' f'(Default: {default.conf_dict["prof_imports"]})', ) + add_argument( + prof_opts, + '--prof-child-procs', + action='store_true', + help='Extend profiling into child Python processes. ' + 'Only works with line profiling (`-l`/`--line-by-line`). ' + f'(EXPERIMENTAL; default: {default.conf_dict["prof_child_procs"]})', + ) out_opts = parser.add_argument_group('output options') if default.conf_dict['outfile']: def_outfile = repr(default.conf_dict['outfile']) @@ -771,6 +789,8 @@ def _add_core_parser_arguments(parser): 'Minimum value (and the value implied if the bare option ' f'is given) is 1 s. (Default: {def_out_int})', ) + # Hidden option for dumping the debug logs to a designated location + add_argument(out_opts, '--debug-log', help=SUPPRESS) def _build_parsers(args=None): @@ -803,8 +823,8 @@ def _build_parsers(args=None): # We've already consumed the `-m `, so we need a dummy # parser for generating the help text; # but the real parser should not consume the `options.script` - # positional arg, and it it got the `--help` option, it should - # hand off the the dummy parser + # positional arg, and if it got the `--help` option, it should + # hand off to the dummy parser real_parser = ArgumentParser(add_help=False, **parser_kwargs) real_parser.add_argument('-h', '--help', action='store_true') help_parser = ArgumentParser(**parser_kwargs) @@ -934,6 +954,15 @@ def main(args=None, *, exit_on_error=True): files created during execution may be deferred to when the interpreter exits. """ + def rmdir_with_pid_guard(pid, dir, *, defer=False, **kwargs): + if os.getpid() != pid: # Leave cleanup to the main process + return + callback = functools.partial(_remove, dir, **kwargs) + if defer: + atexit.register(callback) + else: + callback() + real_parser, help_parser, special_info = _build_parsers(args=args) args = special_info['args'] module = special_info['module'] @@ -968,7 +997,8 @@ def main(args=None, *, exit_on_error=True): cleanup = no_op else: cleanup = functools.partial( - _remove, + rmdir_with_pid_guard, + os.getpid(), tmpdir, recursive=True, missing_ok=True, @@ -985,30 +1015,12 @@ def main(args=None, *, exit_on_error=True): except BaseException: # Defer deletion to after the traceback has been formatted # if needs be - if os.listdir(tmpdir): - atexit.register(cleanup) - else: # Empty tempdir, just delete it - cleanup() + cleanup(defer=bool(os.path.isdir(tmpdir) and os.listdir(tmpdir))) raise else: # Execution succeeded, delete the tempdir ASAP cleanup() -def _touch_tempfile(*args, **kwargs): - """ - Wrapper around :py:func:`tempfile.mkstemp()` which drops and closes - the integer handle (which we don't need and may cause issues on some - platforms). - """ - handle, path = tempfile.mkstemp(*args, **kwargs) - try: - os.close(handle) - except Exception: - os.remove(path) - raise - return path - - def _write_tempfile(source, content, options): """ Called by :py:func:`main()` to handle :command:`kernprof -c` and @@ -1043,104 +1055,33 @@ def _write_tempfile(source, content, options): ) -def _gather_preimport_targets(options, exclude): - """ - Used in _write_preimports - """ - from line_profiler.autoprofile.util_static import modpath_to_modname - from line_profiler.autoprofile.eager_preimports import is_dotted_path - - filtered_targets = [] - recurse_targets = [] - invalid_targets = [] - for target in options.prof_mod: - if is_dotted_path(target): - modname = target - else: - # Paths already normalized by - # `_normalize_profiling_targets()` - if not os.path.exists(target): - invalid_targets.append(target) - continue - if any(os.path.samefile(target, excluded) for excluded in exclude): - # Ignore the script to be run in eager importing - # (`line_profiler.autoprofile.autoprofile.run()` will - # handle it) - continue - modname = modpath_to_modname(target, hide_init=False) - if modname is None: # Not import-able - invalid_targets.append(target) - continue - if modname.endswith('.__init__'): - modname = modname.rpartition('.')[0] - filtered_targets.append(modname) - else: - recurse_targets.append(modname) - if invalid_targets: - invalid_targets = sorted(set(invalid_targets)) - msg = ( - '{} profile-on-import target{} cannot be converted to ' - 'dotted-path form: {!r}'.format( - len(invalid_targets), - '' if len(invalid_targets) == 1 else 's', - invalid_targets, - ) - ) - warnings.warn(msg) - diagnostics.log.warning(msg) - - return filtered_targets, recurse_targets - - -def _write_preimports(prof, options, exclude): +def _write_preimports(prof, options, exclude, keep=False): """ Called by :py:func:`main()` to handle eager pre-imports; not to be invoked on its own. """ - from line_profiler.autoprofile.eager_preimports import ( - write_eager_import_module, - ) - from line_profiler.autoprofile.autoprofile import ( - _extend_line_profiler_for_profiling_imports as upgrade_profiler, - ) + from line_profiler.curated_profiling import ClassifiedPreimportTargets - filtered_targets, recurse_targets = _gather_preimport_targets( - options, exclude - ) - if not (filtered_targets or recurse_targets): - return # We could've done everything in-memory with `io.StringIO` and `exec()`, # but that results in indecipherable tracebacks should anything goes wrong; # so we write to a tempfile and `execfile()` it - upgrade_profiler(prof) temp_mod_path = _touch_tempfile( dir=options.tmpdir, prefix='kernprof-eager-preimports-', suffix='.py' ) - write_module_kwargs = { - 'dotted_paths': filtered_targets, - 'recurse': recurse_targets, - 'static': options.static, - } - temp_file = open(temp_mod_path, mode='w') - if options.debug: - with StringIO() as sio: - write_eager_import_module(stream=sio, **write_module_kwargs) - code = sio.getvalue() - with temp_file as fobj: - print(code, file=fobj) - diagnostics.log.debug( - 'Wrote temporary module for pre-imports to ' - f'{short_string_path(temp_mod_path)!r}' + with open(temp_mod_path, mode='w') as fobj: + preimports = ClassifiedPreimportTargets.from_targets( + options.prof_mod, exclude, ) - else: - with temp_file as fobj: - write_eager_import_module(stream=fobj, **write_module_kwargs) - if not options.dryrun: + preimports.write_preimport_module( + fobj, debug=options.debug, static=options.static, + ) + if preimports and not options.dryrun: ns = {} # Use a fresh namespace execfile(temp_mod_path, ns, ns) # Delete the tempfile ASAP if its execution succeeded - if not diagnostics.KEEP_TEMPDIRS: - _remove(temp_mod_path) + if keep or diagnostics.KEEP_TEMPDIRS: + return temp_mod_path + _remove(temp_mod_path) def _remove(path, *, recursive=False, missing_ok=False): @@ -1154,9 +1095,20 @@ def _remove(path, *, recursive=False, missing_ok=False): path.unlink(missing_ok=missing_ok) -def _dump_filtered_stats(tmpdir, prof, filename): +def _dump_filtered_stats(tmpdir, prof, filename, extra_line_stats=None): import os - import pickle + + if isinstance(prof, ContextualProfile): + # - Not using `line_profiler` + # -> doesn't matter if the source lines can't be retrieved + # -> no need to filter anything + prof.dump_stats(filename) + return + + # Remember to incorporate extra stats where available + line_stats = prof.get_stats() + if extra_line_stats is not None: + line_stats += extra_line_stats # Build list of known temp file paths tempfile_paths = [ @@ -1164,31 +1116,28 @@ def _dump_filtered_stats(tmpdir, prof, filename): for dirpath, _, fnames in os.walk(tmpdir) for fname in fnames ] - - if not tempfile_paths or isinstance(prof, ContextualProfile): + if not tempfile_paths: # - No tempfiles written -> no function lives in tempfiles # -> no need to filter anything - # - Not using `line_profiler` - # -> doesn't matter if the source lines can't be retrieved - # -> no need to filter anything - prof.dump_stats(filename) + line_stats.to_file(filename) return + _dump_filtered_line_stats(line_stats, tempfile_paths, filename) + + +def _dump_filtered_line_stats(stats, exclude, filename): # Filter the filenames to remove data from tempfiles, which will # have been deleted by the time the results are viewed in a # separate process - stats = prof.get_stats() timings = stats.timings for key in set(timings): fname = key[0] try: - if any(os.path.samefile(fname, tmp) for tmp in tempfile_paths): + if any(os.path.samefile(fname, tmp) for tmp in exclude): del timings[key] except OSError: del timings[key] - - with open(filename, 'wb') as f: - pickle.dump(stats, f, protocol=pickle.HIGHEST_PROTOCOL) + stats.to_file(filename) def _format_call_message(func, *args, **kwargs): @@ -1231,13 +1180,110 @@ def _call_with_diagnostics(options, func, *args, **kwargs): return func(*args, **kwargs) -def _pre_profile(options, module, exit_on_error): +class _manage_profiler: """ Prepare the environment to execute profiling with requested options. Note: modifies ``options`` with extra attributes. """ + cache: LineProfilingCache + + def __init__(self, options, module, exit_on_error): + self.options = options + self.module = module + self.exit_on_error = exit_on_error + self._pid = os.getpid() + + def __enter__(self): + from line_profiler.curated_profiling import CuratedProfilerContext + + self.prof = _prepare_profiler( + self.options, self.module, self.exit_on_error, + ) + self._ctx = CuratedProfilerContext( + self.prof, insert_builtin=self.options.builtin, + ) + self._ctx.install() + # Keep the generated pre-imports file to be reused in child + # processes + try: + script_file, preimports_file = _prepare_exec_script( + self.options, self.module, self.prof, + exit_on_error=self.exit_on_error, + keep_preimports_file=self.set_up_child_profiling, + ) + self._set_up_session_cache(script_file, preimports_file) + return self.prof, script_file + except BaseException: + # Make sure that we don't leak the changes made by `._ctx` + # before we've gotten out of `.__enter__()` + self._ctx.uninstall() + raise + + def __exit__(self, *_, **__): + try: + # Guard against when we've forked inside the executed and + # profiled code... + if os.getpid() == self._pid: + extra_stats = None + try: + extra_stats = self._gather_child_prof_stats() + finally: + # Write process-local stats regardless of whether + # child-process stats can be successfully gathered... + # but watch out for when we've forked inside the + # executed code + _post_profile(self.options, self.prof, extra_stats) + finally: + self._ctx.uninstall() + + def _gather_debug_log(self, logfile): + with open(logfile, mode='w') as fobj: + for entry in self.cache._gather_debug_log_entries(): + print(entry.to_text(), file=fobj) + + def _set_up_session_cache(self, script_file, preimports_file): + if not self.set_up_child_profiling: + return + self.cache = cache = _prepare_child_profiling_cache( + self.options, self._ctx, self.prof, preimports_file, script_file, + ) + # Add deferred callbacks for gathering debug logfiles + # (should run right before `.cache.cache_dir` is wiped): + # - Write the debug logs to the `._diagnostics` logger + if cache.debug: + self._ctx.add_cleanup_with_priority( + cache._dump_debug_logs, CLEANUP_PRIORITIES['gather_logs'], + ) + # - Write the debug logs to a specific file + if self.options.debug_log: + self._ctx.add_cleanup_with_priority( + self._gather_debug_log, + CLEANUP_PRIORITIES['gather_logs'], + self.options.debug_log, + ) + + def _gather_child_prof_stats(self): + if not self.set_up_child_profiling: + return None + # Cleaning up here ensures the `multiprocessing` fork-server + # process is rebooted, thus any profiling data on it will be + # properly collected + self.cache.cleanup() + return self.cache.gather_stats() + + @property + def set_up_child_profiling(self): + return bool( + self.options.line_by_line and self.options.prof_child_procs + ) + + +def _prepare_profiler(options, module, exit_on_error): + """ + Set up the appropriate profiler instance. + """ if not options.outfile: extension = 'lprof' if options.line_by_line else 'prof' options.outfile = f'{os.path.basename(options.script)}.{extension}' @@ -1245,7 +1291,10 @@ def _pre_profile(options, module, exit_on_error): f'Using default output destination {short_string_path(options.outfile)!r}' ) - sys.argv = [options.script] + options.args + # Note: we'll restore the content of `sys.argv` with the decorator + # on `main()`, so it's simpler to just `setitem()` instead of + # replacing the whole list object + sys.argv[:] = [options.script] + options.args if module: # Make sure the current directory is on `sys.path` to emulate # `python -m` @@ -1267,24 +1316,26 @@ def _pre_profile(options, module, exit_on_error): execfile(setup_file, ns, ns) if options.line_by_line: - prof = line_profiler.LineProfiler() options.builtin = True + return line_profiler.LineProfiler() elif Profile.__module__ == 'profile': raise RuntimeError( 'non-line-by-line profiling depends on cProfile, ' 'which is not available on this platform' ) else: - prof = ContextualProfile() + return ContextualProfile() - # Overwrite the explicit decorator - global_profiler = line_profiler.profile - install_profiler = global_profiler._kernprof_overwrite - install_profiler(prof) - - if options.builtin: - builtins.__dict__['profile'] = prof +def _prepare_exec_script( + options, module, prof, + *, + exit_on_error=False, + keep_preimports_file=False, +): + """ + Set up the script to be executed among other things. + """ if module: script_file = find_module_script( options.script, static=options.static, exit_on_error=exit_on_error @@ -1304,6 +1355,8 @@ def _pre_profile(options, module, exit_on_error): options.prof_mod = _normalize_profiling_targets(options.prof_mod) if not options.prof_mod: options.preimports = False + + preimports_file = None if options.line_by_line and options.preimports: # We assume most items in `.prof_mod` to be import-able without # significant side effects, but the same cannot be said if it @@ -1311,10 +1364,10 @@ def _pre_profile(options, module, exit_on_error): # even have a `if __name__ == '__main__': ...` guard. So don't # eager-import it. exclude = set() if module else {script_file} - _write_preimports(prof, options, exclude) + preimports_file = _write_preimports( + prof, options, exclude, keep=keep_preimports_file, + ) - options.global_profiler = global_profiler - options.install_profiler = install_profiler if options.output_interval and not options.dryrun: options.rt = RepeatedTimer( max(options.output_interval, 1), prof.dump_stats, options.outfile @@ -1322,7 +1375,53 @@ def _pre_profile(options, module, exit_on_error): else: options.rt = None options.original_stdout = sys.stdout - return script_file, prof + return script_file, preimports_file + + +def _prepare_child_profiling_cache( + options, ctx, prof, preimports_file, script_file +): + """ + Handle the (line-)profiling of spawned/forked child Python + processes. + """ + # Create the cache dir and cache file here; the cache instance will + # be responsible for managing their lifetimes, while derivative + # instances in child processes will merely inherit and use them + cache = LineProfilingCache( + # We don't need to manage this tempdir separately, + # `options.tmpdir` and its contents will be deleted or + # preserved as is appropriate + cache_dir=tempfile.mkdtemp( + dir=options.tmpdir, prefix='line-profiling-cache-', + ), + config=options.config, + profiling_targets=options.prof_mod, + rewrite_module=script_file, + profile_imports=options.prof_imports, + preimports_module=preimports_file, + insert_builtin=options.builtin, + debug=bool(options.debug or options.debug_log), + ) + clean_up = functools.partial(cache.add_cleanup, _remove, missing_ok=True) + clean_up(cache.filename) + + # This file is handed to us at the end of + # `_manage_profiler.__enter__()`; + # normally it is deleted before `.__enter__()` returns, but when + # child-process profiling is used, it is to persist for the lifetime + # of the cache (so that child processes can do the same preimports) + if not (preimports_file is None or diagnostics.KEEP_TEMPDIRS): + clean_up(preimports_file) + + # Handle various setup tasks (see docs thereof) + cache._setup_in_main_process() + cache.profiler = prof + + # Have the context clean up the cache as a failsafe + ctx.add_cleanup(cache.cleanup) + + return cache def _main_profile(options, module=False, exit_on_error=True): @@ -1330,9 +1429,10 @@ def _main_profile(options, module=False, exit_on_error=True): Called by :py:func:`main()` for the actual execution and profiling of code 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: + with _manage_profiler( + options, module, exit_on_error, + ) as (prof, script_file): rmod = functools.partial( run_module, run_name='__main__', alter_sys=True ) @@ -1383,18 +1483,18 @@ def _main_profile(options, module=False, exit_on_error=True): module_ns, module_ns, ) - finally: - _post_profile(options, prof) -def _post_profile(options, prof): +def _post_profile(options, prof, extra_line_stats=None): """ Cleanup setup after executing a main profile """ if options.rt is not None: options.rt.stop() if not options.dryrun: - _dump_filtered_stats(options.tmpdir, prof, options.outfile) + _dump_filtered_stats( + options.tmpdir, prof, options.outfile, extra_line_stats, + ) short_outfile = short_string_path(options.outfile) diagnostics.log.info( ( @@ -1405,9 +1505,15 @@ def _post_profile(options, prof): + f'to {short_outfile!r}' ) if options.verbose > 0 and not options.dryrun: - kwargs = {} - if not isinstance(prof, ContextualProfile): - kwargs.update( + if isinstance(prof, ContextualProfile): + _call_with_diagnostics(options, prof.print_stats) + else: + stats = prof.get_stats() + if extra_line_stats is not None: + stats += extra_line_stats + _call_with_diagnostics( + options, + stats.print, output_unit=options.unit, stripzeros=options.skip_zero, summarize=options.summarize, @@ -1415,7 +1521,6 @@ def _post_profile(options, prof): stream=options.original_stdout, config=options.config, ) - _call_with_diagnostics(options, prof.print_stats, **kwargs) else: py_exe = _python_command() if isinstance(prof, ContextualProfile): @@ -1427,12 +1532,6 @@ def _post_profile(options, prof): f'{quote(py_exe)} -m {show_mod} ' f'{quote(short_outfile)}' ) - # Fully disable the profiler - for _ in range(prof.enable_count): - prof.disable_by_count() - # Restore the state of the global `@line_profiler.profile` - if options.global_profiler: - options.install_profiler(None) if __name__ == '__main__': diff --git a/line_profiler/_child_process_profiling/__init__.py b/line_profiler/_child_process_profiling/__init__.py new file mode 100644 index 00000000..58e0806f --- /dev/null +++ b/line_profiler/_child_process_profiling/__init__.py @@ -0,0 +1,10 @@ +""" +Tooling for profiling child Python processes and gathering their +profiling results. + +Notes: + - THIS IS AN EXPERIMENTAL FEATURE. + + - All contents of this subpackage is to be considered implementation + details. +""" diff --git a/line_profiler/_child_process_profiling/_cache_logging.py b/line_profiler/_child_process_profiling/_cache_logging.py new file mode 100644 index 00000000..94a5c867 --- /dev/null +++ b/line_profiler/_child_process_profiling/_cache_logging.py @@ -0,0 +1,380 @@ +""" +Logging utilities. +""" +from __future__ import annotations + +import os +import re +from collections.abc import Generator +from datetime import datetime +from enum import auto +from itertools import pairwise +from pathlib import Path +from string import Formatter as StringParser +from textwrap import dedent +from typing import TYPE_CHECKING, NamedTuple, TextIO, overload +from typing_extensions import Self + +from .. import _diagnostics as diagnostics +from ..line_profiler_utils import block_indent, StringEnum + + +__all__ = ('CacheLoggingEntry',) + +FILENAME_PATTERN = 'debug_log_{main_pid}_{current_pid}.log' +TIMESTAMP_PATTERN = '[cache-debug-log {timestamp} {level}]' +HEADER_PATTERN = 'PID {current_pid} ({main_pid}): Cache {obj_id:#x}' + +TIMESTAMP_FORMAT = '%Y-%m-%d %H:%M:%S' +TIMESTAMP_MICROSECOND_SEP = ',' +TIMESTAMP_MICROSECOND_PLACES = 3 +TIMESTAMP_SPACING = ' ' + +HEADER_SEP = ': ' +HEADER_MAIN_INDICATOR = 'main process' + + +class LogLevel(StringEnum): + DEBUG = auto() + INFO = auto() + WARNING = auto() + ERROR = auto() + CRITICAL = auto() + + +def get_logger_header(current_pid: int, main_pid: int, obj_id: int) -> str: + """ + Returns: + msg_header (str): + Message header, to be prefixed to messages sent to + :py:data:`line_profiler._diagnostics.log`. + """ + return HEADER_PATTERN.format( + current_pid=current_pid, + main_pid=( + HEADER_MAIN_INDICATOR if main_pid == current_pid else main_pid + ), + obj_id=obj_id, + ) + + +def format_timestamp(ts: datetime) -> str: + """ + Replicate the :py:mod:`logging`'s default formatting for timestamps. + + Example: + >>> ts = datetime(2000, 1, 23, 4, 5, 6, 789000) + >>> as_str = format_timestamp(ts) + >>> print(as_str) + 2000-01-23 04:05:06,789 + >>> assert parse_timestamp(as_str) == ts + """ + return '{}{}{:0{}d}'.format( + ts.strftime(TIMESTAMP_FORMAT), + TIMESTAMP_MICROSECOND_SEP, + int(ts.microsecond / 1000), + TIMESTAMP_MICROSECOND_PLACES, + ) + + +def parse_timestamp(ts: str) -> datetime: + """ + Turn a formatted string timestamp back to a + :py:class:`datetime.datetime` object. + """ + assert TIMESTAMP_MICROSECOND_SEP in ts, ( + f'{ts=!r}, {TIMESTAMP_MICROSECOND_SEP=!r}' + ) + base, _, fractional = ts.rpartition(TIMESTAMP_MICROSECOND_SEP) + # The microsecond field %f must be 6 digits long + if len(fractional) < 6: + fractional = f'{fractional:<06}' + else: + fractional = fractional[:6] + parse_format = f'{TIMESTAMP_FORMAT}{TIMESTAMP_MICROSECOND_SEP}%f' + ts = f'{base}{TIMESTAMP_MICROSECOND_SEP}{fractional}' + return datetime.strptime(ts, parse_format) + + +def add_timestamp( + msg: str, + timestamp: datetime | None = None, + level: str | LogLevel = LogLevel.DEBUG, +) -> str: + """ + Returns: + msg_with_timestamp (str): + (Block-indented) message with timestamp, to be written to + the :py:attr:`LineProfilingCache._debug_log`. + """ + if timestamp is None: + timestamp = datetime.now() + ts_formatted = TIMESTAMP_PATTERN.format( + timestamp=format_timestamp(timestamp), + level=str(level).upper(), + ) + return block_indent(msg, ts_formatted + TIMESTAMP_SPACING) + + +def parse_id(uint: str) -> int: + """ + Example: + >>> n = 123456 + >>> for formatter in str, bin, oct, hex: + ... assert parse_id(formatter(n)) == n + """ + for prefix, base in ('0b', 2), ('0o', 8), ('0x', 16): + if uint.startswith(prefix): + return int(uint[len(prefix):], base=base) + return int(uint) + + +@overload +def fmt_to_regex(fmt: str, /, *auto_numbered_fields: str) -> str: + ... + + +@overload +def fmt_to_regex(fmt: str, /, **named_fields: str) -> str: + ... + + +def fmt_to_regex( + fmt: str, /, *auto_numbered_fields: str, **named_fields: str +) -> str: + """ + Example: + >>> import re + + Simple case: + + >>> pattern = fmt_to_regex( + ... '{func}({args})', func=r'[_\\w][_\\w\\d]+', args='.*', + ... ) + >>> print(pattern) + (?P[_\\w][_\\w\\d]+)\\((?P.*)\\) + >>> regex = re.compile('^' + pattern, re.MULTILINE) + >>> assert not regex.search('0(1)') + >>> match = regex.search(' \\nint(-1.5)') + >>> assert match.group('func', 'args') == ('int', '-1.5') + + Repeated fields: + + >>> palindrome_5l = re.compile(fmt_to_regex( + ... '{first}{second}{third}{second}{first}', + ... first='.', second='.', third='.', + ... )) + >>> print(palindrome_5l.pattern) + (?P.)(?P.)(?P.)(?P=second)(?P=first) + >>> assert not palindrome_5l.match('abbbe') + >>> match = palindrome_5l.match('aBcBa') + >>> assert match.group('first', 'second', 'third') == ( + ... 'a', 'B', 'c', + ... ) + + Auto-numbered fields: + + >>> print(fmt_to_regex( + ... '[{} {}-{}-{} {}:{}:{},{} {}]', + ... # Logger name + ... '.+', + ... # Date + ... r'\\d\\d', r'\\d\\d', r'\\d\\d', + ... # Time + milliseconds + ... r'\\d\\d', r'\\d\\d', r'\\d\\d', r'\\d\\d\\d', + ... # Category + ... 'DEBUG|INFO|WARNING|ERROR|CRITICAL', + ... )) + \\[(.+)\\ (\\d\\d)\\-(\\d\\d)\\-(\\d\\d)\\ \ +(\\d\\d):(\\d\\d):(\\d\\d),(\\d\\d\\d)\\ \ +(DEBUG|INFO|WARNING|ERROR|CRITICAL)\\] + """ + chunks: list[str] = [] + seen_fields: set[str] = set() + for i, (prefix, field, *_) in enumerate(StringParser().parse(fmt)): + chunks.append(re.escape(prefix)) + if field is None: + break # Suffix -> we're done + if field: # Named fields + assert field.isidentifier(), f'{field=!r}' + if field in seen_fields: + chunks.append(f'(?P={field})') + else: + chunks.append(f'(?P<{field}>{named_fields[field]})') + seen_fields.add(field) + else: # Auto-numbered fields + chunks.append(f'({auto_numbered_fields[i]})') + return ''.join(chunks) + + +class CacheLoggingEntry(NamedTuple): + """ + Logging entry written to a log file by + :py:meth:`LineProfilingCache._debug_output`. + + Example: + >>> from datetime import datetime + >>> + >>> + >>> entry = CacheLoggingEntry( + ... datetime(1900, 1, 1, 0, 0, 0, 0), + ... LogLevel.DEBUG, + ... 12345, + ... 12345, + ... 12345678, + ... 'This is a log message;\\nit has multiple lines', + ... ) + >>> print(entry.to_text()) + [cache-debug-log 1900-01-01 00:00:00,000 DEBUG] PID 12345 \ +(main process): Cache 0xbc614e: This is a log message; + it has \ +multiple lines + >>> another_entry = CacheLoggingEntry( + ... datetime(2000, 12, 31, 12, 34, 56, 789000), + ... LogLevel.INFO, + ... 12345, + ... 54321, + ... 87654321, + ... 'FOO BAR BAZ', + ... ) + >>> print(another_entry.to_text()) + [cache-debug-log 2000-12-31 12:34:56,789 INFO] PID 54321 \ +(12345): Cache 0x5397fb1: FOO BAR BAZ + >>> log_text = '\\n'.join([ + ... e.to_text() for e in [entry, another_entry] + ... ]) + >>> assert CacheLoggingEntry.from_text(log_text) == [ + ... entry, another_entry, + ... ] + """ + timestamp: datetime + level: LogLevel + main_pid: int + current_pid: int + cache_id: int + msg: str + + def to_text(self) -> str: + return add_timestamp( + self._get_header() + self.msg, self.timestamp, self.level, + ) + + def _get_header(self) -> str: + return get_logger_header( + self.current_pid, self.main_pid, self.cache_id, + ) + HEADER_SEP + + def write(self, tee: os.PathLike[str] | str | None = None) -> None: + """ + Write the log message using + :py:mod:`line_profiler._diagnostics.log`. If ``tee`` is a path, + also tee thereto with an appropriate timestamp. + """ + log_msg = self._get_header() + self.msg + log_func = getattr(diagnostics.log, self.level.lower()) + log_func(log_msg) + if tee is None: + return + with Path(tee).open(mode='a') as fobj: + full_msg = add_timestamp(log_msg, self.timestamp, self.level) + print(full_msg, file=fobj) + + @classmethod + def new( + cls, main_pid: int, cache_id: int, msg: str, + level: str | LogLevel = LogLevel.DEBUG, + ) -> Self: + return cls( + datetime.now(), LogLevel(level), main_pid, + os.getpid(), cache_id, msg, + ) + + @classmethod + def from_file(cls, file: os.PathLike[str] | str | TextIO) -> list[Self]: + try: + path = Path(file) # type: ignore + except TypeError: # File object + # If we're here, `file` is a file object + if TYPE_CHECKING: + assert isinstance(file, TextIO), f'{file=!r}' + content = file.read() + else: + content = path.read_text() + return cls.from_text(content) + + @classmethod + def from_text(cls, text: str) -> list[Self]: + def gen_timestamps(text: str) -> Generator[re.Match, None, None]: + last_ts_match: re.Match | None = None + while True: + ts_match = timestamp_regex.search( + text, last_ts_match.end() if last_ts_match else 0, + ) + if ts_match: + yield ts_match + last_ts_match = ts_match + else: + return + + def gen_message_blocks(text: str) -> Generator[ + tuple[datetime, LogLevel, re.Match, str], None, None + ]: + timestamps = list(gen_timestamps(text)) + if not timestamps: + return + + # Handle all the entries up till the 2nd-to-last one + for this_match, next_match in pairwise(timestamps): + ts = parse_timestamp(this_match.group('timestamp')) + level = LogLevel(this_match.group('level')) + text_block = text[this_match.start():next_match.start()] + yield (ts, level, this_match, text_block.rstrip('\n')) + # Handle the last entry + last_match = timestamps[-1] + yield ( + parse_timestamp(last_match.group('timestamp')), + LogLevel(last_match.group('level')), + last_match, + text[last_match.start():].rstrip('\n'), + ) + + def get_entries(text: str) -> Generator[Self, None, None]: + for ( + timestamp, level, ts_match, text_block, + ) in gen_message_blocks(text): + # Strip the block indent + ts_text = ts_match.group(0) + assert text_block.startswith(ts_text), ( + f'{text_block=!r}, {ts_text=!r}' + ) + ts_width = len(ts_text) + text_block = dedent(' ' * ts_width + text_block[ts_width:]) + # Strip the header and parse the relevant info from it + header_match = header_regex.match(text_block) + assert header_match, f'{header_regex=!r}, {text_block=!r}' + current_pid = int(header_match.group('current_pid')) + main_pid_ = header_match.group('main_pid') + if main_pid_ == HEADER_MAIN_INDICATOR: + main_pid = current_pid + else: + main_pid = int(main_pid_) + cache_id = parse_id(header_match.group('obj_id')) + # The rest of the block is the message proper + msg = text_block[header_match.end():] + yield cls( + timestamp, level, main_pid, current_pid, cache_id, msg, + ) + + timestamp_pattern = fmt_to_regex( + f'{TIMESTAMP_PATTERN}{TIMESTAMP_SPACING}', + timestamp='.+?', + level='({})'.format('|'.join(LogLevel.__members__)), + ) + timestamp_regex = re.compile('^' + timestamp_pattern, re.MULTILINE) + header_regex = re.compile(fmt_to_regex( + HEADER_PATTERN + HEADER_SEP, + current_pid=r'\d+', + main_pid=r'\d+|' + re.escape(HEADER_MAIN_INDICATOR), + obj_id='.+?', + )) + return list(get_entries(text)) diff --git a/line_profiler/_child_process_profiling/_retrieve_pids.py b/line_profiler/_child_process_profiling/_retrieve_pids.py new file mode 100644 index 00000000..d4c82274 --- /dev/null +++ b/line_profiler/_child_process_profiling/_retrieve_pids.py @@ -0,0 +1,90 @@ +""" +Retrieve active PIDs; uses platform-specific solutions. +""" +import csv +import subprocess +import sys +from collections.abc import Callable, Collection +from shutil import which +from textwrap import indent +from typing import NoReturn + + +__all__ = ( + 'CAN_RETRIEVE_PIDS', 'processes_are_alive', 'get_active_processes', +) + +CAN_RETRIEVE_PIDS: bool +get_active_processes: Callable[[], set[int]] + + +def _find_executable(name: str) -> str: + maybe_path = which(name) + if maybe_path: + return maybe_path + raise RuntimeError(f'cannot find required executable for `{name}`') + + +def processes_are_alive(pids: Collection[int]) -> dict[int, bool]: + """ + Returns: + procs_alive (dict[int, bool]): + Dictionary mapping each PID to whether it is alive + """ + all_pids = get_active_processes() + return {p: p in all_pids for p in set(pids)} + + +def get_active_processes_posix() -> set[int]: + """ + Get the active PIDs using ``ps`` in POSIX(-like) environments. + """ + # `a`: processes from all users + # `u`: excluded, because we don't want usernames clogging up the 1st + # column + # `x`: include detached processes + cmd = [_find_executable('ps'), 'ax', '-o', 'pid'] + ps = subprocess.run(cmd, capture_output=True, text=True, check=True) + _, *procs = ps.stdout.splitlines() # Drop the header line + return {int(p.strip()) for p in procs} + + +def get_active_processes_windows() -> set[int]: + """ + Get the active PIDs using ``tasklist`` on Windows. + """ + # CSV formatted output + cmd = [_find_executable('tasklist'), '/fo', 'csv'] + tasklist = subprocess.run(cmd, capture_output=True, text=True, check=True) + rows = tasklist.stdout.splitlines() + # XXX: this is *probably* locale-safe given that even JP has PID as + # a header... but who knows. + # Might have to test in an RTL language... + try: + return {int(row['pid']) for row in csv.DictReader(rows)} + except Exception as e: # nocover + n = 10 + head = '\n'.join(rows[:n]) + raise RuntimeError( # Context for debugging + 'Error reading PIDs from `tasklist` output ' + f'(first {n} line(s)):\n{indent(head, " ")}' + ) from e + + +def get_active_processes_not_supported() -> NoReturn: # nocover + """ + Unsupported platform; calling this results in a + :py:class:`NotImplementedError`. + """ + raise NotImplementedError(f'unsupported platform: `{sys.platform}`') + + +if sys.platform.startswith('win32'): + CAN_RETRIEVE_PIDS = True + get_active_processes = get_active_processes_windows +elif sys.platform.startswith('wasi'): # nocover + CAN_RETRIEVE_PIDS = False + get_active_processes = get_active_processes_not_supported +else: + CAN_RETRIEVE_PIDS = True + get_active_processes = get_active_processes_posix diff --git a/line_profiler/_child_process_profiling/cache.py b/line_profiler/_child_process_profiling/cache.py new file mode 100644 index 00000000..e217090f --- /dev/null +++ b/line_profiler/_child_process_profiling/cache.py @@ -0,0 +1,1167 @@ +""" +A cache object to be used by for propagating profiling down to child +processes. +""" +from __future__ import annotations + +import atexit +import dataclasses +import os +import site +import sys +import sysconfig +import warnings +try: + import _pickle as pickle +except ImportError: + import pickle # type: ignore[assignment,no-redef] +from collections.abc import ( + Collection, Callable, Generator, Iterable, Mapping, MutableMapping, +) +from functools import partial, cached_property, wraps +from importlib import import_module +from pathlib import Path +from pickle import HIGHEST_PROTOCOL +from textwrap import indent +from types import ModuleType, TracebackType +from typing import Any, ClassVar, Literal, TypeVar, cast, final, overload +from typing_extensions import Concatenate, ParamSpec, Self + +import _line_profiler_hooks as _lp_hooks +from .. import _diagnostics as diagnostics +from ..cleanup import Cleanup, LogLevel, _CALLBACK_REPR_HELPER +from ..curated_profiling import CuratedProfilerContext +from ..line_profiler import LineProfiler, LineStats +from ..toml_config import ConfigSource +from ._cache_logging import CacheLoggingEntry +from ._retrieve_pids import CAN_RETRIEVE_PIDS, processes_are_alive + + +__all__ = ('LineProfilingCache',) + + +T = TypeVar('T') +PS = ParamSpec('PS') + +_THIS_SUBPACKAGE, *_ = (lambda: None).__module__.rpartition('.') +INHERITED_CACHE_ENV_VARNAME_PREFIX = ( + 'LINE_PROFILER_PROFILE_CHILD_PROCESSES_CACHE_DIR' +) +CACHE_FILENAME = 'line_profiler_cache.pkl' +_DEBUG_LOG_FILENAME_PATTERN = 'debug_log_{main_pid}_{current_pid}.log' +_PROFILING_OUTPUT_PREFIX_PATTERN = ( + 'child-prof-output-{main_pid}-{current_pid}-{prof}-' +) +_POSSIBLE_EMPTY_STATS_PREFIX_PATTERN = ( + 'ignore-empty-stats-file-{main_pid}-{current_pid}-' +) + + +def _import_sibling(submodule: str) -> ModuleType: + return import_module(f'{_THIS_SUBPACKAGE}.{submodule}') + + +_private_field = partial(dataclasses.field, init=False, repr=False) + + +class _StatsHelper(Cleanup): + def __init__( + self, + prof: LineProfiler, + outfile: os.PathLike[str] | str, + baseline: LineStats | None = None, + ) -> None: + super().__init__() + self._prof = prof + self.outfile = outfile + self._baseline = baseline + self.add_cleanup(self.dump) + + def __repr__(self) -> str: + name = type(self).__name__ + func = partial(self._prof.dump_stats, self.outfile) + repr_callback = _CALLBACK_REPR_HELPER.repr(func) + return f'<{name} @ {hex(id(self))}: {repr_callback}>' + + def get(self) -> LineStats: + """ + If a ``baseline`` is given (e.g. the stats state inherited + across an :py:func:`os.fork`), subtract it first so that only + work done in *this* process is written; the process the baseline + was inherited from writes the baseline's contents itself. + """ + stats = self._prof.get_stats() + if self._baseline is not None: + stats -= self._baseline + return stats + + def dump(self) -> None: + self.get().to_file(self.outfile) + + __call__ = dump + + def cleanup(self, *args, force: bool = False, **kwargs) -> None: + if force and not any(self._current_context.values()): + self.add_cleanup(self.dump) + super().cleanup(*args, **kwargs) + + +@final +@dataclasses.dataclass +class LineProfilingCache(Cleanup): + """ + Helper object for coordinating a line-profiling session, caching the + info required to make profiling persist into child processes. + """ + cache_dir: os.PathLike[str] | str + config: os.PathLike[str] | str | None = None + profiling_targets: Collection[str] = dataclasses.field( + default_factory=list, + ) + rewrite_module: os.PathLike[str] | str | None = None + profile_imports: bool = False + preimports_module: os.PathLike[str] | str | None = None + main_pid: int = dataclasses.field(default_factory=os.getpid) + # Note: if we're using the line profiler, `kernprof` always sets + # `builtin` to true + insert_builtin: bool = True + debug: bool = diagnostics.DEBUG + + profiler: LineProfiler | None = _private_field(default=None) + _stats_helper: _StatsHelper | None = _private_field(default=None) + # These are unstructured fields; other components can decide on what + # to put in them. They are also pickled by `.dump()`, and are thus + # retrievable in `.load()`-ed instances. + _additional_data: dict[str, Any] = _private_field(default_factory=dict) + + _loaded_instance: ClassVar[LineProfilingCache | None] = None + + def __post_init__(self) -> None: + super().__init__() + + def copy(self, /, **replacements) -> Self: + """ + Make a copy with optionally replaced fields. + + Args: + **replacements (Any): + Optional fields to replace + + Return: + inst (LineProfilingCache): + New instance + """ + init_args: dict[str, Any] = {} + for field, value in self._get_init_args().items(): + init_args[field] = replacements.get(field, value) + return type(self)(**init_args) + + @classmethod + def load(cls) -> Self: + """ + Reconstruct the instance from the environment variables + :env:`LINE_PROFILER_PROFILE_CHILD_PROCESSES_CACHE_PID` and + :env:`LINE_PROFILER_PROFILE_CHILD_PROCESSES_CACHE_DIR_`. + These should have been set from an ancestral Python process. + + Note: + If a previously :py:meth:`.~.load`-ed instance exists, it is + returned instead of a new instance. + """ + # `ty` needs some help here, even if we've marked the class to + # be `@final` + instance = cast(Self | None, cls._loaded_instance) + if instance is None: + pid = os.environ[_lp_hooks.INHERITED_PID_ENV_VARNAME] + cache_varname = f'{INHERITED_CACHE_ENV_VARNAME_PREFIX}_{pid}' + cache_dir = os.environ[cache_varname] + msg = ( + f'PID {os.getpid()} (from {pid}): ' + f'Loading instance from ${{{cache_varname}}} = {cache_dir}' + ) + diagnostics.log.debug(msg) + instance = cls._from_path(cls._get_filename(cache_dir)) + instance._replace_loaded_instance(force=True) + return instance + + def dump(self) -> None: + """ + Serialize the cache instance and dump into the default location + as indicated by :py:attr:`~.cache_dir`, so that they can be + :py:meth:`~.load`-ed by child processes. + + Note: + Cleanup callbacks are not serialized. + """ + content = { + 'init_args': self._get_init_args(), + 'additional_data': self._additional_data, + } + msg = f'Dumping instance data to {self.filename}: {content!r}' + self._debug_output(msg) + with open(self.filename, mode='wb') as fobj: + pickle.dump(content, fobj, protocol=HIGHEST_PROTOCOL) + + def gather_stats( + self, + exclude_pids: Collection[int] | None = None, + *, + on_empty: Literal['error', 'warn', 'ignore'] = 'warn', + on_defective: Literal['error', 'warn', 'ignore'] = 'warn', + ) -> LineStats: + """ + Gather the profiling output files matching ``glob_pattern`` from + :py:attr:`~.cache_dir`, consolidating them into a single + :py:class:`LineStats` object. + + Args: + exclude_pids (Collection[int] | None): + Exclude output from child processes with these PIDs; + the default value :py:const:`None` fetches relevant + PIDs dynamically. + on_empty, on_defective (Literal['error', 'warn', 'ignore']): + Passed to :py:meth:`LineStats.from_files`. + + Returns: + :py:class:`LineStats` instance + """ + def is_empty(path: Path) -> bool: + return not path.stat().st_size + + filter_excludes: Callable[[Iterable[Path]], Iterable[Path]] + if exclude_pids is None: + # NOTE: there is no guarantee that the PID hasn't previously + # been used for another child process that we DID properly + # profile and SHOULD include, so we only filter out empty + # files + exclude_pids = self._get_pids_possibly_lacking_stats() + filter_excludes = partial(filter, is_empty) + else: # User-provided values, who are we to object? + filter_excludes = iter + + fnames_ = set(self._get_profiling_outfiles()) + for pid in exclude_pids: + excludes = filter_excludes(self._get_profiling_outfiles(pid)) + fnames_.difference_update(excludes) + fnames = sorted(fnames_) + self._debug_output( + 'Loading results from {} child profiling file(s): {!r}' + .format(len(fnames), fnames) + ) + if not fnames: + return LineStats.get_empty_instance() + caveat = ( + 'note that empty/malformed profiling files may be produced ' + 'when child processes exits uncleanly (e.g. `.terminate()`-ed ' + 'or `.kill()`-ed by the parent)' + ) + return LineStats.from_files( + *fnames, + on_empty=on_empty, on_defective=on_defective, + _note_on_empty=caveat, _note_on_defective=caveat, + ) + + def _dump_debug_logs(self) -> None: + """ + Gather the debug logfiles in child processes and write their + contents to the logger + (:py:data:`line_profiler._diagnostics.log`). + + Notes: + - The content of each child-process log file is not + re-parsed and is written to the logger as a single + multi-line message. + + - To be called in the main process. + """ + for log in sorted(self._get_debug_logfiles()): + if log == self._debug_log: # Don't double dip + continue + *_, child_pid = log.stem.rpartition('_') + msg = 'Cache log messages from child process {}:\n{}'.format( + child_pid, indent(log.read_text(), ' '), + ) + diagnostics.log.debug(msg) + + def _gather_debug_log_entries( + self, chronological: bool = False, + ) -> list[CacheLoggingEntry]: + """ + Gather and return all entries from debug logfiles sorted by + timestamps. + """ + log_files: Iterable[Path] = self._get_debug_logfiles() + if chronological: # Sorting on the entries -> chronological + to_list: Callable[ + [Iterable[CacheLoggingEntry]], list[CacheLoggingEntry] + ] = sorted + else: + # Otherwise, just sort by filename (entries in each file are + # still chronological) + log_files = sorted(log_files) + to_list = list + return to_list( + entry for log in log_files + for entry in CacheLoggingEntry.from_file(log) + ) + + def _glob(self, *args, **kwargs) -> Iterable[Path]: + return Path(self.cache_dir).glob(*args, **kwargs) + + def _get_debug_logfiles(self) -> Iterable[Path]: + return self._glob(_DEBUG_LOG_FILENAME_PATTERN.format( + main_pid=self.main_pid, current_pid='?*', + )) + + def _get_profiling_outfiles(self, pid: Any = '?*') -> Iterable[Path]: + prefix = _PROFILING_OUTPUT_PREFIX_PATTERN.format( + main_pid=self.main_pid, + current_pid=pid, + # We always format the profiler ID with `hex()`, see + # `._setup_in_child_process()` + prof='0x?*', + ) + return self._glob(prefix + '?*.lprof') + + def inject_env_vars( + self, env: MutableMapping[str, str] | None = None, + ) -> None: + """ + Inject the :py:attr:`~.environ` variables into ``env`` and add + cleanup callbacks to reverse them. + + Args: + env (MutableMapping[str, str] | None): + Dictionary in the format of :py:data:`os.environ`; + default is to use that + """ + self.update_mapping( + os.environ if env is None else env, + self.environ, + _format_debug_msg='Injecting env var ${{{1}}}: {2}'.format, + ) + + def write_pth_hook( + self, *, + prefix: str | None = None, + suffix: str | None = None, + dir: os.PathLike[str] | str | None = None, + # Get rid of the .pth file ASAP so as to be the least disruptive + priority: float = 1, + clean_stale: bool = True, + **kwargs + ) -> Path: + """ + Write a .pth file which allows for setting up profiling in child + Python processes. + + Args: + prefix, suffix (str | None): + Optional filename-stem affixes of the .pth file; default + is to use default values loaded from :py:attr:`.config`. + dir (os.PathLike[str] | str | None): + Optional directory to create the .pth file in; default + is to look up a list of directories where .pth files can + usually be installed to, depending on the environment + and interpreter state. + clean_stale (bool): + Whether to look at the directory where the .pth file + has been written to and prune stale .pth files. + priority, **kwargs: + Passed to :py:meth:`.make_tempfile`. + + Returns: + fpath (Path): + Path to the written .pth file + + Notes: + - Due to normalizations and attachment of extra data, users + should NOT count on the affixes bracketing the filename + stem of the created file. + + - During normal execution, the .pth file should be cleaned + up, as with all other tempfiles created with + :py:meth:`.make_tempfile` created without + ``delete=False``. However, if the Python control flow is + broken (e.g. process killed), cleanup can fail to occur. + Hence the option to ``clean_stale``. + """ + def get_pth_config() -> Mapping[str, Any]: + # Note: the only keys in it should be `prefix` and `suffix` + return ( + self._config_source # Cached + .get_subconfig('child_processes', 'pth_files') + .conf_dict + ) + + if not os.path.exists(self.filename): + self.dump() + assert os.path.exists(self.filename) + + # The string casts are failsafes in case inappropriate values + # (e.g. numbers and booleans) are supplied + if prefix is None: + prefix = str(get_pth_config()['prefix']) + if suffix is None: + suffix = str(get_pth_config()['suffix']) + prefix, suffix_template = self._normalize_pth_affixes(prefix, suffix) + suffix = suffix_template.format(self.main_pid) + + if dir is None: + dirs: list[Path] + dirs = list(self._enumerate_pth_installation_locations()) + else: + dirs = [Path(dir)] + + failures: dict[Path, str] = {} + tempfile_msg_template = 'Created tempfile {0.name!r} at {0.parent}' + for dir in dirs: + try: + fpath = self.make_tempfile( + prefix=prefix, suffix=suffix, + dir=dir, priority=priority, + _format_debug_msg=tempfile_msg_template.format, + **kwargs + ) + except OSError as e: + failures[dir] = self._format_exception(e) + else: + break + else: + raise RuntimeError( + 'cannot create .pth file in any of the following directories ' + f'({{path: error}}): {failures!r}' + ) + + content = self._get_graceful_oneline_call( + _lp_hooks.__name__, + _lp_hooks.load_pth_hook.__name__, + self.main_pid, + ) + try: + fpath.write_text(content) + except Exception: + fpath.unlink(missing_ok=True) + raise + + if clean_stale: + self._cleanup_stale_pth_files( + prefix, suffix_template, fpath.parent, + ) + return fpath + + def _cleanup_stale_pth_files( + self, prefix: str, suffix_template: str, dir: Path, + ) -> None: + def log(msg: str, level: LogLevel = 'debug') -> None: + msg = f'._cleanup_stale_pth_files(): {msg}' + self._debug_output(msg, level) + + def warn(msg: str) -> None: + log(msg, 'warning') + # 3: code calling `._cleanup_stale_pth_files()` + warnings.warn(msg, stacklevel=3) + + if not CAN_RETRIEVE_PIDS: # nocover + warn( + 'lookup for stale .pth files not currently possible ' + f'on this platform (`{sys.platform}`); cleanup aborted', + ) + return + + try: + pth_files = self._find_pth_files(prefix, suffix_template, dir) + is_alive = processes_are_alive(pth_files) + except Exception as e: # nocover + try: + frame = cast(TracebackType, e.__traceback__).tb_frame + context = f'{frame.f_code.co_filename}:{frame.f_lineno}' + except Exception: # E.g. no traceback + context = '' + xc = self._format_exception(e) + if context: + xc = f'{xc} ({context})' + warn( + f'lookup for stale .pth files failed ({xc}); ' + 'cleanup aborted', + ) + return + + nfiles = sum(len(p) for p in pth_files.values()) + nstale = 0 + npruned = 0 + log(f'found {nfiles} matching .pth file(s)') + for ppid, pths in pth_files.items(): + if is_alive[ppid]: + log( + f'skipping over {len(pths)} .pth file(s) ' + f'associated with main PID {ppid} (do not seem stale)' + ) + continue + for pth in pths: + nstale += 1 + try: + pth.unlink() + except Exception as e: + success = False + status = f'failed ({self._format_exception(e)})' + else: + success, status = True, 'succeeded' + npruned += 1 + report: Callable[[str], Any] = log if success else warn + report(f'cleanup for stale .pth file {pth.name!r} {status}') + log(f'summary: {nfiles} found, {nstale} stale, {npruned} pruned') + + @staticmethod + def _find_pth_files( + prefix: str, suffix_template: str, dir: Path, + ) -> dict[int, set[Path]]: + result: dict[int, set[Path]] = {} + assert suffix_template.endswith('ppid-{}.pth') + glob_pattern = f'{prefix}?*{suffix_template.format("?*")}' + for pth in dir.glob(glob_pattern): + *_, suffix = pth.name.rpartition('ppid-') + assert suffix.endswith('.pth') + ppid = int(suffix[:-len('.pth')]) + result.setdefault(ppid, set()).add(pth) + return result + + @staticmethod + def _normalize_pth_affixes(prefix: str, suffix: str) -> tuple[str, str]: + prefix = prefix.rstrip('-') + '-' + # Escape braces because we'll use the suffix as a formatting + # template + suffix = suffix.strip('-').replace('{', '{{').replace('}', '}}') + suffix_chunks: list[str] = ['ppid', '{}'] + if suffix: + suffix_chunks.insert(0, suffix) + suffix_template = ''.join('-' + chunk for chunk in suffix_chunks) + suffix_template += '.pth' + return prefix, suffix_template + + @staticmethod + def _get_graceful_oneline_call( + module: str, func: str, /, *args, **kwargs, + ) -> str: + r""" + Get the content of a one-line .pth file which imports a + function and calls it with the supplied arguments, and fails + gracefully (e.g. due to failed imports) as far as possible. + + Example: + >>> get_pth = LineProfilingCache._get_graceful_oneline_call + >>> run = lambda stmts: exec(stmts, {}) # Isolate side-fxs + + >>> good_stmts = get_pth( + ... 'builtins', 'print', [1, 2], 'b', None, sep='\n', + ... ) + >>> run(good_stmts) + [1, 2] + b + None + + Check that non-frozen modules are correctly handled: + + >>> pprint_stmts = get_pth( + ... 'pprint', 'pprint', [1, 2, 3], width=5, + ... ) + >>> run(pprint_stmts) + [1, + 2, + 3] + + Note that it satisfies the requirements for .pth files: + + >>> assert good_stmts.startswith('import') + >>> assert len(good_stmts.splitlines()) == 1 + + If the import/module-spec lookup fails or the if the + attribute lookup on the module fails, the code is a no-op: + + >>> from contextlib import ExitStack, redirect_stdout + >>> from io import StringIO + + >>> with ExitStack() as stack: + ... fobj = stack.enter_context(StringIO()) + ... stack.enter_context(redirect_stdout(fobj)) + ... bad_module_stmts = get_pth( + ... 'buuiltins', 'print', 'foo', + ... ) + ... run(bad_module_stmts) # Import fails -> no-op + ... bad_func_stmts = get_pth( + ... 'builtins', 'priint', 'foo', + ... ) + ... run(bad_module_stmts) # Attr lookup fails -> no-op + ... assert not (stdout := fobj.getvalue()), stdout + + (Note: when running with :py:mod:`xdoctest`, leaving the + expected output empty does not by default check AGAINST + output; hence the context contraption.) + """ + assert module and not module.isspace() + assert all(chunk.isidentifier() for chunk in module.split()) + assert func.isidentifier() + + call_args = [repr(a) for a in args] + call_args.extend(f'{k}={v!r}' for k, v in kwargs.items()) + call_args = ['mod', repr(func)] + call_args + + statements = [ + 'import importlib.util as iu', + 'dummy = lambda *_, **__: None', + 'call = lambda obj, attr, /, *a, **k: ' + '(func if callable(func := getattr(obj, attr, None)) else dummy)' + '(*a, **k)', + # Import the module + f'spec = iu.find_spec({module!r})', + 'mod = iu.module_from_spec(spec) if spec else None', + "call(spec.loader, 'exec_module', mod) if mod else None", + # Retrieve and call the function + f'call({", ".join(call_args)}) if mod else None', + ] + return '; '.join(statements) + + @staticmethod + def _enumerate_pth_installation_locations() -> Generator[Path, None, None]: + """ + Enumerate locations where the .pth file can potentially be + installed to; the lookup order is: + + - Directory where :py:mod:`_line_profiler_hooks` is installed to + (if among the output of :py:func:`site.getsitepackages`) + + - ``sysconfig.get_path('purelib')`` + + - The output of :py:func:`site.getusersitepackages` + """ + def filter_dirs( + maybe_dirs: Iterable[os.PathLike[str] | str], + ) -> Generator[Path, None, None]: + for path in maybe_dirs: + if os.path.isdir(path): + yield Path(path) + + try: + path = Path(_lp_hooks.__file__).parent + except Exception: + pass + else: + if any( + path.samefile(p) for p in filter_dirs(site.getsitepackages()) + ): + yield path + yield Path(sysconfig.get_path('purelib')) + if site.ENABLE_USER_SITE: + yield Path(site.getusersitepackages()) + + def _debug_output(self, msg: str, /, level: LogLevel = 'debug') -> None: + """ + Beside writing to the logger, also write to the + :py:attr:`~._debug_log`. + """ + entry = CacheLoggingEntry.new(self.main_pid, id(self), msg, level) + try: + entry.write(self._debug_log) + except OSError: # Cache dir may have been rm-ed during cleanup + pass + + def _setup_in_main_process(self, wrap_os_fork: bool = True) -> None: + """ + Set up shop in the main process so that (line-)profiling can + extend into child processes. + + Args: + wrap_os_fork (bool): + Whether to wrap :py:func:`os.fork` which handles + profiling + + Side effects: + + - Instance data written to :py:attr:`~.cache_dir` + + - Environment variables injected + (see :py:meth:`~.inject_env_vars()`) + + - A ``.pth`` file written so that child processes + automatically runs setup code (see + :py:meth:`.write_pth_hook`) + + - :py:func:`os.fork` wrapped so that profiling set up in + forked processes is properly handled (if + ``wrap_os_fork=True``) + + - :py:mod:`multiprocessing` patched so that child processes + managed thereby are properly handled + + - Instance to be returned if :py:func:`~.load()` is called + from now on + """ + self.dump() + self.inject_env_vars() + try: + self.write_pth_hook() + except Exception as e: + xc = self._format_exception(e) + msg = ( + 'cannot write a .pth file for setting up profiling ' + f'in child processes ({xc}); profiling data cannot be ' + 'collected for non-`fork()`ed children' + ) + self._debug_output(msg, 'warning') + warnings.warn(msg, stacklevel=2) + self._setup_common(wrap_os_fork, {'reboot_forkserver': True}) + self._replace_loaded_instance() + + def _setup_in_child_process( + self, + wrap_os_fork: bool = False, + context: str = '', + prof: LineProfiler | None = None, + baseline: LineStats | None = None, + ) -> bool: + """ + Set up shop in a forked/spawned child process so that + (line-)profiling can extend therein. + + Args: + wrap_os_fork (bool): + Whether to wrap :py:func:`os.fork` which handles + profiling; already-forked child processes should set + this to false + context (str): + Optional context from which the function is called, to + be used in log messages + prof (LineProfiler | None): + Optional profiler instance to associate with the cache; + if not provided, an instance is created + baseline (LineStats | None): + Stats which ``prof`` already held when this process + came into existence (only meaningful for forked + processes, which inherit the parent's profiler state); + they are subtracted from every stats dump so that the + pre-existing data isn't double-counted when the parent + gathers and merges the *separately-dumped* child and + parent stats + + Returns: + has_set_up (bool): + False the instance has already been set up prior to + calling this function, true otherwise + """ + def wrap_ctx_debug( + ctx: CuratedProfilerContext, msg: str, /, + level: LogLevel = 'debug', + ) -> None: + self._debug_output(f' Context {id(ctx):#x}: {msg}', level) + + if not context: + context = '...' + self._debug_output(f'Setting up ({context})...') + if self.profiler is not None: # Already set up + self._debug_output(f'Setup aborted ({context})') + return False + + # Create a profiler instance and manage it with + # `CuratedProfilerContext` + if prof is None: + prof = LineProfiler() + self.profiler = prof + ctx = CuratedProfilerContext(prof, insert_builtin=self.insert_builtin) + if self.debug: + self.patch(ctx, '_debug_output', wrap_ctx_debug.__get__(ctx)) + ctx.install() + self.add_cleanup(ctx.uninstall) + self._debug_output(f'Set up `.profiler` at {id(prof):#x}') + + # Do the preimports at `cache.preimports_module` where + # appropriate + if self.preimports_module: + self._debug_output('Loading preimports...') + with open(self.preimports_module, mode='rb') as fobj: + code = compile(fobj.read(), self.preimports_module, 'exec') + exec(code, {}) # Use a fresh, empty namespace + + # - Occupy a tempfile slot in `.cache_dir` + # - Set the profiler up to write thereto when the process + # terminates (with high priority) + # (Also keep a separate reference to the callback for e.g. + # dumping stats ASAP at process exit) + prof_outfile = self.make_tempfile( + prefix=_PROFILING_OUTPUT_PREFIX_PATTERN.format( + main_pid=self.main_pid, + current_pid=os.getpid(), + prof=hex(id(prof)), + ), + suffix='.lprof', + delete=False, + ) + self._stats_helper = sh = _StatsHelper(prof, prof_outfile, baseline) + self.patch( + # If we call `sh.cleanup()` instead of `sh` (e.g. in some + # `multiprocessing` patches), the subsequent debug-log msgs + # are attributed to and handled by this cache instance + sh, '_debug_output', self._debug_output, + cleanup=False, name='._stats_helper', + ) + self.add_cleanup_with_priority(sh, 1) + + # Various setups + self._setup_common(wrap_os_fork, {'reboot_forkserver': False}) + + # Set `.cleanup()` as an atexit hook to handle everything when + # the child process is about to terminate + atexit.register(self._atexit_hook) + + self._debug_output(f'Setup successful ({context})') + return True + + def _setup_common( + self, + wrap_os_fork: bool, + mp_apply_kwargs: dict[str, Any] | None = None, + ) -> None: + if wrap_os_fork: + self._wrap_os_fork() + _import_sibling('multiprocessing_patches').apply( + self, **(mp_apply_kwargs or {}), + ) + + def _wrap_os_fork(self) -> None: + """ + Create a wrapper around :py:func:`os.fork` which handles + profiling. + + Side effects: + + - :py:func:`os.fork` (if available) replaced with the + wrapper + + - :py:meth:`~.cleanup` callback registered undoing that + """ + try: + fork = os.fork + except AttributeError: # Can't fork on this platform + return + + @wraps(fork) + def wrapper() -> int: + ppid = os.getpid() + result = fork() + if result: + return result + # If we're here, we are in the fork + pid = os.getpid() + # Snapshot the profiler state inherited from the parent + # BEFORE any further code runs in the fork; it is used as a + # subtractive baseline for this process's stats dumps so + # that pre-fork data (which the parent dumps itself) isn't + # double-counted at gathering time + if self.profiler is None: + baseline = None + else: + baseline = self.profiler.get_stats() + forked = self.copy() # Ditch inherited cleanups + forked._debug_output(f'Forked: {ppid} -> {pid}') + if forked._replace_loaded_instance(): + forked._debug_output( + 'Superseded cached `.load()`-ed instance in forked process' + ) + # To avoid complications with the inherited cache instance + # `self`: + # - Unregister its `._atexit_hook` + # - Discard its cleanup callbacks + atexit.unregister(self._atexit_hook) + self._contexts.clear() + # Note: we can reuse the profiler instance in the fork, but + # it needs to go through setup so that the separate + # profiling results are dumped into another output file + forked._setup_in_child_process( + False, 'fork', self.profiler, baseline, + ) + return result + + self.patch(os, 'fork', wrapper, name='os') + + def _warn_possible_lack_of_stats( + self, pids: int | Collection[int], + ) -> None: + """ + Register PID(s) which may have created a profiling stats file + without writing to it; when calling :py:meth:`.gather_stats`, + empty stats files associated with those PIDs are ignored instead + of warned against or treated as an error. + """ + if not isinstance(pids, Collection): + pids = pids, + try: + with self._empty_stats_pid_registry.open(mode='a') as fobj: + print(*pids, sep='\n', file=fobj) + except FileNotFoundError: + # At cleanup time, the tempdir backing the cache instance + # may already have ceased to exist; just let that be + pass + + def _get_pids_possibly_lacking_stats(self) -> set[int]: + """ + See also + :py:meth:`._warn_possible_lack_of_stats` + """ + prefix = _POSSIBLE_EMPTY_STATS_PREFIX_PATTERN.format( + main_pid=self.main_pid, + current_pid='?*', # Gather from all child processes + ) + result: set[int] = set() + for registry in self._glob(prefix + '?*.dat'): + from_reg: set[int] = set() + with registry.open() as fobj: + for line in fobj: + try: + from_reg.add(int(line)) + except ValueError: + pass + if from_reg: + self._debug_output( + f'Loaded {len(from_reg)} PID(s) possibly lacking ' + f'profiling output from {registry.name!r}: {from_reg!r}' + ) + result.update(from_reg) + return result + + def make_tempfile(self, **kwargs) -> Path: + """ + Create a fresh tempfile under :py:attr:`~.cache_dir`. The other + arguments are passed as-is to :py:func:`tempfile.mkstemp`. + + Returns: + path (Path): + Path to the created file. + """ + kwargs.setdefault('dir', self.cache_dir) + kwargs.setdefault( + '_format_debug_msg', 'Created tempfile: {0.name!r}'.format, + ) + return super().make_tempfile(**kwargs) + + def _replace_loaded_instance(self, force: bool = False) -> bool: + cls = type(self) + if force or self._consistent_with_loaded_instance: + self.patch(cls, '_loaded_instance', self) + return True + return False + + @classmethod + def _from_path(cls, fname: os.PathLike[str] | str) -> Self: + with open(fname, mode='rb') as fobj: + content = pickle.load(fobj) + instance = cls(**content['init_args']) + instance._additional_data.update(content.get('additional_data', {})) + return instance + + def _get_init_args(self) -> dict[str, Any]: + init_fields = [ + field_obj.name for field_obj in dataclasses.fields(self) + if field_obj.init + ] + return {name: getattr(self, name) for name in init_fields} + + @staticmethod + def _get_filename(cache_dir: os.PathLike[str] | str) -> str: + return os.path.join(cache_dir, CACHE_FILENAME) + + @overload + @classmethod + def _method_wrapper( + cls, + wrapper: Callable[Concatenate[Self, Callable[PS, T], PS], T], + *, + debug: bool | None = None, + wrapper_name: str | None = None, + ) -> Callable[[Callable[PS, T]], Callable[PS, T]]: + ... + + @overload + @classmethod + def _method_wrapper( + cls, wrapper: None = None, *, + debug: bool | None = None, + wrapper_name: str | None = None, + ) -> Callable[ + [Callable[Concatenate[Self, Callable[PS, T], PS], T]], + Callable[[Callable[PS, T]], Callable[PS, T]] + ]: + ... + + @classmethod + def _method_wrapper( + cls, + wrapper: ( + Callable[Concatenate[Self, Callable[PS, T], PS], T] | None + ) = None, + debug: bool | None = None, + wrapper_name: str | None = None, + ) -> ( + Callable[ + [Callable[Concatenate[Self, Callable[PS, T], PS], T]], + Callable[[Callable[PS, T]], Callable[PS, T]] + ] + | Callable[[Callable[PS, T]], Callable[PS, T]] + ): + """ + Convenience wrapper decorator for functions which use the + :py:meth:`load`-ed session instance and wrap another callable. + + Args: + wrapper (Callable[..., T]) + Callable with the call signature + ``(cache, vanilla_impl, *args, **kwargs) -> retval``; + ``*args``, ``**kwargs``, and ``retval`` should be + consistent with that of ``vanilla_impl()``'s. + debug (bool | None) + Whether to format and write debug messages before and + after the call to the ``wrapper`` callable; + if ``debug`` is not set, it will be taken from the + session instance. + wrapper_name (str | None) + Optional name to identify the source of the wrapper, + used in debug messages. + + Returns: + inner_wrapper (Callable[[Callable[PS, T]], Callable[PS, T]]) + Wrapper(-maker) which takes the ``vanilla_impl`` and + return a wrapper around it. + """ + if wrapper is None: + # `ty` doesn't quite support `partial` yet, see issue #1536 + return cast( + Callable[[Callable[PS, T]], Callable[PS, T]], + partial(cls._method_wrapper, debug=debug), + ) + + def inner_wrapper(vanilla_impl: Callable[PS, T]) -> Callable[PS, T]: + @wraps(vanilla_impl) + def wrapped_impl(*args: PS.args, **kwargs: PS.kwargs) -> T: + cache = cls.load() + write = cache._debug_output + debug_: bool | None = debug + call = partial(wrapper, cache, vanilla_impl, *args, **kwargs) + + if debug_ is None: + debug_ = cache.debug + if debug_: + call_fmt = cache._format_call( + vanilla_name, *args, **kwargs, + ) + write( + f'Wrapped call made via `{wrapper_name}`: ' + f'{call_fmt}...', + ) + try: + result = call() + except BaseException as e: + # Note: be more defensive than normal and + # prepared to deal with `BaseException`; this + # decorator is often used for functions invoked + # in child processes which don't cleanly + # terminate + state, outcome = 'failed', cache._format_exception(e) + raise e + else: + state = 'succeeded' + outcome = _CALLBACK_REPR_HELPER.repr(result) + return result + finally: + write( + f'Wrapped call via `{wrapper_name}` {state}: ' + f'{call_fmt} -> {outcome}', + ) + else: + return call() + + vanilla_name = cls._get_name(vanilla_impl) + return wrapped_impl + + if wrapper_name is None: + wrapper_name = cls._get_name(wrapper) + for field in 'name', 'qualname', 'doc': + dunder = f'__{field}__' + value = getattr(wrapper, dunder, None) + if value is not None: + setattr(inner_wrapper, dunder, value) + return inner_wrapper + + @classmethod + def _format_call( + cls, func: Callable[..., Any] | str, /, *args, **kwargs, + ) -> str: + if isinstance(func, partial): + return cls._format_call( + func.func, [*func.args, *args], {**func.keywords, **kwargs}, + ) + call = _CALLBACK_REPR_HELPER.format_call(*args, **kwargs) + if not isinstance(func, str): + func = cls._get_name(func) + return func + call + + @staticmethod + def _format_exception(xc: BaseException) -> str: + formatted = type(xc).__name__ + if str(xc): + formatted = f'{formatted}: {xc}' + return formatted + + @property + def environ(self) -> dict[str, str]: + """ + Environment variables to be injected into and inherited by child + processes. + """ + cache_varname = f'{INHERITED_CACHE_ENV_VARNAME_PREFIX}_{self.main_pid}' + return { + _lp_hooks.INHERITED_PID_ENV_VARNAME: str(self.main_pid), + cache_varname: str(self.cache_dir), + } + + @property + def filename(self) -> str: + return self._get_filename(self.cache_dir) + + @property + def _debug_log(self) -> Path | None: + if not self.debug: + return None + fname = _DEBUG_LOG_FILENAME_PATTERN.format( + main_pid=self.main_pid, current_pid=os.getpid(), + ) + return Path(self.cache_dir) / fname + + @cached_property + def _consistent_with_loaded_instance(self) -> bool: + cls = type(self) + # Note: calling `.load()` can cause a new instance to be created + # and stored at `._loaded_instance`, if there isn't already one + # such instance; guard against that + already_loaded = cls._loaded_instance + try: + return cls.load()._get_init_args() == self._get_init_args() + finally: + cls._loaded_instance = already_loaded + + @cached_property + def _config_source(self) -> ConfigSource: + if self.config is None: + config: str | None = None + else: + config = str(self.config) + return ConfigSource.from_config(config) + + @cached_property + def _empty_stats_pid_registry(self) -> Path: + prefix = _POSSIBLE_EMPTY_STATS_PREFIX_PATTERN.format( + main_pid=self.main_pid, + current_pid=os.getpid(), + ) + return self.make_tempfile(prefix=prefix, suffix='.dat', delete=False) + + @cached_property + def _atexit_hook(self) -> Callable[[], None]: + return partial(self.cleanup, reason='`atexit` callback') diff --git a/line_profiler/_child_process_profiling/multiprocessing_patches/__init__.py b/line_profiler/_child_process_profiling/multiprocessing_patches/__init__.py new file mode 100644 index 00000000..2fd06fcb --- /dev/null +++ b/line_profiler/_child_process_profiling/multiprocessing_patches/__init__.py @@ -0,0 +1,127 @@ +""" +Patch :py:mod:`multiprocessing` so that profiling extends into processes +it creates. + +Notes: + - Based on the implementations in :py:mod:`coverage.multiproc` and + :py:mod:`pytest_autoprofile._multiprocessing`. + + - Results may vary if the process pool is not properly + :py:meth:`multiprocessing.pool.Pool.close`-d and + :py:meth:`multiprocessing.pool.Pool.join`-ed; + see `this caveat `__. +""" +from __future__ import annotations + +import multiprocessing +from collections.abc import Collection +from typing import Literal, get_args + +from ..cache import LineProfilingCache +from ._infrastructure import Registry +from .mp_config import MPConfig + + +__all__ = ('MPConfig', 'Registry', 'apply') + +PublicPatch = Literal['pool', 'process', 'logging'] + +_PATCHED_MARKER = '__line_profiler_patched_multiprocessing__' +_PATCHES = Registry.get_default() + + +def apply( + cache: LineProfilingCache, + reboot_forkserver: bool = True, + patches: Collection[PublicPatch] | None = None, +) -> None: + """ + Set up profiling in :py:mod:`multiprocessing` child processes by + applying patches to the module. + + Args: + cache (LineProfilingCache): + Cache instance governing the profiling run. + reboot_forkserver (bool): + Whether to reboot the global + :py:class`multiprocessing.forkserver.ForkServer` instance + so as to ensure that profiling happens on processes forked + therefrom (see Note). + patches \ +(Collection[Literal['pool', 'process', 'logging'] | None]): + Patches to apply to :py:mod:`multiprocessing`; see the + following section for a description of each; + the default is taken from the TOML config file. + + Patches: + ``'pool'``: + Patch :py:class:`multiprocessing.pool.Pool` and + :py:func:`multiprocessing.pool.worker` so that profiling + output is recorded as each pool-worker child process + completes a task, and written to disk as the pool is + terminated. + ``'process'``: + Patch + :py:meth:`multiprocessing.process.BaseProcess._bootstrap` + so that non-pool-worker child processes write profiling + output on exit. + ``'logging'``: + Patch :py:mod:`multiprocessing.util`'s logging methods (e.g. + ``debug()`` and ``info()``) so that their messages are teed + to the cache's debug log. + + Side effects: + - The aforementioned patches applied + + - If ``reboot_forkserver=True``, fork-server process rebooted: + + - Immediately + + - When ``cache.cleanup()`` is run + + - Cleanup callbacks registered via ``cache.add_cleanup()`` + + Note: + Rebooting the fork server is necessary because its process + statically inherits the environment when it is first spun up + (see :py:func:`multiprocessing.forkserver.ensure_running`). + Thus, without the reboots: + + - If in the same Python process we ever start up two separate + profiling sessions managed by different caches, the child + processes forked from the server will fail to inherit the + updated environment variables injected by the newer cache + instance, leading to the setup code in this subpackage not + being loaded. + + - Since 3.13.8 and 3.14.1, the bug where the ``main_path`` + argument to :py:func:`multiprocessing.forkserver.main` is + unused has been fixed (see ``cpython`` issue `GH-126631`_). + This causes ``sys.modules['__main__']`` to be set up in the + fork-server process, meaning that children forked therefrom + will NOT redo the setup. Thus, the fork-server process itself + will also need to be properly set up for profiling. + + .. _GH-126631: https://github.com/python/cpython/issues/126631 + """ + if getattr(multiprocessing, _PATCHED_MARKER, False): + return + if patches is None: + patches_dict = MPConfig.from_cache(cache).patches + patches_: set[str] = {p for p, use in patches_dict.items() if use} + else: + patches_ = {p.lower() for p in patches} + # Sanity check on `_PATCHES` + for patch in get_args(PublicPatch): + if patch not in _PATCHES: + raise RuntimeError(f'Cannot load patch `{patch}`') + for name, patch in _PATCHES.select(patches_).items(): + if name == '__reboot_forkserver' and not reboot_forkserver: + continue + msg = f'applying `multiprocessing` patch {name!r}' + cache._debug_output(msg.capitalize() + '...') + patch.apply(cache) + cache._debug_output('Done with ' + msg) + # Mark `multiprocessing` as having been patched + cache.patch(multiprocessing, _PATCHED_MARKER, True) diff --git a/line_profiler/_child_process_profiling/multiprocessing_patches/_infrastructure.py b/line_profiler/_child_process_profiling/multiprocessing_patches/_infrastructure.py new file mode 100644 index 00000000..808ff1b3 --- /dev/null +++ b/line_profiler/_child_process_profiling/multiprocessing_patches/_infrastructure.py @@ -0,0 +1,507 @@ +from __future__ import annotations + +import dataclasses +import warnings +from collections.abc import ( + Callable, Collection, Generator, Mapping, Sequence, Set, +) +from functools import partial +from importlib import import_module +from inspect import getattr_static +from operator import attrgetter +from types import MappingProxyType as mappingproxy, ModuleType +from typing import ( + TYPE_CHECKING, + Any, ClassVar, Literal, Protocol, TypeVar, + cast, final, overload, +) +from typing_extensions import Self + +from ... import _diagnostics as diagnostics +from ..cache import LineProfilingCache + + +__all__ = ('Patch', 'SingleModulePatch', 'Registry') + +P = TypeVar('P', bound='Patch') + + +class Patch(Protocol): + """ + Interface for patches. + """ + def apply( + self, + cache: LineProfilingCache, + *, + cleanup: bool = True, + **kwargs + ) -> Any: + """ + Apply the patch. + + Args: + cache (LineProfilingCache): + Session cache + cleanup (bool): + Whether ``cache.cleanup()`` should reverse the patch + **kwargs + Individual implementations should pick the ones they + need and ignore the rest. + + Note: + The patch is responsible for registering the requisite + cleanup callbacks with ``cache`` so that it is reversed when + ``cache.cleanup()`` is called. + """ + ... + + @property + def summary(self) -> Mapping[str, Set[str]]: + """ + A mapping from dotted-path names of objects to the set of + attributes patched thereon. + """ + ... + + @property + def priority(self) -> float | None: + """ + Real number representing how the patch is to be prioritized. A + patch with a HIGHER priority should be applied LATER, that way + it gets to wrap wrappers created by patches applied earlier. + """ + ... + + +@dataclasses.dataclass +class SingleModulePatch: + """ + Patch to apply to a module component in :py:mod:`multiprocessing`. + + Attributes: + submodule (str): + Name of the :py:mod:`multiprocessing` submodule. + targets (dict[str,\ +dict[str, Callable[[Any], Any] | Sequence[Callable[[Any], Any]]]]): + Dictionary mapping (dot-chained) names in said submodule to + a dictionary of patches; said patches dictionary should have + the format of + ``dict[simple_attribute, wrapper | [wrapper1, ...]]``. See + Example for details. + priority (float | None): + Optional numerical value for how to prioritize the patch + (see :py:class:`.Patch`) + + Example: + Consider + ``SingleModulePatch('foo', {'bar.baz': {'foobar': foofoo},\ +'': {'spam': [ham, eggs]}})``. + This instance would perform the following patches on the module + ``multiprocessing.foo``: + + - Replace ``multiprocessing.foo.bar.baz.foobar`` with + ``foofoo(multiprocessing.foo.bar.baz.foobar)`` + + - Replace ``multiprocessing.foo.spam`` with + ``eggs(ham(multiprocessing.foo.spam))``; + note that the two wrappers are applied in order to the + original attribute. + """ + submodule: str + targets: dict[ + str, dict[str, Callable[[Any], Any] | Sequence[Callable[[Any], Any]]] + ] = dataclasses.field(default_factory=dict) + priority: float | None = None + + package: ClassVar[str] = 'multiprocessing' + + def add_target( + self, + target: str, + patches: Mapping[ + str, Callable[[Any], Any] | Sequence[Callable[[Any], Any]] + ], + ) -> Self: + """ + Convenience method for gradually constructing the patch with a + fluent interface. + + Args: + target (str): + Dotted path to the object in :py:attr:`.submodule` + patches (Mapping[str, Callable[[Any], Any] \ +| Sequence[Callable[[Any], Any]]]): + Mapping from patched attrbute names to the wrappers to + apply thereto; sequences of wrappers are applied in + order + + Returns: + This instance + """ + self.targets.setdefault(target, {}).update(patches) + return self + + def add_method( + self, + target: str, + method: str, + wrapper: Callable[[Any], Any], + methodtype: ( + type[classmethod] | type[staticmethod] + | Literal['class', 'static'] | None + ) = None, + ) -> Self: + """ + Convenience method for gradually constructing the patch with a + fluent interface. + + Args: + target (str): + Dotted path to the object in :py:attr:`.submodule` + method (str): + Name of the (class, static, or instance) method to patch + wrapper (Callable[[Any], Any]): + Wrapping callable which takes the method-implementaion + callable and returns a wrapper thereof + methodtype (type[classmethod] | type[staticmethod] | \ +Literal['class', 'static'] | None): + Optional type of the method if not an instance method; + the strings ``'class'`` and ``'static'`` are respective + shorthands for :py:class:`classmethod` and + :py:class:`staticmethod` + + Returns: + This instance + """ + wrappers: Callable[[Any], Any] | list[Callable[[Any], Any]] + if methodtype is None: + wrappers = wrapper + else: + if methodtype == 'class': + methodtype = classmethod + elif methodtype == 'static': + methodtype = staticmethod + wrappers = [attrgetter('__func__'), wrapper, methodtype] + return self.add_target(target, {method: wrappers}) + + def apply( + self, + cache: LineProfilingCache, + *, + cleanup: bool = True, + static: bool = True, + **_ + ) -> list[str]: + """ + Apply the patch. + + Args: + cache (LineProfilingCache): + Session cache + cleanup (bool): + Whether ``cache.cleanup()`` should reverse the patch + static (bool): + Whether to use :py:func:`inspect.getattr_static` to + retrieve to the attributes to be patched on the patch + targets + + Returns: + replacements (list[str]): + Names of entities replaced + """ + submod_name = f'{self.package}.{self.submodule}' + get_attribute = getattr_static if static else getattr + result: list[str] = [] + try: + mod = self.load_module() + except ImportError: # nocover + return [] + + for target in sorted(self.targets, key=len, reverse=True): + if TYPE_CHECKING: + # See `ty` issue #2572 + assert isinstance(target, str) + if target: + try: + obj: Any = attrgetter(target)(mod) + except AttributeError: # nocover + continue + name = f'{submod_name}.{target}' + else: + obj, name = mod, submod_name + replace = partial(cache.patch, obj, cleanup=cleanup, name=name) + for method, method_wrappers in self.targets[target].items(): + if callable(method_wrappers): + method_wrappers = cast( + Sequence[Callable[[Any], Any]], (method_wrappers,), + ) + try: + impl = get_attribute(obj, method) + except AttributeError: + continue + for wrapper in method_wrappers: + impl = wrapper(impl) + replace(method, impl) + result.append(f'{name}.{method}') + return result + + def load_module(self) -> ModuleType: + """ + Returns: + Module object :py:attr:`.module` points to + """ + return import_module(self.module) + + @staticmethod + def _join(s: str, *strs: str, sep: str = '.') -> str: + return sep.join(string for string in (s, *strs) if string) + + @property + def module(self) -> str: + """ + Module where the patches are applied + """ + return self._join(self.package, self.submodule) + + @property + def summary(self) -> mappingproxy[str, frozenset[str]]: + """ + Summary of the dotted paths to the patched objects and their + patched attributes + """ + add_prefix = partial(self._join, self.module) + return mappingproxy({ + add_prefix(target): frozenset(patches) + for target, patches in self.targets.items() + }) + + +@final +class Registry(Mapping[str, Patch]): + """ + Mapping subclass for managing patches. + """ + _default: ClassVar[Registry] + + def __init__(self) -> None: + self._patches: dict[str, tuple[float, Patch]] = {} + + def __repr__(self) -> str: + if self._patches: + patches = ', '.join( + f'{name!r} (priority {priority})' if priority else repr(name) + for name, (priority, _) in self._patches.items() + ) + patches = f'({len(self)} patch(es)): {patches}' + else: + patches = '(0 patches)' + return f'<{type(self).__name__} @ {id(self):#x} {patches}>' + + def __getitem__(self, key: str) -> Patch: + return self._patches[key][1] + + def __iter__(self) -> Generator[str, None, None]: + for name, _ in self._iter_patches(): + yield name + + def __len__(self) -> int: + return len(self._patches) + + def __contains__(self, key: Any) -> bool: + return key in self._patches + + @overload + def register( + self, name: str, patch: P, *, priority: float | None = None, + ) -> P: + ... + + @overload + def register( + self, name: str, patch: None = None, *, priority: float | None = None, + ) -> Patch: + ... + + def register( + self, name: str, patch: Patch | None = None, *, + priority: float | None = None, + ) -> Patch: + """ + Register/look up a patch. + + Args: + name (str): + Name of the patch; patches named with leading double + underscores are considered MANDATORY, and are applied no + matter the user input (e.g. via + ``apply(..., patches=...)`` or the config file). + patch (Patch | None): + Patch object to register; if not provided, look for the + existing patch registered under the name. + priority (float | None): + Optional priority to assign to the patch; default is to + look at :py:attr:`Patch.priority` and to resolve + :py:const:`None` to 0. + + Returns: + Patch object (``patch`` if provided, looked up otherwise) + """ + if patch is not None: + if priority is None: + priority = patch.priority + if not priority: + priority = 0 + old_pri, stored = self._patches.setdefault(name, (priority, patch)) + if stored is not patch: + raise ValueError( + f'name = {name!r}, patch = {patch!r}: ' + f'name already in use by {stored!r}' + ) + if old_pri != priority: # Update priority + self._patches[name] = priority, patch + elif priority is not None: + # Reassign priority of an existing patch + self._patches[name] = priority, self[name] + return self[name] + + def select(self, patches: Collection[str]) -> Registry: + """ + Returns: + New instance with the selected patches + + Note: + Patches whose names are prefixed with double underscores are + considered mandatory and are always selected. + """ + new = Registry() + new._patches.update( + (name, patch_info) for name, patch_info in self._patches.items() + if name.startswith('__') or name in patches + ) + return new + + def _iter_patches( + self, reverse: bool = False, + ) -> Generator[tuple[str, Patch], None, None]: + """ + Iterate over the available patches. + + Note: + Since patches typically consists of function wrappers, and + outer wrappers are both called first and responsible for + calling inner wrappers, patches with a HIGHER priority are + yielded LATER (or earlier if ``reverse=True``. + """ + for _, patches in sorted(self._prioritized.items(), reverse=reverse): + yield from patches.items() + + @property + def summary(self) -> dict[str, frozenset[str]]: + """ + Mapping from the names of the entities affected by the patches + to sets of attributes patched thereon. + """ + summaries = [patch.summary for patch in self.values()] + return { + target: frozenset().union(*(s.get(target, ()) for s in summaries)) + for target in frozenset().union(*summaries) + } + + @property + def _prioritized(self) -> dict[float, dict[str, Patch]]: + """ + Note: + This could've been a static attribute maintained by the + :py:meth:`.register` method, but we don't call this a bunch + so it isn't like we incur a lot of overhead by calculating + it on-the-fly; and it's less error-prone this way. + """ + result: dict[float, dict[str, Patch]] = {} + for name, (priority, patch) in self._patches.items(): + result.setdefault(priority, {})[name] = patch + return result + + @classmethod + def get_default(cls) -> Registry: + """ + Returns: + Instance summarizing the patches loaded from the various + ``.*_patches`` sibling modules + + Note: + This method does NOT create a copy. + + Example: + >>> reg = Registry.get_default() + >>> assert reg.get_default() is reg + + Check for the default plugins that should be installed and + their contents: + + >>> assert 'pool' in reg + >>> assert 'process' in reg + >>> assert 'logging' in reg + + >>> assert ( + ... 'multiprocessing.process.BaseProcess' in reg.summary + ... ) + >>> assert ( + ... 'worker' + ... in reg.summary.get('multiprocessing.pool', set()) + ... ) + """ + def check(patch: P) -> P: + error: str | None = None + if not hasattr(patch, 'priority'): + error = 'expected a `.priority: float | None` field' + elif not isinstance(getattr(patch, 'summary', None), Mapping): + error = 'expected a `.summary: Mapping[str, Set[str]]` field' + elif not callable(getattr(patch, 'apply', None)): + error = ( + 'expected an `.apply(cache: LineProfilingCache, ...)` ' + 'method' + ) + if error: + raise TypeError(f'patch `{patch!r}`: {error}') + return patch + + try: + return cls._default + except AttributeError: + pass + + instance = Registry() + subpkg, *_ = cls.__module__.rpartition('.') + for name, (sibling, patch_loc) in { + '__process_setup': ('_mandatory_patches', 'PROCESS_SETUP_PATCH'), + '__pool_worker_pid': + ('_mandatory_patches', 'POOL_WORKER_PID_PATCH'), + '__reboot_forkserver': + ('_mandatory_patches', 'RebootForkserverPatch'), + '__resource_tracker': + ('_mandatory_patches', 'ResourceTrackerPatch'), + '__spawn_runpy': ('_mandatory_patches', 'RunpyPatch'), + + 'logging': ('_optional_patches', 'LOGGING_PATCH'), + + 'pool': ('_profiling_patches', 'POOL_PATCH'), + 'process': ('_profiling_patches', 'PROCESS_PATCH'), + }.items(): + try: + mod = import_module(f'{subpkg}.{sibling}') + patch = check(cast(Patch, getattr(mod, patch_loc))) + except Exception as e: + error = type(e).__name__ + if str(error): + error = f'{error}: {e}' + msg = ( + f'failed to load patch {name!r} ' + f'from sibling submodule `{subpkg}.{sibling}`: {error}' + ) + diagnostics.log.warning(msg) + warnings.warn(msg) + else: + instance.register(name, patch) + cls._default = instance + return instance diff --git a/line_profiler/_child_process_profiling/multiprocessing_patches/_mandatory_patches.py b/line_profiler/_child_process_profiling/multiprocessing_patches/_mandatory_patches.py new file mode 100644 index 00000000..165738f4 --- /dev/null +++ b/line_profiler/_child_process_profiling/multiprocessing_patches/_mandatory_patches.py @@ -0,0 +1,311 @@ +from __future__ import annotations + +import atexit +import os +import multiprocessing +from collections.abc import Callable +from functools import partial +from multiprocessing.process import BaseProcess +from types import MappingProxyType as mappingproxy, MethodType +from typing import Any, ClassVar, TypeVar, cast +from typing_extensions import Concatenate, ParamSpec + +try: + from multiprocessing import spawn +except ImportError: + _CAN_USE_SPAWN = False +else: + _CAN_USE_SPAWN = True +try: + from multiprocessing import forkserver +except ImportError: + _CAN_USE_FORKSERVER = False +else: + _CAN_USE_FORKSERVER = ( + 'forkserver' in multiprocessing.get_all_start_methods() + ) +try: + from multiprocessing import resource_tracker +except ImportError: + _CAN_USE_RESOURCE_TRACKER = False +else: + _CAN_USE_RESOURCE_TRACKER = True + +from ..cache import LineProfilingCache +from ..runpy_patches import create_runpy_wrapper +from ._infrastructure import SingleModulePatch +from ._pool_patch_helpers import ( + get_per_task_callback_patch, get_worker_finalization_patch, +) + + +__all__ = ( + 'POOL_WORKER_PID_PATCH', 'PROCESS_SETUP_PATCH', + 'RebootForkserverPatch', 'ResourceTrackerPatch', 'RunpyPatch', + 'wrap_bootstrap', +) + +T = TypeVar('T') +P = TypeVar('P', bound=BaseProcess) +PS = ParamSpec('PS') + +# ------------------------------ Helpers ------------------------------- + + +def setup_mp_child( # nocover + cache: LineProfilingCache, proc: BaseProcess, +) -> None: + """ + Perform :py:mod:`multiprocessing`-specific setup in a child process + curated by the module. Currently it does the following: + + - Unregister the :py:mod:`atexit` hook associated with ``cache`` to + avoid possible clashes with the profiling-file writing managed by + this module. + """ + if cache.main_pid == os.getpid(): # Not in a child process + return + xc: Exception | None = None + msg = 'Performing setup for `multiprocessing` child processes...' + cache._debug_output(msg) + setup: Callable[[LineProfilingCache, BaseProcess], Any] + for setup in [_unregister_atexit_hook]: + try: + setup(cache, proc) + except Exception as e: + xc = e + if xc is None: + msg = 'Setup for `multiprocessing` child process succeeded' + cache._debug_output(msg) + else: + xc_str = type(xc).__name__ + if str(xc): + xc_str = f'{xc_str}: {xc}' + cache._debug_output(f'Setup failed: {xc_str}') + raise xc + + +def _unregister_atexit_hook( # nocover + cache: LineProfilingCache, _, +) -> None: + atexit.unregister(cache._atexit_hook) + + +# ----------- `multiprocessing.process.BaseProcess` patches ------------ + + +@LineProfilingCache._method_wrapper # nocover +def wrap_bootstrap( + cache: LineProfilingCache, + vanilla_impl: Callable[Concatenate[BaseProcess, PS], T], + self: BaseProcess, + /, + *args: PS.args, **kwargs: PS.kwargs +) -> T: + """ + Wrap around :py:meth:`.BaseProcess._bootstrap` to perform setups + specific to :py:mod:`multiprocessing`-managed processes. + """ + setup_mp_child(cache, self) + return vanilla_impl(self, *args, **kwargs) + + +PROCESS_SETUP_PATCH = SingleModulePatch('process', priority=1) +PROCESS_SETUP_PATCH.add_method('BaseProcess', '_bootstrap', wrap_bootstrap) + +# ---------------------- PID bookkeeping patches ----------------------- + + +def _get_worker_ntasks(cache: LineProfilingCache, worker: BaseProcess) -> int: + """ + Check if the process has run any tasks; if not, report to the cache. + + Returns: + Number of tasks run by ``worker`` + """ + pid: int | None = getattr(worker, 'pid', None) + ntasks_finalized = _get_ntasks_finalized(cache) + if pid is None: # Dummy process + return 0 + key = id(worker), pid + try: + return ntasks_finalized[key] + except KeyError: + pass + ntasks = _get_ntasks(cache).pop(pid, 0) + msg = 'Worker {0.name!r} (PID: {0.pid}) ran {1} task(s)' + cache._debug_output(msg.format(worker, ntasks)) + if not ntasks: + cache._warn_possible_lack_of_stats(pid) + return ntasks_finalized.setdefault(key, ntasks) + + +def _increment_ntasks(cache: LineProfilingCache, pid: int) -> None: + """ + Take and process the PID of the child process completing the task. + """ + ntasks = _get_ntasks(cache) + ntasks[pid] = ntasks.get(pid, 0) + 1 + + +def _get_ntasks(cache: LineProfilingCache) -> dict[int, int]: + key = 'mp_proc_ntasks' + return cache._additional_data.setdefault(key, cast(dict[int, int], {})) + + +def _get_ntasks_finalized( + cache: LineProfilingCache, +) -> dict[tuple[int, int], int]: + key = 'mp_proc_ntasks_finalized' + return cache._additional_data.setdefault( + key, cast(dict[tuple[int, int], int], {}) + ) + + +def _get_pid(_) -> int: + return os.getpid() + + +POOL_WORKER_PID_PATCH = get_per_task_callback_patch( + _get_pid, _increment_ntasks, '__line_profiler_pool_worker_pid__', +) +get_worker_finalization_patch(_get_worker_ntasks, POOL_WORKER_PID_PATCH) + +# --------------------------- Misc. patches ---------------------------- + + +class RebootForkserverPatch: + """ + Reboot the process backing the global + :py:class:`multiprocessing.forkserver.ForkServer` instance: + + - When the patch is applied, so as to ensure that child processes + forked therefrom actually receives the active patches; and + + - When the session cache is cleaned up, so that child processes + forked therefrom is no longer polluted by the patches. + + Note: + This uses + :py:meth:`multiprocessing.forkserver.ForkServer._stop()` which + is private API, but it's the same hack used in Python's own test + suite -- see the comment to said method. + """ + summary: ClassVar[mappingproxy[str, frozenset[str]]] = mappingproxy({}) + priority: ClassVar[float | None] = None + + @classmethod + def apply(cls, cache: LineProfilingCache, **_) -> None: + if not _CAN_USE_FORKSERVER: + return + cls.reboot() + cache.add_cleanup(cls.reboot) + + @staticmethod + def reboot() -> None: + fs_obj = forkserver._forkserver + stop = getattr(fs_obj, '_stop', None) + # Appease the type-checker since `._stop()` is not public API + if not callable(stop): # nocover + msg = f'ForkServer._stop() (= {stop!r}) not callable' + raise AssertionError(msg) # Shouldn't happen + stop() + + +class ResourceTrackerPatch: + """ + Patch :py:mod:`multiprocessing.resource_tracker` so that + :py:func:`multiprocessing.resource_tracker.ensure_running` and the + eponymous method of + :py:class:`multiprocessing.resource_tracker.ResourceTracker` report + the resource-tracker server PIDs to the session cache. + + Note: + The ``ResourceTracker`` server process is spawned when the first + :py:mod:`multiprocessing` child process is created via the + ``spawn`` or ``forkserver`` start methods. While this server + process does not meaningfully contribute to the profiling result + either way, since it can be created with profiling set up, its + longevity means that :py:meth:`.LineProfilingCache.gather_stats` + often catches empty .lprof files which it has occupied but not + written to. + + To reduce noise while keeping the empty-file warning for other + output files, we report the PIDs used by the server to the + session cache so that they can be ignored if necessary. + """ + if _CAN_USE_RESOURCE_TRACKER: + summary: ClassVar[mappingproxy[str, frozenset[str]]] = mappingproxy({ + 'multiprocessing.resource_tracker': + frozenset({'ensure_running'}), + 'multiprocessing.resource_tracker.ResourceTracker': + frozenset({'ensure_running'}), + }) + else: + summary = mappingproxy({}) + priority: ClassVar[float | None] = None + + @staticmethod + @LineProfilingCache._method_wrapper + def wrap_ensure_running( + cache: LineProfilingCache, + vanilla_impl: Callable[['resource_tracker.ResourceTracker'], None], + self: 'resource_tracker.ResourceTracker', + ) -> None: + """ + Wrap around :py:meth:`multiprocessing.resource_tracker\ +.ResourceTracker.ensure_running` + so that the session cache can keep track of the PIDs used by the + resource-tracer server. + """ + maybe_pids: set[int | None] = {getattr(self, '_pid', None)} + try: + vanilla_impl(self) + finally: + maybe_pids.add(getattr(self, '_pid', None)) + pids = cast(set[int], maybe_pids - {None}) + if pids: + cache._warn_possible_lack_of_stats(pids) + + @classmethod + def apply( + cls, cache: LineProfilingCache, *, cleanup: bool = True, **_, + ) -> list[str]: + if _CAN_USE_RESOURCE_TRACKER: + patch = partial(cache.patch, cleanup=cleanup) + # Patch the method on the class + method = resource_tracker.ResourceTracker.ensure_running + method = cls.wrap_ensure_running(method) + patch(resource_tracker.ResourceTracker, 'ensure_running', method) + # Patch the preexisting bound method on the module + instance = resource_tracker._resource_tracker + bound_method = MethodType(method, instance) + patch(resource_tracker, 'ensure_running', bound_method) + return list(cls.summary) + + +class RunpyPatch: + """ + Patch the copy of :py:mod:`runpy` in the + :py:mod:`multiprocessing.spawn` namespace so that subprocesses can + perform rewrite-based profiling as with + :py:func:`line_profiler.autoprofile.autoprofile.run`. + + See also: + :py:mod:`line_profiler._child_process_profiling.runpy_patches` + """ + summary: ClassVar[mappingproxy[str, frozenset[str]]] + if _CAN_USE_SPAWN and hasattr(spawn, 'runpy'): + summary = mappingproxy({'multiprocessing.spawn': frozenset({'runpy'})}) + else: + summary = mappingproxy({}) + priority: ClassVar[float | None] = None + + @classmethod + def apply( + cls, cache: LineProfilingCache, *, cleanup: bool = True, **_, + ) -> list[str]: + if cls.summary: + patch = partial(cache.patch, cleanup=cleanup) + patch(spawn, 'runpy', create_runpy_wrapper(cache)) + return list(cls.summary) diff --git a/line_profiler/_child_process_profiling/multiprocessing_patches/_optional_patches.py b/line_profiler/_child_process_profiling/multiprocessing_patches/_optional_patches.py new file mode 100644 index 00000000..0fe0f328 --- /dev/null +++ b/line_profiler/_child_process_profiling/multiprocessing_patches/_optional_patches.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +from collections.abc import Callable +from functools import partial +from typing import TypeVar +from typing_extensions import Concatenate, ParamSpec + +from ..cache import LineProfilingCache +from ._infrastructure import SingleModulePatch + + +__all__ = ('LOGGING_PATCH', 'tee_log') + +T = TypeVar('T') +PS = ParamSpec('PS') + +_LOGGERS = ['sub_debug', 'debug', 'info', 'sub_warning', 'warn'] + +# --------------- `multiprocessing.util` logging patches --------------- + + +def _cache_hook( + vanilla_impl: Callable[PS, T], + get_logging_message: Callable[PS, str], + /, + *args: PS.args, + **kwargs: PS.kwargs +) -> T: + msg = get_logging_message(*args, **kwargs) + LineProfilingCache.load()._debug_output(msg) + return vanilla_impl(*args, **kwargs) + + +def tee_log( + marker: str, + vanilla_impl: Callable[Concatenate[str, PS], None], + /, + msg: str, + *args: PS.args, + **kwargs: PS.kwargs +) -> None: + """ + Wrap around logging functions like + :py:func:`multiprocessing.util.debug` so that we can tee log + messages from the package to our own logs. + """ + def get_msg(msg: str, *_, **__) -> str: + return f'`multiprocessing` logging ({marker}): {msg}' + + _cache_hook( + vanilla_impl, get_msg, # type: ignore[arg-type] + msg, *args, **kwargs, + ) + + +LOGGING_PATCH = SingleModulePatch('util').add_target( + # The logging functions exists directly in the module namespace so + # no further attribute access is needed + '', {func: partial(partial, tee_log, func) for func in _LOGGERS}, +) diff --git a/line_profiler/_child_process_profiling/multiprocessing_patches/_pool_patch_helpers.py b/line_profiler/_child_process_profiling/multiprocessing_patches/_pool_patch_helpers.py new file mode 100644 index 00000000..f719aab7 --- /dev/null +++ b/line_profiler/_child_process_profiling/multiprocessing_patches/_pool_patch_helpers.py @@ -0,0 +1,197 @@ +from __future__ import annotations + +import warnings +from collections.abc import Callable +from functools import partial +from multiprocessing.pool import Pool +from multiprocessing.process import BaseProcess +from typing import Any, TypeVar +from typing_extensions import Concatenate, ParamSpec + +from ..cache import LineProfilingCache +from ._infrastructure import SingleModulePatch +from ._queue import Queue, PutWrapper, QuickGetWrapper + + +__all__ = ('get_per_task_callback_patch', 'get_worker_finalization_patch') + +T = TypeVar('T') +P = TypeVar('P', bound=BaseProcess) +PS = ParamSpec('PS') + +_POOL_PATCHES: set[str] = set() + + +def get_per_task_callback_patch( + get_data: Callable[[LineProfilingCache], T], + process_data: Callable[[LineProfilingCache, T], Any], + tag: str, + patch: SingleModulePatch | None = None, +) -> SingleModulePatch: + """ + Create a patch for :py:mod:`multiprocessing.pool` which: + + - Patches :py:func:`multiprocessing.pool.worker` so that extra data + are created in child/worker processes after running EACH task by + ``get_data(cache)``, and pushed back to parent process alongside + said task's result. + + - Patches :py:meth:`multiprocessing.pool.Pool._handle_results` so + that said extra data is, where possible, retrieved from + interprocess communication and processed by the parent's active + :py:class:`.LineProfilingCache` instance. + """ + if tag in _POOL_PATCHES: + raise RuntimeError(f'tag {tag!r} already in use') + _POOL_PATCHES.add(tag) + + wrap_quick_get = partial(QuickGetWrapper, callback=process_data, tag=tag) + + @LineProfilingCache._method_wrapper + def wrap_handle_results( + cache: LineProfilingCache, + vanilla_impl: Callable[ + Concatenate[Queue, Callable[[], tuple[Any, ...] | None], PS], + None + ], + outqueue: Queue, + # Since we patched `outqueue.put()` in the child process, the + # result pushed to the parent is (normally) a `(tag, data, obj)` + # triplet + get: Callable[[], tuple[Any, ...] | None], + *args: PS.args, + **kwargs: PS.kwargs + ) -> None: + """ + Wrap around :py:meth:`multiprocessing.pool.Pool._handle_results` + so that it handles the extra info (result of calling + ``get_data()`` in a child process after each task) included by + ``wrap_worker()`` with ``process_data(cache, data)``. + + Note: + :py:meth:`multiprocessing.pool.Pool._handle_results` is a + static method. + """ + vanilla_impl(outqueue, wrap_quick_get(cache, get), *args, **kwargs) + + @LineProfilingCache._method_wrapper # nocover + def wrap_worker( + cache: LineProfilingCache, + vanilla_impl: Callable[Concatenate[Queue, Queue, PS], None], + inqueue: Queue, + outqueue: Queue, + *args: PS.args, + **kwargs: PS.kwargs + ) -> None: + """ + Wrap around :py:func:`multiprocessing.pool.worker` so that child + processes attach the result of ``get_data()`` as they pass the + task results back to the parent. + + Note: + This is only called in child processes and thus we can't + reliably measure coverage thereon, hence the ``# nocover``. + """ + outqueue = PutWrapper(outqueue, partial(get_data, cache), tag) + return vanilla_impl(inqueue, outqueue, *args, **kwargs) + + if patch is None: + patch = SingleModulePatch('pool') + _check_patch(patch) + patch.add_method('', 'worker', wrap_worker) + patch.add_method('Pool', '_handle_results', wrap_handle_results, 'static') + return patch + + +def get_worker_finalization_patch( + process_worker: Callable[[LineProfilingCache, BaseProcess], Any], + patch: SingleModulePatch | None = None, +) -> SingleModulePatch: + """ + Create a patch for :py:mod:`multiprocessing.pool` which patches + :py:meth:`multiprocessing.pool.Pool._terminate_pool` and + :py:meth:`multiprocessing.pool.Pool._join_exited_workers` so that + when worker processes are finalized a callback is run. + """ + @LineProfilingCache._method_wrapper + def wrap_terminate_pool( + cache: LineProfilingCache, + vanilla_impl: Callable[ + Concatenate[type[Pool], Queue, Queue, Queue, list[P], PS], None + ], + cls: type[Pool], + taskqueue: Queue, + inqueue: Queue, + outqueue: Queue, + pool: list[P], + *args: PS.args, + **kwargs: PS.kwargs + ) -> None: + """ + Wrap around :py:meth:`.Pool._terminate_pool` so that we run + ``process_worker()`` on finished worker processes. + + Note: + :py:meth:`.Pool._terminate_pool` is a class method. + """ + try: + vanilla_impl( + cls, taskqueue, inqueue, outqueue, pool, *args, **kwargs, + ) + finally: + # Guard against dummy pool; see similar code in + # `multiprocessing.pool` + if pool and hasattr(pool[0], 'terminate'): + failures: list[P] = [] + for worker in pool: + if worker.is_alive(): # nocover + failures.append(worker) + continue + process_worker(cache, worker) + if failures: # nocover + msg = ( + f'{len(failures)} worker(s) still alive after ' + f'`Pool.terminate()`: {failures!r}' + ) + cache._debug_output(msg, 'warning') + warnings.warn(msg) + + @LineProfilingCache._method_wrapper + def wrap_join_exited_workers( + cache: LineProfilingCache, + vanilla_impl: Callable[[list[P]], bool], + pool: list[P], + ) -> bool: + """ + Wrap around :py:meth:`.Pool._join_exited_workers` so that we run + ``process_worker()`` on finished worker processes. + + Note: + :py:meth:`.Pool._join_exited_workers` is a static method. + """ + before: dict[int, P] = {id(p): p for p in pool} + if not vanilla_impl(pool): # No workers cleaned up + return False + after: dict[int, P] = {id(p): p for p in pool} + for i, worker in before.items(): + if after.get(i) is not worker: # Worker removed from `pool` + process_worker(cache, worker) + return True + + if patch is None: + patch = SingleModulePatch('pool') + _check_patch(patch) + add = partial(patch.add_method, 'Pool') + add('_terminate_pool', wrap_terminate_pool, 'class') + add('_join_exited_workers', wrap_join_exited_workers, 'static') + return patch + + +def _check_patch(patch: SingleModulePatch) -> None: + if patch.submodule == 'pool': + return + msg = ( + 'patch = {0!r}, .submodule = {0.submodule!r}: ' + 'expected patch target to be `multiprocessing.pool`' + ) + raise AssertionError(msg.format(patch)) diff --git a/line_profiler/_child_process_profiling/multiprocessing_patches/_profiling_patches.py b/line_profiler/_child_process_profiling/multiprocessing_patches/_profiling_patches.py new file mode 100644 index 00000000..6cd8f2fc --- /dev/null +++ b/line_profiler/_child_process_profiling/multiprocessing_patches/_profiling_patches.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +from collections.abc import Callable +from multiprocessing.process import BaseProcess +from pathlib import Path +from typing import TypeVar, cast +from typing_extensions import Concatenate, ParamSpec + +from ...line_profiler import LineStats +from ..cache import LineProfilingCache +from ._infrastructure import SingleModulePatch +from ._pool_patch_helpers import ( + get_per_task_callback_patch, get_worker_finalization_patch, +) + + +__all__ = ( + 'POOL_PATCH', 'PROCESS_PATCH', + 'wrap_bootstrap', 'wrap_process', +) + +T = TypeVar('T') +P = TypeVar('P', bound=BaseProcess) +PS = ParamSpec('PS') + +_POOL_WORKER_MARKER = '__line_profiler_multiprocessing_is_pool_worker__' + +# ------------------------------ Helpers ------------------------------- + + +def dump_stats_quick( + cache: LineProfilingCache, *, reason: str | None = None, +) -> None: + """ + Note: + We don't really care about cleanup in the child process, so just + dump the stats and bail to reduce the chance of end-of-process + shenanigans causing a deadlock... + but do use ``._stats_helper.cleanup()`` instead of + ``.__call__()`` so that we get debugging output (if ``debug`` is + true) + """ + stats_helper = cache._stats_helper + if stats_helper is None: + return + if cache.debug: + stats_helper.cleanup(force=True, reason=reason) + else: + stats_helper() + + +def _mark_worker(worker: P) -> P: + setattr(worker, _POOL_WORKER_MARKER, True) + return worker + + +def _is_marked_worker(proc: BaseProcess) -> bool: + return getattr(proc, _POOL_WORKER_MARKER, False) + + +# ---------------- `multiprocessing.pool.Pool` patches ----------------- + + +@LineProfilingCache._method_wrapper +def wrap_process( + _, vanilla_impl: Callable[PS, P], *args: PS.args, **kwargs: PS.kwargs +) -> P: + """ + Wrap around :py:meth:`.Pool.Process` so that the worker processes + created by the pool are marked and can be distinguished from + processes otherwise managed. + + Notes: + + - :py:meth:`.Pool.Process` is a static method. + + - Technically one can inspect the :py:attr:`.BaseProcess.name` + of the process to see that it is a ``PoolWorker``, but since + said attribute is writable it may be more robust to set up a + separate marker. + """ + return _mark_worker(vanilla_impl(*args, **kwargs)) + + +def _report_stats_and_dest( + cache: LineProfilingCache, +) -> tuple[LineStats, Path] | None: # nocover + """ + Notes: + + - This is only called in child processes and thus we can't + reliably measure coverage thereon; see also + :py:func:`wrap_bootstrap`. + + - In an ideal world, we would have just written profiling output + once as :py:func:`multiprocessing.pool.worker` returns. But: + + - Worker sometimes end up in "dirty" states and deadlock, and + thus has to be terminated. + + - However, terminating a Python process bypasses the + interpreter control flow, meaning that :py:mod:`atexit` + hooks and ``try``-``finally`` blocks aren't executed. + + - On POSIX, this can be mitigated by setting signal handlers, + but signal handling is infamously unreliable on + :py:mod:`multiprocessing` child processes (examples: + `1`_, `2`_, `3`_), causing hangs that are hard to remedy. + + So this is about as good as we can do. + + - Instead of dumping the stats to disk every task, it should be + less overhead for us to just send them back to the parent via + the preexisting connection. + + - Unless the code paths varied significantly between tasks, the + physical size of the stats should not have changed too much – + running the same code more only increment the timing entries, + but do not generate more thereof. Thus, calculating the + delta-stats between tasks and only sending those would have + been a waste of time here; on top of that, the parent would + also have to perform additional processing to accumulate the + deltas. So we just send the entire stats object. + + .. _1: https://github.com/python/cpython/issues/73945 + .. _2: https://github.com/python/cpython/issues/82408 + .. _3: https://github.com/coveragepy/coveragepy/issues/1310 + """ + stats_helper = cache._stats_helper + if stats_helper is None: + return None + return stats_helper.get(), Path(stats_helper.outfile) + + +def _record_stats( + cache: LineProfilingCache, + stats_and_dest: tuple[LineStats, Path] | None, +) -> None: + """ + Record the stats gathered from workers so that they can be dealt + with when the process pool is terminated. + + See also: + :py:func:`_report_stats_and_dest` + """ + if stats_and_dest is None: + return # No-op + stats, dest = stats_and_dest + _get_worker_stats(cache)[dest] = stats + + +def _write_recorded_stats( + cache: LineProfilingCache, worker: BaseProcess, +) -> None: + """ + Write the gathered stats associated with ``worker``. + """ + pid = getattr(worker, 'pid', None) + if pid is None: + return + worker_stats = _get_worker_stats(cache) + xc: Exception | None = None + msg = '{0} (centralized `.dump_stats()`): {1.name!r} (PID: {1.pid}) -> {2}' + for outfile in cache._get_profiling_outfiles(pid): + stats = worker_stats.pop(outfile, None) + if stats is None: + continue + try: + stats.to_file(outfile) + except Exception as e: + xc = e + state, outcome = 'Failed', cache._format_exception(e) + else: + state, outcome = 'Succeeded', repr(outfile.name) + cache._debug_output(msg.format(state, worker, outcome)) + if xc is not None: + raise xc + + +def _get_worker_stats(cache: LineProfilingCache) -> dict[Path, LineStats]: + key = 'mp_pool_worker_stats' + return cache._additional_data.setdefault( + key, cast(dict[Path, LineStats], {}), + ) + + +POOL_PATCH = get_per_task_callback_patch( + _report_stats_and_dest, _record_stats, + '__line_profiler_pool_worker_stats__', +) +get_worker_finalization_patch(_write_recorded_stats, POOL_PATCH) +POOL_PATCH.add_method('Pool', 'Process', wrap_process, 'static') + +# ----------- `multiprocessing.process.BaseProcess` patches ------------ + + +@LineProfilingCache._method_wrapper # nocover +def wrap_bootstrap( + cache: LineProfilingCache, + vanilla_impl: Callable[Concatenate[BaseProcess, PS], T], + self: BaseProcess, + /, + *args: PS.args, **kwargs: PS.kwargs +) -> T: + """ + Wrap around :py:meth:`.BaseProcess._bootstrap` so that profiling + stats are written at the end. + + Notes: + + - This is only invoked in child processes, and + :py:mod:`coverage` seems to be having trouble with them in the + current setup, probably due to issues with .pth file + precendence causing :py:mod:`line_profiler` to be loaded + before it. Hence the ``# nocover``. + + - Since process termination bypasses the Python interpreter (see + notes in :py:func:`wrap_worker`), if a child process is + terminated prematurely (e.g. via + :py:meth:`.BaseProcess.terminate`), profiling data may be + missing. + + - To prevent data corruption/loss, the end-of-function write to + the temporary profiling-stat file only happens for + non-pool-managed :py:class:`BaseProcess` objects, because they + are regularly :py:meth:`.BaseProcess.terminate`-ed by their + managing pool. + """ + try: + return vanilla_impl(self, *args, **kwargs) + finally: + reason = 'exiting `multiprocessing.process.BaseProcess._bootstrap`' + if not _is_marked_worker(self): + dump_stats_quick(cache, reason=reason) + + +PROCESS_PATCH = SingleModulePatch('process') +PROCESS_PATCH.add_method('BaseProcess', '_bootstrap', wrap_bootstrap) diff --git a/line_profiler/_child_process_profiling/multiprocessing_patches/_queue.py b/line_profiler/_child_process_profiling/multiprocessing_patches/_queue.py new file mode 100644 index 00000000..20238ef4 --- /dev/null +++ b/line_profiler/_child_process_profiling/multiprocessing_patches/_queue.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +import warnings +from collections.abc import Callable +from functools import partial +from typing import Any, Generic, Protocol, TypeVar, cast +from typing_extensions import ParamSpec + +from ...cleanup import _CALLBACK_REPR +from ..cache import LineProfilingCache + + +__all__ = ('Queue', 'PutWrapper', 'QuickGetWrapper') + +T = TypeVar('T') +PS = ParamSpec('PS') + +_UNTAGGED_RESULT_WARNING_TEMPLATE = ( + 'received a pool-task result `{result}` without the expected tag {tag!r}; ' + 'the worker process appears not to have been set up for profiling ' + '(e.g. its interpreter never loaded the profiling startup hook), ' + 'so its profiling data will be missing from the output' +) + + +class Queue(Protocol): + """ + Protocol for methods common to e.g. :py:class:`queue.SimpleQueue` + and :py:class:`multiprocessing.queues.SimpleQueue`. + """ + def put(self, obj: Any) -> None: + ... + + def get(self) -> Any: + ... + + +class PutWrapper: + """ + Wrap around a queue (the ``outqueue`` argument to + :py:func:`multiprocessing.pool.worker`) so that each call to its + ``.put()`` is preceded by calling a ``callback()``; if + ``tag`` is given, the object pushed to the parent is replaced with + the triplet ``(tag, callback(), obj)``. + """ + def __init__( + self, + queue: Queue, + callback: Callable[[], Any], + tag: str | None = None, + ) -> None: + self._queue = queue + self._callback = callback + self._tag = tag + + def __getattr__(self, attr: str) -> Any: + return getattr(self._queue, attr) + + def put(self, obj: Any) -> None: + data = self._callback() + if self._tag is not None: + obj = self._tag, data, obj + self._queue.put(obj) + + def get(self) -> Any: + return self._queue.get() + + +class QuickGetWrapper(Generic[PS, T]): + """ + Wrap around a :py:attr:`multiprocessing.pool.Pool._quick_get` to + intercept and process data slipped into the queue by a + :py:class:`PutWrapper`. + + - If the result of the get is :py:const:`None` (i.e. a sentinel + value), it is a no-op. + + - If the result is a 3-tuple consisting of the ``tag``, some + ``data``, followed by the original result, the ``data`` is + processed with + ``callback(cache: LineProfilingCache, data: T) -> Any`` and the + original result is returned. + + - Otherwise, a warning (once per ``tag``) is emitted and the result + is returned as-is. + """ + def __init__( + self, + cache: LineProfilingCache, + get: Callable[PS, tuple[Any, ...] | None], + callback: Callable[[LineProfilingCache, T], Any], + tag: str, + ) -> None: + self._impl = get + self._callback = partial(callback, cache) + self._warn = partial(self._warn_untagged_result_once, cache, tag) + self._tag = tag + + def __call__( + self, /, *args: PS.args, **kwargs: PS.kwargs + ) -> tuple[Any, ...] | None: + """ + Note: + A worker which was never patched (its interpreter didn't run + the profiling startup hook) pushes vanilla un-tagged + results; those are passed through untouched, with a + once-per-session warning, so that a mixed + patched-parent/vanilla-worker setup degrades to missing + profile data instead of killing the pool's result-handler + thread (and thereby deadlocking every + :py:meth:`multiprocessing.pool.AsyncResult.get`). + """ + result = self._impl(*args, **kwargs) + if result is None: + return None + if ( + isinstance(result, tuple) + and len(result) == 3 + and result[0] == self._tag + ): + _, data, orig_result = result + self._callback(cast(T, data)) + return orig_result + self._warn(result) + return result + + @staticmethod + def _warn_untagged_result_once( + cache: LineProfilingCache, tag: str, result: Any, + ) -> None: + key = 'mp_result_handler_untagged_result_warnings' + # No lock: a race just means an extra warning, and this runs on + # the pool's single result-handler thread anyway + has_warned = cache._additional_data.setdefault(key, {}) + if has_warned.get(tag): + return + has_warned[tag] = True + # Log before warning in case the warning is promoted to an error + msg = _UNTAGGED_RESULT_WARNING_TEMPLATE.format( + result=_CALLBACK_REPR(result), tag=tag, + ) + cache._debug_output(msg, 'warning') + warnings.warn(msg) diff --git a/line_profiler/_child_process_profiling/multiprocessing_patches/mp_config.py b/line_profiler/_child_process_profiling/multiprocessing_patches/mp_config.py new file mode 100644 index 00000000..92ef2230 --- /dev/null +++ b/line_profiler/_child_process_profiling/multiprocessing_patches/mp_config.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +import dataclasses +from typing import final +from typing_extensions import Self + +from ...toml_config import ConfigSource +from ..cache import LineProfilingCache + + +__all__ = ('MPConfig',) + + +@final +@dataclasses.dataclass +class MPConfig: + """ + Consolidate the config options into a structured object. + + Notes: + + - This corresponds to the + ``tool.child_processes.multiprocessing`` table in + ``line_profiler.toml``. + + - While the object has currently only a single field, it is + intentionally kept this way to make the config scheme + extensible. + """ + patches: dict[str, bool] + + @classmethod + def from_config(cls, config: ConfigSource) -> Self: + loaded = ( + config + .get_subconfig('child_processes', 'multiprocessing') + .conf_dict + ) + return cls(patches=dict(loaded['patches'])) + + @classmethod + def from_cache(cls, cache: LineProfilingCache) -> Self: + key = 'mp_config' + try: + return cache._additional_data[key] + except KeyError: + config = cls.from_config(cache._config_source) + return cache._additional_data.setdefault(key, config) + + @classmethod + def get_defaults(cls) -> Self: + namespace = globals() + name = '_DEFAULT_CONFIG' + try: + return namespace[name] + except KeyError: + defaults = cls.from_config(ConfigSource.from_default(copy=False)) + return namespace.setdefault(name, defaults) diff --git a/line_profiler/_child_process_profiling/runpy_patches.py b/line_profiler/_child_process_profiling/runpy_patches.py new file mode 100644 index 00000000..22a600ea --- /dev/null +++ b/line_profiler/_child_process_profiling/runpy_patches.py @@ -0,0 +1,133 @@ +""" +Patches for :py:mod:`runpy` to be patched into the namespace of +:py:mod:`multiprocessing.spawn`, so that the rewriting of ``__main__`` +can be continued into child processes. +""" +from __future__ import annotations + +import os +from collections.abc import Callable +from functools import partial +from importlib.util import find_spec +from types import ModuleType +from typing import cast, TypeVar +from typing_extensions import Concatenate, ParamSpec + +from ..autoprofile.ast_tree_profiler import AstTreeProfiler +from ..autoprofile.run_module import AstTreeModuleProfiler +from ..autoprofile.util_static import modname_to_modpath +from ..cleanup import Cleanup +from .cache import LineProfilingCache + + +__all__ = ('create_runpy_wrapper',) + + +PS = ParamSpec('PS') +T = TypeVar('T') + + +THIS_MODULE = (lambda: None).__module__ + + +def _copy_module(name: str) -> ModuleType: + """ + Returns: + module (ModuleType): + Module object, which is a fresh copy of the module named + ``name`` + """ + spec = find_spec(name) + if spec is None: + raise ModuleNotFoundError(name) + assert spec.loader + assert callable(getattr(spec.loader, 'exec_module', None)) + module = ModuleType(spec.name) + for attr, value in { + '__spec__': spec, + '__name__': spec.name, + '__file__': spec.origin, + '__path__': spec.submodule_search_locations, + }.items(): + if value is not None: + setattr(module, attr, value) + spec.loader.exec_module(module) + return module + + +def _exec( + cache: LineProfilingCache, + CodeWriter: type[AstTreeProfiler], + _code, # This represents the first pos arg to `exec()` (ignored) + /, + *args, **kwargs, +) -> None: + """ + To be monkey-patched into :py:mod:`runpy`'s namespace as `exec()` + so that rewritten and autoprofiled code at ``cache.rewrite_module`` + is always executed. + """ + assert cache.rewrite_module + call = cache._format_call('exec', _code, *args, **kwargs) + cache._debug_output(f'Calling via {THIS_MODULE}: `{call}`') + fname = str(cache.rewrite_module) + code_writer = CodeWriter( + fname, + list(cache.profiling_targets), + cache.profile_imports, + ) + code = compile(code_writer.profile(), fname, 'exec') + exec(code, *args, **kwargs) + + +def _run( + cache: LineProfilingCache, + runpy: ModuleType, + func: Callable[Concatenate[str, PS], T], + name: str, + resolve_target_to_path: Callable[[str], str], + CodeWriter: type[AstTreeProfiler], + target: str, + /, + *args: PS.args, **kwargs: PS.kwargs +) -> T: + call = cache._format_call('runpy.' + name, target, *args, **kwargs) + cache._debug_output(f'Calling via {THIS_MODULE}: `{call}`') + if cache.rewrite_module: + try: + filename = resolve_target_to_path(target) + profile = os.path.samefile(filename, cache.rewrite_module) + except Exception as e: + cache._debug_output( + f'{THIS_MODULE}: Failed to check whether code loaded by ' + f'`runpy.{name}(...)` is to be rewritten ' + f'({type(e).__name__}: {e})' + ) + profile = False + else: + profile = False + # If we are about to run the code to be autoprofiled, monkey-patch + # `exec()` into the `runpy` namespace which just rewrites + # `cache.rewrite_module` and executes it + with Cleanup() as cleanup: + if profile: + cleanup.patch(runpy, 'exec', partial(_exec, cache, CodeWriter)) + return func(target, *args, **kwargs) + + +def create_runpy_wrapper(cache: LineProfilingCache) -> ModuleType: + """ + Create a copy of :py:mod:`runpy` which does code rewriting similar + to :py:func:`line_profiler.autoprofile.autoprofile.run` for the + appropriate file as indicated by ``cache``. + """ + runpy = _copy_module('runpy') + for func, resolver, CodeWriter in [ + ('run_path', str, AstTreeProfiler), + ('run_module', modname_to_modpath, AstTreeModuleProfiler), + ]: + impl = getattr(runpy, func) + res = cast(Callable[[str], str], resolver) # Help `mypy` out + wrapper = partial(_run, cache, runpy, impl, func, res, CodeWriter) + setattr(runpy, func, wrapper) + return runpy diff --git a/line_profiler/_threading_patches.py b/line_profiler/_threading_patches.py new file mode 100644 index 00000000..32f13471 --- /dev/null +++ b/line_profiler/_threading_patches.py @@ -0,0 +1,124 @@ +""" +Patch :py:mod:`threading` so that profiling extends consistenly into +processes it creates. +""" +from __future__ import annotations + +import threading +from collections.abc import Callable +from functools import wraps +from typing import TYPE_CHECKING, Any, TypeVar +from typing_extensions import ParamSpec, Concatenate + +from ._line_profiler import ( # type: ignore + USE_LEGACY_TRACE as SHOULD_PATCH_THREADING, +) +from .line_profiler import LineProfiler +from .cleanup import Cleanup + + +__all__ = ('apply', 'SHOULD_PATCH_THREADING') + + +T = TypeVar('T') +PS = ParamSpec('PS') + +_PATCHED_MARKER = '__line_profiler_patched_threading__' + + +def make_syncing_wrapper( + func: Callable[PS, T], prof: LineProfiler, enable_count: int, +) -> Callable[PS, T]: + """ + Wrap the callable ``func`` so that when we spin up a new thread, we + sync the + :py:attr:`line_profiler.line_profiler.LineProfiler.enable_count` of + the active profiler (stored at the cache instance loaded from + :py:meth:`LineProfilingCache.load`) with ``enable_count``. + + Note: + This only seems to work as intended when using the legacy trace + system... + """ + @wraps(func) + def wrapper(*args: PS.args, **kwargs: PS.kwargs) -> T: + if TYPE_CHECKING: + assert hasattr(prof, 'enable_count') + assert isinstance(prof.enable_count, int) + # Note: `prof.enable_count` is most likely to be zero on the new + # thread + thread_enable_count: int = prof.enable_count + for _ in range(enable_count - thread_enable_count): + prof.enable_by_count() + try: + return func(*args, **kwargs) + finally: + # Reset enable counts to avoid problems if the thread id is + # ever reused + for _ in range(prof.enable_count - thread_enable_count): + prof.disable_by_count() + + return wrapper + + +def make_thread_init_wrapper( + prof: LineProfiler, + vanilla_impl: Callable[ + Concatenate[threading.Thread, None, Callable[..., Any] | None, PS], + None + ], +) -> Callable[ + Concatenate[threading.Thread, None, Callable[..., Any] | None, PS], None +]: + """ + Wrap the initializer of :py:class:`threading.Thread` so that the + profiler's :py:attr:`LineProfiler.enable_count` is synced up on + newly spun-up threads. + """ + @wraps(vanilla_impl) + def wrapper( + self: threading.Thread, + group: None = None, + target: Callable[..., Any] | None = None, + *args: PS.args, + **kwargs: PS.kwargs + ) -> None: + enable_count: int | None = getattr(prof, 'enable_count', None) + if target is not None and enable_count: + if TYPE_CHECKING: + assert prof is not None + target = make_syncing_wrapper(target, prof, enable_count) + vanilla_impl(self, group, target, *args, **kwargs) + + return wrapper + + +def apply(cleanup: Cleanup, prof: LineProfiler) -> None: + """ + Set up profiling in threads started by :py:mod:`threading` by + applying patches to the module. + + Args: + cleanup (Cleanup) + Cleanup instance managing the profiling session + + Side effects: + - :py:mod:`threading` marked as having been set up + + - The following methods and functions patched: + + - :py:meth:`threading.Thread.__init__` + + - Cleanup callbacks registered via ``cleanup.add_cleanup()`` + + Note: + This is a no-op when using :py:mod:`sys.monitoring`-based + profiling. + """ + if not SHOULD_PATCH_THREADING: + return + if getattr(threading, _PATCHED_MARKER, False): + return + init_wrapper = make_thread_init_wrapper(prof, threading.Thread.__init__) + cleanup.patch(threading.Thread, '__init__', init_wrapper) + cleanup.patch(threading, _PATCHED_MARKER, True) diff --git a/line_profiler/cleanup.py b/line_profiler/cleanup.py new file mode 100644 index 00000000..b35024ff --- /dev/null +++ b/line_profiler/cleanup.py @@ -0,0 +1,435 @@ +""" +Utilities for cleaning up after ourselves. +""" +from __future__ import annotations + +from collections.abc import ( + Callable, Generator, Iterable, Mapping, MutableMapping, +) +from functools import partial +from inspect import getattr_static +from operator import setitem +from pathlib import Path +from typing import Any, Literal, Protocol, TypeVar, cast, overload +from typing_extensions import Concatenate, ParamSpec, Self + +from .line_profiler_utils import CallbackRepr, make_tempfile +from . import _diagnostics as diagnostics + + +__all__ = ('Cleanup',) + +PS = ParamSpec('PS') +K = TypeVar('K') +V = TypeVar('V') +_Stacks = dict[float, list[Callable[[], Any]]] +_StackContexts = list[_Stacks] +LogLevel = Literal['debug', 'info', 'warning', 'error', 'critical'] + + +_CALLBACK_REPR_HELPER = CallbackRepr(maxother=cast(int, float('inf'))) +_CALLBACK_REPR = _CALLBACK_REPR_HELPER.repr + + +class _LoggingCallback(Protocol): + @overload + def __call__(self, msg: str, /) -> Any: + ... + + @overload + def __call__(self, msg: str, /, level: LogLevel) -> Any: + ... + + def __call__(self, *_, **__): + ... + + +class Cleanup: + """ + Object which holds cleanup callbacks. Also provides convenience + methods for creating tempfiles, updating mappings, and setting + attributes on objects. + """ + def __init__(self, *_, **__) -> None: + self._contexts: _StackContexts = [] + + def __enter__(self) -> Self: + """ + Returns: + The instance + + Note: + This context manager is reentrant; entering the context + create a new set of cleanup stacks, which is then cleaned up + on :py:meth:`~.__exit__`. + + Example: + >>> strings = [] + >>> add = strings.append + >>> with Cleanup() as cleanup: + ... cleanup.add_cleanup(add, 'one') + ... # Increased priority + ... cleanup.add_cleanup_with_priority(add, 1, 'two') + ... add('three') + ... with cleanup: + ... # Decreased priority + ... cleanup.add_cleanup_with_priority( + ... add, -1, 'four', + ... ) + ... cleanup.add_cleanup(add, 'five') + ... add('six') + ... add('seven') + ... # Increased priority + ... cleanup.add_cleanup_with_priority(add, 1, 'eight') + ... + >>> strings # doctest: +NORMALIZE_WHITESPACE + ['three', 'six', 'five', 'four', 'seven', 'eight', 'two', + 'one'] + """ + self._contexts.append({}) + return self + + def __exit__(self, *_, **__) -> Any: + """ + Call ``~.cleanup(1)``, clearing the level of cleanup stacks we + previously :py:meth:`~.__enter__`-ed into. + """ + self.cleanup(1, reason='context exit') + + # Cleanup methods + + def cleanup( + self, levels: int | None = None, *, reason: str | None = None, + ) -> None: + """ + Pop cleanup callbacks from the internal stacks added via + :py:meth:`~.add_cleanup` etc. and call them in order. + + Args: + levels (int | None): + Number of stack levels to clear; passing :py:const`None` + clears the entire stack of callback stacks + reason (str | None): + Optional description of the reason for cleaning up + """ + def pop_all_contexts( + contexts: _StackContexts, + ) -> Generator[_Stacks, None, None]: + while contexts: + yield contexts.pop() + + def pop_n_levels_of_contexts( + contexts: _StackContexts, n: int, + ) -> Generator[_Stacks, None, None]: + for _ in range(n): + try: + yield contexts.pop() + except IndexError: # Ran out of levels + return + + pop_contexts: Iterable[_Stacks] + if levels is None: + pop_contexts = pop_all_contexts(self._contexts) + else: + pop_contexts = pop_n_levels_of_contexts(self._contexts, levels) + cleanup = partial(self._cleanup, self._debug_output, reason=reason) + for stacks in pop_contexts: + cleanup(stacks) + + @staticmethod + def _cleanup( + log: _LoggingCallback, stacks: _Stacks, reason: str | None, + ) -> None: + ncallbacks_total = sum(len(stack) for stack in stacks.values()) + note = f'{ncallbacks_total} callback(s)' + if reason: + note = f'{reason}; {note}' + if not ncallbacks_total: + log(f'Cleanup aborted ({note})') + return + # Bookend the cleanup loop with log messages to help detect if + # child processes are prematurely terminated + log(f'Starting cleanup ({note})...') + ncallbacks_run = 0 + for priority in sorted(stacks, reverse=True): + callbacks = stacks.pop(priority) + while callbacks: + callback = callbacks.pop() + callback_repr = _CALLBACK_REPR(callback) + ncallbacks_run += 1 + try: + callback() + except Exception as e: + success, state = False, 'failed' + msg = f'{callback_repr}: {type(e).__name__}: {e}' + else: + success, state, msg = True, 'succeeded', f'{callback_repr}' + msg = ( + f'- Cleanup {state} ' + f'({ncallbacks_run}/{ncallbacks_total}): {msg}' + ) + log(msg, 'debug' if success else 'warning') + log(f'... cleanup completed ({note})') + + def add_cleanup( + self, callback: Callable[PS, Any], *args: PS.args, **kwargs: PS.kwargs, + ) -> None: + """ + Shorthand for calling :py:meth:`~.add_cleanup_with_priority` + with ``priority=0``, which should be considered the default. + """ + self.add_cleanup_with_priority(callback, 0, *args, **kwargs) + + def add_cleanup_with_priority( + self, callback: Callable[PS, Any], priority: float, /, + *args: PS.args, **kwargs: PS.kwargs, + ) -> None: + """ + Add a cleanup callback to the internal stacks. + + Args: + callback (Callable[..., Any]): + Callback to be called at cleanup + priority (float): + Numeric priority value; callbacks with a HIGHER value + are invoked BEFORE those with lower values + *args, **kwargs: + Arguments ``callback`` should be called with + + Example: + >>> strings = [] + >>> cleanup = Cleanup() + >>> # Default priority + >>> cleanup.add_cleanup(strings.append, 'first') + >>> # Decreased priority + >>> cleanup.add_cleanup_with_priority( + ... strings.append, -1, 'second', + ... ) + >>> # Increased priority + >>> cleanup.add_cleanup_with_priority( + ... strings.append, 1, 'third', + ... ) + >>> cleanup.add_cleanup(strings.append, 'fourth') + >>> assert not strings + >>> cleanup.cleanup() + >>> strings + ['third', 'fourth', 'first', 'second'] + """ + if args or kwargs: + callback = partial(callback, *args, **kwargs) + self._current_context.setdefault(priority, []).append(callback) + header = 'Cleanup callback added' + if priority: + header = f'{header} (priority: {priority})' + self._debug_output(f'{header}: {_CALLBACK_REPR(callback)}') + + # Convenience methods + + def update_mapping( + self, + mapping: MutableMapping[K, V], + updates: Mapping[K, V], + *, + _format_debug_msg: Callable[[Mapping[K, V], K, str], str] = ( + lambda mapping, key, change: 'Update {}[{!r}]: {}'.format( + object.__repr__(mapping), key, change, + ) + ), + ) -> None: + """ + Update a mapping with another and add cleanup callbacks to + reverse them. + + Args: + mapping (MutableMapping[K, V]): + Mapping to be updated + updates (Mapping[K, V]): + Mapping containing the updates + + Example: + >>> d1 = {1: 2, 3: 4} + >>> d2 = d1.copy() + >>> updates = {0: -1, 3: 5} + >>> with Cleanup() as cleanup: + ... cleanup.update_mapping(d1, updates) + ... for key, value in updates.items(): + ... assert d1[key] == value + ... + >>> assert d1 == d2 + """ + for key, value in updates.items(): + try: + old = mapping[key] + except KeyError: + self.add_cleanup(mapping.pop, key, None) + change = f'{value!r} (new)' + else: + self.add_cleanup(setitem, mapping, key, old) + change = f'{old!r} -> {value!r}' + self._debug_output(_format_debug_msg(mapping, key, change)) + mapping[key] = value + + def make_tempfile( + self, *, + delete: bool = True, + priority: float = 0, + _format_debug_msg: Callable[[Path], str] = ( + 'Created tempfile: {}'.format + ), + **kwargs + ) -> Path: + """ + Create a fresh tempfile with :py:func:`tempfile.mkstemp`. + + Args: + delete (bool): + Whether to remove the file on cleanup + priority (float): + Cleanup priority (see + :py:meth:`~.add_cleanup_with_priority`) + **kwargs: + Passed to :py:func:`tempfile.mkstemp` + + Returns: + path (Path): + Path to the created file. + + Example: + >>> prefix, suffix = 'my_file_', '.txt' + >>> with Cleanup() as cleanup: + ... path = cleanup.make_tempfile( + ... prefix=prefix, suffix=suffix, + ... ) + ... assert path.exists() + ... assert path.name.startswith(prefix) + ... assert path.name.endswith(suffix) + ... + >>> assert not path.exists() + """ + path = make_tempfile(**kwargs) + self._debug_output(_format_debug_msg(path)) + if delete: + self.add_cleanup_with_priority( + path.unlink, priority, missing_ok=True, + ) + return path + + def patch( + self, obj: Any, attr: str, value: Any, *, + name: str | None = None, + static: bool = True, + cleanup: bool = True, + priority: float = 0, + ) -> None: + """ + Patch an attribute on an object. + + Args: + obj (Any): + Object to be patched + attr (str): + Name of the attribute + value (Any): + Value to be assigned to said attribute of ``obj`` + name (str | None): + Optional name for ``obj`` to be used in debug messages + static (bool): + Whether to use :py:func:`inspect.getattr_static` to + get the current value of the attribute + cleanup (bool): + Whether to reverse the patch (by resetting or deleting + the attribute) on cleanup + priority (float): + Cleanup priority (see + :py:meth:`~.add_cleanup_with_priority`) + + Example: + >>> class Object(object): + ... pass # Allow setting arbitrary attributes + ... + >>> + >>> obj = Object() + >>> obj.foo = 1 + >>> with Cleanup() as cleanup: + ... cleanup.patch(obj, 'foo', 2) + ... cleanup.patch(obj, 'bar', 3) + ... assert obj.foo == 2 + ... assert obj.bar == 3 + ... + >>> assert obj.foo == 1 + >>> assert not hasattr(obj, 'bar') + """ + if cleanup: + add_cleanup: Callable[ + Concatenate[Callable[..., Any], float, ...], Any + ] = self.add_cleanup_with_priority + else: + # ... yeah gotta disagree with flake8, a lambda makes + # perfect sense here + add_cleanup = lambda *_, **__: None # noqa: E731 + get_attribute = getattr_static if static else getattr + + try: + old = get_attribute(obj, attr) + except AttributeError: + add_cleanup(delattr, priority, obj, attr) + else: + add_cleanup(setattr, priority, obj, attr, old) + setattr(obj, attr, value) + if name is None: + name = self._get_name(obj) + msg = 'Patched `{}.{}` -> `{}`'.format(name, attr, value) + self._debug_output(msg) + + # Helper methods + + @staticmethod + def _get_name(obj: Any, /) -> str: + """ + Get an appropriate name for an arbitrary object. + + Example: + >>> import textwrap + >>> + >>> + >>> Cleanup._get_name(textwrap) + 'textwrap' + >>> Cleanup._get_name(textwrap.dedent) + 'textwrap.dedent' + >>> Cleanup._get_name(str) + 'str' + >>> Cleanup._get_name(print) + 'print' + >>> Cleanup._get_name(object()) # doctest: +ELLIPSIS + '' + """ + if hasattr(obj, '__qualname__'): + name = obj.__qualname__ + elif hasattr(obj, '__name__'): + name = obj.__name__ + else: + return repr(obj) + if hasattr(obj, '__module__'): + if obj.__module__ not in ('builtins', '__builtins__'): + name = f'{obj.__module__}.{name}' + return str(name) + + def _debug_output(self, msg: str, /, level: LogLevel = 'debug') -> None: + """ + Write debugging output. + + Note: + This default implementation just writes to the logger at the + specified level. + """ + log_func = getattr(diagnostics.log, level) + log_func(msg) + + @property + def _current_context(self) -> _Stacks: + try: + return self._contexts[-1] + except IndexError: + ctx: _Stacks = {} + self._contexts.append(ctx) + return ctx diff --git a/line_profiler/curated_profiling.py b/line_profiler/curated_profiling.py new file mode 100644 index 00000000..8f2fd935 --- /dev/null +++ b/line_profiler/curated_profiling.py @@ -0,0 +1,230 @@ +""" +Tools for setting up profiling in a curated environment (e.g. with +the use of :py:mod:`kernprof`). +""" +from __future__ import annotations + +import builtins +import dataclasses +import os +import warnings +from collections.abc import Collection +from io import StringIO +from textwrap import indent +from typing import Any, TextIO, cast +from typing_extensions import Self + +from . import _diagnostics as diagnostics, profile as _global_profiler +from ._threading_patches import apply as apply_threading_patches +from .autoprofile.autoprofile import ( + _extend_line_profiler_for_profiling_imports as upgrade_profiler, +) +from .autoprofile.util_static import modpath_to_modname +from .autoprofile.eager_preimports import ( + is_dotted_path, write_eager_import_module, +) +from .cleanup import Cleanup +from .cli_utils import short_string_path +from .line_profiler import LineProfiler +from .profiler_mixin import ByCountProfilerMixin + + +__all__ = ('ClassifiedPreimportTargets', 'CuratedProfilerContext') + + +@dataclasses.dataclass +class ClassifiedPreimportTargets: + """ + Pre-import targets classified into three bins: ``regular`` targets, + targets to ``recurse`` into, and ``invalid`` targets + """ + regular: list[str] = dataclasses.field(default_factory=list) + recurse: list[str] = dataclasses.field(default_factory=list) + invalid: list[str] = dataclasses.field(default_factory=list) + + def __bool__(self) -> bool: + return bool(self.regular or self.recurse) + + def write_preimport_module( + self, fobj: TextIO, *, debug: bool | None = None, **kwargs + ) -> None: + """ + Convenience interface with + :py:func:`~.write_eager_import_module`, writing a module which + when imported sets up profiling of the targets. + + Args: + fobj (TextIO): + File object to write said module to. + debug (Optional[bool]): + Whether to generate debugging outputs. + kwargs: + Passed to :py:func:`~.write_eager_import_module`. + """ + if self.invalid: + invalid_targets = sorted(set(self.invalid)) + msg = ( + '{} profile-on-import target{} cannot be converted to ' + 'dotted-path form: {!r}'.format( + len(invalid_targets), + '' if len(invalid_targets) == 1 else 's', + invalid_targets, + ) + ) + # Log before warn in case the warning is raised + diagnostics.log.warning(msg) + warnings.warn(msg, stacklevel=2) + + if not self: + return None + # Note: `ty` (but not `mypy`) keeps complaining about the our + # splatting this dict; explicitly use `Any` to tell it to shut + # up. + write_module_kwargs: dict[str, Any] = { + 'dotted_paths': self.regular, + 'recurse': self.recurse, + **kwargs, + } + if diagnostics.DEBUG if debug is None else debug: + with StringIO() as sio: + write_eager_import_module(stream=sio, **write_module_kwargs) + code = sio.getvalue() + print(code, file=fobj) + if hasattr(fobj, 'name'): + fobj_repr = repr(short_string_path(str(fobj.name))) + else: + fobj_repr = repr(fobj) # Fall back + diagnostics.log.debug( + f'Wrote temporary module for pre-imports to {fobj_repr}:\n' + + indent(code, ' ') + ) + else: + write_eager_import_module(stream=fobj, **write_module_kwargs) + + @classmethod + def from_targets( + cls, + targets: Collection[str], + exclude: Collection[os.PathLike[str] | str] = (), + ) -> Self: + """ + Create an instance based on a collection of targets + (like what is supplied to ``kernprof --prof-mod=...``). + + Args: + targets (Collection[str]) + Collection of dotted paths and filenames to profile. + exclude (Collection[str]) + Collections of filenames which are explicitly excluded + from being profiled. + + Return: + New instance. + """ + filtered_targets = [] + recurse_targets = [] + invalid_targets = [] + for target in targets: + if is_dotted_path(target): + modname = target + else: + # Paths already normalized by + # `_normalize_profiling_targets()` + if not os.path.exists(target): + invalid_targets.append(target) + continue + if any( + os.path.samefile(target, excluded) for excluded in exclude + ): + # Ignore the script to be run in eager importing + # (`line_profiler.autoprofile.autoprofile.run()` + # will handle it) + continue + modname = modpath_to_modname(target, hide_init=False) + if modname is None: # Not import-able + invalid_targets.append(target) + continue + if modname.endswith('.__init__'): + modname = modname.rpartition('.')[0] + filtered_targets.append(modname) + else: + recurse_targets.append(modname) + return cls(filtered_targets, recurse_targets, invalid_targets) + + +class CuratedProfilerContext(Cleanup): + """ + Context manager for handling various bookkeeping tasks when setting + up and tearing down profiling: + + - Slipping ``prof`` into the builtin namespace (if + ``insert_builtin`` is true) and :py::deco:`~.profile` + - Patch :py:class:`threading.Thread` so that line-profiling is + enabled on new threads if it is on the spawning threads + - At exit, clearing the ``enable_count`` of ``prof``, properly + disabling it + + Notes: + + - The attributes on this object are to be considered + implementation details, but not its methods and their + signatures. + + - In contrast to the base class (:py:class:`Cleanup`), while + this context manager is still reentrant, reentering in nested + `with: ...` statements is a no-op. + """ + def __init__( + self, + prof: ByCountProfilerMixin, + insert_builtin: bool = False, + builtin_loc: str = 'profile', + ) -> None: + super().__init__() + self.prof = prof + self.insert_builtin = insert_builtin + self.builtin_loc = builtin_loc + self._installed = False + self._kpo = _global_profiler._kernprof_overwrite + + def _global_install(self, prof: ByCountProfilerMixin | None) -> None: + # Wrapper to convince type-checkers it is okay to pass these + # stuff to `._kernprof_overwrite()`. We don't want to patch + # that method's signature because passing non `LineProfiler` + # objects to it should be the exception, not the norm. + self._kpo(cast(LineProfiler, prof)) + + @staticmethod + def _disable_profiler(prof: ByCountProfilerMixin) -> None: + for _ in range(getattr(prof, 'enable_count', 0)): + prof.disable_by_count() + + def install(self) -> None: + if self._installed: + return + # Equip the profiler instance with the + # `.add_imported_function_or_module()` pseudo-method + upgrade_profiler(self.prof) + # Overwrite the explicit profiler (`@line_profiler.profile`) + self._global_install(self.prof) + self.add_cleanup(self._global_install, None) + # Patch `threading` + if isinstance(self.prof, LineProfiler): + apply_threading_patches(self, self.prof) + # Set up hooks to deal with inserting `.prof` as a builtin name + if self.insert_builtin: + self.patch(builtins, self.builtin_loc, self.prof) + # Disable the profiler + self.add_cleanup(self._disable_profiler, self.prof) + + self.patch(self, '_installed', True) + + def uninstall(self) -> None: + self.cleanup(reason='uninstalling profiling context') + + def __enter__(self) -> Self: + self.install() + return self + + def __exit__(self, *_, **__) -> None: + self.uninstall() diff --git a/line_profiler/line_profiler.py b/line_profiler/line_profiler.py index bede6f6c..718c10e4 100755 --- a/line_profiler/line_profiler.py +++ b/line_profiler/line_profiler.py @@ -18,21 +18,12 @@ import tempfile import types import tokenize +import warnings from argparse import ArgumentParser +from collections.abc import Callable, Collection, Mapping, Sequence from datetime import datetime from os import PathLike -from typing import ( - TYPE_CHECKING, - IO, - Callable, - Literal, - Mapping, - Protocol, - Sequence, - TypeVar, - cast, - Tuple, -) +from typing import TYPE_CHECKING, IO, Any, Literal, Protocol, TypeVar, cast try: from ._line_profiler import ( @@ -62,7 +53,7 @@ class _IPythonLike(Protocol): def register_magics(self, magics: type) -> None: ... PS = ParamSpec('PS') - _TimingsMap = Mapping[Tuple[str, int, str], list[Tuple[int, int, int]]] + _TimingsMap = Mapping[tuple[str, int, str], list[tuple[int, int, int]]] T = TypeVar('T') T_co = TypeVar('T_co', covariant=True) @@ -226,6 +217,15 @@ def tokeneater( return super().tokeneater(type, token, srowcol, erowcol, line) +class _EmptyFileError(OSError): + """ + Error raised when trying to read profiling data from an empty file. + """ + def __init__(self, file: PathLike[str] | str) -> None: + super().__init__(str(file)) + self.file = file + + class _WrapperInfo: """ Helper object for holding the state of a wrapper function. @@ -264,8 +264,8 @@ def __eq__(self, other: object) -> bool: Example: >>> from copy import deepcopy >>> stats1 = LineStats( - ... {('foo', 1, 'spam.py'): [(2, 10, 300)], - ... ('bar', 10, 'spam.py'): + ... {('spam.py', 1, 'foo'): [(2, 10, 300)], + ... ('spam.py', 10, 'bar'): ... [(11, 2, 1000), (12, 1, 500)]}, ... 1E-6) >>> stats2 = deepcopy(stats1) @@ -274,7 +274,7 @@ def __eq__(self, other: object) -> bool: >>> assert stats2 != stats1 >>> stats3 = deepcopy(stats1) >>> assert stats1 == stats3 is not stats1 - >>> stats3.timings['foo', 1, 'spam.py'][:] = [(2, 11, 330)] + >>> stats3.timings['spam.py', 1, 'foo'][:] = [(2, 11, 330)] >>> assert stats3 != stats1 """ for attr in 'timings', 'unit': @@ -290,20 +290,20 @@ def __add__(self, other: _StatsLike) -> Self: """ Example: >>> stats1 = LineStats( - ... {('foo', 1, 'spam.py'): [(2, 10, 300)], - ... ('bar', 10, 'spam.py'): + ... {('spam.py', 1, 'foo'): [(2, 10, 300)], + ... ('spam.py', 10, 'bar'): ... [(11, 2, 1000), (12, 1, 500)]}, ... 1E-6) >>> stats2 = LineStats( - ... {('bar', 10, 'spam.py'): + ... {('spam.py', 10, 'bar'): ... [(11, 10, 20000), (12, 5, 1000)], - ... ('baz', 5, 'eggs.py'): [(5, 2, 5000)]}, + ... ('eggs.py', 5, 'baz'): [(5, 2, 5000)]}, ... 1E-7) >>> stats_sum = LineStats( - ... {('foo', 1, 'spam.py'): [(2, 10, 300)], - ... ('bar', 10, 'spam.py'): + ... {('spam.py', 1, 'foo'): [(2, 10, 300)], + ... ('spam.py', 10, 'bar'): ... [(11, 12, 3000), (12, 6, 600)], - ... ('baz', 5, 'eggs.py'): [(5, 2, 500)]}, + ... ('eggs.py', 5, 'baz'): [(5, 2, 500)]}, ... 1E-6) >>> assert stats1 + stats2 == stats2 + stats1 == stats_sum """ @@ -314,20 +314,20 @@ def __iadd__(self, other: _StatsLike) -> Self: """ Example: >>> stats1 = LineStats( - ... {('foo', 1, 'spam.py'): [(2, 10, 300)], - ... ('bar', 10, 'spam.py'): + ... {('spam.py', 1, 'foo'): [(2, 10, 300)], + ... ('spam.py', 10, 'bar'): ... [(11, 2, 1000), (12, 1, 500)]}, ... 1E-6) >>> stats2 = LineStats( - ... {('bar', 10, 'spam.py'): + ... {('spam.py', 10, 'bar'): ... [(11, 10, 20000), (12, 5, 1000)], - ... ('baz', 5, 'eggs.py'): [(5, 2, 5000)]}, + ... ('eggs.py', 5, 'baz'): [(5, 2, 5000)]}, ... 1E-7) >>> stats_sum = LineStats( - ... {('foo', 1, 'spam.py'): [(2, 10, 300)], - ... ('bar', 10, 'spam.py'): + ... {('spam.py', 1, 'foo'): [(2, 10, 300)], + ... ('spam.py', 10, 'bar'): ... [(11, 12, 3000), (12, 6, 600)], - ... ('baz', 5, 'eggs.py'): [(5, 2, 500)]}, + ... ('eggs.py', 5, 'baz'): [(5, 2, 500)]}, ... 1E-6) >>> address = id(stats2) >>> stats2 += stats1 @@ -337,6 +337,114 @@ def __iadd__(self, other: _StatsLike) -> Self: self.timings, self.unit = self._get_aggregated_timings([self, other]) return self + def __sub__(self, other: _StatsLike) -> Self: + """ + Subtract a "baseline" from this instance; the inverse of + :py:meth:`~.__add__`. Entries which reach zero hits and zero + time are dropped. The result is expressed in ``self``'s + :py:attr:`~.unit`. + + Raises: + ValueError: + If ``other`` contains an entry absent from (or larger + than the corresponding entry in) ``self``; + a valid baseline must be a "prefix" of ``self``. + + Example: + >>> baseline = LineStats( + ... {('spam.py', 1, 'foo'): [(2, 10, 300)], + ... ('spam.py', 10, 'bar'): + ... [(11, 2, 1000), (12, 1, 500)]}, + ... 1E-6) + >>> new = LineStats( + ... {('spam.py', 1, 'foo'): [(2, 10, 300)], + ... ('spam.py', 10, 'bar'): + ... [(11, 12, 3000), (12, 6, 600)], + ... ('eggs.py', 5, 'baz'): [(5, 2, 500)]}, + ... 1E-6) + >>> new - baseline + LineStats({('spam.py', 10, 'bar'): [(11, 10, 2000), \ +(12, 5, 100)], ('eggs.py', 5, 'baz'): [(5, 2, 500)]}, 1E-06) + >>> assert (new - baseline) + baseline == new + >>> new - new + LineStats({}, 1E-06) + >>> baseline - new + Traceback (most recent call last): + ... + ValueError: ...not a prefix... + """ + timings, unit = self._get_subtracted_timings(self, other) + return type(self)(timings, unit) + + def __isub__(self, other: _StatsLike) -> Self: + """ + In-place version of :py:meth:`~.__sub__`. + + Example: + >>> baseline = LineStats( + ... {('spam.py', 1, 'foo'): [(2, 10, 300)]}, 1E-6) + >>> stats = LineStats( + ... {('spam.py', 1, 'foo'): [(2, 15, 450)]}, 1E-6) + >>> address = id(stats) + >>> stats -= baseline + >>> assert id(stats) == address + >>> stats + LineStats({('spam.py', 1, 'foo'): [(2, 5, 150)]}, 1E-06) + """ + self.timings, self.unit = self._get_subtracted_timings(self, other) + return self + + @staticmethod + def _get_subtracted_timings(minuend, subtrahend): + """ + Compute ``minuend - subtrahend`` timings, expressed in + ``minuend.unit``; see :py:meth:`~.__sub__`. + """ + def prefix_error(reason): + return ValueError( + 'subtrahend is not a prefix of the minuend ' + f'(cannot subtract): {reason}' + ) + + unit = minuend.unit + factor = subtrahend.unit / unit + timings = { + key: {lineno: (nhits, time) for lineno, nhits, time in entries} + for key, entries in minuend.timings.items() + } + for key, entries in subtrahend.timings.items(): + try: + min_entries = timings[key] + except KeyError: + raise prefix_error(f'{key!r} not in the minuend') from None + for lineno, nhits, time in entries: + try: + prev_nhits, prev_time = min_entries[lineno] + except KeyError: + raise prefix_error( + f'line {lineno} of {key!r} not in the minuend' + ) from None + new_nhits = prev_nhits - nhits + new_time = int(round(prev_time - factor * time, 0)) + if new_nhits < 0 or new_time < 0: + raise prefix_error( + f'line {lineno} of {key!r}: ' + f'({prev_nhits}, {prev_time}) - ({nhits}, {time})' + + ('' if factor == 1 else f' * {factor}') + + f' = ({new_nhits}, {new_time})' + ) + if new_nhits or new_time: + min_entries[lineno] = new_nhits, new_time + else: + del min_entries[lineno] + return { + key: [ + (lineno, nhits, time) + for lineno, (nhits, time) in sorted(entries.items()) + ] + for key, entries in timings.items() if entries + }, unit + def print( self, stream: io.TextIOBase | None = None, @@ -367,17 +475,102 @@ def to_file(self, filename: PathLike[str] | str) -> None: with open(filename, 'wb') as f: pickle.dump(self, f, pickle.HIGHEST_PROTOCOL) + @classmethod + def get_empty_instance(cls) -> Self: + """ + Returns: + instance (LineStats): + New instance without any profiling data. + """ + prof = LineProfiler() + if TYPE_CHECKING: + assert hasattr(prof, 'timer_unit') + return cls({}, cast(float, prof.timer_unit)) + @classmethod def from_files( - cls, file: PathLike[str] | str, /, *files: PathLike[str] | str + cls, + file: PathLike[str] | str, + /, + *files: PathLike[str] | str, + on_empty: Literal['ignore', 'warn', 'error'] = 'warn', + on_defective: Literal['ignore', 'warn', 'error'] = 'error', + _note_on_empty: str | None = None, + _note_on_defective: str | None = None, ) -> Self: """ Utility function to load an instance from the given filenames. + + Args: + file (PathLike[str] | str): + File to load profiling data from + *files (PathLike[str] | str): + Ditto above + on_empty, on_defective (Literal['ignore', 'warn', 'error']): + What to do if some files are empty (resp. otherwise fail + to load): ``'ignore'`` those files, skip them but with a + ``'warn'``-ing, or raise the ``'error'`` as soon as one + is encountered + + Returns: + instance (LineStats): + New instance """ stats_objs = [] - for file in [file, *files]: + failures: dict[str, str] = {} + empty_files: set[str] = set() + all_files = [file, *files] + + for file in all_files: with open(file, 'rb') as f: - stats_objs.append(pickle.load(f)) + try: + if not os.stat(file).st_size: + raise _EmptyFileError(file) + stats_objs.append(pickle.load(f)) + except _EmptyFileError as e: + if on_empty == 'error': + raise + empty_files.add(str(e.file)) + except Exception as e: + if on_defective == 'error': + raise + failure = type(e).__name__ + if str(e): + failure = f'{failure}: {e}' + failures[str(file)] = failure + + problems: Collection[Any] + for problems, description, behavior, note in [ + ( + list(empty_files), + 'is/are empty and thus skipped', + on_empty, + _note_on_empty, + ), + ( + failures, + 'cannot be loaded and thus is/are skipped', + on_defective, + _note_on_defective, + ), + ]: + if not problems: + continue + msg = '{} file(s) out of {} {}: {!r}'.format( + len(problems), len(all_files), description, problems, + ) + if note: + msg = f'{msg}; {note}' + if behavior == 'warn': + # Log before warning because warnings may be promoted to + # errors + diagnostics.log.warning(msg) + warnings.warn(msg, stacklevel=2) + else: # 'ignore' + diagnostics.log.debug(msg) + + if not stats_objs: + return cls.get_empty_instance() return cls.from_stats_objects(*stats_objs) @classmethod @@ -387,23 +580,23 @@ def from_stats_objects( """ Example: >>> stats1 = LineStats( - ... {('foo', 1, 'spam.py'): [(2, 10, 300)], - ... ('bar', 10, 'spam.py'): + ... {('spam.py', 1, 'foo'): [(2, 10, 300)], + ... ('spam.py', 10, 'bar'): ... [(11, 2, 1000), (12, 1, 500)]}, ... 1E-6) >>> stats2 = LineStats( - ... {('bar', 10, 'spam.py'): + ... {('spam.py', 10, 'bar'): ... [(11, 10, 20000), (12, 5, 1000)], - ... ('baz', 5, 'eggs.py'): [(5, 2, 5000)]}, + ... ('eggs.py', 5, 'baz'): [(5, 2, 5000)]}, ... 1E-7) >>> stats_combined = LineStats.from_stats_objects( ... stats1, stats2) >>> assert stats_combined.unit == 1E-6 >>> assert stats_combined.timings == { - ... ('foo', 1, 'spam.py'): [(2, 10, 300)], - ... ('bar', 10, 'spam.py'): + ... ('spam.py', 1, 'foo'): [(2, 10, 300)], + ... ('spam.py', 10, 'bar'): ... [(11, 12, 3000), (12, 6, 600)], - ... ('baz', 5, 'eggs.py'): [(5, 2, 500)]} + ... ('eggs.py', 5, 'baz'): [(5, 2, 500)]} """ timings, unit = cls._get_aggregated_timings([stats, *more_stats]) return cls(timings, unit) @@ -840,7 +1033,7 @@ def show_func( func_name (str): name of profiled function - timings (List[Tuple[int, int, float]]): + timings (list[tuple[int, int, float]]): Measurements for each line (lineno, nhits, time). unit (float): diff --git a/line_profiler/line_profiler_utils.py b/line_profiler/line_profiler_utils.py index 887cdd55..73b83693 100644 --- a/line_profiler/line_profiler_utils.py +++ b/line_profiler/line_profiler_utils.py @@ -5,10 +5,23 @@ from __future__ import annotations import enum -import typing +import os +import sys +from collections.abc import Callable, Collection, Mapping, Sequence +from functools import partial +from pathlib import Path +from reprlib import Repr +from tempfile import mkstemp +from textwrap import indent +from types import MethodType +from typing import TYPE_CHECKING, Any, TypedDict, TypeVar +from typing_extensions import Self, Unpack -if typing.TYPE_CHECKING: - from typing_extensions import Self + +__all__ = ('StringEnum', 'CallbackRepr', 'block_indent', 'make_tempfile') + +# Note: `typing.AnyStr` deprecated since 3.13 +AnyStr = TypeVar('AnyStr', str, bytes) class _StrEnumBase(str, enum.Enum): @@ -49,7 +62,7 @@ def __str__(self) -> str: try: from enum import StrEnum as _StrEnum except ImportError: - if not typing.TYPE_CHECKING: # Don't confuse the typechecker + if not TYPE_CHECKING: # Don't confuse the typechecker _StrEnum = _StrEnumBase @@ -89,3 +102,258 @@ def _missing_(cls, value: object) -> Self | None: for name, instance in cls.__members__.items() } return members.get(value.casefold()) + + +class _ReprAttributes(TypedDict, total=False): + """ + Note: + We use this typed dict instead of directly supplying them in the + :py:meth:`CallbackRepr.__init__()` signature, because we don't + want to bother with the default values there. + """ + maxlevel: int + maxtuple: int + maxlist: int + maxarray: int + maxdict: int + maxset: int + maxfrozenset: int + maxdeque: int + maxstring: int + maxlog: int + maxother: int + fillvalue: str + indent: str | int | None + + +class CallbackRepr(Repr): + """ + :py:class:`reprlib.Repr` subclass to help with representing cleanup + callbacks, special-casing certain relevant object types (see + examples below). + + Example: + >>> from functools import partial + >>> from sys import version_info + + >>> class MyEnviron(dict): + ... def some_method(self) -> None: + ... ... + ... + >>> + >>> class MyRepr(CallbackRepr): + ... # Since we can't instantiate a new `os._Environ`, test + ... # the relevant method with a mock + ... repr_MyEnviron = CallbackRepr.repr__Environ + ... + >>> + >>> r = MyRepr(maxenv=3, maxargs=4, maxstring=15) + + Environ-dict formatting: + + >>> my_env = MyEnviron( + ... foo='1', + ... bar='2', + ... this_varname_is_long_but_isnt_truncated=( + ... "THIS VALUE IS TRUNCATED BECAUSE IT'S TOO LONG" + ... ), + ... baz='4', + ... ) + >>> print(r.repr(my_env)) + environ({'foo': '1', 'bar': '2', \ +'this_varname_is_long_but_isnt_truncated': 'THIS ... LONG', ...}) + + Partial-object formatting: + + >>> r.maxenv = 0 + >>> print(r.repr(my_env.some_method)) + + + Bound-method formatting: + + >>> r.maxargs = 0 + >>> callback_1 = partial(int, base=8) + >>> print(r.repr(callback_1)) + functools.partial(, ...) + + Indentation (Python 3.12+): + + >>> if version_info < (3, 12): + ... from pytest import skip + ... + ... skip( + ... '`Repr.indent` not available on {}.{},{}' + ... .format(*sys.version_info) + ... ) + + >>> r = MyRepr(maxenv=2, maxargs=4) + >>> r.indent = 2 + >>> callback_1 = partial(int, base=8) + >>> print(r.repr(callback_1)) + functools.partial( + , + base=8, + ) + + >>> callback_2 = partial(min, 5, 4, 3, 2, 1) + >>> r.indent = '----' + >>> print(r.repr(callback_2)) + functools.partial( + ----, + ----5, + ----4, + ----3, + ----2, + ----..., + ) + + >>> r.indent = ' ' + >>> r.maxenv = 2 + >>> print(r.repr(my_env.some_method)) + + """ + def __init__( + self, + *, + maxargs: int = 5, + maxenv: int = 3, + **kwargs: Unpack[_ReprAttributes] + ) -> None: + super().__init__() # kwargs are 3.12+ + valid_kwargs = ( + _ReprAttributes.__optional_keys__ + | _ReprAttributes.__required_keys__ + ) + for k, v in kwargs.items(): + if k in valid_kwargs: + setattr(self, k, v) + self.maxargs = maxargs + self.maxenv = maxenv + + def repr__Environ(self, env: os._Environ[AnyStr], level: int) -> str: + """ + Format :py:data:`os.environ` or :py:data:`os.environb`. + """ + get: Callable[[AnyStr], str] = partial(self.repr1, level=level-1) + # Truncate envvar values, but not their names + envvars = ['{!r}: {}'.format(k, get(v)) for k, v in env.items()] + return self._format_items(envvars, ('environ({', '})'), self.maxenv) + + def repr_method(self, method: MethodType, level: int) -> str: + """ + Format a :py:class:`types.MethodType`. + """ + instance = self.repr1(method.__self__, level-1) + func = getattr(method.__func__, '__qualname__', '?') + prefix, suffix = f'' + # Take care of possible multi-line reprs + return block_indent(instance, prefix) + suffix + + def repr_partial(self, ptl: partial, level: int) -> str: + """ + Format a :py:func:`functools.partial`. + """ + name = '{0.__module__}.{0.__qualname__}'.format(type(ptl)) + # The +1 is to account for `ptl.func` + return self._format_call( + level, (name + '(', ')'), self.maxargs + 1, + [ptl.func, *ptl.args], ptl.keywords, + ) + + def format_call(self, /, *args, **kwargs) -> str: + """ + Convenience method for Formatting a call a la + :py:meth:`inspect.BoundArguments.__str__`. + + Example: + >>> r = CallbackRepr(maxargs=3, maxlist=3) + >>> print(r.format_call( + ... [1, 2, 3, 4, 5], 'foo', spam=1, ham=2, + ... )) + ([1, 2, 3, ...], 'foo', spam=1, ...) + """ + return self._format_call( + self.maxlevel, ('(', ')'), self.maxargs, args, kwargs, + ) + + def _format_call( + self, + level: int, + delims: tuple[str, str], + maxargs: int, + args: Sequence[Any], + kwargs: Mapping[str, Any], + ) -> str: + get: Callable[[Any], str] = partial(self.repr1, level=level-1) + args = [get(arg) for arg in args] + args.extend('{}={}'.format(k, get(v)) for k, v in kwargs.items()) + return self._format_items(args, delims, maxargs) + + def _format_items( + self, + items: Collection[str], + delims: tuple[str, str], + maxlen: int | None = None, + ) -> str: + start, end = delims + if maxlen is not None and len(items) > maxlen: + items = list(items)[:maxlen] + ['...'] + indent_prefix: str | None = self._get_indent() + if indent_prefix is None or not items: + return '{}{}{}'.format(start, ', '.join(items), end) + return '\n'.join([ + start, *(indent(item + ',', indent_prefix) for item in items), end, + ]) + + if sys.version_info >= (3, 12): + # Note: `.indent` only available since 3.12 + def _get_indent(self) -> str | None: + indent = self.indent + if indent is None or isinstance(indent, str): + return indent + return ' ' * indent + else: + @staticmethod + def _get_indent() -> None: + return None + + +def block_indent(string: str, prefix: str, fill_char: str = ' ') -> str: + r""" + Example: + >>> string = 'foo\nbar\nbaz' + >>> print(string) + foo + bar + baz + >>> print(block_indent(string, '++++', '-')) + ++++foo + ----bar + ----baz + """ + width = len(prefix) + return prefix + indent(string, fill_char * width)[width:] + + +def make_tempfile(**kwargs) -> Path: + """ + Convenience wrapper around :py:func:`tempfile.mkstemp`, discarding + and closing the integer handle (which if left unattended causes + problems on some platforms). + + Note: + If for whatever reason the handle cannot be closed, the function + errors out and the tempfile is deleted. + """ + handle, fname = mkstemp(**kwargs) + path = Path(fname) + try: + os.close(handle) + return path + except Exception: + path.unlink() + raise diff --git a/line_profiler/rc/line_profiler.toml b/line_profiler/rc/line_profiler.toml index 6680c06a..7a7e3f23 100644 --- a/line_profiler/rc/line_profiler.toml +++ b/line_profiler/rc/line_profiler.toml @@ -94,6 +94,9 @@ preimports = true # - `prof-imports` (bool): # `--prof-imports` (true) or `--no-prof-imports` (false) prof-imports = false +# - `prof-child-procs` (bool): +# `--prof-child-procs` (true) or `--no-prof-child-procs` (false) +prof-child-procs = false # - Misc flags # - `verbose` (count): @@ -206,3 +209,44 @@ hits = 9 time = 12 perhit = 8 percent = 8 + +# XXX: --- Start of implementation details --- +# `line_profiler._child_process_profiling` settings + +[tool.line_profiler.child_processes.pth_files] + +# - `pth_files.prefix`, `.suffix` (str): +# Affixes to use for the stem of the name of the .pth file created +# Notes: +# - May be useful to tweak this in case of issues with .pth file +# precedence. +# - The PID of the main process is always suffixed on top of this +# suffix; this is so that .pth files can be identified by their +# "owning" processes and stale files can be cleaned up. In general +# end users should NOT count on the affixes appearing as-is in the +# final filename. +prefix = '_line_profiler-profiling-hook-' +suffix = '' + +# List of individual patches to apply to `multiprocessing` + +[tool.line_profiler.child_processes.multiprocessing.patches] + +# - `multiprocessing.patches.pool` (bool): +# Whether to patch `multiprocessing.pool.worker()` and `Pool` so that +# child processes report profiling stats to the parent as each task is +# completed +pool = true +# - `multiprocessing.patches.process` (bool): +# Whether to patch `multiprocessing.process.BaseProcess`, so that each +# child process write profiling output before exiting +process = true +# NOTE: for the best result, stick to the default and have both applied + +# - `multiprocessing.patches.logging` (bool): +# Whether to patch logging functions in `multiprocessing.util`, so +# that the internal logs of `multiprocessing` are teed to the session +# cache's debug logs +logging = false + +# XXX: --- End of implementation details --- diff --git a/line_profiler/toml_config.py b/line_profiler/toml_config.py index 781dc60d..236e86da 100644 --- a/line_profiler/toml_config.py +++ b/line_profiler/toml_config.py @@ -106,7 +106,10 @@ def get_subconfig( get_subtable(self.conf_dict, headers, allow_absence=allow_absence), ) new_subtable = [*self.subtable, *headers] - return type(self)(new_dict, self.path, new_subtable) + new_instance = type(self)(new_dict, self.path, new_subtable) + if copy: + new_instance = new_instance.copy() + return new_instance @classmethod def from_default(cls, *, copy: bool = True) -> ConfigSource: @@ -355,7 +358,8 @@ def iter_configs(dir_path): def get_subtable( - table: Mapping[K, Mapping], keys: Sequence[K], *, allow_absence: bool = True + table: Mapping[K, Mapping], keys: Sequence[K], *, + allow_absence: bool = True, ) -> Mapping: """ Arguments: diff --git a/pyproject.toml b/pyproject.toml index 3abd4137..daee1b2d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -126,3 +126,6 @@ docstring-code-format = false unused-ignore-comment = "ignore" unused-type-ignore-comment = "ignore" unresolved-import = "ignore" + +[tool.ty.terminal] +error-on-warning = false # No longer the default since 0.0.52 diff --git a/setup.py b/setup.py index 739d5d96..a0810cc6 100755 --- a/setup.py +++ b/setup.py @@ -314,7 +314,9 @@ def run_cythonize(force=False): setupkw['long_description_content_type'] = 'text/x-rst' setupkw['license'] = 'BSD' setupkw['packages'] = list(setuptools.find_packages()) - setupkw['py_modules'] = ['kernprof', 'line_profiler'] + setupkw['py_modules'] = [ + 'kernprof', 'line_profiler', '_line_profiler_hooks', + ] setupkw['python_requires'] = '>=3.10' setupkw['license_files'] = ['LICENSE.txt', 'LICENSE_Python.txt'] setupkw['package_data'] = {'line_profiler': ['py.typed', '*.pyi', '*.toml']} diff --git a/tests/test_child_procs.py b/tests/test_child_procs.py deleted file mode 100644 index 10be3f78..00000000 --- a/tests/test_child_procs.py +++ /dev/null @@ -1,233 +0,0 @@ -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 diff --git a/tests/test_child_procs/__init__.py b/tests/test_child_procs/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/test_child_procs/_test_child_procs_utils.py b/tests/test_child_procs/_test_child_procs_utils.py new file mode 100644 index 00000000..5bfd740a --- /dev/null +++ b/tests/test_child_procs/_test_child_procs_utils.py @@ -0,0 +1,2376 @@ +from __future__ import annotations + +import dataclasses +import enum +import inspect +import itertools +import multiprocessing.pool +import operator +import os +import pickle +import re +import shlex +import subprocess +import sys +import sysconfig +import threading +import traceback +import warnings +from abc import ABC, abstractmethod +from argparse import ArgumentError +from collections.abc import ( + Callable, Collection, Generator, Iterable, Mapping, + Sequence, Set, +) +from contextlib import ExitStack +from functools import lru_cache, partial, wraps +from io import BytesIO +from importlib import import_module +from multiprocessing.pool import ( # type: ignore + ExceptionWithTraceback as ExceptionHelper, +) +from numbers import Real +from pathlib import Path +from textwrap import dedent, indent +from time import monotonic +from types import MappingProxyType, ModuleType, TracebackType +from typing import ( + TYPE_CHECKING, Any, Generic, IO, Literal, Protocol, TypeVar, + cast, final, overload, +) +from typing_extensions import Concatenate, Self, ParamSpec +from uuid import uuid4 + +import pytest +import ubelt as ub + +from _line_profiler_hooks import load_pth_hook +from kernprof import main as kernprof_main +from line_profiler._child_process_profiling.cache import LineProfilingCache +from line_profiler._child_process_profiling.multiprocessing_patches import ( + _PATCHED_MARKER as MP_PATCHED_MARKER, _PATCHES as _MP_PATCHES, +) +from line_profiler.autoprofile.util_static import modpath_to_modname +from line_profiler.cleanup import Cleanup +from line_profiler.line_profiler import LineProfiler, LineStats +from line_profiler.toml_config import ConfigSource + + +__all__ = ( + 'DEBUG', 'DEFAULT_TIMEOUT', 'NOT_SUPPLIED', 'START_METHODS', + 'PATCH_SUMMARIES', + 'StartMethod', 'ModuleFixture', 'Params', + 'ResultMismatch', 'FuncCallTimeout', + 'preserve_object_attrs', 'preserve_targets', + 'cleanup_extra_pth_files', 'add_timeout', + 'CheckWarnings', + 'mp_patch_is_internal', 'get_mp_patches_toml_text', + 'summarize_mp_patches', 'filter_mp_patch_summary', + 'concat_command_line', 'run_subproc', + 'run_module', 'run_script', 'run_literal_code', 'check_tagged_line_nhits', + 'strip', 'search_cache_logs', +) + + +T = TypeVar('T') +T1 = TypeVar('T1') +T2 = TypeVar('T2') +TCtx_ = TypeVar('TCtx_') +PS = ParamSpec('PS') +C = TypeVar('C', bound=Callable[..., Any]) +StartMethod = Literal['spawn', 'fork', 'forkserver', 'dummy'] + +START_METHODS = set(multiprocessing.get_all_start_methods()) + +DEFAULT_TIMEOUT = 5 # Seconds +DEBUG = True + +# ======================== Misc. helper classes ======================== + + +class _NotSupplied(enum.Enum): + NOT_SUPPLIED = enum.auto() + + +NOT_SUPPLIED = _NotSupplied.NOT_SUPPLIED + + +class _GetAttr(Protocol): + """ + Function signature for functions that behave like + :py:func:`getattr``. + """ + @overload + def __call__(self, obj: Any, attr: str, /) -> Any: + ... + + @overload + def __call__(self, obj: Any, attr: str, default: Any, /) -> Any: + ... + + def __call__(self, *args): + ... + + +@final +class ResultMismatch(ValueError): + def __init__( + self, + expected: Any, + actual: Any | _NotSupplied = NOT_SUPPLIED, + _trunc_tb: int = 0, + ) -> None: + if actual == NOT_SUPPLIED: + msg = f'expected: {expected}' + else: + msg = f'expected {expected}, got {actual}' + super().__init__(msg) + self.expected = expected + self.actual = actual + self._trunc_tb = max(0, _trunc_tb) + + @classmethod + def compare( + cls, expected_: T1, actual_: T2, /, *, + comparator: Callable[[T1, T2], bool] = operator.eq, + expected: str | None = None, + actual: str | None = None, + ) -> None: + if comparator(expected_, actual_): + return + raise cls( + expected_ if expected is None else expected, + actual_ if actual is None else actual, + _trunc_tb=1, + ) + + @classmethod + def pytest_runtest_makereport( + cls, item: pytest.Item, call: pytest.CallInfo, + ) -> Any: + """ + Truncate the tracebacks of instances so that pytest outputs are + more useful and actually stops at the frame where the comparends + are shown. + """ + impl: Callable[..., Any] + impl = item.config.pluginmanager.subset_hook_caller( + 'pytest_runtest_makereport', [cls], + ) + make_report = partial(impl, item=item, call=call) + + xc = call.excinfo + if xc is None: + return make_report() + if not (isinstance(xc.value, cls) and xc.value._trunc_tb): + return make_report() + + tb_stack: list[TracebackType] = [xc.tb] + while tb_stack[-1].tb_next: + tb_stack.append(tb_stack[-1].tb_next) + if len(tb_stack) <= xc.value._trunc_tb: + return make_report() + tb_stack[-(xc.value._trunc_tb + 1)].tb_next = None + + del tb_stack # Help the GC + call.excinfo = xc.from_exception(xc.value.with_traceback(xc.tb)) + return make_report(call=call) + + +class FuncCallTimeout(RuntimeError): + """ + Error raised by the :py:func:`add_timeout` decorator. + """ + pass + + +class _CallableContextManager(ABC, Generic[TCtx_]): + debug: bool + + @abstractmethod + def __enter__(self) -> TCtx_: + ... + + @abstractmethod + def __exit__(self, *a, **k) -> Any: + ... + + def __call__(self, func: Callable[PS, T]) -> Callable[PS, T]: + """ + Wrap ``func()`` so that its calls always happen in the context + of the instance. + """ + @wraps(func) + def wrapper(*args: PS.args, **kwargs: PS.kwargs) -> T: + with self: + return func(*args, **kwargs) + + return wrapper + + def _debug(self, msg: str, **kwargs) -> None: + if not self.debug: + return + header = f'{os.environ["PYTEST_CURRENT_TEST"]}: {type(self).__name__}' + print(f'{header}: {msg}', **kwargs) + + +# ========================== Misc. functions =========================== + + +def strip(s: str) -> str: + return dedent(s).strip('\n') + + +def import_target(target: str) -> Any: + try: + return import_module(target) + except ImportError: # Not a module + assert '.' in target + module, _, attr = target.rpartition('.') + return getattr(import_module(module), attr) + + +def search_cache_logs( + cache: LineProfilingCache, + expecting_logs: bool, + patterns: Mapping[str, bool] | Collection[str], + match_individual_messages: bool = False, + flags: int = 0, +) -> None: + entries = cache._gather_debug_log_entries() + ResultMismatch.compare( + expecting_logs, bool(entries), + expected='logs' if expecting_logs else 'no logs', + actual=repr(entries) if entries else 'nothing', + ) + if not expecting_logs: + return + text_chunks: list[str] = [entry.to_text() for entry in entries] + if not match_individual_messages: + text_chunks = ['\n'.join(text_chunks)] + if isinstance(patterns, Mapping): + to_match: dict[str, bool] = { + str(pat): bool(should_match) + for pat, should_match in patterns.items() + } + else: + to_match = dict.fromkeys(patterns, True) + for pat, should_match in to_match.items(): + pattern = re.compile(pat, flags) + matches = [pattern.search(chunk) for chunk in text_chunks] + if any(matches) == should_match: + continue + error_msg = ( + f'pattern {pattern!r} to {"" if should_match else "not "}match ' + f'{cache!r}\'s logs: {text_chunks!r}' + ) + if should_match: + raise ResultMismatch(error_msg, 'no matches') + else: + raise ResultMismatch(error_msg, [m for m in matches if m]) + + +# ================ `pytest` stuff: fixtures and markers ================ + + +@dataclasses.dataclass +class ModuleFixture: + """ + Convenience wrapper around a Python source file which represents an + importable module. + + Attributes: + path (Path): + File path to the module content. + monkeypatch (pytest.MonkeyPatch): + Monkey-patch object. + dependencies (Collection[ModuleFixture[Any]]): + Other :py:class:`ModuleFixture` objects this instance + depends on. + build_command_line \ +(Callable[Concatenate[bool, StartMethod | None, ...], Sequence[str]] \ +| None): + If the module is supposed to be callable (via + ``python -m ...``), a callable which takes the leading + positional arguments + ``fail: bool, start_method: Literal[...] | None`` followed by + arbitrary arguments, and constructs the command-line + arguments to pass to the module; + note that: + + - ``fail=True`` should result in command-line arguments + causing the executed module to fail. + + - ``start_method`` should select the + :py:mod:`multiprocessing` start method, where ``dummy`` + should cause the analogous APIs at + :py:mod:`multiprocessing.dummy` to be used instead; + if it is :py:const:`None`, the OS default should be used. + get_output (Callable[..., Any] | None): + If the module is supposed to be callable (via + ``python -m ...``), a callable which takes arbitrary + arguments (consistent with those to + ``build_command_line()``) and returns the expected output of + the executed module. + """ + path: Path + monkeypatch: pytest.MonkeyPatch + dependencies: Collection[ModuleFixture] = () + build_command_line: ( + Callable[Concatenate[bool, StartMethod | None, ...], Sequence[str]] + | None + ) = None + get_output: Callable[..., Any] | None = None + + def install( + self, *, + local: bool = False, children: bool = False, deps_only: bool = False, + ) -> None: + """ + Set the module at :py:attr:`~.path` up to be importable. + + Args: + local (bool): + Make it importable for the CURRENT process (via + :py:data:`sys.path`). + children (bool): + Make it importable for CHILD processes (via + ``os.environ['PYTHONPATH']``). + deps_only (bool): + If true, only does the equivalent setup for + dependencies. + """ + for dep in self.dependencies: + dep.install(local=local, children=children) + if deps_only: + return + path = str(self.path.parent) + if local: + self.monkeypatch.syspath_prepend(path) + if children: + self.monkeypatch.setenv('PYTHONPATH', path, prepend=os.pathsep) + + def _import_module_helper(self) -> Generator[ModuleType, None, None]: + def iter_module_names( + module: ModuleFixture, + ) -> Generator[str, None, None]: + yield module.name + for dep in module.dependencies: + yield from iter_module_names(dep) + + self.install(local=True, children=True) + try: + yield import_module(self.name) + finally: + for name in set(iter_module_names(self)): + sys.modules.pop(name, None) + + @staticmethod + def propose_name(prefix: str) -> Generator[str, None, None]: + """ + Propose a valid module name that isn't already occupied. + """ + while True: + name = '_'.join([prefix] + str(uuid4()).split('-')) + if name not in sys.modules: + assert name.isidentifier() + yield name + + @property + def name(self) -> str: + return self.path.stem + + +@final +@dataclasses.dataclass +class Params: + """ + Convenience wrapper around :py:func:`pytest.mark.parametrize`. + """ + params: tuple[str, ...] + values: list[tuple[Any, ...]] + defaults: tuple[Any, ...] + + def __post_init__(self) -> None: + n = len(self.params) + assert all(p.isidentifier() for p in self.params) # Validity + assert len(set(self.params)) == n # Uniqueness + assert len(self.defaults) == n # Consistency + self.values = list(self._unique(self.values)) + assert all(len(v) == n for v in self.values) + + def __mul__(self, other: Self) -> Self: + """ + Form a Cartesian product between the two instances with disjoint + :py:attr:`~.params`, like stacking the + :py:func:`pytest.mark.parametrize `decorators. + + Example: + >>> p1 = Params.new(('a', 'b'), [(0, 0), (1, 2), (3, 4)], + ... defaults=(1, 2)) + >>> p2 = Params.new('c', [0, 5, 6]) + >>> p1 * p2 # doctest: +NORMALIZE_WHITESPACE + Params(params=('a', 'b', 'c'), + values=[(0, 0, 0), (0, 0, 5), (0, 0, 6), + (1, 2, 0), (1, 2, 5), (1, 2, 6), + (3, 4, 0), (3, 4, 5), (3, 4, 6)], + defaults=(1, 2, 0)) + """ + assert not set(self.params) & set(other.params) + return type(self)( + self.params + other.params, + [sv + ov for sv in self.values for ov in other.values], + self.defaults + other.defaults, + ) + + def __add__(self, other: Self) -> Self: + """ + Concatenate two instances: + + - For parameters appearing in both, their lists of values are + concatenated. + + - For parameters appearing in either instance, the missing + values are taken from the other instance's + :py:attr:`~.defaults`. + + Note: + In the case of clashes, the :py:attr:`~.defaults` and the + order of the :py:attr:`~.params` of ``self`` (the left + operand) take precedence. + + Example: + >>> p1 = Params.new(('a', 'b', 'c'), + ... [(0, 0, 0), # defaults + ... (1, 2, 3), (4, 5, 6)]) + >>> p2 = Params.new(('c', 'd'), [(7, 8), (9, 10)], + ... defaults=(-1, -1)) + >>> p1 + p2 # doctest: +NORMALIZE_WHITESPACE + Params(params=('a', 'b', 'c', 'd'), + values=[(0, 0, 0, -1), + (1, 2, 3, -1), + (4, 5, 6, -1), + (0, 0, 7, 8), + (0, 0, 9, 10)], + defaults=(0, 0, 0, -1)) + """ + self_defaults = dict(zip(self.params, self.defaults)) + other_defaults = dict(zip(other.params, other.defaults)) + new_params = tuple(self._unique(self.params + other.params)) + + defaults = {**other_defaults, **self_defaults} + new_defaults_tuple = tuple(defaults[p] for p in new_params) + + new_values: list[tuple[Any, ...]] = [] + for old_values, old_params in [ + (self.values, self.params), (other.values, other.params), + ]: + indices: list[ + tuple[Literal[True], int] | tuple[Literal[False], str] + ] = [ + (True, old_params.index(p)) if p in old_params else (False, p) + for p in new_params + ] + new_values.extend( + tuple( + ( + value[cast(int, index)] + if available else + defaults[cast(str, index)] + ) for available, index in indices + ) + for value in old_values + ) + return type(self)(new_params, new_values, new_defaults_tuple) + + def sorted( + self, + *, + sort_by: Sequence[str] | None = None, + sortable_types: type[Any] | tuple[type[Any], ...] = (Real, str, bytes), + ) -> Self: + """ + Sort by parametrization values. + + Args: + sort_by (Sequence[str] | None): + Column names to sort by; default is to sort by all + sortable params. + sortable_types (type[Any] | tuple[type[Any], ...]): + Type(s) where if a param has all its values being + instances thereof (excl. :py:const:`None`s), said param + is considered sortable. + + Returns: + New instance + """ + def sort_key(obj: Any) -> tuple[bool, str, Any]: + type_name = '{0.__module__}.{0.__qualname__}'.format(type(obj)) + return (obj is None), type_name, obj + + if sort_by is None: + sort_by = self.params + sortable_columns: set[str] = { + param for param, *values in zip(self.params, *self.values) + if all(isinstance(v, sortable_types) or v is None for v in values) + } + sorted_column_indices: tuple[int, ...] = tuple( + i for i, param in enumerate(sort_by) if param in sortable_columns + ) + + if sorted_column_indices: + new_values = sorted( + self.values, key=lambda vtuple: tuple( + sort_key(vtuple[i]) for i in sorted_column_indices + ), + ) + else: # Fallback + new_values = self.values.copy() + return type(self)(self.params, new_values, self.defaults) + + def drop_params(self, params: Collection[str] | str) -> Self: + """ + Return a new instance with the named ``params`` dropped; params + that don't match :py:attr:`.params` are ignored. + + Example: + >>> p = Params.new(('a', 'b'), [(1, 2), (3, 4)]) + >>> p.drop_params('a') + Params(params=('b',), values=[(2,), (4,)], defaults=(2,)) + >>> assert p.drop_params(['c', 'd']) == p + """ + def drop(t: tuple[T, ...]) -> tuple[T, ...]: + return tuple(item for i, item in enumerate(t) if i not in dropped) + + if isinstance(params, str): + params = params, + dropped = {i for i, p in enumerate(self.params) if p in params} + return type(self)( + drop(self.params), + [drop(pvalues) for pvalues in self.values], + drop(self.defaults), + ) + + @overload + def split_on_params( + self, params: tuple[str, ...], *, drop_split_params: bool = True, + ) -> dict[tuple[Any, ...], Self]: + ... + + @overload + def split_on_params( + self, params: str, *, drop_split_params: bool = True, + ) -> dict[Any, Self]: + ... + + def split_on_params( + self, params: tuple[str, ...] | str, *, drop_split_params: bool = True, + ) -> dict[tuple[Any, ...], Self] | dict[Any, Self]: + """ + Return new instances splitting on the values of the named + ``params``; params that don't match :py:attr:`.params` results + in an error. + + Example: + >>> p = Params.new(('a', 'b', 'c'), + ... [(1, 2, True), + ... (1, 2, False), + ... (3, 4, True)]) + + >>> p.split_on_params('a') # doctest: +NORMALIZE_WHITESPACE + {1: Params(params=('b', 'c'), + values=[(2, True), (2, False)], + defaults=(2, True)), + 3: Params(params=('b', 'c'), + values=[(4, True)], + defaults=(2, True))} + + >>> p.split_on_params( # doctest: +NORMALIZE_WHITESPACE + ... ('a', 'b'), + ... ) + {(1, 2): Params(params=('c',), + values=[(True,), (False,)], + defaults=(True,)), + (3, 4): Params(params=('c',), + values=[(True,)], + defaults=(True,))} + + >>> p.split_on_params( # doctest: +NORMALIZE_WHITESPACE + ... 'a', drop_split_params=False, + ... ) + {1: Params(params=('a', 'b', 'c'), + values=[(1, 2, True), (1, 2, False)], + defaults=(1, 2, True)), + 3: Params(params=('a', 'b', 'c'), + values=[(3, 4, True)], + defaults=(1, 2, True))} + + >>> p.split_on_params( # doctest: +NORMALIZE_WHITESPACE + ... ('c', 'd'), + ... ) + Traceback (most recent call last): + ... + ValueError: params = ('c', 'd'): + these params not found: ['d'] + """ + if isinstance(params, str): + params = params, + unpack = True + else: + unpack = False + nonexistent = sorted(set(params) - set(self.params)) + if nonexistent: + raise ValueError( + f'params = {params!r}: these params not found: {nonexistent!r}' + ) + split_params: dict[tuple[Any, ...], list[tuple[Any, ...]]] = {} + indices = tuple(self.params.index(p) for i, p in enumerate(params)) + for pvalues in self.values: + key = tuple(pvalues[i] for i in indices) + split_params.setdefault(key, []).append(pvalues) + new = partial(type(self), params=self.params, defaults=self.defaults) + instances: dict[tuple[Any, ...], Self] = { + key: new(values=values) for key, values in split_params.items() + } + if drop_split_params: + instances = { + key: instance.drop_params(params) + for key, instance in instances.items() + } + if not unpack: + return instances + return {key[0]: instance for key, instance in instances.items()} + + def __call__(self, func: C) -> C: + """ + Mark a callable as with :py:func:`pytest.mark.parametrize`. + """ + # Note: `pytest` automatically assumes single-param values to + # be unpacked, so comply here + if len(self.params) == 1: + marker = pytest.mark.parametrize( + self.params[0], [v[0] for v in self.values], + ) + else: + marker = pytest.mark.parametrize(self.params, self.values) + return marker(func) + + @staticmethod + def _unique(items: Iterable[T]) -> Generator[T, None, None]: + seen: set[T] = set() + for item in items: + if item in seen: + continue + seen.add(item) + yield item + + @overload + @classmethod + def new( + cls, + params: Sequence[str] | str, + values: Sequence[Sequence[Any]], + defaults: Sequence[Any] | _NotSupplied = NOT_SUPPLIED, + ) -> Self: + ... + + @overload + @classmethod + def new( + cls, + params: str, + values: Sequence[Any], + defaults: Any | _NotSupplied = NOT_SUPPLIED, + ) -> Self: + ... + + @classmethod + def new( + cls, + params: Sequence[str] | str, + values: Sequence[Sequence[Any]] | Sequence[Any], + defaults: Sequence[Any] | Any | _NotSupplied = NOT_SUPPLIED, + ) -> Self: + """ + Instantiator more akin to :py:func:`pytest.mark.parametrize`: + + - ``params`` can be provided as a comma-separated string + + - Single parameters can be unpacked (singular param-name string + and param-value sequences) + + - If ``defaults`` are not given, it is implicitly set to the + FIRST item in ``values``. + """ + if isinstance(params, str): + param_list: tuple[str, ...] = tuple( + p.strip() for p in params.split(',') + ) + unpacked = len(param_list) == 1 + else: + param_list = tuple(params) + unpacked = False + if defaults == NOT_SUPPLIED: + defaults, *_ = values + if unpacked: + default_values: tuple[Any, ...] = defaults, + value_tuple_list: list[tuple[Any, ...]] = [(v,) for v in values] + else: + default_values = tuple(defaults) # type: ignore[arg-type] + value_tuple_list = [tuple(v) for v in values] + return cls(param_list, value_tuple_list, default_values) + + +# ============================= Failsafes ============================== + + +class preserve_object_attrs(_CallableContextManager[dict[str, Any]]): + """ + Protect attributes on a single supplied object. + """ + def __init__( + self, obj: Any, attrs: Collection[str], *, + static: bool = True, debug: bool = DEBUG, + ) -> None: + self.obj = obj + self.attrs = set(attrs) + self._callbacks: list[Callable[[], None]] = [] + self.debug = debug + self.static = static + + def __enter__(self) -> dict[str, Any]: + def get_repr(attr: str) -> str: + try: + value = get_attribute(self.obj, attr) + except ValueError: + return '' + else: + return repr(value) + + def delete(attr: str) -> None: + try: + self._debug('Deleted attr `.{} = {}` on `{!r}`'.format( + attr, get_repr(attr), self.obj, + )) + delattr(self.obj, attr) + except AttributeError: + pass + + def reset(attr: str, value: Any) -> None: + self._debug('Reset attr `.{} = {} -> {!r}` on `{!r}`'.format( + attr, get_repr(attr), value, self.obj, + )) + setattr(self.obj, attr, value) + + if self.static: + get_attribute: _GetAttr = inspect.getattr_static + else: + get_attribute = getattr + + result: dict[str, Any] = {} + for attr in self.attrs: + old = get_attribute(self.obj, attr, NOT_SUPPLIED) + if old is NOT_SUPPLIED: + callback = partial(delete, attr) + else: + callback = partial(reset, attr, old) + result[attr] = old + self._callbacks.append(callback) + return result + + def __exit__(self, *_, **__) -> None: + xc: Exception | None = None + for callback in self._callbacks[::-1]: + try: + callback() + except Exception as e: + xc = e + rep_xc = type(xc).__name__ + if str(xc): + rep_xc = f'{rep_xc}: {xc}' + self._debug(f'Callback {callback} failed: {rep_xc}') + if xc is not None: + raise xc + + +class preserve_targets(_CallableContextManager[dict[str, dict[str, Any]]]): + """ + Protect attributes on multiple target objects, which are resolved at + context entry. If ``verify`` is true, check that the targets have + all been restored to their original value on context exit, and raise + an :py:class:`AssertionError` if that is not the case (before + restoring them anyways). + + Example: + >>> import builtins + >>> from functools import partial, wraps + >>> from line_profiler.curated_profiling import ( + ... CuratedProfilerContext, + ... ) + >>> from line_profiler import line_profiler + + >>> assert not hasattr(CuratedProfilerContext, 'foo') + >>> old_main = line_profiler.main + >>> new_context = partial(preserve_targets, debug=False) + + >>> def foo(_) -> None: + ... pass + + >>> @wraps(old_main) + ... def main(*a, **k): + ... return old_main(*a, **k) + + Basic use: + + >>> preserved = { + ... 'line_profiler.curated_profiling' + ... '.CuratedProfilerContext': {'foo'}, + ... 'line_profiler.line_profiler': {'main'}, + ... } + >>> with new_context(preserved) as old: + ... assert old == { + ... 'line_profiler.curated_profiling' + ... '.CuratedProfilerContext': {'foo': NOT_SUPPLIED}, + ... 'line_profiler.line_profiler': {'main': old_main}, + ... } + ... CuratedProfilerContext.foo = foo + ... line_profiler.main = main + ... print('ok') + ... + ok + >>> assert not hasattr(CuratedProfilerContext, 'foo') + >>> assert old_main is \ +old['line_profiler.line_profiler']['main'] + >>> assert old_main is line_profiler.main + >>> assert main is not line_profiler.main + + With ``verify``: + + >>> preserved = {'builtins': {'foo', 'bar'}} + >>> expected = { + ... 'builtins': {'foo': NOT_SUPPLIED, 'bar': NOT_SUPPLIED}, + ... } + >>> ctx = new_context(preserved, verify=True) + >>> with ctx as old: + ... assert old == expected + ... builtins.foo = 1 + ... del builtins.foo + ... print('ok') + ok + >>> with ctx as old: # doctest: +ELLIPSIS + ... assert old == expected + ... builtins.foo = 1 + ... # `bar` not restored -> `.__exit__()` error out + ... builtins.bar = 2 + ... del builtins.foo + ... print('ok, but...') + ok, but... + Traceback (most recent call last): + ... + AssertionError: Compared `builtins.bar` \ +(old:...NOT_SUPPLIED...; new: 2...)... + + >>> assert not hasattr(builtins, 'foo') + >>> assert not hasattr(builtins, 'bar') # Still cleaned up + """ + def __init__( + self, targets: Mapping[str, Collection[str]] | None = None, *, + static: bool = True, debug: bool = DEBUG, verify: bool = False, + ) -> None: + self.targets = targets + self._ctxs: list[tuple[ExitStack, dict[str, dict[str, Any]]]] = [] + self.static = static + self.debug = debug + self.verify = verify + + def __enter__(self) -> dict[str, dict[str, Any]]: + stack = ExitStack() + old: dict[str, Any] = {} + for target, attrs in self.targets.items(): + old[target] = stack.enter_context(preserve_object_attrs( + import_target(target), attrs, + debug=self.debug, static=self.static, + )) + self._ctxs.append((stack, old)) + # Decouple the stored and returned maps for the original values + # in case the user messes around with the latter + old_copy = {key: value.copy() for key, value in old.items()} + return old_copy + + def __exit__(self, *_, **__) -> None: + stack, old = self._ctxs.pop() + try: + if self.verify: + self.compare_with_current_values(old, verbose=self.debug) + finally: + stack.close() + + @staticmethod + def fetch_current_values( + targets: Mapping[str, Collection[str]], static: bool = True, + ) -> dict[str, dict[str, Any]]: + result: dict[str, dict[str, Any]] = {} + na = NOT_SUPPLIED + if static: + get: _GetAttr = inspect.getattr_static + else: + get = getattr + for target, attrs in targets.items(): + obj = import_target(target) + result[target] = {attr: get(obj, attr, na) for attr in attrs} + return result + + @classmethod + def compare_with_current_values( + cls, + old: Mapping[str, Mapping[str, Any]], + comparator: Callable[[Any, Any], bool] = operator.is_, + assert_true: bool | Mapping[str, Mapping[str, bool]] = True, + static: bool = True, + verbose: bool = True, + ) -> dict[str, dict[str, bool]]: + def get_from_mapping(target: str, attr: str) -> bool: + if TYPE_CHECKING: + assert isinstance(assert_true, Mapping) + return assert_true[target][attr] + + def get_from_boolean(*_, **__) -> bool: + return True + + if isinstance(assert_true, Mapping): + get_expected: Callable[[str, str], bool] = get_from_mapping + else: + get_expected = get_from_boolean + + result: dict[str, dict[str, bool]] = {} + new = cls.fetch_current_values(old, static) + failures: list[str] = [] + for target, old_values in old.items(): + new_values = new[target] + cmp_results = result[target] = {} + for attr, old_value in old_values.items(): + if verbose: + print(f'Checking: {target}.{attr}') + new_value = new_values[attr] + cmp_results[attr] = cmp_result = comparator( + new_value, old_value, + ) + format_msg = partial( + '{}: {}'.format, + f'Compared `{target}.{attr}` ' + f'(old: {old_value!r} @ {id(old_value):#x}; ' + f'new: {new_value!r} @ {id(new_value):#x})', + ) + expected_result = get_expected(target, attr) + if assert_true: + if cmp_result == expected_result: + message = format_msg( + f'comparison result with {comparator!r} is ' + f'{cmp_result} (as expected)' + ) + else: + message = format_msg( + f'expected comparison with {comparator!r} to ' + f'return {expected_result}, got {cmp_result}' + ) + failures.append(message) + else: + message = format_msg( + f'comparison result with {comparator!r}: {cmp_result}' + ) + if verbose: + print(message) + assert (not failures), '\n'.join(failures) + return result + + @property + def targets(self) -> dict[str, set[str]]: + if self._targets is None: + # Defer resolution to context entrance + self.targets = PATCH_SUMMARIES['maximal'] + assert self._targets is not None + return self._targets + + @targets.setter + def targets(self, targets: Mapping[str, Collection[str]] | None) -> None: + if targets is None: + self._targets = None + return + self._targets = { + target: set(attrs) for target, attrs in targets.items() + } + + +class cleanup_extra_pth_files(_CallableContextManager[frozenset[str]]): + r""" + Protect the venv where the test is run from temporary .pth files + created by :py:class:`.LineProfilingCache`. + + Example: + >>> from collections.abc import Generator + >>> from os import getpid + >>> from pathlib import Path + >>> from sysconfig import get_path + >>> from uuid import uuid4 + >>> import pytest + + >>> def _propose_alnum() -> Generator[str, None, None]: + ... while True: + ... uuid = str(uuid4()) + ... yield from uuid.split('-') + + >>> def _propose_names( + ... glob_pattern: str, + ... ) -> Generator[str, None, None]: + ... gen_alnum = _propose_alnum() + ... chunks = glob_pattern.split('*') + ... while True: + ... result: list[str] = [chunks[0]] + ... for chunk in chunks[1:]: + ... result.append(next(gen_alnum)) + ... result.append(chunk) + ... yield ''.join(result) + + >>> def get_new_pth(glob_pattern: str) -> Path: + ... gen_names = _propose_names(glob_pattern) + ... while True: + ... path = pth_loc / next(gen_names) + ... if not path.exists(): + ... return path + + >>> pth_loc = Path(get_path('purelib')) + >>> pth_from_cache = get_new_pth( + ... f'_line_profiler-profiling-hook-*-ppid-{getpid()}.pth', + ... ) + >>> pth_unrelated = get_new_pth(f'my-pth-*.pth') + + >>> try: + ... with cleanup_extra_pth_files(config=False, debug=False): + ... pth_from_cache.write_text('import sys') and None + ... pth_unrelated.write_text('import sys') and None + ... assert pth_from_cache.exists() + ... assert pth_unrelated.exists() + ... # This matches the pattern of filenames created by the + ... # cache and is thus dealt with + ... assert not pth_from_cache.exists() + ... # This is unrealated and is thus left alone + ... assert pth_unrelated.exists() + ... except PermissionError: + ... pytest.skip('Can\'t write to the pure-lib directory') + ... finally: # Cleanup + ... for pth in pth_from_cache, pth_unrelated: + ... pth.unlink(missing_ok=True) + """ + def __init__( + self, + config: os.PathLike[str] | str | bool | None = None, + debug: bool = DEBUG, + ) -> None: + self.debug = debug + pattern = self._get_glob_pattern(config) + self._get_files = partial(self.get_pth_files, pattern) + + def __enter__(self) -> frozenset[str]: + self.old = self._get_files() + return self.old + + def __exit__(self, *_, **__) -> None: + for new_pth_file in self._get_files() - self.old: + self._debug(f'Deleting stray .pth file: {new_pth_file!r}') + (self._get_path() / new_pth_file).unlink(missing_ok=True) + del self.old + + @classmethod + def get_pth_files( + cls, pattern: str | None = None, name_only: bool = True, + ) -> frozenset[str]: + if pattern is None: + pattern = cls._get_glob_pattern() + return frozenset( + pth.name if name_only else str(pth) + for pth in cls._get_path().glob(pattern) + ) + + @staticmethod + def _get_path() -> Path: + """ + Note: + This doesn't go through the entire .pth-location search path + (see :py:meth:`.LineProfilingCache.write_pth_hook`), but + suffices for the testing environment. + """ + return Path(sysconfig.get_path('purelib')) + + @staticmethod + def _get_glob_pattern( + config: os.PathLike[str] | str | bool | None = None, + ) -> str: + cmap = ( + ConfigSource.from_config(config) + .get_subconfig('child_processes', 'pth_files') + .conf_dict + ) + prefix, suffix_template = LineProfilingCache._normalize_pth_affixes( + cmap['prefix'], cmap['suffix'], + ) + pattern = '{}*{}'.format(prefix, suffix_template.format(os.getpid())) + assert pattern.endswith('.pth') + return pattern + + +def _cleanup_profiling_in_current_thread() -> None: + """ + Disable all active profiler instances on the current thread, so that + we don't trip over ourselves if a new thread down the road happens + to reuse the thread ID. + + Note: + The thread-local profiler state is supposed to be handled by + :py:mod:`line_profiler._threading_patches` and + :py:class:`.CuratedProfilerContext`, but since a test function + decorated with :py:deco:`.add_timeout` will be isolated in a new + thread BEFORE the fixture setting up such management + (:py:func:`curated_profiler`) is invoked, we don't get that + managing and will have to deal with profilers ourselves. + """ + class Manager(Protocol): + @property + def active_instances(self) -> set[LineProfiler]: + ... + + lp_managers = cast( + dict[int, Manager], getattr(LineProfiler, '_managers', {}), + ) + thread_id = threading.get_ident() + if thread_id not in lp_managers: + return + instances = lp_managers[thread_id].active_instances + for prof in set(instances): + count = cast(int, getattr(prof, 'enable_count', 0)) + for _ in range(count): + prof.disable_by_count() + if count: + print('Disabled', prof, 'in thread', hex(thread_id)) + # Removed after the last `.disable_by_count()` + assert prof not in instances + + +@overload +def add_timeout( + func: Callable[PS, T], *, + timeout: float = DEFAULT_TIMEOUT, + name: str | None = None, +) -> Callable[PS, T]: + ... + + +@overload +def add_timeout( + func: None = None, *, + timeout: float = DEFAULT_TIMEOUT, + name: str | None = None, +) -> Callable[[Callable[PS, T]], Callable[PS, T]]: + ... + + +def add_timeout( + func: Callable[PS, T] | None = None, *, + timeout: float = DEFAULT_TIMEOUT, + name: str | None = None, +) -> Callable[PS, T] | Callable[[Callable[PS, T]], Callable[PS, T]]: + """ + Decorate the test function so that it is run in another thread and + can be timed out; if a ``name`` is provided, it will be set to the + thread. + + Example: + >>> from time import sleep + >>> from threading import enumerate as enum_threads + >>> from uuid import uuid4 + + >>> thread_name = f'my_func-{uuid4()}' + >>> sleep_loop_interval = .0625 + >>> test_is_over = False + + >>> @add_timeout(timeout=.5, name=thread_name) + ... def my_func( + ... n: int, delay: float = 1, error: bool = False, + ... ) -> list[int]: + ... # Sleep in pulses here so the doctest (+ cleanup) + ... # doesn't take 10s if we have `delay=10` + ... nloops, remainder = divmod(delay, sleep_loop_interval) + ... intervals = [sleep_loop_interval] * int(nloops) + ... if remainder: + ... intervals.append(remainder) + ... for interval in intervals: + ... if test_is_over: + ... break + ... sleep(interval) + ... if error: + ... raise RuntimeError('my error message') + ... return list(range(n)) + + Normal execution: + + >>> my_func(3, 0) + [3] + [0, 1, 2, 3] + + Erroring out: + + >>> my_func(3, 0, error=True) + Traceback (most recent call last): + ... + RuntimeError: my error message + + Timing out: + + >>> my_func(4, delay=5) # doctest: +NORMALIZE_WHITESPACE + Traceback (most recent call last): + ... + test_child_procs._test_child_procs_utils.FuncCallTimeout: + my_func(4, delay=5): timed out after 0.5 s + + Note: + When a call is timed out, the background (daemonized) thread + handling it may still be running; it is the caller's + responsibility to locate said thread and SOMEHOW terminate it. + The ``name`` parameter is to facilitate that. + + >>> test_is_over = True # Tell `my_func()` to stop sleeping... + >>> my_thread = next( + ... t for t in enum_threads() if t.name == thread_name + ... ) + >>> my_thread.join() # ... or it may pollute the thread pool + """ + if func is None: + return cast( + Callable[[Callable[PS, T]], Callable[PS, T]], + partial(add_timeout, timeout=timeout, name=name), + ) + + @wraps(func) + def worker( + fobj: IO[bytes], /, *args: PS.args, **kwargs: PS.kwargs + ) -> None: + try: + try: + result = True, func(*args, **kwargs) + except Exception as e: + result = False, ExceptionHelper(e, e.__traceback__) + # Do this instead of directly using `pickle.dump(..., fobj)` + # so that pickling errors and file-handle-related errors are + # handled separately + serialized = pickle.dumps(result, protocol=pickle.HIGHEST_PROTOCOL) + try: + fobj.write(serialized) + fobj.flush() + except Exception: + # Since this is run in a daemon thread, by the time this + # write happens the main thread could've already timed + # out and destroyed `fobj`... in that case just + # gracefully exit + pass + finally: + # See the docstring for that function to see why this is + # needed. + _cleanup_profiling_in_current_thread() + + @wraps(func) + def wrapper(*args: PS.args, **kwargs: PS.kwargs) -> T: + with BytesIO() as bio: + thread = new_thread(args=(bio, *args), kwargs=kwargs) + thread.start() + thread.join(timeout) + if not thread.is_alive(): + successful, result = pickle.loads(bio.getvalue()) + if successful: + return result + assert isinstance(result, Exception) + raise result + args_repr = [repr(a) for a in args] + args_repr.extend(f'{k}={v!r}' for k, v in kwargs.items()) + name = getattr(func, '__name__', repr(func)) + call_repr = f'{name}({", ".join(args_repr)})' + msg = f'{call_repr}: timed out after {timeout:.2g} s' + raise FuncCallTimeout(msg) + + new_thread = partial( + threading.Thread, target=worker, name=name, daemon=True, + ) + + return wrapper + + +# =================== Warning handling/verification ==================== + + +class _WarningInfo(Protocol): + @property + def message(self) -> str | Warning: + ... + + @property + def category(self) -> type[Warning]: + ... + + @property + def filename(self) -> str: + ... + + @property + def lineno(self) -> int: + ... + + @property + def line(self) -> str | None: + ... + + +@dataclasses.dataclass +class _WarningMatcher: + message: str | None = None + category: type[Warning] | None = None + module: str | None = None + lineno: int | None = None + _filters: dict[str, Callable[[Any], Any]] = dataclasses.field( + repr=False, init=False, default_factory=dict, + ) + + def __post_init__(self) -> None: + if self.message is not None: + self._filters['message'] = partial( + self._check_message, re.compile(self.message), + ) + if self.category is not None: + self._filters['category'] = partial( + self._check_category, self.category, + ) + if self.module is not None: + self._filters['filename'] = partial( + self._check_module, re.compile(self.module), + ) + if self.lineno is not None: + self._filters['lineno'] = partial(operator.eq, self.lineno) + + def __repr__(self) -> str: + fields: dict[str, Any] = { + field.name: getattr(self, field.name, None) + for field in dataclasses.fields(self) + if field.repr + } + return '{}({})'.format( + type(self).__name__, + ', '.join( + f'{k}={v!r}' for k, v in fields.items() if v is not None + ), + ) + + def match(self, info: _WarningInfo) -> bool: + for field, check in self._filters.items(): + if not check(getattr(info, field)): + return False + return True + + @staticmethod + def _check_message( + msg_regex: re.Pattern, msg: str | Warning, + ) -> re.Match | None: + if not isinstance(msg, str): + msg = str(msg) + return msg_regex.match(msg) + + @staticmethod + def _check_category(parent: type[Any], maybe_child: type[Any]) -> bool: + try: + return issubclass(maybe_child, parent) + except Exception: + return False + + @staticmethod + def _check_module( + module_regex: re.Pattern, filename: str, + ) -> re.Match | None: + module = modpath_to_modname(filename, hide_main=False, hide_init=False) + return module_regex.match(module) + + +@dataclasses.dataclass +class _WarningContext: + catch_warnings: warnings.catch_warnings = dataclasses.field( + default_factory=partial(warnings.catch_warnings, record=True) + ) + reissue_warnings: bool = False + format_warnings: bool = False + checks: list[tuple[_WarningMatcher, bool]] = dataclasses.field( + default_factory=list, + ) + filters: list[_WarningMatcher] = dataclasses.field(default_factory=list) + _reissue_filtered_warnings: bool | None = dataclasses.field( + init=False, default=None, + ) + + def forbid_warnings(self, *args, **kwargs) -> None: + self.checks.append((self._get_matcher(*args, **kwargs), False)) + + def expect_warnings(self, *args, **kwargs) -> None: + self.checks.append((self._get_matcher(*args, **kwargs), True)) + + def suppress_warnings(self, *args, **kwargs) -> None: + if self._reissue_filtered_warnings is None: + self._reissue_filtered_warnings = False + elif self._reissue_filtered_warnings: + raise RuntimeError( + 'Either `.suppress_warnings()` or `.propagate_warnings()` can ' + 'be called within the same context, but not both' + ) + self.filters.append(self._get_matcher(*args, **kwargs)) + + def propagate_warnings(self, *args, **kwargs) -> None: + if self._reissue_filtered_warnings is None: + self._reissue_filtered_warnings = True + elif not self._reissue_filtered_warnings: + raise RuntimeError( + 'Either `.suppress_warnings()` or `.propagate_warnings()` can ' + 'be called within the same context, but not both' + ) + self.filters.append(self._get_matcher(*args, **kwargs)) + + def check(self, infos: Sequence[_WarningInfo]) -> None: + for matcher, allowed_or_required in self.checks: + matches = [info for info in infos if matcher.match(info)] + if matches and not allowed_or_required: + if self.format_warnings: + matches_repr = ( + f'{len(matches)}:\n{self._format_warnings(matches)}' + ) + else: + matches_repr = f'{len(matches)} ({matches!r})' + raise ResultMismatch( + expected=f'no warnings matching {matcher!r}', + actual=matches_repr, + ) + if not matches and allowed_or_required: + warnings_repr = f'none out of {len(infos)}' + if infos: + if self.format_warnings: + warnings_repr = ( + f'{warnings_repr}:\n{self._format_warnings(infos)}' + ) + else: + warnings_repr = f'{warnings_repr} ({infos!r})' + raise ResultMismatch( + expected=f'warnings matching {matcher!r}', + actual=warnings_repr, + ) + + def reissue(self, infos: Sequence[_WarningInfo]) -> None: + def reissue_(info: _WarningInfo) -> None: + warnings.warn_explicit( + message=info.message, + category=info.category, + filename=info.filename, + lineno=info.lineno, + ) + + # If we haven't called `.suppress_warnings()` or + # `.propagate_warnings()`, handle warnings according to the + # default behavior + if not self.filters: + if self.reissue_warnings: + for info in infos: + reissue_(info) + return + # Otherwise we handle the warnings matching any of the filters + # one way, and the remaining the other way: + # - If we used `.suppress_warnings()`, matching warnings are + # suppressed, and the remainder propagated by default + # - If we used `.propagate_warnings()`, matching warnings are + # propagated, and the remainder suppressed by default + # Note that we are not supposed to have called both methods + reissue_by_default = not self._reissue_filtered_warnings + for info in infos: + should_reissue = reissue_by_default + if any(matcher.match(info) for matcher in self.filters): + should_reissue = not should_reissue + if should_reissue: + reissue_(info) + + @staticmethod + def _format_warnings(infos: Sequence[_WarningInfo]) -> str: + if not infos: + return '' + chunks: list[str] = [] + for info in infos: + text = strip(warnings.formatwarning( + message=info.message, + category=info.category, + filename=info.filename, + lineno=info.lineno, + )) + chunks.append(f'- {info!r}') + chunks.append(indent(text, ' ')) + return '\n'.join(chunks) + + @staticmethod + def _get_matcher( + message: str | None = None, + category: type[Warning] | None = Warning, + module: str | None = None, + lineno: int | None = None, + ) -> _WarningMatcher: + return _WarningMatcher( + message=message, category=category, module=module, lineno=lineno, + ) + + @classmethod + def new( + cls, /, *, + reissue_warnings: bool = False, + format_warnings: bool = False, + **kwargs + ) -> Self: + kwargs['record'] = True + return cls( + warnings.catch_warnings(**kwargs), + reissue_warnings=reissue_warnings, + format_warnings=format_warnings, + ) + + +class CheckWarnings(Sequence[_WarningInfo]): + """ + Helper context for deferring the checking of warnings to until + context exit. + + Example: + >>> import warnings + + Checking for/against warnings: + + >>> cw = CheckWarnings( + ... reissue_warnings=False, format_warnings=False, + ... ) + + >>> with cw: # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE + ... cw.forbid_warnings('foo', UserWarning) + ... warnings.warn('foobar') + ... print('This is printed before the error') + ... + This is printed before the error + Traceback (most recent call last): + ... + test_child_procs._test_child_procs_utils.ResultMismatch: \ +expected no warnings matching + _WarningMatcher(message='foo', + category=), + got 1 ([...]) + + >>> with cw: # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE + ... cw.expect_warnings(category=UserWarning) + ... warnings.warn('foobar', Warning) + ... print('This is printed before the error') + ... + This is printed before the error + Traceback (most recent call last): + ... + test_child_procs._test_child_procs_utils.ResultMismatch: \ +expected warnings matching + _WarningMatcher(category=), + got none out of 1 ([...]) + >>> assert len(cw) == 1 + >>> assert str(cw[0].message) == 'foobar' + + Controlling warning propagation: + + >>> with warnings.catch_warnings(record=True) as cw_outer: + ... with cw: + ... cw.suppress_warnings('foo') + ... warnings.warn('foo') + ... warnings.warn('bar') # Implicitly propagated + ... assert len(cw_outer) == 1 + ... assert str(cw_outer[0].message) == 'bar' + + >>> with warnings.catch_warnings(record=True) as cw_outer: + ... with cw: + ... cw.propagate_warnings('foo') + ... warnings.warn('foo') + ... warnings.warn('bar') # Implicitly suppressed + ... assert len(cw_outer) == 1 + ... assert str(cw_outer[0].message) == 'foo' + """ + def __init__( + self, /, *, + reissue_warnings: bool = True, + format_warnings: bool = True, + **kwargs + ) -> None: + # Note: even with `reissue_warnings=True`, it is not guaranteed + # that all the recorded warnings become visible, depending on + # the warning filters in place before entering the context; so + # we also use `format_warnings=True` by default to include the + # warnings in the error message (if any) + self._new_context: Callable[[], _WarningContext] = partial( + _WarningContext.new, + reissue_warnings=reissue_warnings, + format_warnings=format_warnings, + **kwargs, + ) + self._contexts: list[ + tuple[_WarningContext, Sequence[_WarningInfo]] + ] = [] + self._last_captured: Sequence[_WarningInfo] = [] + + def forbid_warnings(self, *args, **kwargs) -> None: + """ + Equivalent to calling + ``filterwarnings('error', *args, **kwargs)``; + at context exit, if ANY matching warning has been captured, an + error will be raised. + """ + ctx, _ = self._current_context + ctx.forbid_warnings(*args, **kwargs) + + def expect_warnings(self, *args, **kwargs) -> None: + """ + Equivalent to calling + ``filterwarnings('always', *args, **kwargs)``; + at context exit, if NO matching warnings have been captured, an + error will be raised. + """ + ctx, _ = self._current_context + ctx.expect_warnings(*args, **kwargs) + + def suppress_warnings(self, *args, **kwargs) -> None: + """ + Equivalent to calling + ``filterwarnings('never', *args, **kwargs)``; + at context exit, captured warnings are SUPPRESSED if they match + the filter(s) and REISSUED otherwise. + + Note: + Within each context level, one can call EITHER + :py:meth:`.suppress_warnings` or + :py:meth:`.propagate_warnings` but not both. + """ + ctx, _ = self._current_context + ctx.suppress_warnings(*args, **kwargs) + + def propagate_warnings(self, *args, **kwargs) -> None: + """ + Equivalent to calling + ``filterwarnings('always', *args, **kwargs)``; + at context exit, captured warnings are REISSUED if they match + the filter(s) and SUPPRESSED otherwise. + + Note: + Within each context level, one can call EITHER + :py:meth:`.suppress_warnings` or + :py:meth:`.propagate_warnings` but not both. + """ + ctx, _ = self._current_context + ctx.propagate_warnings(*args, **kwargs) + + def __enter__(self) -> Self: + ctx = self._new_context() + infos: Sequence[_WarningInfo] | None + infos = ctx.catch_warnings.__enter__() + assert infos is not None + self._contexts.append((ctx, infos)) + return self + + def __exit__(self, *args, **kwargs) -> None: + ctx, infos = self._contexts.pop() + try: + ctx.check(infos) + finally: + self._last_captured = infos + ctx.catch_warnings.__exit__(*args, **kwargs) + # Now that the warning filters are reset we can reissue + # the warnings without interference + ctx.reissue(infos) + + @overload + def __getitem__(self, i: int, /) -> _WarningInfo: + ... + + @overload + def __getitem__(self, i: slice, /) -> list[_WarningInfo]: + ... + + def __getitem__( + self, i: int | slice, /, + ) -> _WarningInfo | Sequence[_WarningInfo]: + return self._current_warnings[i] + + def __len__(self) -> int: + return len(self._current_warnings) + + @property + def _current_context( + self, + ) -> tuple[_WarningContext, Sequence[_WarningInfo]]: + return self._contexts[-1] + + @property + def _current_warnings(self) -> Sequence[_WarningInfo]: + try: + return self._current_context[1] + except IndexError: + # Outside of contexts, just provide the last captured values + # for convenience + return self._last_captured + + +# ================= `multiprocessing` patch resolution ================= + +_PatchSummary = Mapping[str, Set[str]] + +mp_patch_is_internal: Callable[[str], bool] +mp_patch_is_internal = operator.methodcaller('startswith', '__') + + +def get_patched_attributes( + applied_mp_patches: Collection[str] | None = None, +) -> MappingProxyType[str, frozenset[str]]: + if applied_mp_patches is None: + applied_mp_patches = { + patch for patch, applied in ( + ConfigSource.from_default() + .get_subconfig('child_processes', 'multiprocessing', 'patches') + .conf_dict.items() + ) if applied + } + return _get_patched_attributes(frozenset(applied_mp_patches)) + + +@lru_cache() +def _get_patched_attributes( + applied_mp_patches: frozenset[str], +) -> MappingProxyType[str, frozenset[str]]: + # Get the contents of the individual patches + patches = PATCH_SUMMARIES['minimal'].copy() + iter_summaries = ( + _MP_PATCHES[patch].summary + for patch in applied_mp_patches if patch in _MP_PATCHES + ) + patches = _get_patch_summary_union(patches, *iter_summaries) + return MappingProxyType({ + target: frozenset(attrs) + for target, attrs in filter_mp_patch_summary(patches).items() + }) + + +def get_mp_patches_toml_text(mp_patches: Collection[str]) -> str: + """ + Return the TOML section configuring whether to apply individual + :py:mod:`multiprocessing` patches. + """ + mp_patches_as_dict = { + name: name in mp_patches for name in _MP_PATCHES + if not mp_patch_is_internal(name) + } + return ( + '[tool.line_profiler.child_processes.multiprocessing.patches]\n' + + '\n'.join( + f'{patch} = {str(applied).lower()}' + for patch, applied in mp_patches_as_dict.items() + ) + ) + + +def _get_patch_summary_union( + *summaries: _PatchSummary, +) -> dict[str, frozenset[str]]: + result: dict[str, frozenset[str]] = {} + for summary in summaries: + for target, attrs in summary.items(): + result[target] = result.get(target, frozenset()) | frozenset(attrs) + return result + + +def summarize_mp_patches( + summaries: Collection[tuple[bool, _PatchSummary]] +) -> dict[str, dict[str, bool]]: + """ + Take individual patch summaries and whether each of them is applied, + and return a dictionary representing whether each of the patch + targets is to be patched. + + Example: + >>> summarize_mp_patches([(False, {'foo': {'bar'}})]) + {'foo': {'bar': False}} + >>> summarize_mp_patches([ # doctest: +NORMALIZE_WHITESPACE + ... (False, {'foo': {'bar', 'baz'}}), + ... (True, {'foo': {'baz', 'foobar'}, 'spam': {'ham'}}), + ... (False, {'foo': {'baz'}, 'spam': {'eggs'}}) + ... ]) + {'foo': {'bar': False, 'baz': True, 'foobar': True}, + 'spam': {'eggs': False, 'ham': True}} + """ + def get_all_mentioned(s: Iterable[_PatchSummary]) -> dict[str, set[str]]: + all_items: dict[str, set[str]] = {} + for summary in s: + for target, attrs in summary.items(): + all_items.setdefault(target, set()).update(attrs) + return all_items + + all_items = get_all_mentioned(s for _, s in summaries) + all_patched = get_all_mentioned(s for applied, s in summaries if applied) + result: dict[str, dict[str, bool]] = { + target: + {attr: attr in all_patched.get(target, set()) for attr in attrs} + for target, attrs in all_items.items() + } + # Normalize the order for convenience + return { + target: dict(sorted(attrs.items())) + for target, attrs in sorted(result.items()) + } + + +def filter_mp_patch_summary(summary: _PatchSummary) -> dict[str, set[str]]: + """ + Filter the content of a + :py:attr:`line_profiler._child_process_profiling.multiprocessing\ +._infrastructure.Patch.summary` + so that only patch targets that actually exists remains. + """ + result: dict[str, set[str]] = {} + for target, attrs in summary.items(): + try: + obj = import_target(target) + except ImportError: + continue + present_attrs = {a for a in attrs if hasattr(obj, a)} + # Drop if none of the attributes is present + if present_attrs: + result[target] = present_attrs + return result + + +PATCH_SUMMARIES: dict[ + Literal['minimal', 'maximal', 'default', 'pth_hook'], + dict[str, frozenset[str]] +] = { + 'minimal': {'multiprocessing': frozenset({MP_PATCHED_MARKER})}, +} +# Get patches that are dynamically resolved: while these patches are +# always applied, some of the patch targets are +# platform-/Python-version-specific and may not always exist +_dynamically_resolved_patch_summaries: Iterable[_PatchSummary] = ( + patch.summary for name, patch in _MP_PATCHES.items() + # Basic `multiprocessing` patches are always applied + if mp_patch_is_internal(name) +) +_dynamically_resolved_patch_summaries = itertools.chain( + _dynamically_resolved_patch_summaries, + # some platforms e.g. Windows don't have `fork()` + [{'os': frozenset({'fork'})}], +) +_dynamically_resolved_patch_summaries = cast( # See `ty` issue #3428 + Iterable[_PatchSummary], + map(filter_mp_patch_summary, _dynamically_resolved_patch_summaries), +) +PATCH_SUMMARIES['minimal'] = _get_patch_summary_union( + PATCH_SUMMARIES['minimal'], *_dynamically_resolved_patch_summaries, +) + +# This is only patched if we called +# `_line_profiler_hooks.load_pth_hook()` +PATCH_SUMMARIES['pth_hook'] = { + f'{load_pth_hook.__module__}.{load_pth_hook.__qualname__}': + frozenset({'called'}), +} +# Upper limit of what we could've patched +PATCH_SUMMARIES['maximal'] = _get_patch_summary_union( + PATCH_SUMMARIES['minimal'], + PATCH_SUMMARIES['pth_hook'], + get_patched_attributes([ + name for name in _MP_PATCHES if not mp_patch_is_internal(name) + ]), +) +# Actual patches using the default config +PATCH_SUMMARIES['default'] = _get_patch_summary_union( + PATCH_SUMMARIES['minimal'], get_patched_attributes(), +) + +# ========================= Command execution ========================== + +# `shlex.join()` doesn't work properly on Windows, so use +# `subprocess.list2cmdline()` instead; +# though an "intentionally" undocumented API (cpython issue #10308), +# it's been around since 2.4, seems stable enough, and does exactly what +# is needed +concat_command_line: Callable[[Sequence[str]], str] +if sys.platform.startswith('win32'): + concat_command_line = subprocess.list2cmdline +else: + concat_command_line = shlex.join + + +def run_subproc( + cmd: Sequence[str] | str, /, *args, **kwargs +) -> subprocess.CompletedProcess: + """ + Wrapper around :py:func:`subprocess.run` which writes debugging + output. + """ + class HasStreams(Protocol): + @property + def stdout(self) -> str | bytes | None: + ... + + @property + def stderr(self) -> str | bytes | None: + ... + + if isinstance(cmd, str): + cmd_str = cmd + else: + cmd_str = concat_command_line(cmd) + + print('Command:', cmd_str) + _get_env_deltas(kwargs.get('env')) + print('-- Process start --') + # Note: somehow `mypy` doesn't agree with simply unpacking the + # `*args` into `subprocess.run()`... + status: int | str = '???' + result: HasStreams | None = None + subproc_errors = ( + subprocess.CalledProcessError, subprocess.TimeoutExpired, + ) + time = monotonic() + try: + proc = subprocess.run(cmd, *args, **kwargs) + except Exception as e: + status = 'error' + if isinstance(e, subproc_errors): + result = e + if hasattr(e, 'returncode'): # `CalledProcessError` + status = f'{status} ({e.returncode})' + raise + else: + result, status = proc, proc.returncode + return proc + finally: + time = monotonic() - time + if result is not None: + captured: str | bytes | None + for name, captured, stream in [ + ('stdout', result.stdout, sys.stdout), + ('stderr', result.stderr, sys.stderr), + ]: + if captured is None: + continue + if isinstance(captured, bytes): # `text=False` + captured = captured.decode() + print(f'{name}:\n{indent(captured, " ")}', file=stream) + print( + f'-- Process end (time elapsed: {time:.2f} s / ' + f'return status: {status}) --' + ) + + +# ======================== `kernprof` execution ======================== + + +def _run_as_script( + request: pytest.FixtureRequest, + runner_args: list[str], + test_args: list[str], + test_module: ModuleFixture, + *, + subproc: bool = True, + check_warnings: bool = True, + **kwargs +) -> subprocess.CompletedProcess: + cmd = runner_args + [str(test_module.path)] + test_args + if subproc: + run: Callable[..., subprocess.CompletedProcess] = run_subproc + else: + run = partial(_run_kernprof_main_in_process, request, check_warnings) + test_module.install(children=True, local=not subproc, deps_only=True) + return run(cmd, **kwargs) + + +def _run_as_module( + request: pytest.FixtureRequest, + runner_args: list[str], + test_args: list[str], + test_module: ModuleFixture, + *, + subproc: bool = True, + check_warnings: bool = True, + **kwargs +) -> subprocess.CompletedProcess: + cmd = runner_args + ['-m', test_module.name] + test_args + if subproc: + run: Callable[..., subprocess.CompletedProcess] = run_subproc + else: + run = partial(_run_kernprof_main_in_process, request, check_warnings) + test_module.install(children=True, local=not subproc) + return run(cmd, **kwargs) + + +def _run_as_literal_code( + request: pytest.FixtureRequest, + runner_args: list[str], + test_args: list[str], + test_module: ModuleFixture, + *, + subproc: bool = True, + check_warnings: bool = True, + **kwargs +) -> subprocess.CompletedProcess: + cmd = runner_args + ['-c', test_module.path.read_text()] + test_args + if subproc: + run: Callable[..., subprocess.CompletedProcess] = run_subproc + else: + run = partial(_run_kernprof_main_in_process, request, check_warnings) + test_module.install(children=True, local=not subproc, deps_only=True) + return run(cmd, **kwargs) + + +@cleanup_extra_pth_files() +@preserve_targets(verify=True) +def _run_kernprof_main_in_process( + request: pytest.FixtureRequest, check_warnings: bool, cmd: Sequence[str], + *, + text: bool = False, + capture_output: bool = False, + check: bool = False, + encoding: str | None = None, + env: Mapping[str, str] | None = None, + timeout: float | None = None, + **_kwargs +) -> subprocess.CompletedProcess: + """ + Emulate running ``kernprof`` in a subprocess with in-process + machineries as best as we can, so that we can retrieve more + debugging output when things do go south. Additional safety + railings: + + - `preserve_targets(verify=True)` errors out if any of the patches + is not reversed by :py:func:`kernprof.main` + + - Same goes for the key-value pairs in :py:data:`os.environ`. + + In both cases, if we don't get a match, an + :py:class:`AssertionError` is raised. + """ + def get_streams( + ) -> tuple[str, str] | tuple[bytes, bytes] | tuple[None, None]: + if cap is None: + return None, None + stdout, stderr = cap.readouterr() + if text: + return stdout, stderr + if encoding is None: + encode = operator.methodcaller('encode') + else: + encode = operator.methodcaller('encode', encoding) + return encode(stdout), encode(stderr) + + assert not _kwargs + if cmd[0] == 'kernprof': + args = cmd[1:] + elif any( + list(cmd[:3]) == [python, '-m', 'kernprof'] + for python in ( + 'python', + 'python3', + f'python3.{sys.version_info.minor}', + sys.executable, + ) + ): + args = cmd[3:] + else: + raise AssertionError( + 'Expected commands like `kernprof ...` or ' + f'`python -m kernprof ...`, got {cmd!r}' + ) + + cap: pytest.CaptureFixture | None = None + stdout: str | bytes | None = None + stderr: str | bytes | None = None + if capture_output: + cap = request.getfixturevalue('capsys') + print('Command:', concat_command_line(cmd)) + _get_env_deltas(env) + print('-- Emulated process start --') + get_streams() # Don't include the above in the captured outputs + + status: int | str = '???' + timed_out = False + time = monotonic() + if timeout: + main = add_timeout(kernprof_main, timeout=timeout) + else: + main = kernprof_main + with ExitStack() as stack: + if check_warnings: + # See similar indictments against warnings in + # `test_child_procs.py::_test_apply_mp_patches()` + cw = stack.enter_context(CheckWarnings()) + cw.forbid_warnings('.*resource_tracker', module='multiprocessing') + cw.forbid_warnings( + r'.* file\(s\) .* empty', + category=UserWarning, module='line_profiler', + ) + # The `@add_timeout` decorator spins up a new thread for + # executing `kernprof.main()`; + # explicitly ignore the associated warning when we use + # `start_method='fork'` + if timeout: + cw.suppress_warnings( + r'.*multi-?threaded.*fork\(\)', + category=DeprecationWarning, + ) + try: + try: + cleanup = stack.enter_context(Cleanup()) + if env is not None: + cleanup.update_mapping(os.environ, env) + old_environ = dict(os.environ) + try: + main(args, exit_on_error=False) + finally: + env_diff = _get_env_deltas(os.environ, old_environ, False) + if env_diff: + raise AssertionError( + f'{len(env_diff)} `os.environ` pollution(s): ' + f'{env_diff!r}' + ) + except FuncCallTimeout: # `subprocess` uses `SIGKILL` + returncode, timed_out = -9, True + status = f'error ({returncode})' + except Exception as e: + # Format and output the tracebacks, otherwise we would've + # suppressed them + traceback.print_exception(e) + returncode = 2 if isinstance(e, ArgumentError) else 1 + status = f'error ({returncode})' + else: + status = returncode = 0 + stdout, stderr = get_streams() + + # Turn in-process errors into the corresponding `subprocess` + # errors + if timed_out: + assert timeout is not None # Assure the typechecker + raise subprocess.TimeoutExpired(cmd, timeout, stdout, stderr) + if check and returncode: + raise subprocess.CalledProcessError( + returncode, cmd, stdout, stderr, + ) + return subprocess.CompletedProcess(cmd, returncode, stdout, stderr) + finally: + time = monotonic() - time + captured: str | bytes | None + for name, captured, stream in [ + ('stdout', stdout, sys.stdout), ('stderr', stderr, sys.stderr), + ]: + if captured is None: + continue + if isinstance(captured, bytes): # `text=False` + captured = captured.decode() + print(f'{name}:\n{indent(captured, " ")}', file=stream) + print( + f'-- Emulated process end (time elapsed: {time:.2f} s / ' + f'return status: {status}) --' + ) + + +def _get_env_deltas( + env: Mapping[str, str] | None = None, + original: Mapping[str, str] | None = None, + verbose: bool = True, +) -> list[str]: + if env is None: + return [] + diff: list[str] = [] + if original is None: + original = os.environ + for key in set(original).union(env): + old = original.get(key) + new = env.get(key) + if old is not None is new: + item = f'{old!r} -> (deleted)' + elif old is None is not new: + item = f'{new!r} (added)' + else: + if old == new: + continue + item = f'{old!r} -> {new!r}' + diff.append(f'${{{key}}}: {item}') + if verbose and diff: + print('Env:', indent('\n'.join(diff), ' '), sep='\n') + return diff + + +@cleanup_extra_pth_files() +def _run_test_module( + run_helper: Callable[..., subprocess.CompletedProcess], + request: pytest.FixtureRequest, + test_module: ModuleFixture, + tmp_path_factory: pytest.TempPathFactory, + runner: str | list[str] = 'kernprof', + outfile: str | None = None, + profile: bool = True, + *, + profiled_code_is_tempfile: bool = False, + fail: bool = False, + start_method: StartMethod | None = None, + test_module_args: Sequence[Any] | None = None, + test_module_kwargs: Mapping[str, Any] | None = None, + check: bool = True, + debug_log: str | None = None, + nhits: Mapping[str, int] | None = None, + subproc: bool = True, + **kwargs +) -> tuple[subprocess.CompletedProcess, LineStats | None]: + """ + Returns: + process_running_the_test_module (subprocess.CompletedProcess): + Process object + profiling_stats (LineStats | None): + Line-profiling stats (where available) + """ + if isinstance(runner, str): + runner_cli_args: list[str] = [runner] + else: + runner_cli_args = list(runner) + + if not profile: + nhits = None + + if profile and not profiled_code_is_tempfile: + runner_cli_args.extend(['--prof-mod', str(test_module.path)]) + if nhits is not None: + # We need `kernprof` to write the profiling results immediately + # to preserve data from tempfiles (see note below) + runner_cli_args.append('--view') + + if start_method and start_method not in ('dummy', *START_METHODS): + pytest.skip( + f'`multiprocessing` start method {start_method!r} ' + 'not available on the platform' + ) + if test_module_args is None: + test_module_args = () + if test_module_kwargs is None: + test_module_kwargs = {} + assert test_module.build_command_line + assert test_module.get_output + test_cli_args = test_module.build_command_line( + fail, start_method, *test_module_args, **test_module_kwargs, + ) + + with ub.ChDir(tmp_path_factory.mktemp('mytemp')): + if outfile is not None: + runner_cli_args.extend(['--outfile', outfile]) + if debug_log: + runner_cli_args.extend(['--debug-log', debug_log]) + old_pth_files = cleanup_extra_pth_files.get_pth_files() + try: + proc = run_helper( + request, runner_cli_args, test_cli_args, test_module, + text=True, capture_output=True, check=(check and not fail), + subproc=subproc, + **kwargs + ) + # Checks: + if fail: + # - The process has failed as expected + if check: + assert proc.returncode + else: + # - The result is correctly calculated + expected = test_module.get_output( + *test_module_args, **test_module_kwargs, + ) + output_lines = proc.stdout.splitlines() + ResultMismatch.compare(str(expected), output_lines[0]) + # - Temporary `.pth` file(s) created by + # `LineProfilingCache.write_pth_hook()` has been cleaned up + assert cleanup_extra_pth_files.get_pth_files() == old_pth_files + # - Profiling results are written to the specified file + prof_result: LineStats | None = None + if outfile is None: + assert not list(Path.cwd().iterdir()) + else: + assert os.path.exists(outfile) + assert os.stat(outfile).st_size + if profile: + prof_result = LineStats.from_files(outfile) + # - If we're keeping track, the function is called the + # expected number of times and has run the expected # of + # loops (Note: we do it by parsing the output of + # `kernprof -v` instead of reading the `--outfile`, + # because if the profiled code is in a tempfile the + # profiling data will be dropped in the written outfile) + for tag, num in (nhits or {}).items(): + check_tagged_line_nhits(proc.stdout, tag, num) + except subprocess.TimeoutExpired as e: + # If we're not in running in a subproc, while + # `kernprof.main()` shouldn't have gotten around to write + # the debug logs, we MIGHT be able to retrieve them from the + # active `LineProfilingCache` instance... + try: + if not subproc: + cache = LineProfilingCache.load() + if debug_log is None: + alt_debug_log = 'debug.log' + else: + alt_debug_log = debug_log + '.alt' + with open(alt_debug_log, mode='w') as fobj: + for entry in cache._gather_debug_log_entries(): + print(entry.to_text(), file=fobj) + debug_log = alt_debug_log + except Exception: + pass + finally: + raise e + finally: + if debug_log is not None and os.path.exists(debug_log): + with open(debug_log) as fobj: + print('-- Combined debug logs --', file=sys.stderr) + print(indent(fobj.read(), ' '), end='', file=sys.stderr) + print('-- End of debug logs --', file=sys.stderr) + return proc, prof_result + + +def check_tagged_line_nhits( + output: str, tag: str, nhits: T, *, + comparator: Callable[[T, int], bool] | None = None, +) -> None: + """ + Check the output of :py:meth:`LineStats.print` for the number of + hits on the line tagged with the comment ``# GREP_MARKER[<...>]``. + """ + # The line should be prefixed with 5 numbers: + # lineno, nhits, time, time-per-hit, % time + actual_nhits = 0 + for line in output.splitlines(): + if line.endswith(f'# GREP_MARKER[{tag}]'): + try: + _, n, _, _, _, *_ = line.split() + actual_nhits += int(n) + except Exception: + pass + compare = partial( + ResultMismatch.compare, + nhits, actual_nhits, + expected=f'{nhits} hit(s) on line(s) tagged with {tag!r}', + ) + if comparator is None: + compare() + else: + compare(comparator=comparator) + + +run_module = partial(_run_test_module, _run_as_module) +run_script = partial(_run_test_module, _run_as_script) +run_literal_code = partial( + _run_test_module, _run_as_literal_code, profiled_code_is_tempfile=True, +) diff --git a/tests/test_child_procs/conftest.py b/tests/test_child_procs/conftest.py new file mode 100644 index 00000000..f92009f6 --- /dev/null +++ b/tests/test_child_procs/conftest.py @@ -0,0 +1,422 @@ +from __future__ import annotations + +import os +import shutil +import sysconfig +from collections.abc import Callable, Collection, Generator +from functools import partial +from io import StringIO +from pathlib import Path +from uuid import uuid4 +from tempfile import TemporaryDirectory +from textwrap import indent +from types import ModuleType +from typing import Any + +import pytest + +from line_profiler._child_process_profiling.cache import LineProfilingCache +from line_profiler.curated_profiling import ( + CuratedProfilerContext, ClassifiedPreimportTargets, +) +from line_profiler.autoprofile.util_static import _static_parse +from line_profiler.line_profiler import LineProfiler + +from ._test_child_procs_utils import ( + preserve_object_attrs, ModuleFixture, StartMethod, ResultMismatch, +) + + +__all__ = ( + 'ext_module', 'pool_test_module', 'pool_test_module_clone', + 'process_test_module', + 'ext_module_object', 'pool_test_module_object', + 'process_test_module_object', + 'create_cache', 'curated_profiler', 'another_pid', +) + +# ========================== Module fixtures =========================== + +# Only write the files once per test session + +_EXAMPLES_PATH = Path(__file__).parent / 'multiproc_examples' + + +@pytest.fixture(scope='session') +def _ext_module() -> Generator[Path, None, None]: + name = next(ModuleFixture.propose_name('my_ext_module')) + with TemporaryDirectory() as mydir_str: + my_dir = Path(mydir_str) + my_dir.mkdir(exist_ok=True) + my_module = my_dir / f'{name}.py' + shutil.copy(_EXAMPLES_PATH / 'external_module.py', my_module) + yield my_module + + +@pytest.fixture(scope='session') +def _pool_test_module(_ext_module: Path) -> Generator[Path, None, None]: + yield from _dependent_module_path( + _ext_module, 'pool_test_module.py', 'my_pool_test_module', + ) + + +@pytest.fixture(scope='session') +def _process_test_module(_ext_module: Path) -> Generator[Path, None, None]: + yield from _dependent_module_path( + _ext_module, 'process_test_module.py', 'my_process_test_module', + ) + + +def _dependent_module_path( + _ext_module: Path, basename: str, module_name: str, +) -> Generator[Path, None, None]: + name = next(ModuleFixture.propose_name(module_name)) + body = (_EXAMPLES_PATH / basename).read_text() + body = body.replace('external_module', _ext_module.stem) + with TemporaryDirectory() as mydir_str: + my_dir = Path(mydir_str) + my_dir.mkdir(exist_ok=True) + my_module = my_dir / f'{name}.py' + my_module.write_text(body) + yield my_module + + +def _build_command_line( + fail: bool, start_method: StartMethod | None, *, + use_local_func: bool = False, + nnums: int | None = None, + nprocs: int | None = None, +) -> list[str]: + args: list[str] = [] + if fail: + args.append('--force-failure') + if start_method: + args.extend(['-s', start_method]) + if use_local_func: + args.append('--local') + if nnums: + args.extend(['-l', str(nnums)]) + if nprocs: + args.extend(['-n', str(nprocs)]) + return args + + +def _get_output(nnums_: int, /) -> Callable[..., int]: + def get_output(nnums: int | None = None, **_) -> int: + if nnums is None: + nnums = nnums_ + return nnums * (nnums + 1) // 2 + + return get_output + + +@pytest.fixture +def ext_module( + _ext_module: Path, monkeypatch: pytest.MonkeyPatch, +) -> Generator[ModuleFixture, None, None]: + """ + Yields: + :py:class:`ModuleFixture` helper object containing the code at + ``./multiproc_examples/external_module.py`` + """ + yield ModuleFixture(_ext_module, monkeypatch) + + +@pytest.fixture +def pool_test_module( + _pool_test_module: Path, + ext_module: ModuleFixture, + monkeypatch: pytest.MonkeyPatch, +) -> Generator[ModuleFixture, None, None]: + """ + Yields: + :py:class:`ModuleFixture` helper object containing the code at + ``./multiproc_examples/pool_test_module.py`` + """ + yield from _yield_test_module(_pool_test_module, ext_module, monkeypatch) + + +@pytest.fixture +def pool_test_module_clone( + tmp_path_factory: pytest.TempPathFactory, + monkeypatch: pytest.MonkeyPatch, + pool_test_module: ModuleFixture, + ext_module: ModuleFixture, +) -> Generator[ModuleFixture, None, None]: + """ + Yields: + :py:class:`ModuleFixture` helper object containing the same + code as :py:data:`pool_test_module` + """ + tmpdir = tmp_path_factory.mktemp('my_path') + name = next(ModuleFixture.propose_name('my_cloned_pool_test_module')) + path = tmpdir / f'{name}.py' + path.write_text(pool_test_module.path.read_text()) + yield ModuleFixture( + path, monkeypatch, [ext_module], + pool_test_module.build_command_line, pool_test_module.get_output, + ) + + +@pytest.fixture +def process_test_module( + _process_test_module: Path, + ext_module: ModuleFixture, + monkeypatch: pytest.MonkeyPatch, +) -> Generator[ModuleFixture, None, None]: + """ + Yields: + :py:class:`ModuleFixture` helper object containing the code at + ``./multiproc_examples/process_test_module.py`` + """ + yield from _yield_test_module( + _process_test_module, ext_module, monkeypatch, + ) + + +def _yield_test_module( + test_module: Path, + ext_module: ModuleFixture, + monkeypatch: pytest.MonkeyPatch, +) -> Generator[ModuleFixture, None, None]: + default_nnums = _static_parse('NUM_NUMBERS', test_module) + assert isinstance(default_nnums, int) + yield ModuleFixture( + test_module, monkeypatch, [ext_module], + _build_command_line, _get_output(default_nnums), + ) + + +@pytest.fixture +def ext_module_object( + ext_module: ModuleFixture, +) -> Generator[ModuleType, None, None]: + """ + Yields: + :py:class:`ModuleType` object containing the code at + ``./multiproc_examples/external_module.py``, and is torn down at + the end of the test + """ + yield from ext_module._import_module_helper() + + +@pytest.fixture +def pool_test_module_object( + pool_test_module: ModuleFixture, ext_module_object: ModuleType, +) -> Generator[ModuleType, None, None]: + """ + Yields: + :py:class:`ModuleType` object containing the code at + ``./multiproc_examples/pool_test_module.py``, and is torn down + at the end of the test + """ + yield from pool_test_module._import_module_helper() + + +@pytest.fixture +def process_test_module_object( + process_test_module: ModuleFixture, ext_module_object: ModuleType, +) -> Generator[ModuleType, None, None]: + """ + Yields: + :py:class:`ModuleType` object containing the code at + ``./multiproc_examples/process_test_module.py``, and is torn + down at the end of the test + """ + yield from process_test_module._import_module_helper() + + +# =========================== Misc. fixtures =========================== + + +@pytest.fixture +def create_cache( + tmp_path_factory: pytest.TempPathFactory, + request: pytest.FixtureRequest, +) -> Generator[Callable[..., LineProfilingCache], None, None]: + """ + Wrapper around the :py:class:`LineProfilingCache` instantiator + which: + + - Automatically creates a tempdir and provides it as the + :py:attr:`LineProfilingCache.cache_dir`, + + - Extends the argument ``preimports_module`` to allow for taking + boolean values: + + - ``True``: a temporary preimports module is automatically written + based on ``profiling_targets`` and supplied to the base + constructor. + + - ``False``: equivalent to ``None``. + + - Unless the argument ``_use_curated_profiler: bool = True`` is set + to :py:const:`False`, automatically creates an instance of + :py:class:`LineProfiler` that is curated by a + :py:class:`CuratedProfilerContext` and provides it as the + :py:attr:`LineProfilingCache.profiler`, and + + - At teardown: + + - Removes tempdirs and tempfiles generated. + + - Restores the value of the class' internal reference to the + :py:meth:`LineProfilingCache.load`-ed instance. + + - Calls the `.cleanup()` method of each instance created. + + - Prints these diagnostics for each instance: + + - The stats on the ``.profiler`` associated with each instance + (if any) + + - The stats gathered by + :py:meth:`LineProfilingCache.gather_stats()` + + - The debug logs (if ``.debug`` is true) + """ + def instantiate( + *, + profiling_targets: Collection[str] = (), + preimports_module: os.PathLike[str] | str | bool | None = None, + _use_curated_profiler: bool = True, + **kwargs + ) -> LineProfilingCache: + tmpdir = tmp_path_factory.mktemp('my_cache_dir') + pim: os.PathLike[str] | str | None + if preimports_module in (True, False): + if preimports_module: + targets = ( + ClassifiedPreimportTargets.from_targets(profiling_targets) + ) + if targets: + pim = tmpdir / 'preimports.py' + with pim.open(mode='w') as fobj: + targets.write_preimport_module(fobj) + else: + pim = None + else: + pim = None + else: + # The type checker needs some convincing... + assert not isinstance(preimports_module, bool) + pim = preimports_module + cache = LineProfilingCache( + tmpdir, + profiling_targets=profiling_targets, + preimports_module=pim, + **kwargs, + ) + if _use_curated_profiler: + cache.profiler = request.getfixturevalue('curated_profiler') + instances.append(cache) + return cache + + def print_result( + cache: LineProfilingCache, topic: str, result: str, *notes: str, + ) -> None: + header = '{} ({}):'.format( + topic, '; '.join([f'cache instance {id(cache):#x}', *notes]), + ) + print(header, indent(result, ' '), sep='\n') + + def print_profiler_stats(cache: LineProfilingCache) -> None: + if cache.profiler is None: + result = '' + notes = [] + else: + with StringIO() as sio: + cache.profiler.print_stats(sio) + result = sio.getvalue() + notes = [f'profiler instance {id(cache.profiler):#x}'] + print_result(cache, 'Native profiler stats', result, *notes) + + def print_gathered_stats(cache: LineProfilingCache) -> None: + with StringIO() as sio: + cache.gather_stats().print(sio) + result = sio.getvalue() + print_result(cache, 'Gathered profiler stats', result) + + def print_debug_logs(cache: LineProfilingCache) -> None: + if cache.debug: + result = '\n'.join( + entry.to_text() for entry in cache._gather_debug_log_entries() + ) + else: + result = '' + print_result(cache, 'Gathered debug logs', result) + + instances: list[LineProfilingCache] = [] + handlers: list[Callable[[LineProfilingCache], None]] + handlers = [print_profiler_stats, print_gathered_stats, print_debug_logs] + try: + with preserve_object_attrs(LineProfilingCache, ['_loaded_instance']): + yield instantiate + finally: + for cache in instances: + callbacks: list[Callable[[], Any]] = [cache.cleanup] + callbacks.extend(partial(func, cache) for func in handlers) + for callback in callbacks: + try: + callback() + except Exception: + pass + + +@pytest.fixture +def curated_profiler() -> Generator[LineProfiler, None, None]: + """ + Yields: + Fresh instance of :py:class:`LineProfiler` that is managed by a + :py:class:`CuratedProfilerContext` + """ + prof = LineProfiler() + with CuratedProfilerContext(prof, insert_builtin=True): + yield prof + + +@pytest.fixture +def another_pid() -> int: + """ + Get a PID which is distinct from the current one. + """ + curr_pid = os.getpid() + pid = curr_pid ^ 42 + assert pid != curr_pid + return pid + + +@pytest.fixture(autouse=True) +def _trim_mismatch_traceback(pytestconfig: pytest.Config) -> None: + """ + Truncate the traceback of raised :py:class`ResultMismatch` for more + useful error attribution. + """ + try: + pytestconfig.pluginmanager.register(ResultMismatch) + except ValueError: # Already registered + pass + + +@pytest.fixture() +def check_purelib_dir_writable() -> Generator[None, None, None]: + """ + Check that ``sysconfig.get_path('purelib')`` can be written to; if + not, the test is skipped. + """ + def propose_names() -> Generator[Path, None, None]: + while True: + pth = purelib_path / f'_temp-pth-{uuid4()}.pth' + if not pth.exists(): + yield pth + + purelib_path = Path(sysconfig.get_path('purelib')) + pth = next(propose_names()) + try: + pth.touch() + yield + except OSError: + msg = f"sysconfig.get_path('purelib') ({purelib_path}) not writable" + pytest.skip(msg) + finally: + pth.unlink(missing_ok=True) diff --git a/tests/test_child_procs/multiproc_examples/external_module.py b/tests/test_child_procs/multiproc_examples/external_module.py new file mode 100644 index 00000000..66fce759 --- /dev/null +++ b/tests/test_child_procs/multiproc_examples/external_module.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from itertools import islice + + +def my_external_sum(x: list[int], fail: bool = False) -> int: + result: int = 0 # GREP_MARKER[EXT-INVOCATION] + for item in x: + result += item # GREP_MARKER[EXT-LOOP] + if fail: + raise RuntimeError('forced failure') + return result + + +def split_workload(length: int, n: int) -> list[list[int]]: + """ + Example: + >>> split_workload(10, 3) + [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10]] + >>> split_workload(8, 4) + [[1, 2], [3, 4], [5, 6], [7, 8]] + """ + iter_entries = iter(range(1, length + 1)) + result: list[list[int]] = [] + sublength = length // n + if sublength * n < length: + sublength += 1 + while True: + sublist = list(islice(iter_entries, sublength)) + if not sublist: + break + result.append(sublist) + return result diff --git a/tests/test_child_procs/multiproc_examples/pool_test_module.py b/tests/test_child_procs/multiproc_examples/pool_test_module.py new file mode 100644 index 00000000..71c60654 --- /dev/null +++ b/tests/test_child_procs/multiproc_examples/pool_test_module.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +from argparse import ArgumentParser +from collections.abc import Callable +from multiprocessing import dummy, get_context, Pool +from typing import Literal + +from external_module import my_external_sum +from external_module import split_workload # See issue #433 + + +NUM_NUMBERS = 100 +NUM_PROCS = 4 + + +def my_local_sum(x: list[int], fail: bool = False) -> int: + result: int = 0 # GREP_MARKER[LOCAL-INVOCATION] + for item in x: + result += item # GREP_MARKER[LOCAL-LOOP] + if fail: + raise RuntimeError('forced failure') + return result + + +def sum_in_child_procs( + length: int, n: int, my_sum: Callable[[list[int], bool], int], + start_method: Literal[ + 'fork', 'forkserver', 'spawn', 'dummy' + ] | None = None, + fail: bool = False, +) -> int: + sublists = split_workload(length, n) + if start_method == 'dummy': + pool = dummy.Pool(n) + elif start_method: + pool = get_context(start_method).Pool(n) + else: + pool = Pool(n) + with pool: + subsums = pool.starmap(my_sum, [(sl, fail) for sl in sublists]) + pool.close() + pool.join() + return my_sum(subsums, fail) + + +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) + parser.add_argument( + '-s', '--start-method', + choices=['fork', 'forkserver', 'spawn', 'dummy'], default=None, + ) + parser.add_argument('-f', '--force-failure', action='store_true') + parser.add_argument( + '--local', + action='store_const', + dest='my_sum', + default=my_external_sum, + const=my_local_sum, + ) + options = parser.parse_args(args) + print(sum_in_child_procs( + options.length, options.n, options.my_sum, + start_method=options.start_method, + fail=options.force_failure, + )) + + +if __name__ == '__main__': + main() diff --git a/tests/test_child_procs/multiproc_examples/process_test_module.py b/tests/test_child_procs/multiproc_examples/process_test_module.py new file mode 100644 index 00000000..6416c179 --- /dev/null +++ b/tests/test_child_procs/multiproc_examples/process_test_module.py @@ -0,0 +1,308 @@ +from __future__ import annotations + +import atexit +import os +import pickle +import shutil +import sys +from argparse import ArgumentParser +from collections.abc import Callable, Mapping, Sequence +from contextlib import ExitStack +from dataclasses import dataclass, field +from functools import partial +from itertools import count +from multiprocessing import dummy, get_context +from multiprocessing.process import BaseProcess +from pathlib import Path +from tempfile import mkdtemp, mkstemp, TemporaryDirectory +from time import sleep +from typing import Any, ClassVar, Generic, Literal, TypeVar, cast, final +from typing_extensions import ParamSpec, Self + +from external_module import my_external_sum +from external_module import split_workload # See issue #433 + + +NUM_NUMBERS = 100 +NUM_PROCS = 4 + +T = TypeVar('T') +PS = ParamSpec('PS') +StartMethod = Literal['fork', 'forkserver', 'spawn', 'dummy'] + + +def my_local_sum(x: list[int], fail: bool = False) -> int: + result: int = 0 # GREP_MARKER[LOCAL-INVOCATION] + for item in x: + result += item # GREP_MARKER[LOCAL-LOOP] + if fail: + raise RuntimeError('forced failure') + return result + + +def _format_attrs(attrs: Mapping[str, Any]) -> str: + return ', '.join(f'{attr}={value!r}' for attr, value in attrs.items()) + + +@dataclass(init=False) +class Timeout(RuntimeError): + worker: Worker + timeout: float + + def __init__(self, worker: Worker, timeout: float) -> None: + super().__init__(worker, timeout) + self.worker = worker + self.timeout = timeout + + def _get_repr_attrs(self) -> dict[str, Any]: + return {**self.worker._get_repr_attrs(), 'timeout': self.timeout} + + def __str__(self) -> str: + return 'timed out in {:.2g} s ({})'.format( + self.timeout, _format_attrs(self.worker._get_repr_attrs()), + ) + + def __repr__(self) -> str: + attrs = _format_attrs({ + 'timeout': self.timeout, **self.worker._get_repr_attrs(), + }) + return '<{} @ {:#x} ({})>'.format(type(self).__name__, id(self), attrs) + + +@final +@dataclass +class Worker(Generic[T]): + """ + Example: + >>> import os + >>> from time import sleep + + >>> with Worker.new(sum, args=(range(10),)) as worker: + ... assert os.getpid() != worker.process.pid is not None + ... assert worker.get_result() == 45 + >>> with Worker.new(sleep, args=(10,), daemon=True) as worker: + ... worker.get_result(timeout=.125) # doctest: +ELLIPSIS + Traceback (most recent call last): + ... + process_test_module.Timeout: timed out in 0.12 s (...) + + Note: + For currently unknown reasons, the above doctest results in + pickling errors when run with :py:mod:`xdoctest` but not vanilla + :py:mod:`doctest`; hence the ``SKIP`` directive. + """ + process: BaseProcess + tmpdir: Path + result_callback: Callable[[], T] + _result: T = field(init=False) + _should_interpolate: bool = field(default=False, init=False) + _counter: ClassVar[count] = count() + _module_path: ClassVar[str] = str(Path(__file__).parent) + + def get_result(self, timeout: float | None = None) -> T: + try: + return self._result + except AttributeError: + pass + self.process.join(timeout) + if self.process.is_alive(): + assert timeout is not None + raise Timeout(self, timeout) + self._result = self.result_callback() + return self._result + + def _get_repr_attrs(self) -> dict[str, Any]: + attrs: dict[str, Any] = { + 'name': self.process.name, + 'pid': self.process.pid, + 'exitcode': self.process.exitcode, + } + try: + attrs['result'] = self._result + except AttributeError: + pass + return attrs + + def __repr__(self) -> str: + attrs = _format_attrs(self._get_repr_attrs()) + return '<{} @ {:#x} ({})>'.format(type(self).__name__, id(self), attrs) + + def __enter__(self) -> Self: + # Ensure that this module is import-able; see + # Erotemic/xdoctest#207 + path = self._module_path + self._should_interpolate = interpolate = not any( + os.path.samefile(path, p) for p in sys.path if os.path.exists(p) + ) + if interpolate: + sys.path.insert(0, path) + self.process.start() + return self + + def __exit__(self, *_, **__) -> None: + if self._should_interpolate: + self._should_interpolate = False + try: + sys.path.remove(self._module_path) + except ValueError: + pass + try: + self._end_process('terminate') + finally: + self._end_process('join') + self._end_process('close') + + def _end_process( + self, method: Literal['terminate', 'kill', 'join', 'close'], /, + *args, **kwargs + ) -> None: + if self.process.is_alive(): + getattr(self.process, method)(*args, **kwargs) + sleep(0) + + @classmethod + def new( + cls, + target: Callable[..., T], + *, + args: Sequence[Any] | None = None, + kwargs: Mapping[str, Any] | None = None, + start_method: StartMethod | None = None, + tmpdir: Path | None = None, + **kw + ) -> Worker[T]: + if tmpdir is None: + tmpdir = Path(mkdtemp(prefix='worker-tmpdir-')) + rmdir = partial(shutil.rmtree, tmpdir, ignore_errors=True) + atexit.register(rmdir) + get_process: Callable[..., BaseProcess] + daemon: bool | None = None + use_threads = start_method == 'dummy' + if use_threads: + # Type-checkers don't seem to like a subclass as a + # base-class constructor... + get_process = cast(Callable[..., BaseProcess], dummy.Process) + # Note: `DummyProcess` somehow doesn't take the `daemon` + # parameter, despite `threading.Thread` doing so... + # (see cpython issue GH-#85716, PR GH-#21869) + daemon = kw.pop('daemon', None) + else: + mp_context = get_context(start_method) + assert hasattr(mp_context, 'Process') # Assure `mypy` + get_process = mp_context.Process + handle, fname = mkstemp(suffix='.pkl', dir=tmpdir) + outfile = Path(fname) + os.close(handle) + proc = get_process( + target=cls._write_result, + # Suppress (i.e. don't reraise) errors on threads, or + # `pytest` will complain + args=(outfile, target, use_threads, *(args or ())), + kwargs=(kwargs or {}), + name=f'MyWorker-{cls._get_count()}', + **kw, + ) + if daemon is not None: + # Should've been set except for `DummyProcess` + proc.daemon = daemon + result_callback = cast( + Callable[[], T], partial(cls._read_result, outfile), + ) + return cls(proc, tmpdir, result_callback) + + @classmethod + def _get_count(cls) -> int: + return next(cls._counter) + + @staticmethod + def _read_result(out: Path) -> Any: + with out.open(mode='rb') as fobj: + ok, result = pickle.load(fobj) + if ok: + return result + else: + assert isinstance(result, BaseException) + raise result from None + + @staticmethod + def _write_result( + out: Path, func: Callable[PS, Any], suppress_errors: bool, /, + *args: PS.args, **kwargs: PS.kwargs + ) -> None: + try: + result = True, func(*args, **kwargs) + except BaseException as e: + result = False, type(e)(*e.args) # Truncate traceback + if not suppress_errors: + raise + finally: + with out.open(mode='wb') as fobj: + pickle.dump(result, fobj) + + +def gather_results( + workers: Sequence[Worker[T]], timeout: float | None = None, +) -> list[T]: + """ + Attempt to gather all partial results from all workers even if some + workers errored out. + + Note: + The ``timeout`` is applied on a per-worker basis. + """ + result: list[T] = [] + xc: BaseException | None = None + for worker in workers: + try: + result.append(worker.get_result(timeout)) + except BaseException as e: + xc = e + if xc is None: + return result + raise xc + + +def sum_in_child_procs( + length: int, n: int, my_sum: Callable[[list[int], bool], int], + start_method: StartMethod | None = None, + fail: bool = False, + timeout: float | None = None, +) -> int: + with ExitStack() as stack: + tmpdir = Path(stack.enter_context(TemporaryDirectory())) + new_worker = partial( + Worker.new, start_method=start_method, tmpdir=tmpdir, daemon=True, + ) + workers: list[Worker] = [] + for entries in split_workload(length, n): + worker = new_worker(my_sum, args=(entries, fail)) + workers.append(stack.enter_context(worker)) + return my_sum(gather_results(workers, timeout=timeout), fail) + + +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) + parser.add_argument( + '-s', '--start-method', + choices=['fork', 'forkserver', 'spawn', 'dummy'], default=None, + ) + parser.add_argument('-f', '--force-failure', action='store_true') + parser.add_argument( + '--local', + action='store_const', + dest='my_sum', + default=my_external_sum, + const=my_local_sum, + ) + options = parser.parse_args(args) + print(sum_in_child_procs( + options.length, options.n, options.my_sum, + start_method=options.start_method, + fail=options.force_failure, + )) + + +if __name__ == '__main__': + main() diff --git a/tests/test_child_procs/test_abnormal_execution.py b/tests/test_child_procs/test_abnormal_execution.py new file mode 100644 index 00000000..edb940a3 --- /dev/null +++ b/tests/test_child_procs/test_abnormal_execution.py @@ -0,0 +1,371 @@ +from __future__ import annotations + +import os +import re +import sys +from contextlib import ExitStack +from multiprocessing import get_all_start_methods +from pathlib import Path +# Note: `S_IWRITE` is said to work on Windows but it seems wonky (see +# GitHub issue python/cpython#101675), and it doesn't seem to work on +# Linux either... +from stat import S_IWUSR, S_IWGRP, S_IWOTH +from textwrap import indent +from types import ModuleType +from typing import Literal + +import pytest + +from line_profiler._child_process_profiling.cache import LineProfilingCache + +from ._test_child_procs_utils import ( + run_subproc, strip, check_tagged_line_nhits, +) + + +DEBUG = True + + +class _write_debug_log: + def __init__(self, logfile: os.PathLike[str] | str | None = None) -> None: + self.file = Path(logfile) if logfile else None + + def __enter__(self) -> None: + pass + + def __exit__(self, *_, **__) -> None: + if not (self.file and self.file.exists() and DEBUG): + return + print('-- Combined debug logs --', file=sys.stderr) + print( + indent(self.file.read_text(), ' '), end='', file=sys.stderr, + ) + print('-- End of debug logs --', file=sys.stderr) + + +@pytest.mark.parametrize(('trigger_timeout', 'label1'), + [(True, 'terminate-proc'), (False, 'happy-path')]) +@pytest.mark.parametrize( + ('suppress_error', 'label2'), + [(True, 'suppress-error'), (False, 'propagate-error')]) +def test_prematurely_terminated_process( + tmp_path_factory: pytest.TempPathFactory, + process_test_module_object: ModuleType, + trigger_timeout: bool, + suppress_error: bool, + label1: str, label2: str, +) -> None: + """ + Check that if a :py:class:`multiprocessing.process.BaseProcess` is + explicitly and prematurely ended (e.g. ``.terminate()``-ed or + ``.kill()``-ed, profiling data therein are lost, but it doesn't + crash the session or cause further loss of profiling data. + """ + nlocal = 123 + nchild = 10 + delay = .03125 + nhits = {'INVOCATION': 1, 'LOOP': nlocal} + if trigger_timeout: + # If the timeout isn't enough, exiting the context in the test + # module would have invoked `BaseProcess.terminate()`, impairing + # the collection of profiling data + timeout: float | None = .125 # Would've needed ~ .3 s + else: + timeout = None + nhits['INVOCATION'] += 1 + nhits['LOOP'] += nchild + + tmp = tmp_path_factory.mktemp('mytemp') + test_module = tmp / 'test.py' + out_file = tmp / 'out.lprof' + debug_log: Path | None = None + test_module.write_text(strip(f''' + from __future__ import annotations + + from time import sleep + + from {process_test_module_object.__name__} import Worker + + + def my_sum(n: list[int], delay: float = 0.) -> int: + result: int = 0 # GREP_MARKER[INVOCATION] + for item in range(1, 1 + n): + sleep(delay) + result += item # GREP_MARKER[LOOP] + return result + + + def main() -> None: + my_sum({nlocal}) + try: + with Worker.new(my_sum, args=[{nchild}, {delay}]) as worker: + worker.get_result({timeout}) + except Exception: + if not {suppress_error}: + raise + + + if __name__ == '__main__': + main() + ''')) + + cmd = [ + sys.executable, '-m', 'kernprof', + '--prof-child-procs', + '--line-by-line', + '--view', + f'--prof-mod={test_module}', + f'--outfile={out_file}', + ] + if DEBUG: + debug_log = tmp / 'debug.log' + cmd.append(f'--debug-log={debug_log}') + cmd.append(str(test_module)) + + with _write_debug_log(debug_log): + proc = run_subproc(cmd, capture_output=True, text=True) + + # Check: process termination per-se shouldn't cause `kernprof` + # to error out + assert bool(proc.returncode) == ( + trigger_timeout and not suppress_error + ) + # Check: collection of profiling data is as expected + for tag, num in nhits.items(): + check_tagged_line_nhits(proc.stdout, tag, num) + + +@pytest.mark.parametrize(('corrupt', 'label'), + [(True, 'with-corruption'), (False, 'no-corruption')]) +def test_corrupted_child_stats_file( + tmp_path_factory: pytest.TempPathFactory, + corrupt: bool, + label: str, +) -> None: + """ + Check that if a child's stats file is corrupted, the profiling data + of said child are lost and a warning is issued, but it doesn't crash + the session or cause further loss of profiling data. + """ + def between(range: tuple[int, int], x: int) -> bool: + return min(range) <= x <= max(range) + + nlocal = 123 + nchild = [45, 67, 89, 10] + if corrupt: + # Since we assume that there are 2+ child-profiling-stats files + # and we're corrupting one of them, we should capture SOME but + # not ALL of the child-process data + min_invocations = 2 + max_invocations = len(nchild) + min_loops = nlocal + min(nchild) + max_loops = nlocal + sum(nchild) - min(nchild) + else: + min_invocations = max_invocations = len(nchild) + 1 + min_loops = max_loops = sum(nchild) + nlocal + + tmp = tmp_path_factory.mktemp('mytemp') + test_module = tmp / 'test.py' + out_file = tmp / 'out.lprof' + debug_log: Path | None = None + test_module.write_text(strip(f''' + from __future__ import annotations + + import multiprocessing + + from line_profiler._child_process_profiling.cache import ( + LineProfilingCache, + ) + + + def my_sum(n: list[int]) -> int: + result: int = 0 # GREP_MARKER[INVOCATION] + for item in range(1, 1 + n): + result += item # GREP_MARKER[LOOP] + return result + + + def main() -> None: + # Accrue local data + my_sum({nlocal}) + + # Accrue data in child + with multiprocessing.Pool(2) as pool: + pool.map(my_sum, {nchild[0::2]!r}) + + # Boot up another pool to ensure that we at least have had + # two worker processes while handled tasks + with multiprocessing.Pool(2) as pool: + pool.map(my_sum, {nchild[1::2]!r}) + + # Retrieve one of the non-empty profiling-data files and corrupt + # it + if {corrupt}: + cache = LineProfilingCache.load() + stats, _, *__ = [ + path for path in cache._get_profiling_outfiles() + if path.stat().st_size + ] + stats.write_bytes(b'foo bar baz') + + + if __name__ == '__main__': + main() + ''')) + + cmd = [ + sys.executable, '-m', 'kernprof', + '--prof-child-procs', + '--line-by-line', + '--view', + f'--prof-mod={test_module}', + f'--outfile={out_file}', + ] + if DEBUG: + debug_log = tmp / 'debug.log' + cmd.append(f'--debug-log={debug_log}') + cmd.append(str(test_module)) + + with _write_debug_log(debug_log): + # Check: data corruption shouldn't cause `kernprof` to error out + proc = run_subproc(cmd, capture_output=True, text=True, check=True) + # Check: collection of profiling data is as expected + for tag, nhits_range in [ + ('INVOCATION', (min_invocations, max_invocations)), + ('LOOP', (min_loops, max_loops)), + ]: + check_tagged_line_nhits( + proc.stdout, tag, nhits_range, comparator=between, + ) + # Check: warning for the corrupted file + if corrupt: + assert re.search( + r'UserWarning: .*1 file\(s\) .*cannot be loaded', proc.stderr, + ) + + +@pytest.mark.skipif(sys.platform.startswith('win32'), reason='POSIX-only test') +@pytest.mark.parametrize('start_method', ['spawn', 'fork', 'forkserver']) +@pytest.mark.parametrize( + ('make_unwritable', 'label'), + [(True, 'cannot-write-pth'), (False, 'can-write-pth')]) +@pytest.mark.parametrize(('n', 'nprocs'), [(100, 1)]) +def test_unwritable_purelib_path( + tmp_path_factory: pytest.TempPathFactory, + pool_test_module_object: ModuleType, + make_unwritable: bool, + start_method: Literal['spawn', 'fork', 'forkserver'], + n: int, + nprocs: int, + label: str, +) -> None: + """ + Check that if we can't write a .pth file, the profiling data of + child processes are lost, but it doesn't crash the session or cause + further loss of profiling data. + + Note: + Not being able to write a .pth file is an edge case anyway, so + it's probably alright to skip the test on Windows. + """ + class revoke_write_access: + def __init__(self, path: os.PathLike[str] | str) -> None: + self.path = Path(path) + + def __enter__(self) -> None: + self._perms: int | None = None + if not self.path.is_dir(): + return + if self.path.stat().st_uid != os.getuid(): + return + + mode = self.mode + for bit in S_IWUSR, S_IWGRP, S_IWOTH: + self._make_unwritable(self.path, bit) + if mode == self.mode: + return + self._perms = mode + if DEBUG: + print( + f'Updated {str(self.path)!r}:', + f'{oct(mode)} -> {oct(self.mode)}', + '(disabling write)', + ) + + def __exit__(self, *_, **__) -> None: + if self._perms is None: + return + mode = self.mode + os.chmod(self.path, self._perms) + if DEBUG: + print( + f'Updated {str(self.path)!r}:', + f'{oct(mode)} -> {oct(self.mode)}', + '(enabling write)', + ) + + @staticmethod + def _make_unwritable(path: Path, writable_bytes: int) -> None: + mode = path.stat().st_mode + mask = mode & writable_bytes + if mask: + os.chmod(path, mode - mask) + + @property + def mode(self) -> int: + return self.path.stat().st_mode + + get_pth_locs = LineProfilingCache._enumerate_pth_installation_locations + + # With start methods other than `fork`, we can't extend profiling + # into the worker without writing a .pth + nhits = {'LOCAL-INVOCATION': 1, 'LOCAL-LOOP': nprocs} + if start_method not in get_all_start_methods(): + pytest.skip( + f'start method {start_method!r} not available on ' + f'the platform {sys.platform!r}' + ) + if start_method == 'fork' or not make_unwritable: + nhits['LOCAL-INVOCATION'] += nprocs + nhits['LOCAL-LOOP'] += n + + tmp = tmp_path_factory.mktemp('mytemp') + module_name = pool_test_module_object.__name__ + out_file = tmp / 'out.lprof' + debug_log: Path | None = None + cmd = [ + sys.executable, '-m', 'kernprof', + '--prof-child-procs', + '--line-by-line', + '--view', + f'--prof-mod={module_name}', + f'--outfile={out_file}', + ] + if DEBUG: + debug_log = tmp / 'debug.log' + cmd.append(f'--debug-log={debug_log}') + cmd.extend([ + '-m', + module_name, + f'--start-method={start_method}', + '-l', str(n), + '-n', str(nprocs), + '--local', + ]) + + with _write_debug_log(debug_log): + # Check: even if we can't write a .pth, it shouldn't cause + # `kernprof` to error out + with ExitStack() as stack: + if make_unwritable: + for path in get_pth_locs(): + stack.enter_context(revoke_write_access(path)) + proc = run_subproc(cmd, capture_output=True, text=True, check=True) + # Check: collection of profiling data is as expected + for tag, num in nhits.items(): + check_tagged_line_nhits(proc.stdout, tag, num) + # Check: warning for unwritable .pth locs + if make_unwritable: + assert re.search( + r'UserWarning: .*cannot write a \.pth file', proc.stderr, + ) diff --git a/tests/test_child_procs/test_child_procs.py b/tests/test_child_procs/test_child_procs.py new file mode 100644 index 00000000..0c693df1 --- /dev/null +++ b/tests/test_child_procs/test_child_procs.py @@ -0,0 +1,1273 @@ +from __future__ import annotations + +import dataclasses +import inspect +import operator +import os +import re +import sys +from collections.abc import Callable, Collection, Iterable +from contextlib import ExitStack +from functools import partial +from io import StringIO +from pathlib import Path +from runpy import run_path +from subprocess import CompletedProcess +from textwrap import indent +from types import ModuleType +from typing import Literal, cast + +import pytest + +from _line_profiler_hooks import load_pth_hook +from line_profiler._child_process_profiling.cache import LineProfilingCache +from line_profiler._child_process_profiling.runpy_patches import ( + create_runpy_wrapper, +) +from line_profiler._child_process_profiling.multiprocessing_patches import ( + MPConfig, _PATCHES as MP_PATCHES, +) +from line_profiler.line_profiler import LineStats +from line_profiler.toml_config import ConfigSource + +from ._test_child_procs_utils import ( + DEBUG, DEFAULT_TIMEOUT, NOT_SUPPLIED, START_METHODS, PATCH_SUMMARIES, + ModuleFixture, Params, ResultMismatch, StartMethod, + preserve_object_attrs, preserve_targets, + cleanup_extra_pth_files, + CheckWarnings, + mp_patch_is_internal, get_mp_patches_toml_text, + summarize_mp_patches, filter_mp_patch_summary, + concat_command_line, run_subproc, + run_module, run_script, run_literal_code, check_tagged_line_nhits, + strip, search_cache_logs, +) + + +# ============================= Unit tests ============================= + +# XXX: Tests in this section concern implementation details, and the +# tested APIs and behaviors MUST NOT be relied upon by end-users. + +_DEFAULT_MP_CONFIG = MPConfig.from_config(ConfigSource.from_default()) + + +@pytest.mark.parametrize(('run_profiled_code', 'label1'), + [(True, 'run-profiled'), (False, 'run-unrelated')]) +@pytest.mark.parametrize(('as_module', 'label2'), + [(True, 'run_module'), (False, 'run_path')]) +@pytest.mark.parametrize(('debug', 'label3'), + [(True, 'with-debug'), (False, 'no-debug')]) +def test_runpy_patches( + capsys: pytest.CaptureFixture[str], + ext_module: ModuleFixture, + pool_test_module: ModuleFixture, + pool_test_module_clone: ModuleFixture, + create_cache: Callable[..., LineProfilingCache], + run_profiled_code: bool, + as_module: bool, + debug: bool, + label1: str, label2: str, label3: str, +) -> None: + """ + Test that the :py:mod:`runpy` clone created by + :py:func:`line_profiler._child_process_profiling\ +.create_runpy_wrapper` + correctly sets up profiling when its ``run_*()`` functions are + called. + """ + class restore_argv: + def __enter__(self) -> None: + self.argv = list(sys.argv) + + def __exit__(self, *_, **__) -> None: + sys.argv[:] = self.argv + + cache = create_cache( + rewrite_module=pool_test_module.path, + profiling_targets=[str(ext_module.path)], + profile_imports=True, + debug=debug, + ) + assert cache.profiler is not None + runpy = create_runpy_wrapper(cache) + + nnums = 42 + nprocs = 2 + # If we're running some unrelated code, the profiler should not be + # involved + if run_profiled_code: + module = pool_test_module + num_invocations, num_loops = 1, nprocs + expected_funcs: set[str] = {'my_external_sum', 'split_workload'} + else: + module = pool_test_module_clone + num_invocations, num_loops = 0, 0 + expected_funcs = set() + if as_module: + first_arg = module.name + runner = partial(runpy.run_module, alter_sys=True) + called_func = 'run_module' + else: + first_arg = str(module.path) + runner = runpy.run_path + called_func = 'run_path' + + # Check that the code is run + module.install(local=True, deps_only=not as_module) + with restore_argv(): + sys.argv[:] = [first_arg, f'--length={nnums}', '-n', str(nprocs)] + runner(first_arg, run_name='__main__') + stdout = capsys.readouterr().out + assert stdout.rstrip('\n') == str(nnums * (nnums + 1) // 2) + + # Check that profiler has received the appropriate targets + funcs = {func.__name__ for func in getattr(cache.profiler, 'functions')} + assert funcs == expected_funcs + + # Check that calls in the current process are profiled iif the + # correct file is executed + with StringIO() as sio: + cache.profiler.print_stats(sio) + stats = sio.getvalue() + check_tagged_line_nhits(stats, 'EXT-INVOCATION', num_invocations) + check_tagged_line_nhits(stats, 'EXT-LOOP', num_loops) + + # Check the debug-log entries are correctly gathered + search_cache_logs( + cache, + debug, + { + rf'calling .*{called_func}\(': True, + r'calling .*exec\(': run_profiled_code, + }, + match_individual_messages=True, + flags=re.IGNORECASE, + ) + + +def test_cache_dump_load( + create_cache: Callable[..., LineProfilingCache], +) -> None: + """ + Test that: + + - We can round-trip the cache via :py:meth:`LineProfilingCache.dump` + and :py:meth:`LineProfilingCache.load` + + - The same instance is :py:meth:`LineProfilingCache.load`-ed in + subsequent calls + """ + original = create_cache( + profiling_targets=['foo', 'bar', 'baz'], main_pid=123456, + ) + cache_instances: list[LineProfilingCache] = [original] + envvars: set[str] = set(os.environ) + try: + original.inject_env_vars() # Needed for `.load()` + # Also test slipping stuff into the `._additional_data` + original._additional_data['foo'] = [1, 'string', None] + try: + # Env vars should be inserted + assert set(os.environ) == envvars.union(original.environ) > envvars + original.dump() + loaded = original.load() + cache_instances.append(loaded) + reloaded = original.load() + cache_instances.append(reloaded) + assert original is not loaded is reloaded + # Compare init fields + for field in dataclasses.fields(LineProfilingCache): + if not field.init: + continue + assert ( + getattr(original, field.name) + == getattr(loaded, field.name) + ) + # Compare `._additional_data` + assert original._additional_data == loaded._additional_data + finally: # Explicitly cleanup + for cache in cache_instances: + cache.cleanup() + finally: # Env vars restored after cleanup + assert set(os.environ) == envvars + + +@(Params.new(('wrap_os_fork', 'label1'), + [(True, 'with-wrap-fork'), (False, 'no-wrap-fork')]) + + Params.new(('debug', 'label2'), + [(True, 'with-debug'), (False, 'no-debug')]) + + Params.new(('patch_pool', 'patch_process', 'intercept_logs', 'label3'), + [(True, True, True, 'all-patches'), + (True, True, False, 'pool-and-process'), + (True, False, True, 'pool-and-logging'), + (True, False, False, 'pool-only'), + (False, True, True, 'process-and-logging'), + (False, True, False, 'process-only'), + (False, False, True, 'logging-only'), + (False, False, False, 'no-patches')])).sorted() +@pytest.mark.usefixtures('check_purelib_dir_writable') +def test_cache_setup_main_process( + tmp_path_factory: pytest.TempPathFactory, + create_cache: Callable[..., LineProfilingCache], + wrap_os_fork: bool, + debug: bool, + patch_pool: bool, + patch_process: bool, + intercept_logs: bool, + label1: str, label2: str, label3: str, +) -> None: + """ + Test that :py:meth:`LineProfilingCache._setup_in_main_process` works + as expected. + """ + mp_patches: set[str] = set() + if patch_pool: + mp_patches.add('pool') + if patch_process: + mp_patches.add('process') + if intercept_logs: + mp_patches.add('logging') + + config = tmp_path_factory.mktemp('myconfig') / 'mytoml.toml' + config.write_text(get_mp_patches_toml_text(mp_patches)) + cache = create_cache(debug=debug, config=config) + + # Check that only the requested patches are applied + patches = summarize_mp_patches([ + (True, PATCH_SUMMARIES['minimal']), + *( + (name in mp_patches, filter_mp_patch_summary(patch.summary)) + for name, patch in MP_PATCHES.items() + if not mp_patch_is_internal(name) + ), + ]) + try: + patches['os']['fork'] = wrap_os_fork + except KeyError: + # `os.fork()` pruned because it doesn't exist on e.g. Windows + assert not hasattr(os, 'fork') + + with ExitStack() as stack: + patched = stack.enter_context(preserve_targets(patches)) + compare_patched = partial( + preserve_targets.compare_with_current_values, patched, + ) + original_pths = stack.enter_context(cleanup_extra_pth_files()) + cache._setup_in_main_process(wrap_os_fork=wrap_os_fork) + # There should be exactly one extra `.pth` file + new_pth_hook, = cleanup_extra_pth_files.get_pth_files() - original_pths + # Check whether the patches are applied + compare_patched(operator.is_not, assert_true=patches) + # Check that the instance is set as the `.load()`-ed one + assert cache is cache.load() + # Check whether the patches are reversed + cache.cleanup() + compare_patched() + + # Check the debug-log output + patterns: dict[str, bool] = dict.fromkeys( + [ + r'\(main process\)', + r'Injecting env var.*\$\{LINE_PROFILER_\w+\}', + re.escape(new_pth_hook), + ], + True, + ) + for target, maybe_patches in patches.items(): + patterns.update( + ('Patched.*' + re.escape(f'{target}.{attr}'), is_patched) + for attr, is_patched in maybe_patches.items() + ) + search_cache_logs(cache, debug, patterns) + + +@pytest.mark.parametrize(('wrap_os_fork', 'label1'), + [(True, 'with-wrap-fork'), (False, 'no-wrap-fork')]) +@pytest.mark.parametrize(('preimports', 'label2'), + [(True, 'with-preimports'), (False, 'no-preimports')]) +@pytest.mark.parametrize(('new_profiler', 'label3'), + [(True, 'no-profiler'), (False, 'with-profiler')]) +@pytest.mark.parametrize(('debug', 'label4'), + [(True, 'with-debug'), (False, 'no-debug')]) +@pytest.mark.parametrize('n', [100]) +@preserve_targets() +def test_cache_setup_child( + create_cache: Callable[..., LineProfilingCache], + ext_module_object: ModuleType, + another_pid: int, + wrap_os_fork: bool, + preimports: bool, + new_profiler: bool, + debug: bool, + n: int, + label1: str, label2: str, label3: str, label4: str, +) -> None: + """ + Test that :py:meth:`LineProfilingCache._setup_in_child_process` + works as expected. + """ + def list_profiled_funcs() -> list[str]: + return [ + f'{func.__module__}.{func.__qualname__}' + for func in getattr(cache.profiler, 'functions', []) + ] + + func = ext_module_object.my_external_sum + cache = create_cache( + profiling_targets=[f'{func.__module__}.{func.__qualname__}'], + preimports_module=preimports, + _use_curated_profiler=not new_profiler, + main_pid=another_pid, + debug=debug, + ) + assert (cache.profiler is None) == new_profiler + + seen_funcs = list_profiled_funcs() + if preimports: + preimport_targets = list(cache.profiling_targets) + else: + preimport_targets = [] + + with preserve_object_attrs(os, ['fork']) as preserved: + old_fork = preserved['fork'] + # Check that we're only setting up if there isn't already a + # profiler + assert cache._setup_in_child_process( + wrap_os_fork=wrap_os_fork, context='test_cache_setup_child', + ) == new_profiler + assert cache.profiler + if not new_profiler: + return + + # Check that the profiler has been presented with the profiling + # target + assert list_profiled_funcs() == (seen_funcs + preimport_targets) + + # Check that on cache cleanup: + # - Profiling data is collected + # - `os.fork()` is restored + # - The warning for empty profiling files is only issued when + # expected + assert func(range(1, n + 1)) == n * (n + 1) // 2 + stats = cache.profiler.get_stats() + for callback, has_nonempty_file, has_stats, fork_patched in [ + (lambda: None, False, False, wrap_os_fork), + (cache.cleanup, True, preimports, False), + ]: + callback() + with CheckWarnings() as cw: + handle_warning: list[Callable[..., None]] = [] + if has_nonempty_file: + handle_warning.append(cw.forbid_warnings) + else: # Check for the warning but don't reissue it + handle_warning.extend([ + cw.expect_warnings, cw.suppress_warnings, + ]) + for handle in handle_warning: + handle(r'.* file\(s\) .* empty', module='line_profiler') + gathered = cache.gather_stats() + assert any(gathered.timings.values()) == has_stats, gathered + if hasattr(os, 'fork'): + assert (os.fork is not old_fork) == fork_patched + else: # E.g. Windows + assert old_fork == NOT_SUPPLIED + # Check that after cleaning up the profiler has been disabled + assert not getattr(cache.profiler, 'enable_count', 0) + + # Check that profiling results have been written to the cache + # directory + stats_file, = Path(cache.cache_dir).glob('*.lprof') + assert LineStats.from_files(stats_file) == stats == gathered + + # Check the debug-log output + patterns = { + f'Set up .*profiler.* {id(cache.profiler):#x}': True, + 'Loading preimports': preimports, + 'Created .*' + re.escape(stats_file.name): True, + '[Ss]ucceeded.*dump_stats': True, + 'Loading results .*' + re.escape(stats_file.name): True, + } + search_cache_logs(cache, debug, patterns) + + +@pytest.mark.parametrize('ppid_should_match', [True, False, None]) +@preserve_targets() +def test_load_pth_hook( + create_cache: Callable[..., LineProfilingCache], + another_pid: int, + ppid_should_match: bool | None, +) -> None: + """ + Simulate calling :py:func:`_line_profiler_hooks.load_pth_hook()` in + a child process. + + Notes: + + - The function is CALLED in the .pth file, but we don't actually + NEED a .pth file to call and test it. + + - The counterpart :py:meth:`line_profiler\ +._child_process_profiling.cache.LineProfilingCache.write_pth_hook()` + is implicitly tested in + :py:func:`test_cache_setup_main_process()`. + """ + # This test is mostly here to hack coverage; since the function is + # only to be called in child processes, `coverage` seems to have + # trouble getting data on it... + + # We basically only need this cache instance to set up the + # environment variables and the requisite files... + cache = create_cache(main_pid=another_pid) + if ppid_should_match is not None: + cache.inject_env_vars() + if ppid_should_match: + call_ppid = another_pid + else: # On a PPID mismatch, the function bails after checking + call_ppid = another_pid + 10 + else: + # Without the requisite envvars, the hook should bail very + # quickly (due to the `environ` lookup erroring out), regardless + # of the provided PPID + call_ppid = 0 + cache.dump() + + compare = preserve_targets.compare_with_current_values + patches = { + target: frozenset(attr for attr, patched in attrs.items() if patched) + for target, attrs in summarize_mp_patches([ + (True, PATCH_SUMMARIES['default']), + (True, PATCH_SUMMARIES['pth_hook']), + ]).items() + } + with preserve_targets(patches) as patched: + try: + # NOTE: this creates a cache instance that isn't + # automatically cleaned up by the `create_cache()` + # fixture!!! Hence the try-finally + load_pth_hook(call_ppid) + # Check that the patches are applied where appropriate + assert ( + getattr(load_pth_hook, 'called', False) + == bool(ppid_should_match) + ) + if ppid_should_match: + compare(patched, operator.is_not) + else: # no-op + compare(patched) + return + # Check that calling `load_pth_hook()` again is a no-op + with preserve_targets(patches) as re_patched: + load_pth_hook(call_ppid) + compare(re_patched) + finally: + try: + current_cache = LineProfilingCache.load() + except Exception: + pass + else: + current_cache.cleanup() + # Check that the patches are reversed + compare(patched) + + +@cleanup_extra_pth_files() +@preserve_targets(verify=True) +def _test_apply_mp_patches_inner( + tmp_path_factory: pytest.TempPathFactory, + create_cache: Callable[..., LineProfilingCache], + ext_module_object: ModuleType, + test_module_object: ModuleType, + start_method: StartMethod, + mp_patches: Collection[str], + fail: bool, + n: int, + nprocs: int, +) -> None: + def is_valid_stats_file(path: os.PathLike[str] | str) -> bool: + try: + LineStats.from_files(path, on_empty='error', on_defective='error') + except Exception: + return False + return True + + def get_lineno(path: os.PathLike[str] | str, query: str) -> int: + with Path(path).open() as fobj: + for i, line in enumerate(fobj): + if query in line: + return 1 + i + raise RuntimeError( + f'Did not find line containing {query!r} in {path!r}', + ) + + config = tmp_path_factory.mktemp('myconfig') / 'mytoml.toml' + intercept_logs = 'logging' in mp_patches + assert {'pool', 'process'} & set(mp_patches) # Needed for profiling + config.write_text(get_mp_patches_toml_text(mp_patches)) + + # Note: no need to test the case for `my_local_sum()` separately, + # with `preimports_module=True`, both are just imported and added + # to the profiler, so the code paths are the same + profiled_func = ext_module_object.my_external_sum + called_func = partial( + test_module_object.sum_in_child_procs, + n=nprocs, + my_sum=profiled_func, + start_method=start_method, + fail=fail, + ) + + func_name = f'{profiled_func.__module__}.{profiled_func.__qualname__}' + cache = create_cache( + profiling_targets=[func_name], + preimports_module=True, + config=config, + debug=True, + ) + # Note: + # - The reversibility of the patches have already been tested in + # `test_cache_setup_main_process()`, so we just actually test the + # patched-in components themselves here. + # - `._setup_in_main_process()` doesn't include actually doing the + # preimports. To may the results more consistent between + # `start_method='dummy'` and the others, manually do them below. + cache._setup_in_main_process() # This calls `apply()` + assert cache.profiler is not None + assert cache.preimports_module is not None + run_path(str(cache.preimports_module), {'profile': cache.profiler}) + + timing_key = ( + inspect.getfile(profiled_func), + inspect.getsourcelines(profiled_func)[1], + profiled_func.__qualname__, + ) + assert ext_module_object.__file__ + loop_line = get_lineno(ext_module_object.__file__, 'EXT-LOOP') + + nloops_expected = n + if not fail: + # Counts from the one final sum over the parallel results + nloops_expected += nprocs + + fail_msg = 'forced failure' + try: + result = called_func(n) + except RuntimeError as e: + if not (fail and str(e) == fail_msg): + raise + else: + if fail: + msg = f"expected `RuntimeError({fail_msg!r})`, no error raised" + raise ValueError(msg) + else: # Check correctness of the results + assert result == n * (n + 1) // 2 + + # Check that calls in children are traced + cache.cleanup() + stats = cache.profiler.get_stats() + stats += cache.gather_stats() + entries = stats.timings[timing_key] + nloops = sum(nhits for ln, nhits, _ in entries if ln == loop_line) + ResultMismatch.compare(nloops_expected, nloops) + + # Check the debug logs to see if we have done everything right, esp. + # the logging interception part not covered by other tests + patterns: dict[str, bool] = {} + iter_stats: Iterable[Path] = Path(cache.cache_dir).glob('*.lprof') + iter_stats = cast( # See `ty` issue #3428 + Iterable[Path], filter(is_valid_stats_file, iter_stats), + ) + pat = '[Ss]ucceeded.*dump_stats.*{}' + patterns.update({ + pat.format(re.escape(path.name)): True for path in iter_stats + }) + logger_pat = '{} {}'.format( + re.escape('`multiprocessing` logging'), + r'\((sub_)?debug|info|sub_warning|warn\)', + ) + patterns[logger_pat] = intercept_logs + search_cache_logs(cache, True, patterns) + + +def _test_apply_mp_patches( + patch_process: bool | None = None, + patch_pool: bool | None = None, + intercept_logs: bool | None = None, + *, + start_method: StartMethod, + **kwargs +) -> None: + if start_method not in ('dummy', *START_METHODS): + pytest.skip( + f'`multiprocessing` start method {start_method!r} ' + 'not available on the platform' + ) + + patches = cast(dict[str, bool], _DEFAULT_MP_CONFIG.patches.copy()) + for name, applied in { + 'pool': patch_pool, 'process': patch_process, + 'logging': intercept_logs, + }.items(): + if applied is not None: + patches[name] = applied + mp_patches = [name for name, applied in patches.items() if applied] + + with CheckWarnings() as cw: + cw.forbid_warnings('.*resource_tracker', module='multiprocessing') + cw.forbid_warnings( + r'.* file\(s\) .* empty', + category=UserWarning, module='line_profiler', + ) + _test_apply_mp_patches_inner( + mp_patches=mp_patches, start_method=start_method, **kwargs, + ) + + +@(Params.new('start_method', ['fork', 'forkserver', 'spawn', 'dummy'], + defaults='dummy') + # We only need to check if `intercept_logs = logging` work, the other + # parametrizations don't matter + + Params.new(('intercept_logs', 'label1'), + [(True, 'with-logging'), (False, 'no-logging')], + defaults=(None, 'default-logging'))).sorted() +@pytest.mark.parametrize( + ('test_module', 'patch_process', 'patch_pool', 'label2'), + [('pool_test_module', True, True, 'patch-pool-and-process'), + ('pool_test_module', False, True, 'patch-pool-only'), + ('process_test_module', True, True, 'patch-pool-and-process'), + ('process_test_module', True, False, 'patch-process-only')]) +@pytest.mark.parametrize(('n', 'nprocs'), [(100, 2)]) +@pytest.mark.usefixtures('check_purelib_dir_writable') +def test_apply_mp_patches_success( + request: pytest.FixtureRequest, + tmp_path_factory: pytest.TempPathFactory, + create_cache: Callable[..., LineProfilingCache], + ext_module_object: ModuleType, + test_module: str, + start_method: StartMethod, + patch_process: bool, + patch_pool: bool, + intercept_logs: bool | None, + n: int, + nprocs: int, + label1: str, + label2: str, +) -> None: + """ + Test that :py:func:`line_profiler._child_process_profiling\ +.multiprocessing_patches.apply` + works as expected for Python code using parallelism based on + :py:class:`multiprocessing.pool.Pool` and + :py:class:`multiprocessing.process.BaseProcess`, when the parallel + workload DOES NOT error out. + + See also: + :py:func:`test_apply_mp_patches_failure` + """ + test_module_object = request.getfixturevalue(test_module + '_object') + assert isinstance(test_module_object, ModuleType) + _test_apply_mp_patches( + patch_process, + patch_pool, + intercept_logs, + tmp_path_factory=tmp_path_factory, + create_cache=create_cache, + ext_module_object=ext_module_object, + test_module_object=test_module_object, + start_method=start_method, + fail=False, + n=n, + nprocs=nprocs, + ) + + +@pytest.mark.parametrize('start_method', + ['fork', 'forkserver', 'spawn', 'dummy']) +@pytest.mark.parametrize( + ('test_module', 'patch_process', 'patch_pool', 'label'), + [('pool_test_module', True, True, 'patch-pool-and-process'), + ('pool_test_module', False, True, 'patch-pool-only'), + ('process_test_module', True, True, 'patch-pool-and-process'), + ('process_test_module', True, False, 'patch-process-only')]) +@pytest.mark.parametrize(('n', 'nprocs'), [(100, 2)]) +@pytest.mark.usefixtures('check_purelib_dir_writable') +def test_apply_mp_patches_failure( + request: pytest.FixtureRequest, + tmp_path_factory: pytest.TempPathFactory, + create_cache: Callable[..., LineProfilingCache], + ext_module_object: ModuleType, + test_module: str, + start_method: StartMethod, + patch_process: bool, + patch_pool: bool, + n: int, + nprocs: int, + label: str, +) -> None: + """ + Test that :py:func:`line_profiler._child_process_profiling\ +.multiprocessing_patches.apply` + works as expected for Python code using parallelism based on + :py:class:`multiprocessing.pool.Pool` and + :py:class:`multiprocessing.process.BaseProcess`, when the parallel + workload DOES error out. + + See also: + :py:func:`test_apply_mp_patches_success` + """ + test_module_object = request.getfixturevalue(test_module + '_object') + assert isinstance(test_module_object, ModuleType) + _test_apply_mp_patches( + patch_process, + patch_pool, + tmp_path_factory=tmp_path_factory, + create_cache=create_cache, + ext_module_object=ext_module_object, + test_module_object=test_module_object, + start_method=start_method, + fail=True, + n=n, + nprocs=nprocs, + ) + + +# XXX: End of tests for implementation details + +# ========================= Integration tests ========================== + + +def _get_mp_start_method_fuzzer(label_name: str | None) -> Params: + """ + Returns: + :py:class:`Params` object which does a full Cartesian-product + fuzz between ``fail`` (true or false) and ``start_method`` + ('fork', 'forkserver', and 'spawn'; default :py:const:`None`) + """ + if label_name is None: + label_name, drop_label = '_', True + else: + drop_label = False + fuzz_fail = Params.new(('fail', label_name), + [(True, 'failure'), (False, 'success')], + defaults=(False, 'success')) + if drop_label: + fuzz_fail = fuzz_fail.drop_params(label_name) + fuzz_start = Params.new('start_method', ['fork', 'forkserver', 'spawn'], + defaults=None) + return fuzz_fail * fuzz_start + + +@(Params.new('test_module', ['pool_test_module', 'process_test_module']) + * Params.new(('run_func', 'label1'), + [(run_module, 'module'), (run_script, 'script')]) + * Params.new(('use_local_func', 'label2'), + [(True, 'local'), (False, 'ext')]) + # Python can't pickle things unless they resided in a retrievable + # location (so not the script supplied by `python -c`); this also + # means that `process_test_module` cannot be `python -c`-ed at all, + # because `Worker` is locally-defined + + Params.new( + ('test_module', 'run_func', 'label1', 'use_local_func', 'label2'), + [('pool_test_module', run_literal_code, 'literal-code', False, 'ext')]) + # Also fuzz the parallelization-related stuff, esp. check what + # happens if an exception is raised inside the parallelly-run func + + _get_mp_start_method_fuzzer('label3') + + Params.new(('nnums', 'nprocs'), [(200, None), (None, 3)], + defaults=(None, None))).sorted() +def test_multiproc_script_sanity_check( + run_func: Callable[..., CompletedProcess], + request: pytest.FixtureRequest, + test_module: str, + tmp_path_factory: pytest.TempPathFactory, + use_local_func: bool, + fail: bool, + start_method: StartMethod | None, + nnums: int | None, + nprocs: int | None, + # Dummy arguments to make `pytest` output more legible + label1: str, label2: str, label3: str, +) -> None: + """ + Sanity check that the test modules function as expected when run + with vanilla Python. + """ + module_fixture = request.getfixturevalue(test_module) + assert isinstance(module_fixture, ModuleFixture) + run_func( + request, module_fixture, tmp_path_factory, + runner=sys.executable, profile=False, + fail=fail, + start_method=start_method, + test_module_kwargs=dict( + use_local_func=use_local_func, nnums=nnums, nprocs=nprocs, + ), + ) + + +@pytest.mark.parametrize( + ('run_func', 'label1'), + [(run_module, 'module'), + (run_script, 'script'), + (run_literal_code, 'literal-code')] +) +@pytest.mark.parametrize( + ('runner', 'outfile', 'profile', + 'label2'), # Dummy argument to make `pytest` output more legible + # This is essentially a no-op since it doesn't actually do + # line-profiling, but we check that code path for completeness + [([sys.executable, '-m', 'kernprof', '-q', '--no-line'], + 'out.prof', False, 'cProfile')] + # Run line profiling with and w/o profiling targets + + [([sys.executable, '-m', 'kernprof', '-q', '-l'], 'out.lprof', False, + 'line_profiler-inactive'), + ([sys.executable, '-m', 'kernprof', '-q', '-l'], 'out.lprof', True, + 'line_profiler-active')], +) +def test_running_multiproc_script( + run_func: Callable[..., CompletedProcess], + request: pytest.FixtureRequest, + pool_test_module: ModuleFixture, + tmp_path_factory: pytest.TempPathFactory, + runner: str | list[str], + outfile: str | None, + profile: bool, + # Dummy arguments to make `pytest` output more legible + label1: str, label2: str, +) -> None: + """ + Check that `kernprof` can RUN a test module in various contexts + (`kernprof [...] `, `kernprof [...] -m `, and + `kernprof [...] -c "code"`). + + Notes: + - See issue #422 for the original motivation. + + - This test does not test the actual profiling, just the + execution of the code and presence of profiling data + thereafter. + + - Because this is mostly a sanity check, we only run one of the + test modules. + """ + run_func( + request, pool_test_module, tmp_path_factory, runner, outfile, profile, + ) + + +_fuzz_prof_mp_run_func = Params.new(('run_func', 'label1'), + [(run_module, 'module'), + (run_script, 'script'), + (run_literal_code, 'literal-code')], + defaults=(run_script, 'script')) +_fuzz_prof_mp_markers = ( + (_fuzz_prof_mp_run_func + + Params.new(('prof_child_procs', 'label2'), + [(True, 'with-child-prof'), (False, 'no-child-prof')]) + + _get_mp_start_method_fuzzer(None)) + # Test all `multiproc` start methods with both locally- and + # externally-defined profiling targets + * Params.new(('use_local_func', 'label3'), + [(True, 'local'), (False, 'external')], + defaults=(False, 'external')) + # Test all of the above with both test modules + * Params.new('test_module', ['pool_test_module', 'process_test_module']) + # Add missing params + * Params.new(('preimports', 'label4'), [(False, 'no-preimports')]) + * Params.new(('subproc', 'label5'), [(False, 'in-proc')]) + # The 'with-preimports' case is already tested rather thoroughly in + # `test_apply_mp_patches()`, so exclude these from the above "main" + # param matrix and just test the different `kernprof` modes via the + # `run_func()`s + + (_fuzz_prof_mp_run_func + * Params.new(('preimports', 'label4'), [(True, 'with-preimports')])) + # Similar to the above, we also run at least 1 subtest with + # `subproc=True` for each `(start_method, fail)` + + (_get_mp_start_method_fuzzer(None) + * Params.new(('subproc', 'label5'), [(True, 'subproc')])) +).sorted().split_on_params('fail') + + +def _test_profiling_multiproc_script( + run_func: Callable[..., CompletedProcess], + request: pytest.FixtureRequest, + test_module: ModuleFixture, + ext_module: ModuleFixture, + tmp_path_factory: pytest.TempPathFactory, + prof_child_procs: bool, + preimports: bool, + use_local_func: bool, + fail: bool, + start_method: StartMethod | None, + nnums: int, + nprocs: int, + **kwargs +) -> None: + # How many calls do we expect? + nhits = dict.fromkeys( + ['EXT-INVOCATION', 'EXT-LOOP', 'LOCAL-INVOCATION', 'LOCAL-LOOP'], 0, + ) + # Make sure we're profiling the right function + tag = 'LOCAL' if use_local_func else 'EXT' + tag_call = tag + '-INVOCATION' + tag_loop = tag + '-LOOP' + if not fail: + # The final sum in the parent process should always be profiled + # unless the child processes failed and we never returned from + # `Pool.starmap()` (or `gather_results()`) + nhits[tag_call] += 1 + nhits[tag_loop] += nprocs + if prof_child_procs: + # When profiling extends into child processes, each of them + # invokes the sum function once and when combined they loop thru + # all the items + nhits[tag_call] += nprocs + nhits[tag_loop] += nnums + + runner = [sys.executable, '-m', 'kernprof', '-l'] + runner.extend([ + '--{}prof-child-procs'.format('' if prof_child_procs else 'no-'), + '--{}preimports'.format('' if preimports else 'no-'), + ]) + if not use_local_func: + # Also make sure to include the external module in `--prof-mod` + runner.append(f'--prof-mod={ext_module.name}') + kwargs.setdefault('timeout', DEFAULT_TIMEOUT) + if prof_child_procs and DEBUG: + kwargs.setdefault('debug_log', 'debug.log') + run_func( + request, test_module, tmp_path_factory, + runner=runner, + outfile='out.lprof', + profile=True, + fail=fail, + start_method=start_method, + nhits=nhits, + test_module_kwargs=dict( + use_local_func=use_local_func, + nnums=nnums, + nprocs=nprocs, + ), + **kwargs, + ) + + +@(_fuzz_prof_mp_markers[False]) +@pytest.mark.parametrize( + # XXX: should we explicitly test the single-proc case? We already + # have quite a lot of subtests tho... + ('nnums', 'nprocs'), [(2000, 3)], +) +@pytest.mark.usefixtures('check_purelib_dir_writable') +def test_profiling_multiproc_script_success( + run_func: Callable[..., CompletedProcess], + request: pytest.FixtureRequest, + test_module: str, + ext_module: ModuleFixture, + tmp_path_factory: pytest.TempPathFactory, + prof_child_procs: bool, + preimports: bool, + use_local_func: bool, + start_method: StartMethod | None, + nnums: int, + nprocs: int, + subproc: bool, + # Dummy arguments to make `pytest` output more legible + label1: str, label2: str, label3: str, label4: str, label5: str, +) -> None: + """ + Check that `kernprof` can PROFILE the test modules in various + contexts when the parallel workload runs without errors, optionally + extending profiling into child processes. + + Note: + This test function is heavily parametrized. Here is why that is + necessary: + + - ``run_func`` tests the different ``kernprof`` modes (see + :py:func:`~.test_running_multiproc_script`). + + - ``test_module`` chooses what kind of parallelism the test + module should use (:py:class:`multiprocessing.pool.Pool` vs + :py:class:`multiprocessing.process.BaseProcess`). + + - ``preimports`` tests that both mechanisms for setting up + profiling targets work: + + - :py:const:`True`: child processes import the module + generated by + :py:mod:`line_profiler.autoprofile.eager_preimports`, like + the main :py:mod:`kernprof` process does. + + - :py:const:`False`: child processes rewrite the executed code + before passing it to :py:mod:`runpy`, similar to what + :py:mod:`line_profiler.autoprofile.autoprofile` does. + + These code paths go through different + :py:mod:`multiprocessing` components that we have patched and + thus needs separate testing. + + - ``use_local_func`` tests that we can consistently set up + profiling in both functions locally-defined in the profiled + code and imported by it. + + - ``fail`` tests that our patches and hook doesn't choke when + exceptions occur in child processes, and profiling data can + still be collected. + + - ``start_method`` tests whether all available + :py:mod:`multiprocessing` start methods are covered. + + - ``prof_child_procs`` of course toggles whether to do the + patches to set up profiling in child processes. + + See also: + :py:func:`test_profiling_multiproc_script_failure` + """ + module_fixture = request.getfixturevalue(test_module) + assert isinstance(module_fixture, ModuleFixture) + _test_profiling_multiproc_script( + run_func=run_func, + request=request, + test_module=module_fixture, + ext_module=ext_module, + tmp_path_factory=tmp_path_factory, + prof_child_procs=prof_child_procs, + preimports=preimports, + use_local_func=use_local_func, + fail=False, + start_method=start_method, + nnums=nnums, + nprocs=nprocs, + subproc=subproc, + ) + + +@(_fuzz_prof_mp_markers[True]) +@pytest.mark.parametrize(('nnums', 'nprocs'), [(2000, 3)]) +@pytest.mark.usefixtures('check_purelib_dir_writable') +def test_profiling_multiproc_script_failure( + run_func: Callable[..., CompletedProcess], + request: pytest.FixtureRequest, + test_module: str, + ext_module: ModuleFixture, + tmp_path_factory: pytest.TempPathFactory, + prof_child_procs: bool, + preimports: bool, + use_local_func: bool, + start_method: StartMethod | None, + nnums: int, + nprocs: int, + subproc: bool, + # Dummy arguments to make `pytest` output more legible + label1: str, label2: str, label3: str, label4: str, label5: str, +) -> None: + """ + Check that `kernprof` can PROFILE the test modules in various + contexts when the parallel workload errors out, optionally + extending profiling into child processes. + + See also: + :py:func:`test_profiling_multiproc_script_success` + """ + module_fixture = request.getfixturevalue(test_module) + assert isinstance(module_fixture, ModuleFixture) + _test_profiling_multiproc_script( + run_func=run_func, + request=request, + test_module=module_fixture, + ext_module=ext_module, + tmp_path_factory=tmp_path_factory, + prof_child_procs=prof_child_procs, + preimports=preimports, + use_local_func=use_local_func, + fail=True, + start_method=start_method, + nnums=nnums, + nprocs=nprocs, + subproc=subproc, + ) + + +_TestBarePythonCase = Literal['subprocess.run', 'os.system', 'os.fork'] +_fuzz_bare = ( + Params.new('case', ['subprocess.run', 'os.system', 'os.fork']) + * Params.new(('prof_child_procs', 'label'), + [(True, 'with-child-prof'), (False, 'no-child-prof')]) + * Params.new('n', [200]) +) + + +def _test_profiling_bare_python( + tmp_path_factory: pytest.TempPathFactory, + ext_module: ModuleFixture, + case: _TestBarePythonCase, + prof_child_procs: bool, + fail: bool, + n: int, +) -> None: + ext_module.install(children=True) + temp_dir = tmp_path_factory.mktemp('mytemp') + + out_file = temp_dir / 'out.lprof' + debug_log_file = temp_dir / 'debug.log' + write_debug = DEBUG and prof_child_procs + cmd = [ + sys.executable, '-m', 'kernprof', '-lv', '--preimports', + f'--prof-mod={ext_module.name}', + f'--outfile={out_file}', + '--{}prof-child-procs'.format('' if prof_child_procs else 'no-'), + ] + if write_debug: + cmd.append(f'--debug-log={debug_log_file}') + + if case == 'os.fork': + if not callable(getattr(os, 'fork', None)): + pytest.skip(f'Cannot fork on {sys.platform}') + + code = strip(""" + import os + + from {EXT_MODULE} import my_external_sum + + + def main(n: int, fail: bool) -> None: + numbers = list(range(1, 1 + n)) + child_pid = os.fork() + offset = int(not child_pid) # Literal[0, 1] + try: + my_external_sum(numbers[offset::2], fail) + finally: + if child_pid: # Always join the forked child + os.waitpid(child_pid, 0) + + + if __name__ == '__main__': + main({N}, {FAIL}) + """.format( + EXT_MODULE=ext_module.name, + N=n, + FAIL=fail, + )) + cmd.extend(['-c', code]) + nhits = {'EXT-INVOCATION': 2, 'EXT-LOOP': n} + if not prof_child_procs: + for k in nhits: # We still have the data in the main proc + nhits[k] //= 2 + else: + script_path = temp_dir / 'my-script.py' + script_content = strip(""" + from {EXT_MODULE} import my_external_sum + + + if __name__ == '__main__': + numbers = list(range(1, 1 + {N})) + result = my_external_sum(numbers, {FAIL}) + """.format( + EXT_MODULE=ext_module.name, + N=n, + FAIL=fail, + )) + script_path.write_text(script_content) + + sub_cmd = [sys.executable, str(script_path)] # FIXME + if case == 'subprocess.run': + code = strip(f""" + import subprocess + + + subprocess.run({sub_cmd!r}, check=True) + """) + elif case == 'os.system': + code = strip(""" + import os + + + if os.system({!r}): + raise RuntimeError('called process failed') + """.format(concat_command_line(sub_cmd))) + else: + msg = f'case = {case!r}: expected {_TestBarePythonCase}' + raise AssertionError(msg) + cmd.extend(['-c', code]) + nhits = {'EXT-INVOCATION': 1, 'EXT-LOOP': n} + if not prof_child_procs: + for k in nhits: + nhits[k] = 0 + + proc = run_subproc( + cmd, text=True, capture_output=True, timeout=DEFAULT_TIMEOUT, + ) + + try: + # Check that the code errors out when expected + assert bool(fail) == bool(proc.returncode) + # Check that the profiling output is as expected + for tag, num in nhits.items(): + check_tagged_line_nhits(proc.stdout, tag, num) + finally: + if write_debug: + print('-- Combined debug logs --', file=sys.stderr) + print( + indent(debug_log_file.read_text(), ' '), + end='', file=sys.stderr, + ) + print('-- End of debug logs --', file=sys.stderr) + + +@_fuzz_bare +@pytest.mark.usefixtures('check_purelib_dir_writable') +def test_profiling_bare_python_success( + tmp_path_factory: pytest.TempPathFactory, + ext_module: ModuleFixture, + case: _TestBarePythonCase, + prof_child_procs: bool, + n: int, + # Dummy argument to make `pytest` output more legible + label: str, +) -> None: + """ + Check that `kernprof` can profile the target functions if the code + invokes another bare Python process (via either + :py:func:`os.system`, :py:func:`subprocess.run`, or forking) that + calls them and exits WITHOUT errors. + + See also: + :py:func:`test_profiling_bare_python_failure` + """ + _test_profiling_bare_python( + tmp_path_factory=tmp_path_factory, + ext_module=ext_module, + case=case, + prof_child_procs=prof_child_procs, + fail=False, + n=n, + ) + + +@_fuzz_bare +@pytest.mark.usefixtures('check_purelib_dir_writable') +def test_profiling_bare_python_failure( + tmp_path_factory: pytest.TempPathFactory, + ext_module: ModuleFixture, + case: _TestBarePythonCase, + prof_child_procs: bool, + n: int, + label: str, +) -> None: + """ + Check that `kernprof` can profile the target functions if the code + invokes another bare Python process (via either + :py:func:`os.system`, :py:func:`subprocess.run`, or forking) that + calls them and exits WITH errors. + + See also: + :py:func:`test_profiling_bare_python_success` + """ + _test_profiling_bare_python( + tmp_path_factory=tmp_path_factory, + ext_module=ext_module, + case=case, + prof_child_procs=prof_child_procs, + fail=True, + n=n, + ) diff --git a/tests/test_child_procs/test_fork_stats_attribution.py b/tests/test_child_procs/test_fork_stats_attribution.py new file mode 100644 index 00000000..cbf2901c --- /dev/null +++ b/tests/test_child_procs/test_fork_stats_attribution.py @@ -0,0 +1,165 @@ +""" +Regression tests: profiling data must be attributed to the process that +produced it exactly once. + +Forked children inherit the parent profiler's accumulated stats; if they +dump them verbatim, every line the parent executed before the fork is +counted once per forked child when the results are merged (the parent +dumps the same data itself). See ``LineProfilingCache._wrap_os_fork()`` +and ``_StatsHelper.dump()`` for the subtractive-baseline fix these tests +pin down. +""" +import multiprocessing +import os +import subprocess +import sys +from pathlib import Path +from textwrap import indent +from typing import Literal + +import pytest + +from line_profiler import LineStats + +from ._test_child_procs_utils import strip as code_block + + +StartMethod = Literal['spawn', 'forkserver', 'fork'] +API = Literal['process', 'pool', 'cfut'] + +LOOP_COUNT = 5000 +NUM_TASKS = 4 +TIMEOUT = 120 + +_API_SNIPPETS: dict[API, str] = { + 'process': code_block(""" + procs = [ctx.Process(target=child_work, args=(i,)) + for i in range(num_children)] + for p in procs: + p.start() + for p in procs: + p.join() + """), + 'pool': code_block(""" + with ctx.Pool(2) as pool: + results = pool.map(child_work, range(num_tasks)) + assert results == [i * 2 for i in range(num_tasks)] + """), + 'cfut': code_block(""" + import concurrent.futures as cf + + with cf.ProcessPoolExecutor(max_workers=2, mp_context=ctx) as ex: + results = list(ex.map(child_work, range(num_tasks))) + assert results == [i * 2 for i in range(num_tasks)] + """), +} + + +def _write_workload( + path: Path, method: StartMethod, api: API, num_children: int = 1, +) -> dict[str, int]: + blocks = { + 'import': 'import multiprocessing as mp', + 'child_work': code_block(""" + def child_work(x): + doubled = x * 2 + return doubled + """), + 'main': code_block(f""" + def main(): + num_tasks = {NUM_TASKS} + num_children = {num_children} + acc = 0 + for i in range({LOOP_COUNT}): + acc += i + ctx = mp.get_context({method!r}) + """), + 'guard': code_block(""" + if __name__ == '__main__': + main() + """), + } + blocks['main'] = '{}\n{}'.format( + blocks['main'], indent(_API_SNIPPETS[api], ' '), + ) + source = '\n\n\n'.join(blocks.values()) + path.write_text(source) + print( + 'Current test: {}'.format(os.environ.get('PYTEST_CURRENT_TEST')), + 'Workload script:\n{}'.format(indent(source, ' ')), + sep='\n\n', + ) + lines = source.splitlines() + return { + 'acc': 1 + lines.index(' acc += i'), + 'child_work': 1 + lines.index(' doubled = x * 2'), + } + + +def _run_kernprof(tmp_path: Path, script: Path) -> LineStats: + """ + Run the real CLI in a real subprocess (the .pth/env machinery only + fully engages for a fresh interpreter) and load the merged stats. + """ + outfile = tmp_path / 'out.lprof' + proc = subprocess.run( + [sys.executable, '-m', 'kernprof', + '-l', f'--prof-mod={script}', '--prof-child-procs', + f'--outfile={outfile}', str(script)], + cwd=tmp_path, capture_output=True, text=True, timeout=TIMEOUT, + ) + assert proc.returncode == 0, (proc.stdout, proc.stderr) + return LineStats.from_files(outfile) + + +def _nhits(stats: LineStats, func_name: str, lineno: int) -> int: + matches = [ + entry + for (_, _, func), entries in stats.timings.items() + if func == func_name + for entry in entries + if entry[0] == lineno + ] + assert matches, ( + f'no timings recorded for line {lineno} of {func_name}(); ' + f'keys = {sorted(stats.timings)!r}' + ) + return sum(nhits for _, nhits, _ in matches) + + +@pytest.mark.parametrize('api', sorted(_API_SNIPPETS)) +@pytest.mark.parametrize('method', ['fork', 'forkserver', 'spawn']) +@pytest.mark.usefixtures('check_purelib_dir_writable') +def test_exact_stats_across_start_methods( + tmp_path: Path, method: StartMethod, api: API, +) -> None: + """ + The merged profile must show the parent's loop exactly once and the + children's work exactly ``NUM_TASKS`` times, no matter how the + children came into existence. + """ + if method not in multiprocessing.get_all_start_methods(): + pytest.skip(f'start method {method!r} unavailable') + script = tmp_path / 'workload.py' + expected_child_hits = 1 if api == 'process' else NUM_TASKS + linenos = _write_workload(script, method, api) + stats = _run_kernprof(tmp_path, script) + assert _nhits(stats, 'main', linenos['acc']) == LOOP_COUNT + assert ( + _nhits(stats, 'child_work', linenos['child_work']) + == expected_child_hits + ) + + +@pytest.mark.usefixtures('check_purelib_dir_writable') +def test_stats_do_not_scale_with_fork_children(tmp_path: Path) -> None: + """ + Pre-fork parent data must not be re-contributed once per child. + """ + if 'fork' not in multiprocessing.get_all_start_methods(): + pytest.skip("start method 'fork' unavailable") + script = tmp_path / 'workload.py' + linenos = _write_workload(script, 'fork', 'process', num_children=3) + stats = _run_kernprof(tmp_path, script) + assert _nhits(stats, 'main', linenos['acc']) == LOOP_COUNT + assert _nhits(stats, 'child_work', linenos['child_work']) == 3 diff --git a/tests/test_child_procs/test_pool_protocol.py b/tests/test_child_procs/test_pool_protocol.py new file mode 100644 index 00000000..4db4d0e7 --- /dev/null +++ b/tests/test_child_procs/test_pool_protocol.py @@ -0,0 +1,165 @@ +""" +Tests for the tagged pool result protocol. + +Patched pool workers push ``(tag, pid, result)`` triplets so th patched +parent can e.g. attribute results to worker PIDs. A worker that was +never patched (e.g. its interpreter never loaded the profiling startup +hook, or it runs via ``multiprocessing.set_executable()`` pointing at a +different Python) pushes vanilla results; the parent must pass those +through with a warning rather than crash its result-handler thread — +the crash presents to the user as ``pool.map()`` hanging forever. +""" +import math +import multiprocessing +import warnings +from collections.abc import Callable +from typing import Any, Literal +from typing_extensions import Never + +import pytest + +from line_profiler._child_process_profiling.cache import LineProfilingCache +from line_profiler._child_process_profiling.multiprocessing_patches import ( + _queue, apply as mp_apply, +) + +from ._test_child_procs_utils import CheckWarnings + + +GET_TIMEOUT = 60 +TEST_TAG = '__test_pool_protocol_test_tag__' + + +def test_put_wrapper_tags_results() -> None: + """ + Check that :py:class:`line_profiler._child_process_profiling.\ +multiprocessing_patches._queue.PutWrapper` + behaves as expected, inserting a tag and the return value of a + callback into the pushed tuple. + """ + pushed: list[Any] = [] + + class FakeQueue: + @staticmethod + def put(obj: Any) -> None: + pushed.append(obj) + + @staticmethod + def get() -> Never: # Just here to comply with the interface + raise NotImplementedError + + wrapper = _queue.PutWrapper(FakeQueue(), lambda: 1234, tag=TEST_TAG) + wrapper.put(('job', 0, 'result')) + assert pushed == [(TEST_TAG, 1234, ('job', 0, 'result'))] + + +@pytest.mark.parametrize( + 'incoming, expected_result, expected_data, is_tagged, is_sentinel', + [ + ( + (TEST_TAG, 1234, ('job', 0, 'ok')), + ('job', 0, 'ok'), 1234, True, False, + ), + # vanilla worker + (('job', 0, 'ok'), ('job', 0, 'ok'), None, False, False), + (None, None, None, True, True), # queue sentinel, never warns + ], +) +def test_quick_get_wrapper_unwrapping_results( + create_cache: Callable[..., LineProfilingCache], + incoming: tuple[Any, ...] | None, + expected_result: tuple[Any, ...] | None, + expected_data: Any, + is_tagged: bool, + is_sentinel: bool, +) -> None: + """ + Check that :py:class:`line_profiler._child_process_profiling.\ +multiprocessing_patches._queue.QuickGetWrapper` + behaves as expected, pulling a tag and the extra data from the + pulled tuple (where appropriate) and returning the original result. + """ + def get() -> tuple[Any, ...] | None: + return incoming + + def store(_, data: int) -> None: + stored.append(data) + + cache = create_cache(_use_curated_profiler=False) + stored: list[Any] = [] + quick_get = _queue.QuickGetWrapper(cache, get, store, TEST_TAG) + checked_warning = { + 'category': UserWarning, + 'message': f'.*without the expected tag.*{TEST_TAG}', + } + if is_tagged: + with CheckWarnings() as cw: + cw.forbid_warnings(**checked_warning) + assert quick_get() == expected_result + else: + with CheckWarnings(reissue_warnings=False) as cw: + cw.expect_warnings(**checked_warning) + assert quick_get() == expected_result + # ... but only once per session + with CheckWarnings() as cw: + cw.forbid_warnings(**checked_warning) + assert quick_get() == expected_result + if is_tagged and not is_sentinel: + assert stored == [expected_data] + else: + assert not stored + + +def _run_pool_map( + method: Literal['spawn', 'fork', 'forkserver'], +) -> list[float]: + import os + + ctx = multiprocessing.get_context(method) + with ctx.Pool(1) as pool: + result = pool.map_async(math.sqrt, [0, 1, 4, 9]) + # A hard timeout so that a protocol regression fails fast + # instead of hanging the suite (the historical failure mode) + values = result.get(timeout=GET_TIMEOUT) + # Guard against environments where the pool silently degrades: + # the values must really come from another process + worker_pid = pool.apply_async(os.getpid).get(timeout=GET_TIMEOUT) + assert worker_pid != os.getpid() + return values + + +def test_patched_parent_with_vanilla_spawn_worker( + create_cache: Callable[..., LineProfilingCache], +) -> None: + """ + The parent is patched, but the spawn children know nothing of the + profiling session (no env vars are injected, no .pth hook exists + for them), so they push vanilla un-tagged results: the map must + still complete, with a warning. + """ + if 'spawn' not in multiprocessing.get_all_start_methods(): + pytest.skip("start method 'spawn' unavailable") + cache = create_cache(_use_curated_profiler=False) + # Make the wrapped methods resolve `LineProfilingCache.load()` to + # this instance without injecting env vars (children stay vanilla) + cache._replace_loaded_instance(force=True) + mp_apply(cache) + with pytest.warns(UserWarning, match='without the expected'): + assert _run_pool_map('spawn') == [0.0, 1.0, 2.0, 3.0] + + +def test_patched_parent_with_patched_fork_worker( + create_cache: Callable[..., LineProfilingCache], +) -> None: + """ + Control case: fork children inherit the parent's patches, results + arrive tagged, and no warning fires. + """ + if 'fork' not in multiprocessing.get_all_start_methods(): + pytest.skip("start method 'fork' unavailable") + cache = create_cache() + cache._replace_loaded_instance(force=True) + mp_apply(cache) + with warnings.catch_warnings(): + warnings.simplefilter('error') + assert _run_pool_map('fork') == [0.0, 1.0, 2.0, 3.0]