fix(controller): preserve timeout reason across report retries

agentlightning/controller/local_reconciler.py

@@ -186,6 +186,9 @@ async def _finish_proc(self, rollout: Rollout, item: Proc) -> bool:
             patched = await self._patch(rollout.rollout_id, RolloutState.RUNNING, last_attempt_id=item.attempt_id)
             if not patched:
                 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(

tests/controller/test_local_reconciler.py

@@ -3,14 +3,17 @@
 """Unit tests for local subprocess reconciliation."""
 
 import asyncio
-from unittest.mock import AsyncMock
+import signal
+import time
+from typing import cast
+from unittest.mock import AsyncMock, Mock
 
 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 +25,15 @@ 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)
@@ -41,6 +48,26 @@ def _reconciler(*, state: RolloutState = RolloutState.QUEUING) -> tuple[LocalRec
     return LocalReconciler(api, config), api
 
 
+class _Process:
+    def __init__(self, returncode: int | None) -> None:
+        self.returncode = returncode
+        self.pid = 1234
+        self.wait = AsyncMock()
+
+
+def _proc(*, returncode: int | None, killed: bool = False, attempt_id: str = "attempt-1") -> tuple[Proc, _Process]:
+    process = _Process(returncode)
+    return (
+        Proc(
+            attempt_id=attempt_id,
+            proc=cast(asyncio.subprocess.Process, process),
+            spawned_at=time.monotonic(),
+            killed=killed,
+        ),
+        process,
+    )
+
+
 @pytest.mark.asyncio
 async def test_shutdown_does_not_spawn_queued_rollouts(monkeypatch: pytest.MonkeyPatch) -> None:
     reconciler, api = _reconciler()
@@ -100,3 +127,159 @@ 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("failed_terminal_patches", [0, 1, 2])
+@pytest.mark.parametrize("returncode", [-9, 0])
+async def test_timeout_reconciliation_retries_and_retains_process_record(
+    monkeypatch: pytest.MonkeyPatch,
+    failed_terminal_patches: int,
+    returncode: int,
+) -> None:
+    reconciler, api = _reconciler(state=RolloutState.RUNNING, timeout_seconds=1)
+    rollout = Rollout.model_validate(api.get.return_value.json()[0])
+    item, process = _proc(returncode=None)
+    reconciler._rid_to_proc[rollout.rollout_id] = item
+    item.spawned_at = 0.0
+    api.patch.side_effect = [
+        *[httpx.ConnectError("server unavailable") for _ in range(failed_terminal_patches)],
+        _response({}),
+    ]
+    killpg = Mock()
+    monkeypatch.setattr("agentlightning.controller.local_reconciler.os.killpg", killpg)
+
+    async def complete_wait() -> None:
+        process.returncode = returncode
+
+    process.wait.side_effect = complete_wait
+    spawn_for = AsyncMock(return_value=True)
+    monkeypatch.setattr(reconciler, "_spawn_for", spawn_for)
+
+    for _ in range(failed_terminal_patches + 1):
+        await reconciler._reconcile_once()
+
+    assert item.killed
+    assert reconciler._rid_to_proc[rollout.rollout_id] is item
+    assert item.attempt_id == "attempt-1"
+    process.wait.assert_awaited_once_with()
+    killpg.assert_called_once_with(process.pid, signal.SIGKILL)
+    spawn_for.assert_not_awaited()
+    assert [call.kwargs["json"]["status"] for call in api.patch.await_args_list] == [
+        {"state": "failed", "error_message": "local subprocess timed out"}
+    ] * (failed_terminal_patches + 1)
+
+    api.get.return_value = _response([])
+    await reconciler._reconcile_once()
+    assert reconciler._rid_to_proc[rollout.rollout_id] is item
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("returncode", "expected"),
+    [
+        (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"}),
+    ],
+)
+@pytest.mark.parametrize("failed_terminal_patches", [0, 1])
+async def test_unmarked_terminal_process_retries_and_retains_record(
+    returncode: int,
+    expected: dict[str, str],
+    failed_terminal_patches: int,
+) -> None:
+    reconciler, api = _reconciler(state=RolloutState.RUNNING)
+    rollout = Rollout.model_validate(api.get.return_value.json()[0])
+    item, _ = _proc(returncode=returncode)
+    reconciler._rid_to_proc[rollout.rollout_id] = item
+    api.patch.side_effect = [
+        *[httpx.ConnectError("server unavailable") for _ in range(failed_terminal_patches)],
+        _response({}),
+    ]
+
+    for _ in range(failed_terminal_patches + 1):
+        await reconciler._reconcile_once()
+
+    assert reconciler._rid_to_proc[rollout.rollout_id] is item
+    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
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("patch_fails", [False, True])
+async def test_run_stops_then_shutdown_kills_after_final_reconcile(
+    monkeypatch: pytest.MonkeyPatch,
+    patch_fails: bool,
+) -> None:
+    reconciler, api = _reconciler(state=RolloutState.RUNNING)
+    item, process = _proc(returncode=None)
+    reconciler._rid_to_proc["rollout-1"] = item
+    events: list[str] = []
+
+    async def get(*args: object, **kwargs: object) -> httpx.Response:
+        del args, kwargs
+        events.append("reconcile")
+        reconciler.stop()
+        return _response([Rollout.model_validate(api.get.return_value.json()[0]).model_dump(mode="json")])
+
+    async def complete_wait() -> None:
+        process.returncode = -9
+
+    process.wait.side_effect = complete_wait
+    killpg = Mock(side_effect=lambda *args: events.append("kill"))
+    monkeypatch.setattr("agentlightning.controller.local_reconciler.os.killpg", killpg)
+    api.get.side_effect = get
+    if patch_fails:
+        api.patch.side_effect = httpx.ConnectError("server unavailable")
+
+    await reconciler.run()
+
+    assert events == ["reconcile", "reconcile", "kill"]
+    assert item.killed
+    assert reconciler._rid_to_proc["rollout-1"] is item
+    assert api.patch.await_args.kwargs["json"]["status"] == {
+        "state": "failed",
+        "error_message": "local controller shutdown",
+    }
+
+
+@pytest.mark.asyncio
+async def test_timeout_during_run_retries_original_reason_during_shutdown_without_rekill(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    reconciler, api = _reconciler(state=RolloutState.RUNNING, timeout_seconds=1)
+    rollout = Rollout.model_validate(api.get.return_value.json()[0])
+    item, process = _proc(returncode=None)
+    item.spawned_at = 0.0
+    reconciler._rid_to_proc[rollout.rollout_id] = item
+    events: list[str] = []
+
+    async def get(*args: object, **kwargs: object) -> httpx.Response:
+        del args, kwargs
+        events.append("reconcile")
+        reconciler.stop()
+        return _response([rollout.model_dump(mode="json")])
+
+    async def complete_wait() -> None:
+        process.returncode = -9
+
+    process.wait.side_effect = complete_wait
+    killpg = Mock(side_effect=lambda *args: events.append("kill"))
+    monkeypatch.setattr("agentlightning.controller.local_reconciler.os.killpg", killpg)
+    api.get.side_effect = get
+    api.patch.side_effect = [httpx.ConnectError("server unavailable"), _response({})]
+
+    await reconciler.run()
+
+    assert events == ["reconcile", "kill", "reconcile"]
+    assert reconciler._rid_to_proc[rollout.rollout_id] is item
+    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"},
+    ]