fix(asap-tools): keep ProcessMonitor samples when a monitored process dies

asap-tools/experiments/classes/process_monitor.py

@@ -176,6 +176,34 @@ def update_pid_monitor_map(self, p) -> List[ProcessMetricSnapshot]:
 
         return iteration_info
 
+    def _sample_counts(self) -> dict:
+        # All list fields of a pid entry grow in lockstep, so one length per pid suffices.
+        return {
+            pid: len(m[self.monitors[0]]) for pid, m in self.pid_monitor_map.items()
+        }
+
+    def _terminate_early(self, dead_pid: int, counts_before_round: dict):
+        """
+        Roll back the partial sampling round so every series stays aligned, and
+        mark the dead seed pid so consumers can tell the run was cut short.
+        """
+        for pid in list(self.pid_monitor_map):
+            if pid not in counts_before_round:
+                del self.pid_monitor_map[pid]
+                continue
+            n = counts_before_round[pid]
+            for value in self.pid_monitor_map[pid].values():
+                if isinstance(value, list):
+                    del value[n:]
+        self.pid_monitor_map[dead_pid][
+            constants.PROCESS_MONITOR_EXITED_AT_SAMPLE_KEY
+        ] = counts_before_round[dead_pid]
+        print(
+            f"Monitored process {dead_pid} "
+            f"({self.pid_monitor_map[dead_pid]['keyword']}) exited; "
+            f"stopping monitor after {counts_before_round[dead_pid]} samples"
+        )
+
     def run(self):
         # NOTE: Possibility of init() (and close()) being called more than once if multiple
         #       processes get started up that were passed the same reference
@@ -201,27 +229,41 @@ def run(self):
 
                 iteration_info = []
                 stop_requested = False
+                seed_died = False
+                counts_before_round = self._sample_counts()
                 for pid, p in self.psutil_handles.items():
                     if self.pipe.poll(0):
                         stop_requested = True
                         break
-                    iteration_info += self.update_pid_monitor_map(p)
+                    try:
+                        iteration_info += self.update_pid_monitor_map(p)
+                        children = (
+                            p.children(recursive=True) if self.include_children else []
+                        )
+                    except psutil.NoSuchProcess:
+                        self._terminate_early(pid, counts_before_round)
+                        seed_died = True
+                        break
                     if (
                         self.thread_attribution_keyword is not None
                         and self.pid_monitor_map[pid]["keyword"]
                         == self.thread_attribution_keyword
                     ):
                         self._compute_thread_group_cpu(pid, elapsed)
                     if self.include_children:
-                        for child in p.children(recursive=True):
+                        for child in children:
                             if self.pipe.poll(0):
                                 stop_requested = True
                                 break
                             if child.pid not in self.pid_monitor_map:
                                 self.add_child_pid_to_map(pid, child.pid)
                                 self.child_handles[child.pid] = child
                             handle = self.child_handles[child.pid]
-                            iteration_info += self.update_pid_monitor_map(handle)
+                            try:
+                                iteration_info += self.update_pid_monitor_map(handle)
+                            except psutil.NoSuchProcess:
+                                # Children come and go; the dead child's series just ends here.
+                                continue
                             if (
                                 self.thread_attribution_keyword is not None
                                 and self.pid_monitor_map[handle.pid]["keyword"]
@@ -231,7 +273,7 @@ def run(self):
                         if stop_requested:
                             break
 
-                if stop_requested:
+                if stop_requested or seed_died:
                     break
 
                 self.update_hooks(iteration_info)

asap-tools/experiments/constants.py

@@ -36,6 +36,8 @@
 INGEST_MONITOR_SHUTDOWN_POLL_INTERVAL_SECONDS = 0.5
 PROCESS_MONITOR_STOP_TIMEOUT_SECONDS = 30
 PROCESS_MONITOR_JOIN_TIMEOUT_SECONDS = 10
+# monitor_output.json: set on a seed pid's entry when it exited mid-run; value is its sample count.
+PROCESS_MONITOR_EXITED_AT_SAMPLE_KEY = "exited_at_sample"
 
 PROMETHEUS_CONFIG_DIR = "prometheus_config"
 PROMETHEUS_CONFIG_FILE = "prometheus.yml"

asap-tools/experiments/post_experiment/single_experiment/compare_costs.py

@@ -358,6 +358,26 @@ def main(args):
         with open(monitor_info_file, "r") as f:
             monitor_info = json.load(f)
 
+        exited = [
+            f"{pid} ({entry['keyword']}) at sample "
+            f"{entry[constants.PROCESS_MONITOR_EXITED_AT_SAMPLE_KEY]}"
+            for pid, entry in monitor_info.items()
+            if constants.PROCESS_MONITOR_EXITED_AT_SAMPLE_KEY in entry
+        ]
+        if exited:
+            msg = f"Mode {experiment_mode}: monitored process exited mid-run: {', '.join(exited)}"
+            if not args.allow_partial_runs:
+                print(
+                    f"ERROR: {msg}. Skipping this mode; "
+                    "pass --allow_partial_runs to include it anyway.",
+                    file=sys.stderr,
+                )
+                continue
+            print(
+                f"WARNING: {msg}; costs cover only the samples before exit.",
+                file=sys.stderr,
+            )
+
         resources_across_pids = defaultdict(list)
 
         pids = list(monitor_info.keys())
@@ -608,6 +628,14 @@ def main(args):
             "5th-percentile-of-this-run heuristic."
         ),
     )
+    parser.add_argument(
+        "--allow_partial_runs",
+        action="store_true",
+        help=(
+            "Include modes where a monitored process exited mid-run, computing "
+            "costs over the samples before the exit (otherwise such modes are skipped)."
+        ),
+    )
     parser.add_argument(
         "--machine-readable",
         action="store_true",

asap-tools/experiments/remote_monitor.py

@@ -334,6 +334,12 @@ def main(args):
 
         if prometheus_client_service.use_container:
             while prometheus_client_service.is_healthy():
+                # The sampler only sends data unprompted when a monitored process died.
+                if control_pipe.poll(0):
+                    logger.error(
+                        "A monitored process exited mid-run; stopping prometheus client early"
+                    )
+                    break
                 logger.debug(
                     "Waiting for prometheus client container to stop running..."
                 )
@@ -380,7 +386,9 @@ def _request_shutdown(signum, _frame):
             signal.signal(signal.SIGINT, signal.SIG_DFL)
     elif args.execution_mode == "timed":
         logger.debug(f"Running for {args.time_to_run} seconds")
-        time.sleep(args.time_to_run)
+        # Returns early if the sampler sends data because a monitored process died.
+        if control_pipe.poll(args.time_to_run):
+            logger.error("A monitored process exited mid-run; stopping early")
 
     if qe_flamegraph_procs:
         logger.debug("Stopping profiling for query engine pids")
@@ -406,6 +414,13 @@ def _request_shutdown(signum, _frame):
         if os.path.exists(monitor_output_file):
             os.remove(monitor_output_file)
     else:
+        for pid, entry in monitor_info.items():
+            if constants.PROCESS_MONITOR_EXITED_AT_SAMPLE_KEY in entry:
+                logger.error(
+                    f"Monitored process {pid} ({entry['keyword']}) exited after "
+                    f"{entry[constants.PROCESS_MONITOR_EXITED_AT_SAMPLE_KEY]} samples; "
+                    f"{args.monitor_output_file} covers only the run up to that point"
+                )
         with open(monitor_output_file, "w") as f:
             json.dump(monitor_info, f)
 

asap-tools/experiments/requirements.txt

@@ -8,3 +8,4 @@ msgpack==1.1.0
 numpy>=1.26,<2
 pandas==2.0.3
 plotnine==0.12.4
+psutil==5.9.8

asap-tools/experiments/tests/test_process_monitor.py

@@ -0,0 +1,105 @@
+"""Tests for ProcessMonitor behavior when monitored processes exit mid-run."""
+
+import multiprocessing
+import subprocess
+import threading
+import time
+import unittest
+from collections import namedtuple
+from unittest import mock
+
+import psutil
+
+import constants
+from classes import process_monitor
+
+MONITORS = ["memory_info", "cpu_percent"]
+INTERVAL_SECONDS = 0.05
+_Mem = namedtuple("_Mem", ["rss"])
+
+
+def _fake_process(pid, children=()):
+    p = mock.Mock()
+    p.pid = pid
+    p.as_dict.return_value = {"memory_info": _Mem(rss=1), "cpu_percent": 0.0}
+    p.children.return_value = list(children)
+    return p
+
+
+class SeedProcessExitTest(unittest.TestCase):
+    def test_seed_exit_returns_aligned_samples_with_exit_marker(self):
+        # A seed pid dying used to crash the sampler and lose every pid's samples.
+        survivor = subprocess.Popen(["sleep", "60"])
+        victim = subprocess.Popen(["sleep", "60"])
+        self.addCleanup(survivor.kill)
+        self.addCleanup(victim.kill)
+
+        monitor, control_pipe, monitor_pipe = process_monitor.start_monitor(
+            [survivor.pid, victim.pid],
+            ["survivor", "victim"],
+            INTERVAL_SECONDS,
+            MONITORS,
+            include_children=True,
+            hooks=None,
+        )
+        time.sleep(INTERVAL_SECONDS * 5)
+        victim.kill()
+        victim.wait()
+
+        self.assertTrue(control_pipe.poll(5), "sampler should send data unprompted")
+        info = process_monitor.stop_monitor(
+            monitor, control_pipe, monitor_pipe, timeout=5
+        )
+
+        self.assertIsNotNone(info)
+        n = len(info[survivor.pid]["cpu_percent"])
+        self.assertGreater(n, 0)
+        # The survivor is sampled before the victim in each round, so this also
+        # checks that the partial round was rolled back.
+        for entry in info.values():
+            self.assertEqual(len(entry["cpu_percent"]), n)
+            self.assertEqual(len(entry["memory_info"]), n)
+        self.assertEqual(
+            info[victim.pid][constants.PROCESS_MONITOR_EXITED_AT_SAMPLE_KEY], n
+        )
+        self.assertNotIn(
+            constants.PROCESS_MONITOR_EXITED_AT_SAMPLE_KEY, info[survivor.pid]
+        )
+
+
+class ChildProcessExitTest(unittest.TestCase):
+    def test_child_exit_does_not_stop_monitor(self):
+        dead_child = _fake_process(2)
+        dead_child.as_dict.side_effect = psutil.NoSuchProcess(2)
+        seed = _fake_process(1, children=[dead_child])
+
+        with mock.patch.object(psutil, "Process", return_value=seed):
+            control_pipe, monitor_pipe = multiprocessing.Pipe()
+            monitor = process_monitor.MyMonitor(
+                [1],
+                ["seed"],
+                monitor_pipe,
+                INTERVAL_SECONDS,
+                MONITORS,
+                hooks=None,
+                include_children=True,
+            )
+        # Run in a thread so the mocks stay in effect.
+        thread = threading.Thread(target=monitor.run)
+        thread.start()
+        self.assertEqual(control_pipe.recv(), "ready")
+
+        # A crash closes the pipe (poll -> True); a surviving monitor stays quiet until stopped.
+        self.assertFalse(control_pipe.poll(INTERVAL_SECONDS * 5))
+        control_pipe.send("stop")
+        self.assertTrue(control_pipe.poll(5))
+        info = control_pipe.recv()
+        thread.join(5)
+
+        self.assertGreater(len(info[1]["cpu_percent"]), 1)
+        self.assertNotIn(constants.PROCESS_MONITOR_EXITED_AT_SAMPLE_KEY, info[1])
+        self.assertEqual(info[2]["cpu_percent"], [])
+
+
+if __name__ == "__main__":
+    unittest.main()