From ba8abc741068f88987714e9cd29220d96594f67b Mon Sep 17 00:00:00 2001 From: Anvarbek Ziyodov <55726523+Mr-Perfectuz@users.noreply.github.com> Date: Sat, 17 Oct 2026 22:49:14 +0900 Subject: [PATCH 1/2] fix(verl): prevent stale/failed rollouts from breaking batch completion Reject leftover server results, allow successful retries to overwrite empty failed placeholders, and soften the completed==queued assert so training does not crash when counts drift (#355). --- agentlightning/verl/daemon.py | 119 ++++++++++++++-- tests/verl/__init__.py | 1 + tests/verl/conftest.py | 25 ++++ tests/verl/test_daemon_completed_rollouts.py | 138 +++++++++++++++++++ 4 files changed, 271 insertions(+), 12 deletions(-) create mode 100644 tests/verl/__init__.py create mode 100644 tests/verl/conftest.py create mode 100644 tests/verl/test_daemon_completed_rollouts.py diff --git a/agentlightning/verl/daemon.py b/agentlightning/verl/daemon.py index 873e23b67..4337c0bcc 100644 --- a/agentlightning/verl/daemon.py +++ b/agentlightning/verl/daemon.py @@ -349,6 +349,15 @@ def run_server(): async def _async_set_up(self, data: Dict[str, Any], server_addresses: List[str], is_train: bool = True): """Async helper to set up data and resources on the server.""" self.clear_data_and_server() + # Drain leftover completions from the previous step so they cannot be + # mixed into the new batch (see #355 / stale server result cache). + if self.mode == "v0": + try: + stale = await self.server.retrieve_completed_rollouts() + if stale: + print(f"Warning: Drained {len(stale)} stale completed rollout(s) from server before setup.") + except Exception as exc: # pragma: no cover - best-effort cleanup + print(f"Warning: Failed to drain stale completed rollouts: {exc}") if server_addresses != self.backend_llm_server_addresses: self.backend_llm_server_addresses = server_addresses if self.mode == "v1" and not self.llm_proxy.is_running(): @@ -504,30 +513,116 @@ async def _validate_data_v1(self, rollout: Rollout) -> RolloutLegacy: return result_rollout + @staticmethod + def _is_failed_rollout(rollout: RolloutLegacy) -> bool: + """Return True when a rollout has no usable reward or triplets. + + Backend errors (e.g. vLLM HTTP 400 on over-long prompts) often complete + a rollout with ``final_reward is None`` and empty/missing triplets. The + store may retry the same ``rollout_id``; a later successful result must + be allowed to overwrite this placeholder. + """ + return rollout.final_reward is None and not rollout.triplets + + def _drop_stale_completed_rollouts(self) -> int: + """Drop completed rollouts that are not part of the current batch. + + Stale server/store results from a previous step can otherwise inflate + ``_completed_rollouts_v0`` past ``_total_tasks_queued`` (see #355). + """ + stale_ids = [ + rollout_id + for rollout_id in list(self._completed_rollouts_v0) + if rollout_id not in self._task_id_to_original_sample + ] + for rollout_id in stale_ids: + print(f"Warning: Dropping stale completed rollout {rollout_id} not in current batch.") + del self._completed_rollouts_v0[rollout_id] + return len(stale_ids) + + def _ensure_completed_matches_queued(self) -> None: + """Reconcile completed rollouts with the current batch before asserts.""" + self._drop_stale_completed_rollouts() + completed = len(self._completed_rollouts_v0) + queued = self._total_tasks_queued + if completed == queued: + return + print( + f"Warning: completed rollouts ({completed}) != queued tasks ({queued}); " + "continuing with rollouts that belong to the current batch." + ) + + def _pending_rollout_ids(self) -> List[str]: + """Rollout IDs that still need a successful result (missing or failed).""" + pending: List[str] = [] + for rollout_id in self._task_id_to_original_sample: + existing = self._completed_rollouts_v0.get(rollout_id) + if existing is None or self._is_failed_rollout(existing): + pending.append(rollout_id) + return pending + async def _async_run_until_finished(self, verbose: bool = True): """Async helper to wait for all tasks to complete.""" - while len(self._completed_rollouts_v0) < self._total_tasks_queued: + # After every queued ID has some result, keep polling briefly so a store + # retry can overwrite failed placeholders (empty reward/triplets). + failed_grace_polls_remaining: Optional[int] = None + while True: if self.mode == "v0": completed_batch = await self.server.retrieve_completed_rollouts() else: - completed_batch = await self.store.wait_for_rollouts( - rollout_ids=list(self._task_id_to_original_sample.keys()), timeout=0 - ) + pending_ids = self._pending_rollout_ids() + if not pending_ids and failed_grace_polls_remaining is None: + break + wait_ids = pending_ids or list(self._task_id_to_original_sample.keys()) + completed_batch = await self.store.wait_for_rollouts(rollout_ids=wait_ids, timeout=0) for rollout in completed_batch: - if rollout.rollout_id in self._completed_rollouts_v0: - # Already processed, skip - continue if isinstance(rollout, Rollout): rollout = await self._validate_data_v1(rollout) else: self._validate_data(rollout) if rollout.rollout_id not in self._task_id_to_original_sample: + # Reject unregistered / leftover IDs from earlier batches (#355). print(f"Warning: Received unknown rollout ID {rollout.rollout_id}, skipping.") - else: - self._completed_rollouts_v0[rollout.rollout_id] = rollout + continue + + existing = self._completed_rollouts_v0.get(rollout.rollout_id) + if existing is not None and not self._is_failed_rollout(existing): + # Keep the first successful result for this rollout_id. + continue + if existing is not None and self._is_failed_rollout(rollout): + # No upgrade from one failed placeholder to another. + continue + + # Accept first result, or overwrite a failed placeholder with a + # successful retry for the same rollout_id. + self._completed_rollouts_v0[rollout.rollout_id] = rollout + self._drop_stale_completed_rollouts() if verbose: print(f"Completed {len(self._completed_rollouts_v0)}/{self._total_tasks_queued} tasks...") - await asyncio.sleep(5) + + missing = [ + rollout_id + for rollout_id in self._task_id_to_original_sample + if rollout_id not in self._completed_rollouts_v0 + ] + failed = [ + rollout_id + for rollout_id, rollout in self._completed_rollouts_v0.items() + if self._is_failed_rollout(rollout) + ] + if missing: + failed_grace_polls_remaining = None + await asyncio.sleep(5) + continue + if failed: + if failed_grace_polls_remaining is None: + # A few extra polls (~15s) for in-flight retries to land. + failed_grace_polls_remaining = 3 + elif failed_grace_polls_remaining > 0: + failed_grace_polls_remaining -= 1 + await asyncio.sleep(5) + continue + break print("All tasks finished.") @@ -556,7 +651,7 @@ def run_until_all_finished(self, verbose: bool = True): def get_test_metrics(self): """Calculates and returns metrics for a validation run.""" assert not self.is_train, "This method should only be called during validation." - assert len(self._completed_rollouts_v0) == self._total_tasks_queued + self._ensure_completed_matches_queued() sample_stat_list: List[Dict[str, Any]] = [] sample_stat_list_by_source: Dict[str, List[Dict[str, Any]]] = defaultdict( @@ -638,7 +733,7 @@ def get_train_data_batch(self, max_prompt_length: int, max_response_length: int, truncation, and tensor creation for the PPO training loop. """ assert self.is_train, "This method should only be called during training." - assert len(self._completed_rollouts_v0) == self._total_tasks_queued + self._ensure_completed_matches_queued() # 1. Reconstruct the `finished_id_to_sample_info` structure from completed rollouts finished_id_to_sample_info: Dict[str, Dict[str, Any]] = {} diff --git a/tests/verl/__init__.py b/tests/verl/__init__.py new file mode 100644 index 000000000..2a50eae89 --- /dev/null +++ b/tests/verl/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Microsoft. All rights reserved. diff --git a/tests/verl/conftest.py b/tests/verl/conftest.py new file mode 100644 index 000000000..76a6d04cc --- /dev/null +++ b/tests/verl/conftest.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Test fixtures for VERL daemon unit tests. + +``agentlightning.verl.daemon`` imports ``verl.DataProto`` at module load time. +These tests only exercise daemon bookkeeping helpers, so we stub the VERL +package when it is not installed. +""" + +from __future__ import annotations + +import sys +import types + + +def pytest_configure() -> None: + if "verl" in sys.modules: + return + verl_stub = types.ModuleType("verl") + + class DataProto: # noqa: D401 - minimal stub for import + """Placeholder for verl.DataProto.""" + + verl_stub.DataProto = DataProto + sys.modules["verl"] = verl_stub diff --git a/tests/verl/test_daemon_completed_rollouts.py b/tests/verl/test_daemon_completed_rollouts.py new file mode 100644 index 000000000..9b8d42f12 --- /dev/null +++ b/tests/verl/test_daemon_completed_rollouts.py @@ -0,0 +1,138 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Regression tests for AgentModeDaemon completed-rollout bookkeeping (#355).""" + +from __future__ import annotations + +import asyncio +import importlib.util +import sys +import types +from pathlib import Path +from typing import List, Optional +from unittest.mock import AsyncMock + +import pytest + +from agentlightning.types import RolloutLegacy, Triplet + + +def _load_daemon_module(): + """Load daemon.py without importing agentlightning.verl package __init__.""" + if "verl" not in sys.modules: + verl_stub = types.ModuleType("verl") + + class DataProto: + pass + + verl_stub.DataProto = DataProto + sys.modules["verl"] = verl_stub + + daemon_path = Path(__file__).resolve().parents[2] / "agentlightning" / "verl" / "daemon.py" + spec = importlib.util.spec_from_file_location("agentlightning_verl_daemon_under_test", daemon_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +_daemon_module = _load_daemon_module() +AgentModeDaemon = _daemon_module.AgentModeDaemon + + +def _bare_daemon(*, mode: str = "v1") -> AgentModeDaemon: + """Build a daemon instance without starting servers or proxy threads.""" + daemon = object.__new__(AgentModeDaemon) + daemon.mode = mode + daemon._total_tasks_queued = 0 + daemon._completed_rollouts_v0 = {} + daemon._task_id_to_original_sample = {} + daemon.is_train = True + return daemon + + +def _failed(rollout_id: str) -> RolloutLegacy: + return RolloutLegacy(rollout_id=rollout_id, final_reward=None, triplets=None) + + +def _success(rollout_id: str, reward: float = 1.0) -> RolloutLegacy: + return RolloutLegacy( + rollout_id=rollout_id, + final_reward=reward, + triplets=[Triplet(prompt={"token_ids": [1]}, response={"token_ids": [2]})], + ) + + +def test_is_failed_rollout_detects_empty_backend_failures() -> None: + assert AgentModeDaemon._is_failed_rollout(_failed("r1")) is True + assert AgentModeDaemon._is_failed_rollout(RolloutLegacy(rollout_id="r1", final_reward=None, triplets=[])) is True + assert AgentModeDaemon._is_failed_rollout(_success("r1")) is False + assert ( + AgentModeDaemon._is_failed_rollout(RolloutLegacy(rollout_id="r1", final_reward=0.0, triplets=None)) is False + ) + + +def test_drop_stale_completed_rollouts_removes_unregistered_ids() -> None: + daemon = _bare_daemon() + daemon._task_id_to_original_sample = {"r1": {"data_id": "d1"}} + daemon._total_tasks_queued = 1 + daemon._completed_rollouts_v0 = { + "r1": _success("r1"), + "stale-from-previous-batch": _success("stale-from-previous-batch"), + } + + dropped = daemon._drop_stale_completed_rollouts() + + assert dropped == 1 + assert list(daemon._completed_rollouts_v0.keys()) == ["r1"] + + +def test_ensure_completed_matches_queued_drops_extras_without_asserting() -> None: + daemon = _bare_daemon() + daemon._task_id_to_original_sample = {"r1": {"data_id": "d1"}, "r2": {"data_id": "d2"}} + daemon._total_tasks_queued = 2 + daemon._completed_rollouts_v0 = { + "r1": _success("r1"), + "r2": _success("r2"), + "extra": _success("extra"), + } + + daemon._ensure_completed_matches_queued() + + assert set(daemon._completed_rollouts_v0) == {"r1", "r2"} + + +@pytest.mark.asyncio +async def test_async_run_until_finished_rejects_unknown_and_overwrites_failed(monkeypatch: pytest.MonkeyPatch) -> None: + daemon = _bare_daemon(mode="v1") + daemon._task_id_to_original_sample = {"r1": {"data_id": "d1"}, "r2": {"data_id": "d2"}} + daemon._total_tasks_queued = 2 + + polls: List[List[RolloutLegacy]] = [ + # First poll: one failure, one unknown stale id from a previous batch. + [_failed("r1"), _success("stale-old")], + # Second poll: successful retry for r1 and a success for r2. + [_success("r1", reward=0.5), _success("r2")], + ] + + async def fake_wait_for_rollouts(*, rollout_ids: List[str], timeout: Optional[float] = None) -> List[RolloutLegacy]: + del timeout + assert set(rollout_ids).issubset({"r1", "r2"}) + if not polls: + return [] + return polls.pop(0) + + daemon.store = AsyncMock() + daemon.store.wait_for_rollouts = fake_wait_for_rollouts + daemon._validate_data = lambda rollout: None # type: ignore[method-assign] + + async def immediate_sleep(_: float) -> None: + return None + + monkeypatch.setattr(asyncio, "sleep", immediate_sleep) + + await daemon._async_run_until_finished(verbose=False) + + assert set(daemon._completed_rollouts_v0) == {"r1", "r2"} + assert daemon._completed_rollouts_v0["r1"].final_reward == 0.5 + assert AgentModeDaemon._is_failed_rollout(daemon._completed_rollouts_v0["r1"]) is False From 937836c35867bdedf3d1adb69438f2e7e4c10f31 Mon Sep 17 00:00:00 2001 From: Mr-Perfectuz <55726523+Mr-Perfectuz@users.noreply.github.com> Date: Sun, 18 Oct 2026 00:15:37 +0900 Subject: [PATCH 2/2] fix(verl): keep grace-polling after failed-only batch results Co-authored-by: Cursor --- agentlightning/verl/daemon.py | 4 +- tests/verl/test_daemon_completed_rollouts.py | 42 ++++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/agentlightning/verl/daemon.py b/agentlightning/verl/daemon.py index 4337c0bcc..5570b9305 100644 --- a/agentlightning/verl/daemon.py +++ b/agentlightning/verl/daemon.py @@ -618,7 +618,9 @@ async def _async_run_until_finished(self, verbose: bool = True): if failed_grace_polls_remaining is None: # A few extra polls (~15s) for in-flight retries to land. failed_grace_polls_remaining = 3 - elif failed_grace_polls_remaining > 0: + # Independent of init: always sleep/poll while grace remains so a + # retry that arrives on the next poll can replace the failure. + if failed_grace_polls_remaining > 0: failed_grace_polls_remaining -= 1 await asyncio.sleep(5) continue diff --git a/tests/verl/test_daemon_completed_rollouts.py b/tests/verl/test_daemon_completed_rollouts.py index 9b8d42f12..ba1173ba0 100644 --- a/tests/verl/test_daemon_completed_rollouts.py +++ b/tests/verl/test_daemon_completed_rollouts.py @@ -136,3 +136,45 @@ async def immediate_sleep(_: float) -> None: assert set(daemon._completed_rollouts_v0) == {"r1", "r2"} assert daemon._completed_rollouts_v0["r1"].final_reward == 0.5 assert AgentModeDaemon._is_failed_rollout(daemon._completed_rollouts_v0["r1"]) is False + + +@pytest.mark.asyncio +async def test_async_run_until_finished_grace_polls_when_only_failure_present( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """When every queued ID has a result but it failed, keep polling for retries. + + Unlike the mixed missing+failed case, this hits the grace-init branch with + ``missing == []`` so a successful retry on the next poll can replace it. + """ + daemon = _bare_daemon(mode="v1") + daemon._task_id_to_original_sample = {"r1": {"data_id": "d1"}} + daemon._total_tasks_queued = 1 + + polls: List[List[RolloutLegacy]] = [ + [_failed("r1")], + [_success("r1", reward=0.7)], + ] + + async def fake_wait_for_rollouts(*, rollout_ids: List[str], timeout: Optional[float] = None) -> List[RolloutLegacy]: + del timeout + assert set(rollout_ids) == {"r1"} + if not polls: + return [] + return polls.pop(0) + + daemon.store = AsyncMock() + daemon.store.wait_for_rollouts = fake_wait_for_rollouts + daemon._validate_data = lambda rollout: None # type: ignore[method-assign] + + async def immediate_sleep(_: float) -> None: + return None + + monkeypatch.setattr(asyncio, "sleep", immediate_sleep) + + await daemon._async_run_until_finished(verbose=False) + + assert set(daemon._completed_rollouts_v0) == {"r1"} + assert daemon._completed_rollouts_v0["r1"].final_reward == 0.7 + assert AgentModeDaemon._is_failed_rollout(daemon._completed_rollouts_v0["r1"]) is False + assert polls == []