fix(controller): release completed local process records

agentlightning/controller/local_reconciler.py

@@ -166,7 +166,8 @@ async def _reconcile_once(self, *, spawn_queued: bool = True) -> None:
                     await self._patch(rollout.rollout_id, RolloutState.RUNNING, last_attempt_id=item.attempt_id)
                 continue
 
-            await self._finish_proc(rollout, item)
+            if await self._finish_proc(rollout, item):
+                self._rid_to_proc.pop(rollout.rollout_id)
 
         now = time.monotonic()
         for rollout_id, item in list(self._rid_to_proc.items()):
@@ -178,14 +179,17 @@ async def _reconcile_once(self, *, spawn_queued: bool = True) -> None:
                 timeout is not None
                 and (now - item.spawned_at) > timeout
                 and await self._kill_process_group(rollout_id, item)
+                and await self._patch(rollout_id, RolloutState.FAILED, "local subprocess timed out")
             ):
-                await self._patch(rollout_id, RolloutState.FAILED, "local subprocess timed out")
+                self._rid_to_proc.pop(rollout_id)
 
     async def _finish_proc(self, rollout: Rollout, item: Proc) -> bool:
         if rollout.status.state == RolloutState.QUEUING:
             patched = await self._patch(rollout.rollout_id, RolloutState.RUNNING, last_attempt_id=item.attempt_id)
             if not patched:
                 return False
+        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(

tests/controller/test_local_reconciler.py

@@ -3,14 +3,15 @@
 """Unit tests for local subprocess reconciliation."""
 
 import asyncio
-from unittest.mock import AsyncMock
+import time
+from unittest.mock import AsyncMock, MagicMock
 
 import httpx
 import pytest
 from omegaconf import OmegaConf
 
 from agentlightning.client import AgentLightningAsyncClient
-from agentlightning.controller.local_reconciler import LocalReconciler
+from agentlightning.controller.local_reconciler import LocalReconciler, Proc
 from agentlightning.schemas import Rollout, RolloutConfig, RolloutLifecycleStatus, RolloutState
 
 
@@ -22,11 +23,13 @@ def _response(json: object) -> httpx.Response:
     )
 
 
-def _reconciler(*, state: RolloutState = RolloutState.QUEUING) -> tuple[LocalReconciler, AsyncMock]:
+def _reconciler(
+    *, state: RolloutState = RolloutState.QUEUING, timeout_seconds: int = 3600
+) -> tuple[LocalReconciler, AsyncMock]:
     rollout = Rollout(
         rollout_id="rollout-1",
         input={"question": "1 + 1"},
-        config=RolloutConfig(),
+        config=RolloutConfig(timeout_seconds=timeout_seconds),
         status=RolloutLifecycleStatus(state=state, created_at=1.0, updated_at=1.0),
     )
     api = AsyncMock(spec=AgentLightningAsyncClient)
@@ -100,3 +103,127 @@ async def test_shutdown_still_fails_running_rollout_without_local_process() -> N
         "state": "failed",
         "error_message": "local subprocess is not running",
     }
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("returncode", "expected_status"),
+    [
+        (0, {"state": "succeeded", "last_attempt_id": "attempt-1"}),
+        (1, {"state": "failed", "error_message": "subprocess exited with code 1"}),
+        (-9, {"state": "failed", "error_message": "subprocess exited with code -9"}),
+    ],
+)
+async def test_reconcile_removes_completed_process_after_terminal_patch(
+    returncode: int, expected_status: dict[str, str]
+) -> None:
+    reconciler, api = _reconciler(state=RolloutState.RUNNING)
+    proc = MagicMock(spec=asyncio.subprocess.Process)
+    proc.returncode = returncode
+    reconciler._rid_to_proc["rollout-1"] = Proc("attempt-1", proc, spawned_at=time.monotonic())
+
+    await reconciler._reconcile_once()
+
+    assert "rollout-1" not in reconciler._rid_to_proc
+    assert api.patch.await_args.kwargs["json"]["status"] == expected_status
+
+
+@pytest.mark.asyncio
+async def test_reconcile_keeps_completed_process_when_terminal_patch_fails_then_retries() -> None:
+    reconciler, api = _reconciler(state=RolloutState.RUNNING)
+    proc = MagicMock(spec=asyncio.subprocess.Process)
+    proc.returncode = 0
+    reconciler._rid_to_proc["rollout-1"] = Proc("attempt-1", proc, spawned_at=time.monotonic())
+    api.patch.side_effect = [
+        httpx.Response(500, request=httpx.Request("PATCH", "http://server/api/rollouts/rollout-1")),
+        _response({}),
+    ]
+
+    await reconciler._reconcile_once()
+
+    assert "rollout-1" in reconciler._rid_to_proc
+
+    await reconciler._reconcile_once()
+
+    assert "rollout-1" not in reconciler._rid_to_proc
+    assert api.patch.await_count == 2
+
+
+@pytest.mark.asyncio
+async def test_reconcile_keeps_running_process() -> None:
+    reconciler, api = _reconciler(state=RolloutState.RUNNING)
+    proc = MagicMock(spec=asyncio.subprocess.Process)
+    proc.returncode = None
+    reconciler._rid_to_proc["rollout-1"] = Proc("attempt-1", proc, spawned_at=time.monotonic())
+
+    await reconciler._reconcile_once()
+
+    assert reconciler._rid_to_proc["rollout-1"].proc is proc
+    api.patch.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_reconcile_removes_timed_out_process_after_failure_patch(monkeypatch: pytest.MonkeyPatch) -> None:
+    reconciler, api = _reconciler(state=RolloutState.RUNNING, timeout_seconds=1)
+    proc = MagicMock(spec=asyncio.subprocess.Process)
+    proc.returncode = None
+    reconciler._rid_to_proc["rollout-1"] = Proc("attempt-1", proc, spawned_at=time.monotonic() - 2)
+    kill_process_group = AsyncMock(return_value=True)
+    monkeypatch.setattr(reconciler, "_kill_process_group", kill_process_group)
+
+    await reconciler._reconcile_once()
+
+    assert "rollout-1" not in reconciler._rid_to_proc
+    kill_process_group.assert_awaited_once()
+    assert api.patch.await_args.kwargs["json"]["status"] == {
+        "state": "failed",
+        "error_message": "local subprocess timed out",
+    }
+
+
+@pytest.mark.asyncio
+async def test_reconcile_keeps_timed_out_process_when_failure_patch_fails_then_retries(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    reconciler, api = _reconciler(state=RolloutState.RUNNING, timeout_seconds=1)
+    proc = MagicMock(spec=asyncio.subprocess.Process)
+    proc.returncode = None
+    reconciler._rid_to_proc["rollout-1"] = Proc("attempt-1", proc, spawned_at=time.monotonic() - 2)
+    async def kill_process_group(_: str, item: Proc) -> bool:
+        item.killed = True
+        proc.returncode = -9
+        return True
+
+    monkeypatch.setattr(reconciler, "_kill_process_group", kill_process_group)
+    api.patch.side_effect = [
+        httpx.Response(500, request=httpx.Request("PATCH", "http://server/api/rollouts/rollout-1")),
+        _response({}),
+    ]
+
+    await reconciler._reconcile_once()
+
+    assert "rollout-1" in reconciler._rid_to_proc
+
+    await reconciler._reconcile_once()
+
+    assert "rollout-1" not in reconciler._rid_to_proc
+    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_reconcile_finishes_queued_completed_process_after_running_transition() -> None:
+    reconciler, api = _reconciler()
+    proc = MagicMock(spec=asyncio.subprocess.Process)
+    proc.returncode = 0
+    reconciler._rid_to_proc["rollout-1"] = Proc("attempt-1", proc, spawned_at=time.monotonic())
+
+    await reconciler._reconcile_once()
+
+    assert "rollout-1" not in reconciler._rid_to_proc
+    assert [call.kwargs["json"]["status"] for call in api.patch.await_args_list] == [
+        {"state": "running", "last_attempt_id": "attempt-1"},
+        {"state": "succeeded", "last_attempt_id": "attempt-1"},
+    ]