Skip to content

Commit

Permalink
Make hooks work in bench_command subprocesses (#197)
Browse files Browse the repository at this point in the history
* Make hooks work in bench_command subprocesses

* Fix tests
  • Loading branch information
mdboom authored Aug 27, 2024
1 parent bff934a commit a833224
Show file tree
Hide file tree
Showing 4 changed files with 58 additions and 15 deletions.
11 changes: 8 additions & 3 deletions pyperf/_command.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import functools
import json
import os.path
import subprocess
import sys
Expand All @@ -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


Expand Down
13 changes: 13 additions & 0 deletions pyperf/_hooks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
38 changes: 35 additions & 3 deletions pyperf/_process_time.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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:]

Expand All @@ -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__":
Expand Down
11 changes: 2 additions & 9 deletions pyperf/_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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())

Expand Down

0 comments on commit a833224

Please sign in to comment.