Add generic stage-aware perf hook infrastructure (#683)

benchpress/plugins/hooks/perf.py

@@ -8,6 +8,7 @@
 
 import logging
 import os
+import threading
 import traceback
 
 from benchpress.lib import open_source
@@ -76,6 +77,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
@@ -91,8 +113,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.
@@ -102,34 +296,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"

benchpress/plugins/hooks/perf_monitors/memstat.py

@@ -22,8 +22,8 @@
 
 
 class MemStat(Monitor):
-    def __init__(self, interval, job_uuid, additional_counters=()):
-        super(MemStat, self).__init__(interval, "mem-stat", job_uuid)
+    def __init__(self, interval, job_uuid, additional_counters=(), subdir=None):
+        super(MemStat, self).__init__(interval, "mem-stat", job_uuid, subdir=subdir)
         counters = {"MemTotal", "MemFree", "MemAvailable", "SwapTotal", "SwapFree"}
         self.counters = counters.union(set(additional_counters))
         self.run_collector = False

benchpress/plugins/hooks/perf_monitors/mpstat.py

@@ -12,8 +12,8 @@
 
 
 class MPStat(Monitor):
-    def __init__(self, interval, job_uuid):
-        super(MPStat, self).__init__(interval, "mpstat", job_uuid)
+    def __init__(self, interval, job_uuid, subdir=None):
+        super(MPStat, self).__init__(interval, "mpstat", job_uuid, subdir=subdir)
         self.headers = []
 
     def run(self):

benchpress/plugins/hooks/perf_monitors/netstat.py

@@ -14,8 +14,8 @@
 
 
 class NetStat(Monitor):
-    def __init__(self, interval, job_uuid, additional_counters=()):
-        super(NetStat, self).__init__(interval, "net-stat", job_uuid)
+    def __init__(self, interval, job_uuid, additional_counters=(), subdir=None):
+        super(NetStat, self).__init__(interval, "net-stat", job_uuid, subdir=subdir)
         counters = {"rx_bytes", "rx_packets", "tx_bytes", "tx_packets"}
         self.counters = counters.union(set(additional_counters))
         self.run_collector = False

benchpress/plugins/hooks/perf_monitors/perfstat.py

@@ -39,8 +39,10 @@ def unpack_perf_stat_line(line, delim=","):
 
 
 class PerfStat(Monitor):
-    def __init__(self, interval, job_uuid, additional_events=(), delim=","):
-        super(PerfStat, self).__init__(interval, "perf-stat", job_uuid)
+    def __init__(
+        self, interval, job_uuid, additional_events=(), delim=",", subdir=None
+    ):
+        super(PerfStat, self).__init__(interval, "perf-stat", job_uuid, subdir=subdir)
         self.events = ["instructions", "cycles"] + list(additional_events)
         self.delim = delim