diff --git a/verifiers/utils/process_utils.py b/verifiers/utils/process_utils.py index 3fdc8dbd74..c148ff3eda 100644 --- a/verifiers/utils/process_utils.py +++ b/verifiers/utils/process_utils.py @@ -1,9 +1,12 @@ """Process lifecycle utilities.""" +import contextlib import logging import os import signal +import subprocess import threading +import time from multiprocessing.connection import Connection from multiprocessing.process import BaseProcess @@ -83,6 +86,153 @@ def terminate_processes( t.join() +def _proc_session_processes(session: int) -> dict[int, bool] | None: + """Read session membership from procfs, or return None when unavailable.""" + try: + entries = os.scandir("/proc") + except OSError: + return None + + members: dict[int, bool] = {} + with entries: + for entry in entries: + if not entry.name.isdigit(): + continue + try: + with open(f"/proc/{entry.name}/stat") as stat: + fields = stat.read().rsplit(")", 1)[1].split() + if int(fields[3]) == session: + members[int(entry.name)] = fields[0] != "Z" + except (IndexError, OSError, ValueError): + # Processes can exit between scandir and opening their stat file. A live + # member is seen on the next fixed-point scan below. + continue + return members + + +def _session_processes(session: int) -> dict[int, bool] | None: + """Map current POSIX session PIDs to whether each process is non-zombie.""" + if (members := _proc_session_processes(session)) is not None: + return members + + try: + listing = subprocess.run( + ["ps", "-A", "-o", "pid=,sid=,stat="], + check=True, + capture_output=True, + text=True, + timeout=5, + ).stdout + except (OSError, subprocess.SubprocessError): + logger.warning("Could not read POSIX session %d", session) + return None + + members = {} + for line in listing.splitlines(): + try: + raw_pid, raw_sid, state = line.split() + pid, sid = int(raw_pid), int(raw_sid) + except ValueError: + continue + if sid == session: + members[pid] = not state.startswith("Z") + return members + + +def kill_process_session(process: BaseProcess, timeout: float = 10.0) -> None: + """Kill a worker-owned POSIX session, including its separate process groups. + + Pool workers are session leaders, and framework subprocess runtimes create their own + process groups inside that session. Session membership therefore survives worker + exit and does not rely on a stale parent-PID snapshot. Members are stopped until the + session is stable, preventing new forks between discovery and ``SIGKILL``. + + Call this only after giving the process a chance to terminate gracefully. It is a + POSIX helper; if the process did not become its own session leader, only its PID is + killed so an unrelated caller session is never targeted. + """ + root = process.pid + assert root is not None + # The caller waits on the process sentinel without joining first. An exited worker + # therefore remains a zombie and keeps its PID/SID reserved until this cleanup joins it. + members = _session_processes(root) + if members is None: + logger.error( + "Cannot discover worker session %d; killing only the worker process", root + ) + with contextlib.suppress(Exception): + process.kill() + process.join(timeout=timeout) + return + if root not in members: + with contextlib.suppress(Exception): + process.kill() + process.join(timeout=timeout) + return + + live = {pid for pid, is_live in members.items() if is_live} + discovery_failures = 0 + deadline = time.monotonic() + timeout + old_mask = signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGINT, signal.SIGTERM}) + try: + try: + while live and time.monotonic() < deadline: + for pid in live: + with contextlib.suppress(ProcessLookupError, PermissionError): + os.kill(pid, signal.SIGSTOP) + discovered = _session_processes(root) + if discovered is None: + discovery_failures += 1 + if discovery_failures < 3: + continue + logger.error("Could not verify stopped worker session %d", root) + break + discovery_failures = 0 + current = {pid for pid, is_live in discovered.items() if is_live} + if current <= live: + live = current + break + live = current + else: + if live: + logger.warning("Timed out freezing worker session %d", root) + finally: + # Drain rather than kill one snapshot: a member that forked just before SIGSTOP or + # the freeze deadline may have added a process that only the next scan can see. + # Every successful SIGKILL round removes possible forkers, so this converges without + # letting an unfreezable member keep the bounded freeze phase alive forever. + try: + kill_deadline = time.monotonic() + timeout + while live and time.monotonic() < kill_deadline: + signaled = False + for pid in live: + with contextlib.suppress(ProcessLookupError, PermissionError): + if os.getsid(pid) == root: + os.kill(pid, signal.SIGKILL) + signaled = True + if not signaled: + logger.error("Could not kill worker session %d", root) + break + current = _session_processes(root) + if current is None: + logger.error("Could not verify killed worker session %d", root) + break + live = {pid for pid, is_live in current.items() if is_live} + if live: + logger.error( + "Worker session %d still has %d live process(es) after SIGKILL", + root, + len(live), + ) + finally: + with contextlib.suppress(BaseException): + process.kill() + with contextlib.suppress(BaseException): + process.join(timeout=timeout) + finally: + signal.pthread_sigmask(signal.SIG_SETMASK, old_mask) + + def use_threading_tqdm_lock() -> None: """Pin tqdm to a threading lock so it never lazily creates an `mp.RLock` a killed worker would leak (a `resource_tracker` semaphore warning). No-op if tqdm isn't installed.""" diff --git a/verifiers/v1/runtimes/subprocess.py b/verifiers/v1/runtimes/subprocess.py index daf6483621..7235ab1feb 100644 --- a/verifiers/v1/runtimes/subprocess.py +++ b/verifiers/v1/runtimes/subprocess.py @@ -54,14 +54,14 @@ async def run(self, argv: list[str], env: dict[str, str]) -> ProgramResult: cwd=self.workdir, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, - start_new_session=True, # own process group, so we can reap the whole tree + process_group=0, # own group inside the pool worker's cleanup session ) try: stdout, stderr = await proc.communicate() finally: # If the await didn't finish, the caller cancelled it (e.g. the rollout's # scoring_timeout / harness_timeout fired): communicate() leaves the process - # running, so SIGKILL its whole group (start_new_session => pgid == pid) — otherwise + # running, so SIGKILL its whole group (process_group=0 => pgid == pid) — otherwise # a hung child (a wedged uv/sympy verify) outlives the rollout and leaks CPU. A # no-op once it has exited on its own. if proc.returncode is None: @@ -88,7 +88,7 @@ async def run_background( cwd=self.workdir, stdout=f, stderr=asyncio.subprocess.STDOUT, - start_new_session=True, # own process group, so cleanup() reaps the whole tree + process_group=0, # own group inside the pool worker's cleanup session ) self._background.append( proc @@ -104,7 +104,7 @@ async def write(self, path: str, data: bytes) -> None: def cleanup(self) -> None: for proc in self._background: - # Kill the whole group (start_new_session => pgid == pid), not just proc.pid, so a + # Kill the whole group (process_group=0 => pgid == pid), not just proc.pid, so a # background server's children (sh -> uv -> python) are reaped too. with contextlib.suppress(ProcessLookupError, PermissionError): os.killpg(os.getpgid(proc.pid), signal.SIGTERM) diff --git a/verifiers/v1/serve/pool.py b/verifiers/v1/serve/pool.py index faa8b5f503..4f75a2ad63 100644 --- a/verifiers/v1/serve/pool.py +++ b/verifiers/v1/serve/pool.py @@ -28,8 +28,10 @@ import os import signal import threading +import time import uuid from collections.abc import Callable +from multiprocessing.connection import wait import msgpack import zmq @@ -38,6 +40,7 @@ from verifiers.v1.env import EnvConfig from verifiers.v1.serve.server import EnvServer from verifiers.v1.serve.types import HealthResponse, RunGroupRequest +from verifiers.utils.process_utils import kill_process_session logger = logging.getLogger(__name__) @@ -117,7 +120,7 @@ def _spawn_worker(self) -> None: address = f"ipc://{self._worker_path(i)}" parent_conn, child_conn = self._mpctx.Pipe() proc = self._mpctx.Process( - target=serve_env, + target=_serve_pool_worker, kwargs=dict( max_workers=1, address=address, @@ -241,24 +244,43 @@ async def run(self) -> None: self._shutdown() def _shutdown(self) -> None: - for w in self.workers: - with contextlib.suppress(Exception): - w["pipe"].close() - with contextlib.suppress(Exception): - w["process"].terminate() - for w in self.workers: - with contextlib.suppress(Exception): - w["process"].join(timeout=10) - if w["process"].is_alive(): + old_handlers = { + sig: signal.signal(sig, signal.SIG_IGN) + for sig in (signal.SIGINT, signal.SIGTERM) + } + try: + for w in self.workers: + with contextlib.suppress(Exception): + w["pipe"].close() + with contextlib.suppress(Exception): + w["process"].terminate() + + # Wait on sentinels rather than join/is_alive: an exited worker remains a zombie, + # reserving the PID that identifies its cleanup session until that session is reaped. + pending = {w["process"].sentinel for w in self.workers} + deadline = time.monotonic() + 10 + while pending and (remaining := deadline - time.monotonic()) > 0: + pending.difference_update(wait(pending, timeout=remaining)) + + for w in self.workers: with contextlib.suppress(Exception): - w["process"].kill() - with contextlib.suppress(Exception): - w["dealer"].close() - with contextlib.suppress(OSError): - os.unlink(self._worker_path(w["index"])) - self.frontend.close() - self.ctx.term() - logger.info("EnvServerPool down") + kill_process_session(w["process"]) + with contextlib.suppress(Exception): + w["dealer"].close() + with contextlib.suppress(OSError): + os.unlink(self._worker_path(w["index"])) + self.frontend.close() + self.ctx.term() + logger.info("EnvServerPool down") + finally: + for sig, handler in old_handlers.items(): + signal.signal(sig, handler) + + +def _serve_pool_worker(**kwargs) -> None: + """Run a pool worker as the leader of the session that owns its subprocesses.""" + os.setsid() + serve_env(**kwargs) def env_config_data(config) -> dict: