Hold only a run that trained every rollout to the run-end store check

The run-end sweep failed every synchronous --ci-test run that left an object, and raised on a
failed sweep in any run. A run that stopped before training a rollout still holds that
rollout's objects, so a crash ended in the sweep's AssertionError, with the real error only as
its context. A --partial-rollout run ends with aborted groups back in the data buffer, so a
correct run failed. And a Mooncake failure at teardown turned a finished non-CI run into a
failed one, though a failed delete only logs.

get() now files every rollout until its release, so dispose can tell a run that stopped early.
The sweep holds a CI run to leaving nothing, and to a working sweep, only when no rollout is
left unreleased; a partial-rollout run only logs what it left, like a fully-async one; outside
CI a failed sweep only logs. A dispose test checks the sweep runs once, after the rest of
teardown.

Author: guapisolo

docs/advanced/mooncake-rollout-transfer.md

@@ -287,10 +287,11 @@ Requirements:
 
 Known defect: an object stored by a request attempt that was then retried, or whose delete
 failed (Miles logs it), leaks and stays until the run ends. When the run ends, Miles removes
-every object left under the run's key prefix and logs how many keys it removed; under
-`--ci-test`, a synchronous run that left any fails. A fully-async run ends with requests in
-flight, and an engine can still store an object for one after that. A periodic sweep during the
-run is planned (TODO).
+every object left under the run's key prefix and logs how many keys it removed, or that the
+removal failed. Under `--ci-test`, a run that trained every rollout it generated fails if the
+removal failed or found any object, unless it is fully async or uses `--partial-rollout`: those
+end with groups still buffered. A fully-async run ends with requests in flight, and an engine
+can still store an object for one after that. A periodic sweep during the run is planned (TODO).
 
 Rejected at launch: no replay array enabled, `--cluster-backend kubernetes`, PD
 disaggregation, external rollout engines, the legacy rollout, a custom convert function,

miles/ray/rollout/rollout_executor.py

@@ -97,6 +97,7 @@ def __init__(
         self._inference_controller_provider = inference_controller_provider
         self._output_snapshotter = _RolloutExecutorOutputSnapshotter(args=args)
         # Output store refs of each rollout's trained and dropped samples, deleted once it has trained.
+        # Every rollout get() began is here until its release, so one left at dispose was never trained.
         self._unreleased_store_refs: dict[int, dict[str, Any]] = {}
         self._delete_store_refs = None
 
@@ -189,11 +190,23 @@ async def dispose(self) -> None:
 
     def _sweep_output_store(self) -> None:
         """Remove what the run left in the output store, which nothing reads once it ends."""
-        if not (left := sweep_output_store(self.args)):
+        # Only a run that trained every rollout it generated answers to CI: one that stopped early (a crash, a debug
+        # exit) still holds that rollout's objects, and a crash already has its own error.
+        strict = self.args.ci_test and not self._unreleased_store_refs
+        try:
+            left = sweep_output_store(self.args)
+        except Exception:
+            if strict:
+                raise
+            # Like a failed delete, a failed sweep only leaks objects nothing reads.
+            logger.exception("Failed to sweep the output store; the run's objects stay in Mooncake")
+            return
+        if not left:
             return
         message = f"Swept {left} output store keys the run left behind"
-        # A fully-async run ends with groups buffered or in flight; a synchronous one has released every rollout.
-        if self.args.ci_test and not self.args.fully_async:
+        # A fully-async run ends with groups buffered or in flight, a partial rollout with aborted groups back in
+        # the data buffer; any other run that trained every rollout it generated has released them all.
+        if strict and not (self.args.fully_async or self.args.partial_rollout):
             raise AssertionError(f"{message}: a retried request, a failed delete or a missed release")
         logger.warning(message)
 
@@ -202,6 +215,7 @@ def _sweep_output_store(self) -> None:
     @event_logger_context(lambda _self, rollout_id, trainer_model_id=None: dict(rollout_id=rollout_id))
     async def get(self, rollout_id: int, trainer_model_id: str | None = None) -> RolloutDataPack:
         self.last_get_rollout_id_of_model_id[trainer_model_id] = rollout_id
+        self._unreleased_store_refs[rollout_id] = {}
         event_analyzer.run_sample_ownership_analysis(args=self.args)
         replay = self._output_snapshotter.get(trainer_model_id=trainer_model_id, rollout_id=rollout_id)
         if replay is None:
@@ -216,8 +230,7 @@ async def get(self, rollout_id: int, trainer_model_id: str | None = None) -> Rol
         else:
             data, metadata = replay.data, replay.metadata
         # The dp schedule and a custom convert may leave rows out of the train data: collect from the samples.
-        if store_refs := store_refs_of(data) | take_dropped_refs():
-            self._unreleased_store_refs[rollout_id] = store_refs
+        self._unreleased_store_refs[rollout_id] = store_refs_of(data) | take_dropped_refs()
 
         with SampleOwnershipRecorder.suppress_drop_logging() if replay is not None else nullcontext():
             train_data = convert_samples_to_train_data(

tests/fast/ray/rollout/test_rollout_executor.py

@@ -174,6 +174,33 @@ async def test_synchronous_train_and_eval_rollout_disposers_are_accepted(
 
         assert disposed == ["train", "eval"]
 
+    async def test_a_run_with_the_output_store_sweeps_it_once_after_the_rest_of_dispose(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Last, so the producers have stopped writing and a failed CI check raises after the other teardown."""
+        disposed: list[str] = []
+        executor = RolloutExecutor.__new__(RolloutExecutor)
+        executor._disposed = False
+        executor.use_legacy_rollout_v1 = False
+        executor.generate_rollout = _SynchronousDisposable(disposed, "train")
+        executor.eval_generate_rollout = _SynchronousDisposable(disposed, "eval")
+        executor.data_source = object()
+        executor.args = make_args(ci_test=True, fully_async=False, partial_rollout=False)
+        executor._metric_checker = None
+        executor._delete_store_refs = lambda refs: None
+        executor._unreleased_store_refs = {}
+        monkeypatch.setattr(rollout_executor_module, "CheckpointEvalFn", _SynchronousDisposable)
+        monkeypatch.setattr(rollout_executor_module.event_analyzer, "run_analysis_from_args", lambda _args: None)
+        monkeypatch.setattr(
+            rollout_executor_module.event_analyzer, "run_sample_ownership_analysis", lambda *, args: None
+        )
+        monkeypatch.setattr(rollout_executor_module, "sweep_output_store", lambda args: disposed.append("sweep") or 0)
+
+        await executor.dispose()
+        await executor.dispose()
+
+        assert disposed == ["train", "eval", "sweep"]
+
     async def test_the_final_ownership_check_runs_before_the_run_ends(self, monkeypatch: pytest.MonkeyPatch) -> None:
         """Violations that only mature after the last get would otherwise never be checked."""
         checked: list[Namespace] = []
@@ -320,6 +347,7 @@ def _make_executor(tmp_path: Path, rollout_fn: _CountingRolloutFn) -> RolloutExe
     executor._train_parallel_configs_of_model_id = {None: make_train_parallel_config()}
     executor._weight_versions_of_model_id = {}
     executor.last_get_rollout_id_of_model_id = {}
+    executor._unreleased_store_refs = {}
     executor.custom_convert_samples_to_train_data_func = None
     executor.custom_reward_post_process_func = None
     executor._output_snapshotter = _RolloutExecutorOutputSnapshotter(args=executor.args)

tests/fast/ray/rollout/test_rollout_executor_multi_policy.py

@@ -69,6 +69,7 @@ def _make_executor() -> RolloutExecutor:
     executor._train_parallel_configs_of_model_id = {}
     executor._rollouts_since_publish_of_model_id = defaultdict(int)
     executor.last_get_rollout_id_of_model_id = {}
+    executor._unreleased_store_refs = {}
     executor._eval_lock = asyncio.Lock()
     executor._eval_fleet = None
     executor._metric_checker = None

tests/fast/ray/rollout/test_rollout_executor_store_refs.py

@@ -110,14 +110,23 @@ async def test_refs_without_a_store_client_fail_loudly():
         await executor.release(5)
 
 
+def run_end_args(*, ci_test=True, fully_async=False, partial_rollout=False):
+    return make_args(ci_test=ci_test, fully_async=fully_async, partial_rollout=partial_rollout)
+
+
 @pytest.mark.parametrize(
-    ("ci_test", "fully_async", "fails"),
-    [(True, False, True), (True, True, False), (False, False, False)],
+    ("args", "fails"),
+    [
+        (run_end_args(), True),
+        (run_end_args(fully_async=True), False),
+        (run_end_args(partial_rollout=True), False),
+        (run_end_args(ci_test=False), False),
+    ],
 )
-def test_keys_left_at_the_end_fail_only_a_synchronous_ci_run(monkeypatch, caplog, ci_test, fully_async, fails):
-    """A fully-async run ends with groups buffered or in flight, so only a synchronous one must leave none."""
+def test_keys_left_at_the_end_fail_only_a_synchronous_ci_run(monkeypatch, caplog, args, fails):
+    """A fully-async run ends with groups buffered or in flight, a partial rollout with aborted groups in its buffer."""
     executor, _ = make_executor({})
-    executor.args = make_args(ci_test=ci_test, fully_async=fully_async)
+    executor.args = args
     monkeypatch.setattr(rollout_executor_module, "sweep_output_store", lambda args: 4)
 
     if fails:
@@ -130,9 +139,57 @@ def test_keys_left_at_the_end_fail_only_a_synchronous_ci_run(monkeypatch, caplog
 
 def test_a_run_that_left_nothing_logs_nothing(monkeypatch, caplog):
     executor, _ = make_executor({})
-    executor.args = make_args(ci_test=True, fully_async=False)
+    executor.args = run_end_args()
     monkeypatch.setattr(rollout_executor_module, "sweep_output_store", lambda args: 0)
 
     executor._sweep_output_store()
 
     assert "Swept" not in caplog.text
+
+
+def _fail_the_generation(rollout_id):
+    raise RuntimeError("engine died")
+
+
+@pytest.mark.parametrize("stopped_in", ["train", "generation"])
+async def test_a_run_that_stopped_before_training_a_rollout_only_logs_what_it_left(monkeypatch, caplog, stopped_in):
+    """Its error is the crash itself; the untrained rollout's objects are expected, not a leak."""
+    if stopped_in == "train":
+        executor, _ = make_executor({5: [sample_owning("a")]})
+        await executor.get(rollout_id=5)
+    else:
+        executor, _ = make_executor({5: [sample_owning("a")]}, during_get=_fail_the_generation)
+        with pytest.raises(RuntimeError, match="engine died"):
+            await executor.get(rollout_id=5)
+    executor.args = run_end_args()
+    monkeypatch.setattr(rollout_executor_module, "sweep_output_store", lambda args: 4)
+
+    executor._sweep_output_store()
+
+    assert "Swept 4 output store keys" in caplog.text
+
+
+def _fail_the_sweep(args):
+    raise RuntimeError("Mooncake store setup failed: -1")
+
+
+@pytest.mark.parametrize(
+    ("args", "stopped_early", "fails"),
+    [(run_end_args(), False, True), (run_end_args(ci_test=False), False, False), (run_end_args(), True, False)],
+)
+async def test_a_failed_sweep_fails_only_a_ci_run_that_trained_every_rollout(
+    monkeypatch, caplog, args, stopped_early, fails
+):
+    """Like a failed delete, a failed sweep only leaks objects nothing reads; CI still wants to see it."""
+    executor, _ = make_executor({5: [sample_owning("a")]})
+    if stopped_early:
+        await executor.get(rollout_id=5)
+    executor.args = args
+    monkeypatch.setattr(rollout_executor_module, "sweep_output_store", _fail_the_sweep)
+
+    if fails:
+        with pytest.raises(RuntimeError, match="setup failed"):
+            executor._sweep_output_store()
+    else:
+        executor._sweep_output_store()
+        assert "Failed to sweep the output store" in caplog.text