Merge pull request #230 from maurycy/tachyon

Author: pablogsal

pyperf/_hooks.py

@@ -5,8 +5,10 @@
 
 import abc
 import importlib.metadata
+import os
 import os.path
 import shlex
+import signal
 import subprocess
 import sys
 import tempfile
@@ -168,3 +170,84 @@ def exec_perf_cmd(self, cmd):
         self.ctl_fd.write(f"{cmd}\n")
         self.ctl_fd.flush()
         self.ack_fd.readline()
+
+
+class tachyon(HookBase):
+    """Profile the benchmark using sampling profiler (Tachyon).
+
+    The value of the `PYPERF_TACHYON_OPTS` environment variable is
+    appended to the `profiling.sampling attach` command line.
+
+    This hook does not generate output filenames. Use -o or --output in
+    `PYPERF_TACHYON_OPTS` to control output. For most formats, -o can
+    point at an existing directory and the profiler will auto-generate a
+    filename inside it.
+
+    Configuration environment variables:
+        PYPERF_TACHYON_OPTS: Extra arguments passed to
+            `python -m profiling.sampling attach`.
+    """
+
+    def __init__(self):
+        if sys.platform == "win32":
+            raise HookError("tachyon hook is not supported on Windows")
+
+        if sys.version_info < (3, 15):
+            raise HookError(
+                "tachyon hook requires Python 3.15+, "
+                "current version: %s.%s"
+                % (sys.version_info.major, sys.version_info.minor)
+            )
+
+        try:
+            import profiling.sampling  # noqa: F401
+        except ImportError:
+            raise HookError("profiling.sampling module not available")
+
+        self.extra_opts = os.environ.get("PYPERF_TACHYON_OPTS", "")
+
+        self._proc = None
+
+    def __enter__(self):
+        if self._proc is not None:
+            self._stop_profiler()
+
+        cmd = [
+            sys.executable,
+            "-m", "profiling.sampling",
+            "attach",
+            str(os.getpid()),
+        ]
+        cmd += shlex.split(self.extra_opts)
+
+        self._proc = subprocess.Popen(
+            cmd,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+
+    def __exit__(self, _exc_type, _exc_value, _traceback):
+        self._stop_profiler()
+
+    def _stop_profiler(self):
+        if not self._proc:
+            return
+
+        if self._proc.poll() is None:
+            self._proc.send_signal(signal.SIGINT)
+            try:
+                self._proc.wait(timeout=30)
+            except subprocess.TimeoutExpired:
+                self._proc.terminate()
+                try:
+                    self._proc.wait(timeout=5)
+                except subprocess.TimeoutExpired:
+                    self._proc.kill()
+                    self._proc.wait()
+
+        self._proc = None
+
+    def teardown(self, metadata):
+        self._stop_profiler()
+        if self.extra_opts:
+            metadata["tachyon_extra_opts"] = self.extra_opts

pyperf/_utils.py

@@ -264,6 +264,7 @@ def create_environ(inherit_environ, locale, copy_all):
                 "PYTHONPATH", "PYTHON_CPU_COUNT", "PYTHON_GIL",
                 # Pyperf specific variables
                 "PYPERF_PERF_RECORD_DATA_DIR", "PYPERF_PERF_RECORD_EXTRA_OPTS",
+                "PYPERF_TACHYON_OPTS",
                 ]
     if locale:
         copy_env.extend(('LANG', 'LC_ADDRESS', 'LC_ALL', 'LC_COLLATE',

pyproject.toml

@@ -59,6 +59,7 @@ pyperf = "pyperf.__main__:main"
 [project.entry-points."pyperf.hook"]
 perf_record = "pyperf._hooks:perf_record"
 pystats = "pyperf._hooks:pystats"
+tachyon = "pyperf._hooks:tachyon"
 _test_hook = "pyperf._hooks:_test_hook"
 
 [tool.setuptools]