Add support for profiling benchmarks using perf-record (#214)

* Add support for profiling benchmarks using `perf-record`

Add a hook that profiles benchmarks using `perf-record` when enabled

* Rename hook to perf_record

* Fix docstring

* Fix formatting

* Remove assignment to unused variable

* Update docs

* Inherit perf record env vars

Author: mpage

doc/run_benchmark.rst

@@ -206,3 +206,17 @@ The `Tools/scripts/summarize_stats.py <https://github.com/python/cpython/blob/ma
 
 Statistics are not cleared between runs.
 If you need to delete statistics from a previous run, remove the files in ``/tmp/py_stats`` (Unix) or ``C:\temp\py_stats`` (Windows).
+
+Profiling benchmarks using ``perf record``
+==========================================
+``pyperf`` supports profiling benchmark execution using ``perf
+record``. ``perf`` is only enabled while the benchmark is running to avoid
+profiling unrelated parts of ``pyperf`` itself.
+
+One profile data file is generated for each benchmark run.  These files have
+the basename of ``perf.data.<uuid>`` and are written to the current directory
+by default. The directory can be overridden by setting the
+``PYPERF_PERF_RECORD_DATA_DIR`` environment variable.
+
+The value of the ``PYPERF_PERF_RECORD_EXTRA_OPTS`` environment variable is
+appended to the command line of ``perf record`` if it is provided.

pyperf/_hooks.py

@@ -5,7 +5,12 @@
 
 import abc
 import importlib.metadata
+import os.path
+import shlex
+import subprocess
 import sys
+import tempfile
+import uuid
 
 
 def get_hooks():
@@ -108,3 +113,58 @@ def __enter__(self):
 
     def __exit__(self, _exc_type, _exc_value, _traceback):
         sys._stats_off()
+
+
+class perf_record(HookBase):
+    """Profile the benchmark using perf-record.
+
+    Profile data is written to the current directory directory by default, or
+    to the value of the `PYPERF_PERF_RECORD_DATA_DIR` environment variable, if
+    it is provided.
+
+    Profile data files have a basename of the form `perf.data.<uuid>`
+
+    The value of the `PYPERF_PERF_RECORD_EXTRA_OPTS` environment variable is
+    appended to the command line of perf-record, if provided.
+    """
+
+    def __init__(self):
+        self.tempdir = tempfile.TemporaryDirectory()
+        self.ctl_fifo = self.mkfifo(self.tempdir.name, "ctl_fifo")
+        self.ack_fifo = self.mkfifo(self.tempdir.name, "ack_fifo")
+        perf_data_dir = os.environ.get("PYPERF_PERF_RECORD_DATA_DIR", "")
+        perf_data_basename = f"perf.data.{uuid.uuid4()}"
+        cmd = ["perf", "record",
+               "--pid", str(os.getpid()),
+               "--output", os.path.join(perf_data_dir, perf_data_basename),
+               "--control", f"fifo:{self.ctl_fifo},{self.ack_fifo}"]
+        extra_opts = os.environ.get("PYPERF_PERF_RECORD_EXTRA_OPTS", "")
+        cmd += shlex.split(extra_opts)
+        self.perf = subprocess.Popen(
+            cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+        self.ctl_fd = open(self.ctl_fifo, "w")
+        self.ack_fd = open(self.ack_fifo, "r")
+
+    def __enter__(self):
+        self.exec_perf_cmd("enable")
+
+    def __exit__(self, _exc_type, _exc_value, _traceback):
+        self.exec_perf_cmd("disable")
+
+    def teardown(self, metadata):
+        try:
+            self.exec_perf_cmd("stop")
+            self.perf.wait(timeout=120)
+        finally:
+            self.ctl_fd.close()
+            self.ack_fd.close()
+
+    def mkfifo(self, tmpdir, basename):
+        path = os.path.join(tmpdir, basename)
+        os.mkfifo(path)
+        return path
+
+    def exec_perf_cmd(self, cmd):
+        self.ctl_fd.write(f"{cmd}\n")
+        self.ctl_fd.flush()
+        self.ack_fd.readline()

pyperf/_utils.py

@@ -270,7 +270,10 @@ def create_environ(inherit_environ, locale, copy_all):
     env = {}
     copy_env = ["PATH", "HOME", "TEMP", "COMSPEC", "SystemRoot", "SystemDrive",
                 # Python specific variables
-                "PYTHONPATH", "PYTHON_CPU_COUNT", "PYTHON_GIL"]
+                "PYTHONPATH", "PYTHON_CPU_COUNT", "PYTHON_GIL",
+                # Pyperf specific variables
+                "PYPERF_PERF_RECORD_DATA_DIR", "PYPERF_PERF_RECORD_EXTRA_OPTS",
+                ]
     if locale:
         copy_env.extend(('LANG', 'LC_ADDRESS', 'LC_ALL', 'LC_COLLATE',
                          'LC_CTYPE', 'LC_IDENTIFICATION', 'LC_MEASUREMENT',

pyproject.toml

@@ -56,6 +56,7 @@ dev = [
 pyperf = "pyperf.__main__:main"
 
 [project.entry-points."pyperf.hook"]
+perf_record = "pyperf._hooks:perf_record"
 pystats = "pyperf._hooks:pystats"
 _test_hook = "pyperf._hooks:_test_hook"