diff --git a/agentlightning/controller/local_reconciler.py b/agentlightning/controller/local_reconciler.py index 847bfe736..fcbd3a3ab 100644 --- a/agentlightning/controller/local_reconciler.py +++ b/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( diff --git a/tests/controller/test_local_reconciler.py b/tests/controller/test_local_reconciler.py index de0d1bda4..1a67ac456 100644 --- a/tests/controller/test_local_reconciler.py +++ b/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"}, + ]