fix(omp-rpc): fixed rpc process cleanup and moved status data off the event loop
- Updated `RpcClient` to spawn `omp` in a new session, cache its process group, and terminate descendants on stop. - Fixed `stop()` to signal the cached process group with SIGTERM then SIGKILL so leaked grandchildren are terminated. - Moved status-query reads in `create_app` to `asyncio.to_thread`, preventing FastAPI event-loop stalls. - Added regression coverage for stopping a spawned grandchild process in the client tests.
This commit is contained in:
@@ -3,6 +3,7 @@ from __future__ import annotations
|
||||
import json
|
||||
import os
|
||||
import queue
|
||||
import signal
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
@@ -112,6 +113,69 @@ _DEFAULT_ERROR_HISTORY_LIMIT = 128
|
||||
_TODO_STATUS_VALUES = frozenset({"pending", "in_progress", "completed", "abandoned"})
|
||||
|
||||
|
||||
def _process_group_id(process: subprocess.Popen[Any]) -> int | None:
|
||||
"""Process-group id of `process`, or `None` when groups are unavailable.
|
||||
|
||||
Captured right after spawn so teardown can signal the whole group even
|
||||
after the leader is reaped — POSIX `os.getpgid` fails on a reaped pid.
|
||||
"""
|
||||
getpgid = getattr(os, "getpgid", None)
|
||||
if getpgid is None:
|
||||
return None
|
||||
try:
|
||||
return getpgid(process.pid)
|
||||
except OSError:
|
||||
return None
|
||||
|
||||
|
||||
def _terminate_process_group(process: subprocess.Popen[Any], pgid: int | None) -> None:
|
||||
"""Terminate the subprocess *and* every descendant sharing its group.
|
||||
|
||||
omp is spawned with `start_new_session=True`, so it leads a session/group
|
||||
that also contains children spawned by the agent's `bash` tool (e.g. a
|
||||
`bun test` run). Signalling only the leader pid would orphan those
|
||||
grandchildren: they reparent to the container init and keep running
|
||||
untracked — how a runaway test ballooned to tens of GB of RAM. Signal the
|
||||
whole group, escalating SIGTERM -> SIGKILL, so descendants die with the
|
||||
task even when the leader has already exited on its own (the graceful
|
||||
stdin-close path).
|
||||
|
||||
`pgid` is captured at spawn; `os.killpg` is POSIX-only, so without it
|
||||
(Windows) we fall back to terminating the leader process alone.
|
||||
"""
|
||||
killpg = getattr(os, "killpg", None)
|
||||
if pgid is None or killpg is None:
|
||||
if process.poll() is None:
|
||||
process.terminate()
|
||||
try:
|
||||
process.wait(timeout=1.0)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
try:
|
||||
process.wait(timeout=1.0)
|
||||
except subprocess.TimeoutExpired:
|
||||
pass
|
||||
return
|
||||
|
||||
def _signal_group(sig: int) -> None:
|
||||
try:
|
||||
killpg(pgid, sig)
|
||||
except OSError:
|
||||
# ESRCH: the group is already empty. Teardown is best-effort.
|
||||
pass
|
||||
|
||||
_signal_group(signal.SIGTERM)
|
||||
try:
|
||||
process.wait(timeout=1.0)
|
||||
except subprocess.TimeoutExpired:
|
||||
pass
|
||||
_signal_group(signal.SIGKILL)
|
||||
try:
|
||||
process.wait(timeout=1.0)
|
||||
except subprocess.TimeoutExpired:
|
||||
pass
|
||||
|
||||
|
||||
def _clone_json_value(value: object) -> JsonValue:
|
||||
if value is None or isinstance(value, (str, int, float, bool)):
|
||||
return cast(JsonValue, value)
|
||||
@@ -330,6 +394,7 @@ class RpcClient:
|
||||
)
|
||||
|
||||
self._process: subprocess.Popen[str] | None = None
|
||||
self._pgid: int | None = None
|
||||
self._stdout_thread: threading.Thread | None = None
|
||||
self._stderr_thread: threading.Thread | None = None
|
||||
self._ready = threading.Event()
|
||||
@@ -429,8 +494,10 @@ class RpcClient:
|
||||
encoding="utf-8",
|
||||
errors="replace",
|
||||
bufsize=1,
|
||||
start_new_session=True,
|
||||
)
|
||||
self._process = process
|
||||
self._pgid = _process_group_id(process)
|
||||
|
||||
self._stdout_thread = threading.Thread(
|
||||
target=self._read_stdout_loop, name="omp-rpc-stdout", daemon=True
|
||||
@@ -486,13 +553,7 @@ class RpcClient:
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
if process.poll() is None:
|
||||
process.terminate()
|
||||
try:
|
||||
process.wait(timeout=1.0)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
process.wait(timeout=1.0)
|
||||
_terminate_process_group(process, self._pgid)
|
||||
finally:
|
||||
if process.stdout is not None:
|
||||
try:
|
||||
@@ -516,6 +577,7 @@ class RpcClient:
|
||||
self._pending_host_tool_calls.clear()
|
||||
self._pending_host_uri_requests.clear()
|
||||
self._process = None
|
||||
self._pgid = None
|
||||
if self._stdout_thread is not None:
|
||||
self._stdout_thread.join(timeout=1.0)
|
||||
if self._stderr_thread is not None:
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import signal
|
||||
import sys
|
||||
import tempfile
|
||||
import textwrap
|
||||
import threading
|
||||
import time
|
||||
@@ -1125,5 +1129,94 @@ class StopUnblocksPromptAndWaitTests(unittest.TestCase):
|
||||
client.stop()
|
||||
|
||||
|
||||
class TerminatesProcessGroupTests(unittest.TestCase):
|
||||
"""Regression: stop() must reap descendants the agent spawned, not only
|
||||
the omp leader.
|
||||
|
||||
A `bun test` launched by the agent's `bash` tool runs as a grandchild of
|
||||
the omp process. Before the fix, stop() signalled only the leader pid, so
|
||||
such grandchildren reparented to the container init and kept running —
|
||||
once ballooning to tens of GB of RAM. omp is now spawned in its own
|
||||
session and stop() tears down the whole process group.
|
||||
"""
|
||||
|
||||
@unittest.skipUnless(hasattr(os, "killpg"), "POSIX process groups only")
|
||||
def test_stop_kills_grandchild_spawned_by_server(self) -> None:
|
||||
work = tempfile.mkdtemp()
|
||||
self.addCleanup(shutil.rmtree, work, ignore_errors=True)
|
||||
pid_file = os.path.join(work, "gc.pid")
|
||||
beat_file = os.path.join(work, "gc.beat")
|
||||
gc_script = os.path.join(work, "gc.py")
|
||||
with open(gc_script, "w", encoding="utf-8") as handle:
|
||||
handle.write(
|
||||
textwrap.dedent(
|
||||
f"""
|
||||
import os, time
|
||||
with open({pid_file!r}, "w") as f:
|
||||
f.write(str(os.getpid()))
|
||||
while True:
|
||||
with open({beat_file!r}, "w") as f:
|
||||
f.write(str(time.time()))
|
||||
time.sleep(0.02)
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
def _reap_leaked_grandchild() -> None:
|
||||
try:
|
||||
with open(pid_file, encoding="utf-8") as f:
|
||||
os.kill(int(f.read()), signal.SIGKILL)
|
||||
except (OSError, ValueError):
|
||||
pass
|
||||
|
||||
self.addCleanup(_reap_leaked_grandchild)
|
||||
|
||||
# Fake omp server: spawn the long-lived grandchild, signal ready, then
|
||||
# idle until torn down (sleep past stdin EOF so the group is still
|
||||
# alive when stop() fires).
|
||||
server = textwrap.dedent(
|
||||
f"""
|
||||
import json, subprocess, sys, time
|
||||
subprocess.Popen([sys.executable, {gc_script!r}])
|
||||
print(json.dumps({{"type": "ready"}}), flush=True)
|
||||
for _line in sys.stdin:
|
||||
pass
|
||||
time.sleep(30)
|
||||
"""
|
||||
)
|
||||
|
||||
client = RpcClient(
|
||||
command=[sys.executable, "-u", "-c", server],
|
||||
startup_timeout=2.0,
|
||||
request_timeout=2.0,
|
||||
)
|
||||
client.start()
|
||||
try:
|
||||
deadline = time.time() + 2.0
|
||||
while time.time() < deadline and not os.path.exists(pid_file):
|
||||
time.sleep(0.02)
|
||||
self.assertTrue(os.path.exists(pid_file), "grandchild never started")
|
||||
with open(pid_file, encoding="utf-8") as f:
|
||||
os.kill(int(f.read()), 0) # alive before teardown
|
||||
finally:
|
||||
client.stop()
|
||||
|
||||
# The grandchild writes `time.time()` every 20ms. Once the group is
|
||||
# killed it stops writing, so the file contents stay frozen. Compare
|
||||
# contents (not mtime) to stay independent of filesystem timestamp
|
||||
# resolution.
|
||||
time.sleep(0.2)
|
||||
with open(beat_file, encoding="utf-8") as f:
|
||||
first = f.read()
|
||||
time.sleep(0.3)
|
||||
with open(beat_file, encoding="utf-8") as f:
|
||||
second = f.read()
|
||||
self.assertEqual(
|
||||
second,
|
||||
first,
|
||||
"grandchild kept running after stop() — process group leaked",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user