fix(verl): stop admitting batches at the epoch limit

agentlightning/verl/trainer.py

@@ -55,6 +55,10 @@
 RolloutManagerT = TypeVar("RolloutManagerT", bound=AglRolloutManagerBase)
 
 
+class _TrainingDataExhausted(Exception):
+    """The configured training epochs have no more input batches."""
+
+
 def _batch_dict_len(batch: dict[str, Any] | None) -> int:
     """Leading-dim length of a dataloader batch dict (0 if absent/empty)."""
     if batch is None or not batch:
@@ -274,10 +278,13 @@ def _next_train_batch_dict_for_rollout(self) -> dict[str, Any]:
         async_train_batch_size = async_cfg.async_train_batch_size
         n_carry_over = len({rollout.data_id for rollout in self._carry_over_rollouts})
         n_new = int(async_train_batch_size) - n_carry_over
-        return self._next_train_batch_dict_with_size(n_new)
+        batch = self._next_train_batch_dict_with_size(n_new)
+        if not batch and not self._carry_over_rollouts:
+            raise _TrainingDataExhausted
+        return batch
 
     def _next_train_batch_dict(self) -> dict[str, Any]:
-        while True:
+        while self.epoch <= self.config.trainer.total_epochs:
             if self._train_dataloader_iter is None:
                 self._train_dataloader_iter = iter(self.train_dataloader)
 
@@ -286,6 +293,7 @@ def _next_train_batch_dict(self) -> dict[str, Any]:
 
             self.epoch += 1
             self._train_dataloader_iter = None
+        raise _TrainingDataExhausted
 
     def _next_train_batch_dict_with_size(self, size: int) -> dict[str, Any]:
         if size <= 0:
@@ -340,15 +348,12 @@ def concat_batches(batches: list[dict[str, Any]]) -> dict[str, Any]:
                 remaining = 0
 
         while remaining > 0:
-            if self._train_dataloader_iter is None:
-                self._train_dataloader_iter = iter(self.train_dataloader)
-
             try:
-                batch_dict = next(self._train_dataloader_iter)
-            except StopIteration:
-                self.epoch += 1
-                self._train_dataloader_iter = None
-                continue
+                batch_dict = self._next_train_batch_dict()
+            except _TrainingDataExhausted:
+                # The last refill may be partial. An empty refill still lets
+                # the async runner finish groups admitted by an earlier step.
+                break
 
             batch_size = _batch_dict_len(batch_dict)
             if batch_size <= remaining:
@@ -360,7 +365,7 @@ def concat_batches(batches: list[dict[str, Any]]) -> dict[str, Any]:
                 self._train_dataloader_buf = tail
                 remaining = 0
 
-        return concat_batches(collected)
+        return concat_batches(collected) if collected else {}
 
     def _rollout(self, gen_batch: DataProto, is_train: bool) -> tuple[DataProto, dict[str, Any]]:
         """Run Agent Lightning rollouts and return the resulting DataProto plus metrics."""
@@ -381,11 +386,14 @@ def _rollout(self, gen_batch: DataProto, is_train: bool) -> tuple[DataProto, dic
             rollout_manager.delete_model()
             rollout_manager.register_model(server_addresses)
             previous_carry_over_rollouts = list(self._carry_over_rollouts)
+            available_groups = _batch_dict_len(data_dict) + len(
+                {rollout.data_id for rollout in previous_carry_over_rollouts}
+            )
             completed_rollouts, new_carry_over_rollouts = rollout_manager.enqueue_and_wait_until_group_completed(
                 data_dict,
                 previous_carry_over_rollouts,
                 is_train=True,
-                target_finished_group_num=self.config.data.train_batch_size,
+                target_finished_group_num=min(self.config.data.train_batch_size, available_groups),
             )
             self._carry_over_rollouts = new_carry_over_rollouts
             async_rollout_metrics = self._compute_async_rollout_metrics(
@@ -680,6 +688,7 @@ def fit(self):
         self.global_steps = 0
         self.epoch = 0
         self._train_dataloader_iter = None
+        self._train_dataloader_buf = None
         self._carry_over_rollouts = []
         self._load_checkpoint()
         # Push loaded weights before the first rollout or validation.
@@ -698,9 +707,10 @@ def fit(self):
         self.global_steps += 1
         self.epoch += 1
         last_val_metrics = None
+        last_update_step: int | None = None
 
         while True:
-            if self.global_steps > self.total_training_steps or self.epoch > self.config.trainer.total_epochs:
+            if self.global_steps > self.total_training_steps:
                 progress_bar.close()
                 return
 
@@ -712,15 +722,30 @@ def fit(self):
                 else False
             )
 
-            with marked_timer("step", timing_raw):
-                result = self._train_step(timing_raw, curr_step_profile)
+            try:
+                with marked_timer("step", timing_raw):
+                    result = self._train_step(timing_raw, curr_step_profile)
+            except _TrainingDataExhausted:
+                # Exhaustion is discovered when requesting the next input, so
+                # finalize the last update without counting an extra step.
+                if last_update_step is not None:
+                    self.global_steps = last_update_step
+                    if self.config.trainer.test_freq > 0 and self.global_steps % self.config.trainer.test_freq != 0:
+                        last_val_metrics = self._validate()
+                        logger.log(data=last_val_metrics, step=self.global_steps)
+                    if self.config.trainer.save_freq > 0 and self.global_steps % self.config.trainer.save_freq != 0:
+                        self._save_checkpoint()
+                pprint(f"Final validation metrics: {last_val_metrics}")
+                progress_bar.close()
+                return
             if result is None:
                 print("AgentLightningRayPPOTrainer: train step returned no batch; skipping update.")
                 # SC experiment budgets count optimizer updates, not empty collections.
                 if self.config.actor_rollout_ref.actor.policy_loss.loss_mode != "score_centering":
                     self.global_steps += 1
                 continue
             metrics, step_batch = result
+            last_update_step = self.global_steps
 
             # Compute timing and throughput after the step timer records its total duration.
             metrics.update(compute_timing_metrics(batch=step_batch, timing_raw=timing_raw))
@@ -735,8 +760,7 @@ def fit(self):
             ):
                 with marked_timer("validate", timing_raw, color="green"):
                     val_metrics = self._validate()
-                    if is_last_step:
-                        last_val_metrics = val_metrics
+                    last_val_metrics = val_metrics
                 metrics.update(val_metrics)
 
             if self.config.trainer.save_freq > 0 and (

docs/20-trainer-configuration.md

@@ -193,6 +193,19 @@ Transition mode does not use `trajectory_max_prompt_length` or `trajectory_max_r
 
 Use transition mode when every request-response call should remain a separate training sample.
 
+## Training limits
+
+`trainer.total_epochs` limits how many times the trainer reads the training dataloader.
+After the final epoch, synchronous collection stops before starting another rollout batch.
+Asynchronous collection stops admitting new inputs and finishes the groups already in flight;
+its last refill may contain fewer inputs than `async_train_batch_size` requests. The configured
+`trainer.total_training_steps` remains an independent upper bound and can stop training earlier.
+
+When the epoch limit ends training, enabled validation and checkpointing run at the last
+successful training step, unless they already ran there on their regular schedule. No extra
+optimizer step is taken to discover that the data is exhausted. Final partial collections still
+follow the PPO mini-batch alignment rules described below.
+
 ## Algorithm correctness
 
 The following settings control how rollout data contributes to optimization:

tests/verl/test_training_epoch_limit.py

@@ -0,0 +1,206 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Exercise epoch termination through fit without launching GPU workers."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any, cast
+
+import numpy as np
+import pytest
+import torch
+from omegaconf import OmegaConf
+from verl import DataProto
+
+from agentlightning.verl import trainer as trainer_module
+
+
+def _trainer(
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    rows: int = 4,
+    epochs: int = 1,
+    async_size: int | None = None,
+    total_steps: int = 100,
+    frequency: int = 10,
+    skip_updates: bool = False,
+) -> Any:
+    trainer = cast(Any, object.__new__(trainer_module.AgentLightningRayPPOTrainer))
+    trainer.config = OmegaConf.create(
+        {
+            "trainer": {
+                "project_name": "test",
+                "experiment_name": "epochs",
+                "logger": [],
+                "val_before_train": False,
+                "total_epochs": epochs,
+                "test_freq": frequency,
+                "save_freq": frequency,
+            },
+            "global_profiler": {"steps": None},
+            "actor_rollout_ref": {"actor": {"policy_loss": {"loss_mode": "score_centering"}}, "rollout": {"n": 1}},
+            "agentlightning": {
+                "async_rollout": {"async_train_batch_size": async_size},
+                "trace_aggregator": {},
+                "reward_fillna_value": 0.0,
+            },
+            "data": {"train_batch_size": 2, "max_prompt_length": 1, "max_response_length": 1},
+        }
+    )
+    trainer.is_async = async_size is not None
+    trainer.total_training_steps = total_steps
+    trainer.train_dataloader = [
+        {"item": np.arange(start, min(start + 2, rows)), "fake_ids": torch.ones(min(2, rows - start), 1)}
+        for start in range(0, rows, 2)
+    ]
+    trainer._load_checkpoint = lambda: None
+    trainer.checkpoint_manager = SimpleNamespace(update_weights=lambda _step: None)
+    trainer.resource_pool_manager = SimpleNamespace(get_n_gpus=lambda: 0)
+    trainer.tokenizer = SimpleNamespace(pad_token_id=0)
+    trainer.async_rollout_manager = SimpleNamespace(server_addresses=["test"])
+    trainer._resume_all_rollout_generation = lambda: None
+    trainer._resume_gateway = lambda: None
+    trainer._pause_and_drain_gateway = lambda **_kwargs: {}
+    trainer._compute_async_rollout_metrics = lambda **_kwargs: {}
+    trainer._rollout_lifecycle_metrics = lambda _rollouts: {}
+    trainer.fetched = []
+    trainer.completed = []
+    trainer.updates = []
+    trainer.validations = []
+    trainer.checkpoints = []
+    trainer.targets = []
+    trainer._save_checkpoint = lambda: trainer.checkpoints.append(trainer.global_steps)
+
+    def validate():
+        trainer.validations.append(trainer.global_steps)
+        return {"validation/reward": 1.0}
+
+    trainer._validate = validate
+    monkeypatch.setattr(trainer_module, "Tracking", lambda **_kwargs: SimpleNamespace(log=lambda **_kwargs: None))
+    monkeypatch.setattr(trainer_module, "compute_timing_metrics", lambda **_kwargs: {})
+
+    class Manager:
+        next_id = 0
+
+        def delete_model(self):
+            pass
+
+        def register_model(self, _addresses):
+            pass
+
+        def enqueue_and_wait_until_group_completed(self, data, carry, *, is_train, target_finished_group_num):
+            assert is_train
+            active = list(carry)
+            for item in data.get("item", []):
+                active.append(SimpleNamespace(data_id=str(self.next_id), item=int(item)))
+                self.next_id += 1
+            # A target above the available groups would hang the real manager.
+            assert 0 < target_finished_group_num <= len(active)
+            trainer.targets.append(target_finished_group_num)
+            return active[:target_finished_group_num], active[target_finished_group_num:]
+
+    class Adapter:
+        def __init__(self, **_kwargs):
+            pass
+
+        def get_train_data_batch(self, completed, **_kwargs):
+            trainer.completed.extend(row.item for row in completed)
+            return DataProto.from_dict(tensors={"fake_ids": torch.ones(len(completed), 1)}), {}
+
+    manager = Manager()
+    trainer._make_rollout_manager = lambda _cls: manager
+    monkeypatch.setattr(trainer_module, "RolloutAdapter", Adapter)
+
+    def train_step(_timing, _profile):
+        batch_dict = trainer._next_train_batch_dict_for_rollout()
+        trainer.fetched.extend(batch_dict.get("item", []))
+        batch = DataProto.from_single_dict(batch_dict)
+        if trainer.is_async:
+            batch, _ = trainer._rollout(batch, is_train=True)
+        if skip_updates:
+            return None
+        trainer.updates.append(trainer.global_steps)
+        return {}, batch
+
+    trainer._train_step = train_step
+    return trainer
+
+
+@pytest.mark.parametrize("epochs", [1, 2])
+def test_fit_stops_at_epoch_limit_before_an_extra_update(monkeypatch, epochs):
+    trainer = _trainer(monkeypatch, epochs=epochs)
+
+    trainer.fit()
+
+    assert trainer.fetched == [0, 1, 2, 3] * epochs
+    assert trainer.updates == list(range(1, 2 * epochs + 1))
+    assert trainer.validations == [2 * epochs]
+    assert trainer.checkpoints == [2 * epochs]
+
+
+@pytest.mark.parametrize("rows,async_size", [(4, 4), (6, 5), (3, 5)])
+@pytest.mark.parametrize("epochs", [1, 2])
+def test_async_fit_drains_admitted_groups_without_refilling_past_epoch_limit(monkeypatch, rows, async_size, epochs):
+    trainer = _trainer(monkeypatch, rows=rows, epochs=epochs, async_size=async_size)
+
+    trainer.fit()
+
+    assert trainer.fetched == list(range(rows)) * epochs
+    assert trainer.completed == trainer.fetched
+    assert trainer._carry_over_rollouts == []
+    assert sum(trainer.targets) == rows * epochs
+    assert len(trainer.updates) == (rows * epochs + 1) // 2
+    assert trainer.validations == [trainer.updates[-1]]
+    assert trainer.checkpoints == [trainer.updates[-1]]
+
+
+def test_epoch_completion_does_not_repeat_scheduled_validation_or_checkpoint(monkeypatch):
+    trainer = _trainer(monkeypatch, frequency=2)
+
+    trainer.fit()
+
+    assert trainer.validations == [2]
+    assert trainer.checkpoints == [2]
+
+
+def test_explicit_step_limit_still_stops_training_first(monkeypatch):
+    trainer = _trainer(monkeypatch, epochs=2, total_steps=1)
+
+    trainer.fit()
+
+    assert trainer.fetched == [0, 1]
+    assert trainer.updates == trainer.validations == trainer.checkpoints == [1]
+
+
+def test_epoch_limit_stops_empty_collections_without_advancing_optimizer_steps(monkeypatch):
+    trainer = _trainer(monkeypatch, skip_updates=True)
+
+    trainer.fit()
+
+    assert trainer.fetched == [0, 1, 2, 3]
+    assert trainer.updates == trainer.validations == trainer.checkpoints == []
+
+
+def test_actual_train_step_accepts_an_empty_refill_for_carry_over(monkeypatch):
+    trainer = _trainer(monkeypatch, async_size=4)
+    trainer.epoch = 2
+    trainer.global_steps = 2
+    trainer._train_dataloader_iter = None
+    trainer._carry_over_rollouts = [SimpleNamespace(data_id="pending")]
+
+    class ReachedRollout(Exception):
+        pass
+
+    def rollout(gen_batch, *, is_train):
+        assert is_train
+        assert gen_batch.non_tensor_batch == {}
+        assert len(gen_batch) == 0
+        assert gen_batch.meta_info["global_steps"] == 2
+        raise ReachedRollout
+
+    trainer._rollout = rollout
+    # Exercise DataProto construction and the inherited _get_gen_batch, not the
+    # simplified train-step stub used to observe the fit loop above.
+    with pytest.raises(ReachedRollout):
+        trainer_module.AgentLightningRayPPOTrainer._train_step(trainer, {}, False)