Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion agentlightning/controller/local_reconciler.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,7 +138,7 @@ async def _reconcile_once(self, *, spawn_queued: bool = True) -> None:
params = httpx.QueryParams()
params = params.add("state_in", RolloutState.QUEUING.value)
params = params.add("state_in", RolloutState.RUNNING.value)
params = params.add("limit", 50)
params = params.add("limit", max(50, self._pool_size))
response = await self._api.get("/api/rollouts", params=params)
response.raise_for_status()
rollouts = [Rollout.model_validate(item) for item in response.json()]
Expand Down
54 changes: 54 additions & 0 deletions tests/controller/test_local_reconciler.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,60 @@ async def test_normal_reconcile_still_spawns_queued_rollouts(monkeypatch: pytest
spawn_for.assert_awaited_once()


@pytest.mark.asyncio
@pytest.mark.parametrize(
("pool_size", "expected_limit"),
[(25, 50), (50, 50), (60, 60), (128, 128)],
)
async def test_reconcile_fetches_enough_rollouts_to_fill_pool(
monkeypatch: pytest.MonkeyPatch,
pool_size: int,
expected_limit: int,
) -> None:
rollouts = [
Rollout(
rollout_id=f"rollout-{index}",
input={"index": index},
config=RolloutConfig(),
status=RolloutLifecycleStatus(created_at=1.0, updated_at=1.0),
)
for index in range(128)
]
requested_limits: list[int] = []
api = AsyncMock(spec=AgentLightningAsyncClient)

async def get(*args: object, **kwargs: object) -> httpx.Response:
del args
params = kwargs["params"]
assert isinstance(params, httpx.QueryParams)
limit = int(params["limit"])
requested_limits.append(limit)
return _response([rollout.model_dump(mode="json") for rollout in rollouts[:limit]])

api.get.side_effect = get
api.patch.return_value = _response({})
config = OmegaConf.create(
{
"runner_type": "local",
"local_runner": {"maximum_size": pool_size, "poll_interval": 0.01},
}
)
reconciler = LocalReconciler(api, config)

async def spawn(rollout: Rollout) -> bool:
item, _ = _proc(returncode=None)
reconciler._rid_to_proc[rollout.rollout_id] = item
return True

monkeypatch.setattr(reconciler, "_spawn_for", spawn)

await reconciler._reconcile_once()
await reconciler._reconcile_once()

assert len(reconciler._rid_to_proc) == pool_size
assert requested_limits == [expected_limit, expected_limit]


@pytest.mark.asyncio
async def test_reconcile_does_not_spawn_after_stop_requested_during_poll(monkeypatch: pytest.MonkeyPatch) -> None:
reconciler, api = _reconciler()
Expand Down