Per-sub-benchmark PMU + sysstat for WDL prod_set

Summary:
The WDL bench `prod_set` job runs ~25 individual sub-benchmarks
(memcpy_benchmark, openssl, lzbench, ProtocolBench, ...) back-to-back in a
single benchpress run. Today the perf hook spans the whole run, so PMU +
sysstat data ends up smeared across all sub-benchmarks in one set of CSVs.
Distinguishing IPC, topdown breakdown, mpstat etc. per sub-benchmark
requires teasing apart timestamps after the fact, which is brittle.

This change adds an opt-in stage-aware mode to the perf hook so each
sub-benchmark gets its own folder of perf data:

  benchmark_metrics_<uuid>/
      memcpy_benchmark/
          mpstat.csv
          mem-stat.csv
          perf-stat.csv
          topdown-... .csv
          (etc -- one set per enabled perf monitor)
      hash_hash_benchmark/
          ...
      ...

The mechanism:

  1. The perf hook accepts a new option `stage_aware: true`. When set,
     `before_job` does NOT start any monitors. Instead it creates a FIFO
     under benchmark_metrics_<uuid>/perf_stage.fifo, advertises its path
     via the env var BENCHPRESS_PERF_STAGE_FIFO, and spawns a
     coordinator thread.

  2. The coordinator reads commands from the FIFO. Each
     "START <stage_name>" allocates a fresh set of perf monitors with
     that stage name as a sub-folder; each "STOP" terminates the
     monitors and writes their CSVs. Multiple START/STOP cycles are
     supported. `after_job` writes a final __EXIT__ to drain the
     coordinator.

  3. Every existing perf monitor (mpstat, memstat, netstat, perfstat,
     vmstat, cpufreq*, power, topdown -- including IntelPerfSpect/3,
     BasePerfUtil, AMDPerfUtil, ARMPerfUtil, NVPerfUtil,
     NeoVerseV3PerfUtil) gains a `subdir` constructor arg that the base
     `Monitor.gen_path` joins under benchmark_metrics_<uuid>/. Existing
     callers that don't pass `subdir` keep their flat layout, so default
     mode is byte-for-byte unchanged.

  4. `packages/wdl_bench/run_prod.sh` emits "START <bench>" before each
     sub-benchmark and "STOP" after. The emit is gated on the
     BENCHPRESS_PERF_STAGE_FIFO env var, so running prod_set without
     stage-aware mode (or without the perf hook at all) is unchanged.

  5. A new job entry `prod_set_per_stage_perf` in jobs_wdl.yml
     documents the wiring; it has the same shape as `prod_set`. To
     activate, run:

         ./benchpress run prod_set_per_stage_perf \\
             -k perf -a '{"perf": {"stage_aware": true}}'

Differential Revision: D108110315

Author: charles-typ

benchpress/config/jobs_wdl.yml

@@ -183,9 +183,44 @@
   args:
     - '--type {type}'
     - '--output {output}'
+    - '--name {name}'
+  vars:
+    - 'type=prod'
+    - 'output=wdl_bench_results.txt'
+    - 'name=none'
+  hooks:
+    - hook: cpu-mpstat
+      options:
+        args:
+          - '-u'   # utilization
+          - '1'    # second interval
+    - hook: copymove
+      options:
+        is_move: true
+        after:
+          - 'benchmarks/wdl_bench/wdl_bench_results.txt'
+          - 'benchmarks/wdl_bench/out_*.json'
+
+# Stage-aware variant of prod_set. Runs the same set of WDL sub-benchmarks
+# but signals the perf hook before/after each sub-benchmark via a FIFO so
+# every sub-benchmark gets its own folder of PMU + sysstat data under
+# benchmark_metrics_<uuid>/<sub_benchmark_name>/. To use, run
+#   ./benchpress run prod_set_per_stage_perf -k perf -a '{"perf": {"stage_aware": true}}'
+# (or with whatever monitor overrides you want).
+- name: prod_set_per_stage_perf
+  benchmark: wdl_bench
+  description: >
+    Same as prod_set, but pairs with the perf hook's stage-aware mode to
+    collect a separate set of PMU + sysstat data per WDL sub-benchmark.
+    Pass `-k perf -a '{"perf": {"stage_aware": true}}'` to enable.
+  args:
+    - '--type {type}'
+    - '--output {output}'
+    - '--name {name}'
   vars:
     - 'type=prod'
     - 'output=wdl_bench_results.txt'
+    - 'name=none'
   hooks:
     - hook: cpu-mpstat
       options:

benchpress/plugins/hooks/perf.py

@@ -8,6 +8,7 @@
 
 import logging
 import os
+import threading
 import traceback
 
 from benchpress.lib import open_source
@@ -75,6 +76,27 @@
 logger = logging.getLogger(__name__)
 
 
+# Env var used to advertise the stage FIFO path to a child benchmark process.
+# When set, a stage-aware benchmark (e.g. WDL prod_set's run_prod.sh) writes
+# stage markers like ``START <stage_name>`` and ``STOP`` to that FIFO so
+# this hook can rotate the perf monitor set per sub-benchmark and emit one
+# folder of perf data per stage.
+PERF_STAGE_FIFO_ENV = "BENCHPRESS_PERF_STAGE_FIFO"
+
+
+def _resolve_subdir_option(opts):
+    """Return whether the user opted into stage-aware mode.
+
+    Stage-aware mode is requested either by setting ``stage_aware: true`` in
+    the perf hook options, or by leaving the default (False). Today the only
+    benchmark that emits stage markers is WDL prod_set, but the wiring is
+    generic.
+    """
+    if not isinstance(opts, dict):
+        return False
+    return bool(opts.get("stage_aware", False))
+
+
 class Perf(Hook):
     def before_job(self, opts, job):
         self.opts = DEFAULT_OPTIONS
@@ -90,8 +112,180 @@ def before_job(self, opts, job):
         if not os.path.isdir(self.benchmark_metrics_dir):
             os.mkdir(self.benchmark_metrics_dir)
 
+        self._stage_aware = _resolve_subdir_option(opts)
+        self._stage_thread = None
+        self._stage_fifo_path = None
+        self._stage_lock = threading.Lock()
+        self._current_stage = None
+        self._current_monitors = []
+
+        if self._stage_aware:
+            self._start_stage_coordinator(job)
+        else:
+            # Legacy / default path: start one set of monitors that span the
+            # entire benchmark.
+            self.monitors = self._build_and_run_monitors(job, subdir=None)
+
+    def after_job(self, opts, job):
+        if self._stage_aware:
+            self._stop_stage_coordinator()
+            return
+        for monitor in self.monitors:
+            monitor.terminate()
+        for monitor in self.monitors:
+            monitor.write_csv()
+
+    # ------------------------------------------------------------------
+    # Stage-aware mode
+    # ------------------------------------------------------------------
+
+    def _start_stage_coordinator(self, job):
+        """Open a FIFO under the benchmark metrics dir and spawn a thread
+        that reads ``START <stage>`` / ``STOP`` commands. Each START gets
+        a fresh monitor set whose CSVs land in
+        ``benchmark_metrics_<uuid>/<stage>/``.
+
+        The FIFO path is exposed via the ``BENCHPRESS_PERF_STAGE_FIFO``
+        env var so the child benchmark process can find it. We deliberately
+        DO NOT start any monitors yet -- the coordinator launches them
+        when the first STAGE command arrives.
+        """
+        self._stage_fifo_path = os.path.join(
+            self.benchmark_metrics_dir, "perf_stage.fifo"
+        )
+        # Recreate the FIFO each run so a stale one from a previous run
+        # never confuses us.
+        if os.path.exists(self._stage_fifo_path):
+            os.unlink(self._stage_fifo_path)
+        os.mkfifo(self._stage_fifo_path, 0o600)
+        os.environ[PERF_STAGE_FIFO_ENV] = self._stage_fifo_path
+        self._stage_stop_event = threading.Event()
+        self._stage_thread = threading.Thread(
+            target=self._stage_loop,
+            args=(job,),
+            name="perf-stage-coordinator",
+            daemon=True,
+        )
+        self._stage_thread.start()
+        logger.info(
+            f"Perf hook: stage-aware mode active, FIFO at {self._stage_fifo_path}"
+        )
+
+    def _stop_stage_coordinator(self):
+        # Tell the coordinator to exit and unblock the FIFO read by writing
+        # a final STOP. A standalone writer also lets us flush any lingering
+        # monitor set without depending on the benchmark script having sent
+        # its own STOP.
+        self._stage_stop_event.set()
+        try:
+            with open(self._stage_fifo_path, "w") as f:
+                f.write("STOP\n")
+                f.write("__EXIT__\n")
+        except Exception as e:
+            logger.warning(f"Perf hook: failed to nudge stage FIFO: {e}")
+        if self._stage_thread is not None:
+            self._stage_thread.join(timeout=30)
+        # Clean up the FIFO and the env var so a subsequent benchmark in the
+        # same process doesn't accidentally re-use stale state.
+        try:
+            os.unlink(self._stage_fifo_path)
+        except OSError:
+            pass
+        os.environ.pop(PERF_STAGE_FIFO_ENV, None)
+
+    def _stage_loop(self, job):
+        """Read commands from the FIFO until __EXIT__. Each line is one of:
+
+          START <stage_name>
+          STOP
+          __EXIT__
+
+        START allocates fresh monitors and runs them; STOP terminates them
+        and writes their CSVs. Multiple START/STOP cycles are supported.
+
+        Implementation note: each writer that closes the FIFO causes EOF on
+        the reader, so we have to re-open the FIFO after every burst rather
+        than holding a single file handle. Each ``open(..., "r")`` blocks
+        until a writer is available again -- which is exactly the behavior
+        we want for a coordinator that's awaiting the next stage marker.
+        """
+        try:
+            while not self._stage_stop_event.is_set():
+                with open(self._stage_fifo_path, "r") as fifo:
+                    for line in fifo:
+                        line = line.strip()
+                        if not line:
+                            continue
+                        if line == "__EXIT__":
+                            self._end_current_stage(job)
+                            return
+                        if line.startswith("START "):
+                            stage = line[len("START ") :].strip()
+                            self._begin_stage(job, stage)
+                            continue
+                        if line == "STOP":
+                            self._end_current_stage(job)
+                            continue
+                        logger.warning(
+                            f"Perf hook: ignoring unknown stage command: {line!r}"
+                        )
+        except Exception as e:
+            logger.warning(
+                f"Perf hook: stage coordinator crashed: {type(e).__name__}: {e}"
+            )
+
+    def _begin_stage(self, job, stage_name):
+        with self._stage_lock:
+            if self._current_stage is not None:
+                # Implicit stop of the previous stage -- the script forgot to
+                # close it. Don't lose data; flush before starting the next.
+                logger.warning(
+                    f"Perf hook: implicit STOP of stage "
+                    f"{self._current_stage!r} before START {stage_name!r}"
+                )
+                self._end_current_stage_locked(job)
+            # Sanitize the stage name into a filesystem-friendly subdir.
+            sanitized = _sanitize_subdir(stage_name)
+            logger.info(f"Perf hook: starting stage {sanitized!r}")
+            self._current_stage = sanitized
+            self._current_monitors = self._build_and_run_monitors(job, subdir=sanitized)
+
+    def _end_current_stage(self, job):
+        with self._stage_lock:
+            self._end_current_stage_locked(job)
+
+    def _end_current_stage_locked(self, job):
+        if self._current_stage is None:
+            return
+        logger.info(f"Perf hook: stopping stage {self._current_stage!r}")
+        for monitor in self._current_monitors:
+            try:
+                monitor.terminate()
+            except Exception as e:
+                logger.warning(
+                    f"Perf hook: terminating monitor {monitor.name} failed: {e}"
+                )
+        for monitor in self._current_monitors:
+            try:
+                monitor.write_csv()
+            except Exception as e:
+                logger.warning(f"Perf hook: write_csv on {monitor.name} failed: {e}")
+        self._current_stage = None
+        self._current_monitors = []
+
+    # ------------------------------------------------------------------
+    # Shared monitor setup
+    # ------------------------------------------------------------------
+
+    def _build_and_run_monitors(self, job, subdir):
+        """Instantiate every enabled monitor (with the given subdir) and
+        start it. Returns the list of started monitors.
+
+        Refactored out of the original ``before_job`` body so both default
+        mode and stage-aware mode share the same monitor-bring-up logic.
+        """
         should_run_perf_stat = True
-        self.monitors = []
+        monitors = []
         for mon_name in AVAIL_MONITORS.keys():
             # `enabled` is a perf-hook-level flag, not a monitor constructor
             # arg. Pop it out before passing the rest to the monitor class.
@@ -101,34 +295,45 @@ def before_job(self, opts, job):
                 continue
             try:
                 MonitorClass = AVAIL_MONITORS[mon_name]
-                monitor = MonitorClass(job_uuid=job.uuid, **init_args)
+                monitor = MonitorClass(job_uuid=job.uuid, subdir=subdir, **init_args)
                 # We should disable PerfStat (and not run anything that uses PMU)
                 # if IntelPerfSpect3 is enabled.
                 if isinstance(monitor, topdown.IntelPerfSpect3) and monitor.supported:
                     logger.info(
                         "Disabling PerfStat to avoid conflict with IntelPerfSpect3"
                     )
                     should_run_perf_stat = False
-                self.monitors.append(monitor)
+                monitors.append(monitor)
             except Exception as e:
                 logger.warning(
                     f"Failed to load the perf monitor {mon_name} due to the following exception:"
                 )
                 logger.warning(traceback.print_exception(type(e), e, e.__traceback__))
 
-        for monitor in self.monitors:
+        for monitor in monitors:
             try:
                 if isinstance(monitor, perfstat.PerfStat) and not should_run_perf_stat:
                     continue
                 monitor.run()
             except Exception as e:
                 logger.warning(
-                    f"Could not run perf monitor {mon_name} due to the following exception:"
+                    f"Could not run perf monitor {monitor.name} due to the following exception:"
                 )
                 logger.warning(traceback.print_exception(type(e), e, e.__traceback__))
+        return monitors
 
-    def after_job(self, opts, job):
-        for monitor in self.monitors:
-            monitor.terminate()
-        for monitor in self.monitors:
-            monitor.write_csv()
+
+def _sanitize_subdir(name):
+    """Make a stage name safe to use as a directory component.
+
+    Strip path separators and trim. Keep alphanumerics, dashes, underscores;
+    replace anything else with underscore.
+    """
+    out = []
+    for ch in name.strip():
+        if ch.isalnum() or ch in "-_.":
+            out.append(ch)
+        else:
+            out.append("_")
+    cleaned = "".join(out).strip("._")
+    return cleaned or "unnamed_stage"

benchpress/plugins/hooks/perf_monitors/__init__.py

@@ -21,21 +21,44 @@
 
 class Monitor:
     def gen_path(self, filename):
-        return os.path.join(
-            get_artifacts_dir(), f"benchmark_metrics_{self.job_uuid}", filename
-        )
+        """Build a path under the job's benchmark_metrics directory.
+
+        When ``self.subdir`` is set (per-stage perf collection in stage-aware
+        mode), the path is nested one level further so each stage gets its
+        own sub-folder. Without ``subdir`` the layout is unchanged from
+        legacy behavior:
 
-    def __init__(self, interval, name, job_uuid):
-        """Initialize some common parameters and storage variables"""
+            benchmark_metrics_<uuid>/<filename>            # no subdir
+            benchmark_metrics_<uuid>/<subdir>/<filename>   # stage-aware mode
+        """
+        base = os.path.join(get_artifacts_dir(), f"benchmark_metrics_{self.job_uuid}")
+        if getattr(self, "subdir", None):
+            base = os.path.join(base, self.subdir)
+        return os.path.join(base, filename)
+
+    def __init__(self, interval, name, job_uuid, subdir=None):
+        """Initialize some common parameters and storage variables.
+
+        Args:
+            subdir: Optional sub-folder under benchmark_metrics_<uuid>/. Used
+                by the perf hook's stage-aware mode to give each WDL prod_set
+                sub-benchmark its own slice of perf data. None preserves the
+                original flat layout.
+        """
         self.name = name
         self.interval = interval
         # Reserved for result processing
         self.res = []
         # Reserved for original output of the monitoring process
         self.output = ""
         self.job_uuid = job_uuid
+        self.subdir = subdir
         self.logpath = self.gen_path(f"{name}.log")
         self.csvpath = self.gen_path(f"{name}.csv")
+        # Make sure the (possibly nested) directory exists before opening
+        # the log file. Stage-aware mode creates per-stage dirs lazily so we
+        # can't rely on the perf hook to have made them.
+        os.makedirs(os.path.dirname(self.logpath), exist_ok=True)
         self.logfile = open(self.logpath, "w", buffering=1)  # noqa: P201
 
     def __del__(self):

benchpress/plugins/hooks/perf_monitors/cpufreq_cpuinfo.py

@@ -14,8 +14,10 @@
 
 
 class CPUFreq(Monitor):
-    def __init__(self, interval, job_uuid):
-        super(CPUFreq, self).__init__(interval, "cpufreq_cpuinfo", job_uuid)
+    def __init__(self, interval, job_uuid, subdir=None):
+        super(CPUFreq, self).__init__(
+            interval, "cpufreq_cpuinfo", job_uuid, subdir=subdir
+        )
         self.run_freq_collector = False
         self.supported = os.path.exists(
             "/sys/devices/system/cpu/cpu0/cpufreq/cpuinfo_cur_freq"

benchpress/plugins/hooks/perf_monitors/cpufreq_scaling.py

@@ -14,8 +14,10 @@
 
 
 class CPUFreq(Monitor):
-    def __init__(self, interval, job_uuid):
-        super(CPUFreq, self).__init__(interval, "cpufreq_scaling", job_uuid)
+    def __init__(self, interval, job_uuid, subdir=None):
+        super(CPUFreq, self).__init__(
+            interval, "cpufreq_scaling", job_uuid, subdir=subdir
+        )
         self.run_freq_collector = False
         self.supported = os.path.exists(
             "/sys/devices/system/cpu/cpu0/cpufreq/scaling_cur_freq"