diff --git a/src/dstack/_internal/cli/services/presets/agent.py b/src/dstack/_internal/cli/services/presets/agent.py index 64cedcfa6..b92535db5 100644 --- a/src/dstack/_internal/cli/services/presets/agent.py +++ b/src/dstack/_internal/cli/services/presets/agent.py @@ -9,7 +9,7 @@ from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path -from typing import Any, AsyncIterator, Callable, Literal, Optional, Sequence, get_args +from typing import Any, AsyncIterator, Callable, Iterable, Literal, Optional, Sequence, get_args import psutil from pydantic import ValidationError @@ -314,10 +314,16 @@ def agent_alive() -> bool: offset_store=offset_store, ) ) + descendants = _process_descendants(proc.pid) + descendant_watcher = asyncio.create_task(_watch_process_descendants(proc.pid, descendants)) try: returncode = await proc.wait() output = await collect_task finally: + descendant_watcher.cancel() + with suppress(asyncio.CancelledError): + await descendant_watcher + _terminate_processes(descendants.values()) # Never leave the collector orphaned when an await above raises # (cancellation, interrupt, or a wait failure). if not collect_task.done(): @@ -585,6 +591,7 @@ async def _terminate_process(proc: asyncio.subprocess.Process) -> None: await asyncio.to_thread(_terminate_windows_process_tree, proc.pid) await proc.wait() return + _terminate_processes(_process_descendants(proc.pid).values()) # The Windows branch returns above; Pyright still checks these POSIX-only APIs on Windows. if hasattr(os, "killpg"): with suppress(ProcessLookupError): @@ -613,6 +620,7 @@ def terminate_agent_process(agent: Optional[PresetSessionProcess]) -> None: if IS_WINDOWS: _terminate_windows_process_tree(agent_pid) return + _terminate_processes(_process_descendants(agent_pid).values()) with suppress(OSError): os.killpg(agent_pid, signal.SIGTERM) # pyright: ignore[reportAttributeAccessIssue] for _ in range(_TERMINATE_GRACE_SECONDS * 10): @@ -623,6 +631,34 @@ def terminate_agent_process(agent: Optional[PresetSessionProcess]) -> None: os.killpg(agent_pid, signal.SIGKILL) # pyright: ignore[reportAttributeAccessIssue] +def _process_descendants(pid: int) -> dict[int, psutil.Process]: + try: + processes = psutil.Process(pid).children(recursive=True) + except psutil.NoSuchProcess: + return {} + return {process.pid: process for process in processes} + + +async def _watch_process_descendants(pid: int, descendants: dict[int, psutil.Process]) -> None: + while True: + descendants.update(_process_descendants(pid)) + if not pid_running(pid): + return + await asyncio.sleep(0.05) + + +def _terminate_processes(processes: Iterable[psutil.Process]) -> None: + processes = list(processes) + for process in processes: + with suppress(psutil.NoSuchProcess): + process.terminate() + _, alive = psutil.wait_procs(processes, timeout=_TERMINATE_GRACE_SECONDS) + for process in alive: + with suppress(psutil.NoSuchProcess): + process.kill() + psutil.wait_procs(alive, timeout=_TERMINATE_GRACE_SECONDS) + + def _terminate_windows_process_tree(pid: int) -> None: try: root = psutil.Process(pid) diff --git a/src/tests/_internal/cli/services/presets/test_agent.py b/src/tests/_internal/cli/services/presets/test_agent.py index f101258c5..ee4622e24 100644 --- a/src/tests/_internal/cli/services/presets/test_agent.py +++ b/src/tests/_internal/cli/services/presets/test_agent.py @@ -18,6 +18,7 @@ ClaudeAuth, _build_claude_command, _prepare_subprocess_command, + _run_claude_process, _terminate_process, build_preset_agent_env, get_claude_auth, @@ -481,6 +482,57 @@ async def test_terminates_windows_process_tree(self): assert proc.returncode is not None assert not psutil.pid_exists(child_pid) + @pytest.mark.skipif(IS_WINDOWS, reason="exercises POSIX process groups") + @pytest.mark.asyncio + async def test_terminates_detached_process_tree(self): + proc = await asyncio.create_subprocess_exec( + sys.executable, + "-c", + ( + "import subprocess,sys,time; " + "p=subprocess.Popen([sys.executable,'-c','import time; time.sleep(60)'], " + "start_new_session=True); print(p.pid,flush=True); time.sleep(60)" + ), + start_new_session=True, + stdout=asyncio.subprocess.PIPE, + ) + assert proc.stdout is not None + child_pid = int((await proc.stdout.readline()).decode()) + + await _terminate_process(proc) + + assert proc.returncode is not None + assert not psutil.pid_exists(child_pid) + + @pytest.mark.skipif(IS_WINDOWS, reason="exercises POSIX process groups") + @pytest.mark.asyncio + async def test_cleans_detached_process_after_agent_exits(self, tmp_path): + child_pid_path = tmp_path / "child.pid" + script = tmp_path / "agent.py" + script.write_text( + "import pathlib, subprocess, sys, time; " + "child = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(60)'], " + "start_new_session=True); " + "pathlib.Path(sys.argv[1]).write_text(str(child.pid)); " + 'print(\'{"type": "result", "structured_output": {"ok": true}}\', flush=True); ' + "time.sleep(0.2)" + ) + workspace, session = _agent_setup(tmp_path) + + output, returncode = await _run_claude_process( + command=[sys.executable, str(script), str(child_pid_path)], + prompt="prompt", + env=_subprocess_env(), + workspace=workspace, + redacted_values=(), + session=session, + offset_store=open_session_offsets(session), + ) + + assert returncode == 0 + assert output.report_data == {"ok": True} + assert not psutil.pid_exists(int(child_pid_path.read_text())) + class TestRecordMirror: def test_mirrors_complete_lines_redacted(self, tmp_path):