From 9651340c3064dcb7a893349f46a14c2f3e380d28 Mon Sep 17 00:00:00 2001 From: Blue <3067670134@qq.com> Date: Sat, 3 Oct 2026 19:19:31 +0800 Subject: [PATCH] fix: release completed local subprocesses while retaining launch history --- agentlightning/controller/local_reconciler.py | 59 +++++++-- tests/controller/test_local_reconciler.py | 117 ++++++++++++++++-- 2 files changed, 156 insertions(+), 20 deletions(-) diff --git a/agentlightning/controller/local_reconciler.py b/agentlightning/controller/local_reconciler.py index fcbd3a3ab..3f92b3436 100644 --- a/agentlightning/controller/local_reconciler.py +++ b/agentlightning/controller/local_reconciler.py @@ -60,6 +60,16 @@ class Proc: killed: bool = False +@dataclass(frozen=True, slots=True) +class CompletedProc: + """Compact history for a subprocess whose terminal status was reported.""" + + attempt_id: str + returncode: int + spawned_at: float + killed: bool + + def _build_env_from_map(task_input: object, env_map: dict[str, str]) -> dict[str, str]: env: dict[str, str] = {} for name, path in env_map.items(): @@ -106,6 +116,7 @@ def __init__( self._pool_size = int(self._runner_config.maximum_size) self._tick_interval = float(self._runner_config.poll_interval) self._rid_to_proc: dict[str, Proc] = {} + self._completed_procs: dict[str, CompletedProc] = {} self._stop = asyncio.Event() async def run(self) -> None: @@ -146,6 +157,8 @@ async def _reconcile_once(self, *, spawn_queued: bool = True) -> None: live_count = sum(1 for item in self._rid_to_proc.values() if item.proc.returncode is None) for rollout in rollouts: + if rollout.rollout_id in self._completed_procs: + continue item = self._rid_to_proc.get(rollout.rollout_id) if item is None: @@ -179,7 +192,9 @@ async def _reconcile_once(self, *, spawn_queued: bool = True) -> None: and (now - item.spawned_at) > timeout and await self._kill_process_group(rollout_id, item) ): - await self._patch(rollout_id, RolloutState.FAILED, "local subprocess timed out") + patched = await self._patch(rollout_id, RolloutState.FAILED, "local subprocess timed out") + if patched: + self._archive_proc(rollout_id, item) async def _finish_proc(self, rollout: Rollout, item: Proc) -> bool: if rollout.status.state == RolloutState.QUEUING: @@ -188,14 +203,38 @@ async def _finish_proc(self, rollout: Rollout, item: Proc) -> bool: return False # Normal timeouts reach here; shutdown kills happen after final reconciliation. if item.killed: - return await self._patch(rollout.rollout_id, RolloutState.FAILED, "local subprocess timed out") - if item.proc.returncode == 0: - return await self._patch(rollout.rollout_id, RolloutState.SUCCEEDED, last_attempt_id=item.attempt_id) - return await self._patch( - rollout.rollout_id, - RolloutState.FAILED, - f"subprocess exited with code {item.proc.returncode}", + patched = await self._patch(rollout.rollout_id, RolloutState.FAILED, "local subprocess timed out") + elif item.proc.returncode == 0: + patched = await self._patch( + rollout.rollout_id, + RolloutState.SUCCEEDED, + last_attempt_id=item.attempt_id, + ) + else: + patched = await self._patch( + rollout.rollout_id, + RolloutState.FAILED, + f"subprocess exited with code {item.proc.returncode}", + ) + if patched: + self._archive_proc(rollout.rollout_id, item) + return patched + + def _archive_proc(self, rollout_id: str, item: Proc) -> None: + """Replace a reported terminal process with compact history.""" + if self._rid_to_proc.get(rollout_id) is not item: + return + returncode = item.proc.returncode + if returncode is None: + log.warning("Refusing to archive running subprocess", rollout_id=rollout_id, pid=item.proc.pid) + return + self._completed_procs[rollout_id] = CompletedProc( + attempt_id=item.attempt_id, + returncode=returncode, + spawned_at=item.spawned_at, + killed=item.killed, ) + del self._rid_to_proc[rollout_id] async def _kill_process_group(self, rollout_id: str, item: Proc) -> bool: """SIGKILL the worker process group and wait for exit.""" @@ -270,7 +309,9 @@ async def _shutdown(self) -> None: for rollout_id, item in list(self._rid_to_proc.items()): if item.proc.returncode is None and await self._kill_process_group(rollout_id, item): - await self._patch(rollout_id, RolloutState.FAILED, "local controller shutdown") + patched = await self._patch(rollout_id, RolloutState.FAILED, "local controller shutdown") + if patched: + self._archive_proc(rollout_id, item) async def _patch( self, diff --git a/tests/controller/test_local_reconciler.py b/tests/controller/test_local_reconciler.py index 1a67ac456..26a07a4fd 100644 --- a/tests/controller/test_local_reconciler.py +++ b/tests/controller/test_local_reconciler.py @@ -3,8 +3,10 @@ """Unit tests for local subprocess reconciliation.""" import asyncio +import gc import signal import time +import weakref from typing import cast from unittest.mock import AsyncMock, Mock @@ -132,7 +134,7 @@ async def test_shutdown_still_fails_running_rollout_without_local_process() -> N @pytest.mark.asyncio @pytest.mark.parametrize("failed_terminal_patches", [0, 1, 2]) @pytest.mark.parametrize("returncode", [-9, 0]) -async def test_timeout_reconciliation_retries_and_retains_process_record( +async def test_timeout_reconciliation_retries_then_archives_process_record( monkeypatch: pytest.MonkeyPatch, failed_terminal_patches: int, returncode: int, @@ -156,12 +158,20 @@ async def complete_wait() -> None: spawn_for = AsyncMock(return_value=True) monkeypatch.setattr(reconciler, "_spawn_for", spawn_for) - for _ in range(failed_terminal_patches + 1): + for _ in range(failed_terminal_patches): await reconciler._reconcile_once() + assert reconciler._rid_to_proc[rollout.rollout_id] is item + assert rollout.rollout_id not in reconciler._completed_procs + + await reconciler._reconcile_once() assert item.killed - assert reconciler._rid_to_proc[rollout.rollout_id] is item - assert item.attempt_id == "attempt-1" + assert rollout.rollout_id not in reconciler._rid_to_proc + history = reconciler._completed_procs[rollout.rollout_id] + assert history.attempt_id == "attempt-1" + assert history.returncode == returncode + assert history.spawned_at == 0.0 + assert history.killed process.wait.assert_awaited_once_with() killpg.assert_called_once_with(process.pid, signal.SIGKILL) spawn_for.assert_not_awaited() @@ -171,7 +181,7 @@ async def complete_wait() -> None: api.get.return_value = _response([]) await reconciler._reconcile_once() - assert reconciler._rid_to_proc[rollout.rollout_id] is item + assert rollout.rollout_id not in reconciler._rid_to_proc @pytest.mark.asyncio @@ -184,7 +194,7 @@ async def complete_wait() -> None: ], ) @pytest.mark.parametrize("failed_terminal_patches", [0, 1]) -async def test_unmarked_terminal_process_retries_and_retains_record( +async def test_unmarked_terminal_process_retries_then_archives_record( returncode: int, expected: dict[str, str], failed_terminal_patches: int, @@ -198,16 +208,80 @@ async def test_unmarked_terminal_process_retries_and_retains_record( _response({}), ] - for _ in range(failed_terminal_patches + 1): + for _ in range(failed_terminal_patches): await reconciler._reconcile_once() + assert reconciler._rid_to_proc[rollout.rollout_id] is item + assert rollout.rollout_id not in reconciler._completed_procs + + await reconciler._reconcile_once() - assert reconciler._rid_to_proc[rollout.rollout_id] is item + assert rollout.rollout_id not in reconciler._rid_to_proc + history = reconciler._completed_procs[rollout.rollout_id] + assert history.attempt_id == "attempt-1" + assert history.returncode == returncode + assert history.spawned_at == item.spawned_at + assert not history.killed assert [call.kwargs["json"]["status"] for call in api.patch.await_args_list] == [expected] * ( failed_terminal_patches + 1 ) api.get.return_value = _response([]) await reconciler._reconcile_once() - assert reconciler._rid_to_proc[rollout.rollout_id] is item + assert rollout.rollout_id not in reconciler._rid_to_proc + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stale_state", [RolloutState.QUEUING, RolloutState.RUNNING]) +async def test_completed_rollout_ignores_stale_active_state( + monkeypatch: pytest.MonkeyPatch, + stale_state: RolloutState, +) -> None: + reconciler, api = _reconciler(state=RolloutState.RUNNING) + rollout = Rollout.model_validate(api.get.return_value.json()[0]) + item, _ = _proc(returncode=0) + reconciler._rid_to_proc[rollout.rollout_id] = item + + await reconciler._reconcile_once() + api.patch.reset_mock() + rollout.status.state = stale_state + api.get.return_value = _response([rollout.model_dump(mode="json")]) + spawn_for = AsyncMock(return_value=True) + monkeypatch.setattr(reconciler, "_spawn_for", spawn_for) + + await reconciler._reconcile_once() + + assert rollout.rollout_id not in reconciler._rid_to_proc + assert rollout.rollout_id in reconciler._completed_procs + spawn_for.assert_not_awaited() + api.patch.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_live_process_remains_active() -> None: + reconciler, api = _reconciler(state=RolloutState.RUNNING) + item, _ = _proc(returncode=None) + reconciler._rid_to_proc["rollout-1"] = item + + await reconciler._reconcile_once() + reconciler._archive_proc("rollout-1", item) + + assert reconciler._rid_to_proc["rollout-1"] is item + assert "rollout-1" not in reconciler._completed_procs + api.patch.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_archived_process_does_not_keep_process_object_alive() -> None: + reconciler, api = _reconciler(state=RolloutState.RUNNING) + rollout = Rollout.model_validate(api.get.return_value.json()[0]) + item, process = _proc(returncode=0) + process_ref = weakref.ref(process) + reconciler._rid_to_proc[rollout.rollout_id] = item + + await reconciler._reconcile_once() + del item, process + gc.collect() + + assert process_ref() is None @pytest.mark.asyncio @@ -241,7 +315,12 @@ async def complete_wait() -> None: assert events == ["reconcile", "reconcile", "kill"] assert item.killed - assert reconciler._rid_to_proc["rollout-1"] is item + if patch_fails: + assert reconciler._rid_to_proc["rollout-1"] is item + assert "rollout-1" not in reconciler._completed_procs + else: + assert "rollout-1" not in reconciler._rid_to_proc + assert reconciler._completed_procs["rollout-1"].killed assert api.patch.await_args.kwargs["json"]["status"] == { "state": "failed", "error_message": "local controller shutdown", @@ -277,9 +356,25 @@ async def complete_wait() -> None: await reconciler.run() assert events == ["reconcile", "kill", "reconcile"] - assert reconciler._rid_to_proc[rollout.rollout_id] is item + assert rollout.rollout_id not in reconciler._rid_to_proc + assert reconciler._completed_procs[rollout.rollout_id].killed process.wait.assert_awaited_once_with() assert [call.kwargs["json"]["status"] for call in api.patch.await_args_list] == [ {"state": "failed", "error_message": "local subprocess timed out"}, {"state": "failed", "error_message": "local subprocess timed out"}, ] + + +@pytest.mark.asyncio +async def test_shutdown_retains_live_process_when_kill_wait_times_out(monkeypatch: pytest.MonkeyPatch) -> None: + reconciler, api = _reconciler(state=RolloutState.RUNNING) + item, process = _proc(returncode=None) + reconciler._rid_to_proc["rollout-1"] = item + process.wait.side_effect = TimeoutError + monkeypatch.setattr("agentlightning.controller.local_reconciler.os.killpg", Mock()) + + await reconciler._shutdown() + + assert reconciler._rid_to_proc["rollout-1"] is item + assert "rollout-1" not in reconciler._completed_procs + api.patch.assert_not_awaited()