Merge pull request #113 from rsasaki0109/agent/fix-local-env-and-benchmark-cleanup

Fix local environment and benchmark cleanup

Author: rsasaki0109

CMakeLists.txt

@@ -789,6 +789,14 @@ if(BUILD_TESTING)
     NAME bootstrap_colcon_workspace_shell
     COMMAND bash -n ${CMAKE_CURRENT_SOURCE_DIR}/scripts/bootstrap_colcon_workspace.sh
   )
+  add_test(
+    NAME setup_local_env
+    COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test/test_setup_local_env.py
+  )
+  add_test(
+    NAME benchmark_runner_cleanup
+    COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test/test_benchmark_runner_cleanup.py
+  )
   add_test(
     NAME koide_hard_imu_deskew_smoke_shell
     COMMAND bash -n ${CMAKE_CURRENT_SOURCE_DIR}/scripts/run_koide_hard_imu_deskew_smoke.sh

scripts/benchmark_runner

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import argparse
+import atexit
 import csv
 import json
 import math
@@ -147,6 +148,46 @@ class ResourceMonitor(threading.Thread):
                 last_wall = current_wall
 
 
+class ManagedProcessCleanup:
+    """Best-effort process-group cleanup for normal and exceptional exits."""
+
+    def __init__(self, timeout_sec: float) -> None:
+        self._timeout_sec = timeout_sec
+        self._processes: List[ManagedProcess] = []
+        self._monitor: Optional[ResourceMonitor] = None
+        self._lock = threading.Lock()
+        self._cleaned = False
+
+    def add(self, managed: ManagedProcess) -> ManagedProcess:
+        with self._lock:
+            self._processes.append(managed)
+        return managed
+
+    def set_monitor(self, monitor: ResourceMonitor) -> None:
+        with self._lock:
+            self._monitor = monitor
+
+    def stop_all(self) -> None:
+        with self._lock:
+            if self._cleaned:
+                return
+            self._cleaned = True
+            processes = list(reversed(self._processes))
+            monitor = self._monitor
+        if monitor is not None:
+            monitor.stop()
+            monitor.join(timeout=2.0)
+        for managed in processes:
+            try:
+                stop_process(managed, self._timeout_sec)
+            except (OSError, subprocess.SubprocessError):
+                try:
+                    managed.stdout_stream.close()
+                    managed.stderr_stream.close()
+                except OSError:
+                    pass
+
+
 def build_stats(samples: List[Dict[str, float]]) -> Dict[str, Optional[float]]:
     cpu_values = [sample["cpu_percent"] for sample in samples]
     rss_values = [sample["rss_mb"] for sample in samples]
@@ -177,11 +218,17 @@ def launch_process(
     env = os.environ.copy()
     if extra_env:
         env.update(extra_env)
+
+    def prepare_child_process() -> None:
+        os.setsid()
+        signal.pthread_sigmask(
+            signal.SIG_UNBLOCK, {signal.SIGINT, signal.SIGTERM})
+
     process = subprocess.Popen(
         ["bash", "-lc", command],
         stdout=stdout_stream,
         stderr=stderr_stream,
-        preexec_fn=os.setsid,
+        preexec_fn=prepare_child_process,
         env=env,
     )
     return ManagedProcess(
@@ -195,22 +242,66 @@ def launch_process(
     )
 
 
+def launch_registered_process(
+    cleanup: ManagedProcessCleanup,
+    name: str,
+    command: str,
+    output_dir: str,
+    extra_env: Optional[Dict[str, str]] = None,
+) -> ManagedProcess:
+    """Launch a process group and immediately register it for cleanup."""
+    termination_signals = {signal.SIGINT, signal.SIGTERM}
+    previous_mask = signal.pthread_sigmask(signal.SIG_BLOCK, termination_signals)
+    try:
+        managed = launch_process(name, command, output_dir, extra_env)
+        return cleanup.add(managed)
+    finally:
+        signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask)
+
+
+def process_group_exists(process_group_id: int) -> bool:
+    try:
+        os.killpg(process_group_id, 0)
+    except ProcessLookupError:
+        return False
+    except PermissionError:
+        return True
+    return True
+
+
+def wait_for_process_group_exit(
+    managed: ManagedProcess, process_group_id: int, timeout_sec: float
+) -> bool:
+    deadline = time.monotonic() + timeout_sec
+    while process_group_exists(process_group_id) and time.monotonic() < deadline:
+        managed.process.poll()
+        time.sleep(0.02)
+    managed.process.poll()
+    return not process_group_exists(process_group_id)
+
+
 def stop_process(managed: ManagedProcess, timeout_sec: float) -> int:
+    process_group_id = managed.process.pid
+    for shutdown_signal in (signal.SIGINT, signal.SIGTERM, signal.SIGKILL):
+        if not process_group_exists(process_group_id):
+            break
+        try:
+            os.killpg(process_group_id, shutdown_signal)
+        except ProcessLookupError:
+            break
+        if wait_for_process_group_exit(managed, process_group_id, timeout_sec):
+            break
+
     if managed.process.poll() is None:
-        os.killpg(os.getpgid(managed.process.pid), signal.SIGINT)
         try:
             managed.process.wait(timeout=timeout_sec)
         except subprocess.TimeoutExpired:
-            os.killpg(os.getpgid(managed.process.pid), signal.SIGTERM)
-            try:
-                managed.process.wait(timeout=timeout_sec)
-            except subprocess.TimeoutExpired:
-                os.killpg(os.getpgid(managed.process.pid), signal.SIGKILL)
-                managed.process.wait(timeout=timeout_sec)
+            managed.process.kill()
+            managed.process.wait(timeout=timeout_sec)
 
     managed.stdout_stream.close()
     managed.stderr_stream.close()
-    return int(managed.process.returncode)
+    return int(managed.process.returncode if managed.process.returncode is not None else -1)
 
 
 def detect_process_died_positions(log_path: str, target_process_pattern: str) -> Dict[str, int]:
@@ -369,9 +460,17 @@ def wait_for_pid(pattern: str, timeout_sec: float) -> Optional[int]:
     return None
 
 
+def exit_on_termination_signal(signum: int, _frame: object) -> None:
+    """Convert SIGTERM into normal interpreter unwinding so atexit cleanup runs."""
+    raise SystemExit(128 + signum)
+
+
 def main() -> int:
     args = parse_args()
     os.makedirs(args.output_dir, exist_ok=True)
+    cleanup = ManagedProcessCleanup(args.shutdown_timeout)
+    atexit.register(cleanup.stop_all)
+    signal.signal(signal.SIGTERM, exit_on_termination_signal)
 
     pose_csv_path = os.path.join(args.output_dir, "pose_trace.csv")
     diagnostic_csv_path = os.path.join(args.output_dir, "alignment_status.csv")
@@ -400,10 +499,14 @@ def main() -> int:
 
     started_at = now_iso()
     settle_start_monotonic = time.monotonic()
-    system = launch_process("system", system_command, args.output_dir, launch_env)
-    recorder = launch_process("pose_recorder", recorder_command, args.output_dir, launch_env)
+    system = launch_registered_process(
+        cleanup, "system", system_command, args.output_dir, launch_env)
+    recorder = launch_registered_process(
+        cleanup, "pose_recorder", recorder_command, args.output_dir, launch_env)
     diagnostic_recorder = (
-        launch_process("diagnostic_recorder", diagnostic_recorder_command, args.output_dir, launch_env)
+        launch_registered_process(
+            cleanup, "diagnostic_recorder", diagnostic_recorder_command,
+            args.output_dir, launch_env)
         if args.diagnostic_topic
         else None
     )
@@ -412,13 +515,15 @@ def main() -> int:
     monitored_pid = wait_for_pid(args.target_process_pattern, args.settle_seconds)
     if monitored_pid is not None:
         monitor = ResourceMonitor(monitored_pid, resource_csv_path, args.monitor_interval)
+        cleanup.set_monitor(monitor)
         monitor.start()
 
     settle_elapsed = time.monotonic() - settle_start_monotonic
     time.sleep(max(0.0, args.settle_seconds - settle_elapsed))
 
     bag_started_monotonic = time.monotonic()
-    bag = launch_process("bag_play", bag_command, args.output_dir, launch_env)
+    bag = launch_registered_process(
+        cleanup, "bag_play", bag_command, args.output_dir, launch_env)
     bag_stopped_by_runner = False
     if args.bag_duration > 0.0:
         deadline = bag_started_monotonic + args.bag_duration
@@ -524,4 +629,8 @@ def main() -> int:
 
 
 if __name__ == "__main__":
-    sys.exit(main())
+    try:
+        exit_code = main()
+    except KeyboardInterrupt:
+        exit_code = 130
+    sys.exit(exit_code)

scripts/setup_local_env.sh

@@ -15,7 +15,14 @@ else
 fi
 unset _lidarloc_ws_root_default
 _lidarloc_prefix="${LIDAR_LOCALIZATION_LOCAL_PREFIX:-${_lidarloc_ws_root}/local_prefix}"
-_lidarloc_overlay="${LIDAR_LOCALIZATION_OVERLAY:-${_lidarloc_ws_root}/build_ws/install/setup.bash}"
+if [[ -n "${LIDAR_LOCALIZATION_OVERLAY:-}" ]]; then
+  _lidarloc_overlay="${LIDAR_LOCALIZATION_OVERLAY}"
+elif [[ -f "${_lidarloc_ws_root}/install/setup.bash" ]]; then
+  _lidarloc_overlay="${_lidarloc_ws_root}/install/setup.bash"
+else
+  # Keep compatibility with the older nested workspace layout.
+  _lidarloc_overlay="${_lidarloc_ws_root}/build_ws/install/setup.bash"
+fi
 
 _lidarloc_ros_distro=""
 for _lidarloc_candidate in "${LIDAR_LOCALIZATION_ROS_DISTRO:-}" "${ROS_DISTRO:-}" humble jazzy; do

test/test_benchmark_runner_cleanup.py

@@ -0,0 +1,86 @@
+#!/usr/bin/env python3
+
+from pathlib import Path
+import signal
+import subprocess
+import tempfile
+import time
+import unittest
+
+
+REPO = Path(__file__).resolve().parents[1]
+RUNNER = REPO / "scripts" / "benchmark_runner"
+
+
+def child_pids(pid: int) -> list[int]:
+    path = Path(f"/proc/{pid}/task/{pid}/children")
+    try:
+        return [int(value) for value in path.read_text().split()]
+    except FileNotFoundError:
+        return []
+
+
+def cmdline(pid: int) -> str:
+    try:
+        return Path(f"/proc/{pid}/cmdline").read_bytes().replace(b"\0", b" ").decode()
+    except FileNotFoundError:
+        return ""
+
+
+class BenchmarkRunnerCleanupTest(unittest.TestCase):
+    def run_interrupted_benchmark(self, interrupt_signal: signal.Signals, expected_code: int):
+        with tempfile.TemporaryDirectory() as tmp:
+            output = Path(tmp) / "output"
+            process = subprocess.Popen(
+                [
+                    str(RUNNER),
+                    "--bag-path", str(Path(tmp) / "unused_bag"),
+                    "--output-dir", str(output),
+                    "--system-command", "exec -a benchmark_cleanup_system sleep 60",
+                    "--bag-command", "exec -a benchmark_cleanup_bag sleep 60",
+                    "--diagnostic-topic", "",
+                    "--target-process-pattern", "benchmark-target-that-does-not-exist",
+                    "--settle-seconds", "0.1",
+                    "--post-roll-seconds", "0",
+                    "--shutdown-timeout", "0.2",
+                ],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+            try:
+                deadline = time.monotonic() + 5.0
+                named_children = []
+                while time.monotonic() < deadline:
+                    named_children = [
+                        pid for pid in child_pids(process.pid)
+                        if "benchmark_cleanup_" in cmdline(pid)]
+                    if len(named_children) == 2:
+                        break
+                    time.sleep(0.02)
+                self.assertEqual(len(named_children), 2)
+                tracked = child_pids(process.pid)
+
+                process.send_signal(interrupt_signal)
+                self.assertEqual(process.wait(timeout=5.0), expected_code)
+                deadline = time.monotonic() + 2.0
+                while (
+                    any(Path(f"/proc/{pid}").exists() for pid in tracked)
+                    and time.monotonic() < deadline
+                ):
+                    time.sleep(0.02)
+                survivors = [pid for pid in tracked if Path(f"/proc/{pid}").exists()]
+                self.assertFalse(survivors, f"managed processes survived cleanup: {survivors}")
+            finally:
+                if process.poll() is None:
+                    process.kill()
+                    process.wait(timeout=5.0)
+
+    def test_sigint_stops_all_managed_process_groups(self):
+        self.run_interrupted_benchmark(signal.SIGINT, 130)
+
+    def test_sigterm_stops_all_managed_process_groups(self):
+        self.run_interrupted_benchmark(signal.SIGTERM, 128 + signal.SIGTERM)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/test_setup_local_env.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+
+from pathlib import Path
+import os
+import subprocess
+import tempfile
+import unittest
+
+
+REPO = Path(__file__).resolve().parents[1]
+SCRIPT = REPO / "scripts" / "setup_local_env.sh"
+
+
+class SetupLocalEnvTest(unittest.TestCase):
+    def run_setup(self, workspace: Path, explicit: Path | None = None) -> str:
+        env = os.environ.copy()
+        env["LIDAR_LOCALIZATION_WS_ROOT"] = str(workspace)
+        env["LIDAR_LOCALIZATION_ROS_DISTRO"] = os.environ.get("ROS_DISTRO", "jazzy")
+        env.pop("LIDAR_LOCALIZATION_OVERLAY", None)
+        if explicit is not None:
+            env["LIDAR_LOCALIZATION_OVERLAY"] = str(explicit)
+        result = subprocess.run(
+            ["bash", "-c", f'source "{SCRIPT}" && printf %s "$LIDAR_TEST_OVERLAY_SOURCED"'],
+            check=True, capture_output=True, text=True, env=env)
+        return result.stdout
+
+    def test_conventional_workspace_install_is_sourced(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            workspace = Path(tmp)
+            setup = workspace / "install" / "setup.bash"
+            setup.parent.mkdir(parents=True)
+            setup.write_text("export LIDAR_TEST_OVERLAY_SOURCED=conventional\n")
+            self.assertEqual(self.run_setup(workspace), "conventional")
+
+    def test_explicit_overlay_has_priority(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            workspace = Path(tmp)
+            conventional = workspace / "install" / "setup.bash"
+            explicit = workspace / "explicit_setup.bash"
+            conventional.parent.mkdir(parents=True)
+            conventional.write_text("export LIDAR_TEST_OVERLAY_SOURCED=conventional\n")
+            explicit.write_text("export LIDAR_TEST_OVERLAY_SOURCED=explicit\n")
+            self.assertEqual(self.run_setup(workspace, explicit), "explicit")
+
+
+if __name__ == "__main__":
+    unittest.main()