diff --git a/pyperf/_command.py b/pyperf/_command.py index 465d971a..592bf5c0 100644 --- a/pyperf/_command.py +++ b/pyperf/_command.py @@ -1,4 +1,5 @@ import functools +import json import os.path import subprocess import sys @@ -23,18 +24,22 @@ def bench_command(command, task, loops): % proc.returncode) rss = None + metadata = {} try: lines = output.splitlines() timing = float(lines[0]) - if len(lines) >= 2: - rss = int(lines[1]) + rss = int(lines[1]) + metadata = json.loads(lines[2]) except ValueError: raise ValueError("failed to parse script output: %r" % output) - if rss: + if rss and rss > 0: # store the maximum max_rss = task.metadata.get('command_max_rss', 0) task.metadata['command_max_rss'] = max(max_rss, rss) + + task.metadata.update(metadata) + return timing diff --git a/pyperf/_hooks.py b/pyperf/_hooks.py index cb621047..226e150f 100644 --- a/pyperf/_hooks.py +++ b/pyperf/_hooks.py @@ -31,6 +31,19 @@ def get_selected_hooks(hook_names): yield hook_mapping[hook_name] +def instantiate_selected_hooks(hook_names): + hook_managers = {} + for hook in get_selected_hooks(hook_names): + try: + hook_managers[hook.name] = hook.load()() + except HookError as e: + print(f"ERROR setting up hook '{hook.name}':", file=sys.stderr) + print(str(e), file=sys.stderr) + sys.exit(1) + + return hook_managers + + class HookError(Exception): pass diff --git a/pyperf/_process_time.py b/pyperf/_process_time.py index 88bceb4e..199dc885 100644 --- a/pyperf/_process_time.py +++ b/pyperf/_process_time.py @@ -14,6 +14,8 @@ If resource.getrusage() is available: compute the maximum RSS memory in bytes per process and writes it into stdout as a second line. """ +import contextlib +import json import os import subprocess import sys @@ -91,6 +93,27 @@ def bench_process(loops, args, kw, profile_filename=None): return (dt, max_rss) +def load_hooks(metadata): + hook_names = [] + while "--hook" in sys.argv: + hook_idx = sys.argv.index("--hook") + hook_name = sys.argv[hook_idx + 1] + hook_names.append(hook_name) + del sys.argv[hook_idx] + del sys.argv[hook_idx] + + if len(hook_names): + # Only import pyperf if we know we have hooks + import pyperf._hooks + + hook_managers = pyperf._hooks.instantiate_selected_hooks(hook_names) + metadata["hooks"] = ", ".join(hook_managers.values()) + else: + hook_managers = {} + + return hook_managers + + def main(): # Make sure that the pyperf module wasn't imported if 'pyperf' in sys.modules: @@ -111,6 +134,9 @@ def main(): else: profile_filename = None + metadata = {} + hook_managers = load_hooks(metadata) + loops = int(sys.argv[1]) args = sys.argv[2:] @@ -125,15 +151,21 @@ def main(): kw['stdout'] = devnull kw['stderr'] = subprocess.STDOUT - dt, max_rss = bench_process(loops, args, kw, profile_filename) + with contextlib.ExitStack() as stack: + for hook in hook_managers.values(): + stack.enter_context(hook) + dt, max_rss = bench_process(loops, args, kw, profile_filename) if devnull is not None: devnull.close() + for hook in hook_managers.values(): + hook.teardown(metadata) + # Write timing in seconds into stdout print(dt) - if max_rss: - print(max_rss) + print(max_rss or -1) + print(json.dumps(metadata)) if __name__ == "__main__": diff --git a/pyperf/_worker.py b/pyperf/_worker.py index 5d423e8b..1c2b31dd 100644 --- a/pyperf/_worker.py +++ b/pyperf/_worker.py @@ -6,7 +6,7 @@ import pyperf from pyperf._formatter import (format_number, format_value, format_values, format_timedelta) -from pyperf._hooks import get_selected_hooks, HookError +from pyperf._hooks import instantiate_selected_hooks from pyperf._utils import MS_WINDOWS, MAC_OS, percentile, median_abs_dev @@ -60,14 +60,7 @@ def _compute_values(self, values, nvalue, task_func = self.task_func - hook_managers = {} - for hook in get_selected_hooks(args.hook): - try: - hook_managers[hook.name] = hook.load()() - except HookError as e: - print(f"ERROR setting up hook '{hook.__name__}:'", file=sys.stderr) - print(str(e), file=sys.stderr) - sys.exit(1) + hook_managers = instantiate_selected_hooks(args.hook) if len(hook_managers): self.metadata["hooks"] = ", ".join(hook_managers.keys())