rpx_benchmark.cleanup

  1"""Graceful-shutdown plumbing for long-running CLI scripts.
  2
  3Problem this solves
  4-------------------
  5Top-level scripts (``run_depth.py``, ``visualize_rerun.py``,
  6``dataset_hub.cli``, ...) spawn subprocesses: HF download workers, the
  7Rerun viewer, DataLoader workers, ffmpeg encoders. When the user hits
  8Ctrl+C or the parent receives SIGTERM, Python's default behaviour
  9raises ``KeyboardInterrupt`` only in the main thread — child processes
 10keep running and accumulate as orphans (visible as background CPU /
 11memory the user has to hunt down later).
 12
 13What this module does
 14---------------------
 15:func:`install_signal_cleanup` installs three guarantees in one call:
 16
 171. The current process becomes its **own process group leader** via
 18   ``os.setpgrp()``. Every child / grandchild spawned afterwards
 19   inherits the group id by default.
 202. ``SIGINT`` and ``SIGTERM`` deliver a ``SIGTERM`` to the **whole
 21   group**, wait up to ``timeout`` seconds for graceful exit, then send
 22   ``SIGKILL`` to anything still alive. The signalled process exits via
 23   ``sys.exit(128 + signum)`` so its own exit code reflects the signal.
 243. A double-tap (Ctrl+C twice within the timeout) skips the wait and
 25   goes straight to SIGKILL of the group.
 26
 27Idempotent: calling it twice is a no-op. Linux-only — on macOS / BSD
 28the process-group dance still works, on Windows the call is a no-op
 29(no equivalent group semantics, no harm in importing).
 30"""
 31
 32from __future__ import annotations
 33
 34import atexit
 35import os
 36import signal
 37import sys
 38import time
 39from types import FrameType
 40from typing import Optional
 41
 42_INSTALLED = False
 43_KILL_TIMEOUT = 5.0
 44_DOUBLE_TAP_AT: Optional[float] = None
 45
 46
 47def _signal_to_group(sig: int) -> None:
 48    """Send ``sig`` to our process group, ignoring failures (group already
 49    empty, no permission, race with a child exiting normally)."""
 50    try:
 51        os.killpg(os.getpgrp(), sig)
 52    except (ProcessLookupError, PermissionError, OSError):
 53        pass
 54
 55
 56def _on_signal(signum: int, _frame: Optional[FrameType]) -> None:
 57    global _DOUBLE_TAP_AT
 58    now = time.monotonic()
 59    name = signal.Signals(signum).name if signum in signal.Signals._value2member_map_ else str(signum)
 60
 61    # Double-tap (e.g. two Ctrl+C in quick succession) → escalate immediately.
 62    if _DOUBLE_TAP_AT is not None and now - _DOUBLE_TAP_AT < _KILL_TIMEOUT:
 63        sys.stderr.write(f"\n[cleanup] {name} (double-tap) — SIGKILL process group\n")
 64        _signal_to_group(signal.SIGKILL)
 65        os._exit(128 + signum)
 66    _DOUBLE_TAP_AT = now
 67
 68    sys.stderr.write(
 69        f"\n[cleanup] {name} received — SIGTERM process group, "
 70        f"waiting up to {_KILL_TIMEOUT:.0f}s for graceful exit "
 71        "(press Ctrl+C again to SIGKILL)\n"
 72    )
 73    # SIGTERM the children. We don't SIGTERM ourselves — we want to
 74    # complete shutdown and exit normally.
 75    pgid = os.getpgrp()
 76    for pid in _list_children(pgid):
 77        if pid == os.getpid():
 78            continue
 79        try:
 80            os.kill(pid, signal.SIGTERM)
 81        except (ProcessLookupError, PermissionError):
 82            pass
 83
 84    # Wait for the group to drain.
 85    deadline = now + _KILL_TIMEOUT
 86    while time.monotonic() < deadline and _children_still_alive(pgid):
 87        time.sleep(0.1)
 88
 89    if _children_still_alive(pgid):
 90        sys.stderr.write("[cleanup] timeout — SIGKILL stragglers\n")
 91        for pid in _list_children(pgid):
 92            if pid == os.getpid():
 93                continue
 94            try:
 95                os.kill(pid, signal.SIGKILL)
 96            except (ProcessLookupError, PermissionError):
 97                pass
 98
 99    sys.exit(128 + signum)
100
101
102def _list_children(pgid: int) -> list[int]:
103    """Return PIDs in our process group (Linux /proc walk; falls back to []
104    on platforms without /proc)."""
105    try:
106        entries = os.listdir("/proc")
107    except OSError:
108        return []
109    out: list[int] = []
110    for name in entries:
111        if not name.isdigit():
112            continue
113        try:
114            with open(f"/proc/{name}/stat", "rb") as f:
115                fields = f.read().split()
116            # /proc/<pid>/stat field 5 (1-indexed) is pgid
117            if len(fields) >= 5 and int(fields[4]) == pgid:
118                out.append(int(name))
119        except (OSError, ValueError):
120            continue
121    return out
122
123
124def _children_still_alive(pgid: int) -> bool:
125    return any(pid != os.getpid() for pid in _list_children(pgid))
126
127
128def _on_atexit() -> None:
129    """On normal interpreter shutdown, SIGTERM any lingering descendants.
130
131    Don't escalate to SIGKILL here — normal exit usually means children
132    have already cleaned themselves up; this is a safety net.
133    """
134    pgid = os.getpgrp()
135    for pid in _list_children(pgid):
136        if pid == os.getpid():
137            continue
138        try:
139            os.kill(pid, signal.SIGTERM)
140        except (ProcessLookupError, PermissionError):
141            pass
142
143
144def install_signal_cleanup(timeout: float = 5.0) -> None:
145    """Install signal handlers + process group setup.
146
147    Safe to call multiple times — only the first call takes effect.
148    No-op on Windows (no process-group semantics).
149
150    Parameters
151    ----------
152    timeout
153        Seconds to wait between sending SIGTERM and escalating to
154        SIGKILL. The default (5s) gives child PyTorch / HF workers
155        enough time to flush state without making Ctrl+C feel slow.
156    """
157    global _INSTALLED, _KILL_TIMEOUT
158    if _INSTALLED:
159        return
160    if sys.platform.startswith("win"):
161        _INSTALLED = True
162        return
163    _KILL_TIMEOUT = float(timeout)
164    # New process group → killing the leader kills every descendant.
165    try:
166        os.setpgrp()
167    except OSError:
168        # Already a group leader, or unsupported on this platform — fine.
169        pass
170    signal.signal(signal.SIGINT, _on_signal)
171    signal.signal(signal.SIGTERM, _on_signal)
172    atexit.register(_on_atexit)
173    _INSTALLED = True