fix(controller): release completed subprocesses while retaining launch history

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,

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()