diff --git a/pyperf/_hooks.py b/pyperf/_hooks.py index 8c9e7cdf..9d910327 100644 --- a/pyperf/_hooks.py +++ b/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 diff --git a/pyperf/_utils.py b/pyperf/_utils.py index 540e7eed..c8261522 100644 --- a/pyperf/_utils.py +++ b/pyperf/_utils.py @@ -259,6 +259,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', diff --git a/pyproject.toml b/pyproject.toml index 04efe88f..caa49942 100644 --- a/pyproject.toml +++ b/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]