[ananthsub]

Branch ananthsub/token-capture-store-retire, on #3938 in the partial-rollout checkpoint stack, so this diff shows only the complete-record store change. It doesn't need #3938: the same commit applied to main alone is branch ananthsub/token-capture-store-retire-independent, which also carries identical copies of the two helpers #3938 adds (validate_rollout_ids and RolloutRemoval), and passes the unit suite on its own. Related and independent: #3937 (striped locks; a small conflict around the removed sweeps) and #4268 (bounded ledger cache).

What changed and why

The complete-record token store (TokenCaptureStore, behind TokenSource) could retire a rollout only through the conditional TokenSource.drop, which needs the consumed snapshot. That leaves three gaps:

• Masked and failed captures are never retired: rollout_collection keeps them for diagnosis, and a trainer had no operation to remove them later.
• Abandoned attempts have no snapshot, so drop can't retire them at all.
• Fences (the retired state drop leaves behind, which the store's docs called a tombstone) were only removable by sweep_retired, an age-based sweep that nothing called, which is a timer in disguise.
• Abandoned captures were only removable by sweep_stale (feat(token-capture): sweep abandoned captures safely #3969), the same kind of age-based sweep. It guesses that a capture is abandoned once all its files are older than a cutoff, which the caller must keep above the longest possible session. Nothing in Gym called it either.

This PR gives the store the same two batched operations #3938 gives the capture ledger, with the same names and meaning:

Operation	Capture ledger (#3938)	Complete-record token store (this PR)
Delete: remove the data and the fence	CaptureLedger.delete(rollout_ids)	TokenSource.delete(rollout_ids)
Retire: remove the data, keep a fence	CaptureLedger.retire(rollout_ids)	TokenSource.retire(rollout_ids): drop without the snapshot check; keeps the retired fence

• TokenCaptureStore.retire and delete are async like the rest of TokenSource, with retire_now and delete_now as synchronous forms, following freeze and freeze_now.
• drop and retire share one retire step, so both leave the same fence, which refuses later writes, begin_call, mark_incomplete, and freeze.
• A retired rollout refuses them with TokenCaptureRetiredError, a new subclass of TokenCaptureFrozenError, instead of a plain RuntimeError. The sink already drops a late write to a frozen rollout quietly. Before this, a late call of an abandoned attempt, which is exactly what retire is for, was logged with a traceback, counted as a capture failure, and the sink then tried to mark the retired rollout incomplete. Now it is dropped like a late write after a freeze. A late call that failed to build, produced no token ids, or arrived under a mismatched identity is dropped quietly the same way: the sink tries to mark the rollout first, and when the rollout turns out to be frozen or retired it skips its warning and the capture-failure counter.
• Fewer fsyncs. retire syncs the directory once per batch instead of once per rollout, and drop syncs it once instead of twice: the directory fsync after the state file is replaced already covers the deletions before it. Retiring takes one file fsync per rollout instead of three. delete syncs every non-empty batch, even when it finds nothing to remove, for the same reason as the capture ledger's delete: an overlapping delete may not have synced yet, and a reused rollout ID must not see the old fence or records come back after a crash.
• Retire works in two synced phases. retire_now and drop first write and sync the fences, then remove the records and sync again. A rollout is therefore never left without records and without a fence, and an overlapping retire of the same rollout cannot return before its fence is durable: the first sync also runs when only records are left to remove. A retry removes records left behind by a retire that failed part way or by a crash. A batch that fails part way still syncs what it changed; a failing sync there is logged, not raised, so it never hides the original error. Every non-empty batch syncs its fences before returning, even when it wrote none, so an overlapping retire can never report success before the fence is durable. The second phase re-checks each fence under the rollout's lock, so a rollout ID that a delete cleared and a new attempt reused between the phases keeps its new records.
• The store's docstrings and the external agent harness page now call that retired state a fence, the capture ledger's word for it, instead of a tombstone.
• delete is now batched; it was single-ID and synchronous. Its one caller, clear_token_captures_for_rollouts, now makes one batched call.
• sweep_retired and sweep_stale are removed. A capture's end of life is an explicit call from its owner, not its age: retire when the training framework is done with a rollout, including an abandoned attempt, and delete once nothing can still write.
• Both operations validate the whole batch first, reject a bare string, and report removed and absent like the ledger's, typed as the same RolloutRemovalPayload TypedDict.
• The conformance suite checks both operations, so an external TokenSource has to implement them. Its retirement check requires a late put, mark_incomplete, and begin_call to fail with TokenCaptureRetiredError, as TokenSource.retire documents, rather than accepting any exception, and lets any other freeze error fail the check.
• rollout_collection behavior is unchanged: it still keeps masked and failed captures. Callers can now remove them.
• Docs: the external agent harness page describes when to retire and delete.

How it works

flowchart TD
  A["Absent<br/>no records, no fence"] -- "first put, begin_call, or mark_incomplete" --> L["Live<br/>records on disk"]
  L -- "freeze" --> FZ["Frozen<br/>snapshot taken"]
  FZ -- "drop (exact snapshot)" --> R["Retired<br/>records gone, fence kept"]
  L -- "retire" --> R
  FZ -- "retire" --> R
  A -- "retire" --> R
  R -- "late put, begin_call, or mark_incomplete<br/>(TokenCaptureRetiredError)" --> R
  R -- "delete" --> A
  L -- "delete" --> A
  FZ -- "delete" --> A

Loading

drop stays the way to retire a consumed snapshot, because it refuses when the rollout changed after the snapshot. retire is for everything else the consumer is done with; delete removes fences once nothing can still write, and resets an ID before reuse.

What the caller must do

Situation	Call
Built and durably delivered	retire_rollout_token_capture (unchanged: conditional drop)
Masked or failed, evidence no longer needed	retire([ids])
Abandoned attempt	retire([ids]); its agent may still be running, and the fence makes its later writes fail
Retired attempts once nothing of them can still write (end of run, or after restarting every Gym server)	delete([ids]) to remove the fences
Reusing a rollout ID	delete([ids]) first (rollout_collection already does this before re-dispatch)

Always pass a list; both operations are idempotent, so a failed batch can be retried as is.

Compatibility

• TokenSource gains retire and delete. Gym never checks a TokenSource against the protocol at runtime, but the conformance suite now requires both, so an external source must add them to keep passing.
• TokenCaptureStore.delete(rollout_id) was synchronous and single-ID; it is now async delete(rollout_ids), with delete_now(rollout_ids) as the synchronous form. A caller still writing store.delete("id") gets an un-awaited coroutine and deletes nothing, so update such calls to delete_now(["id"]). Gym and NeMo RL have no such callers outside the one updated here.
• Writes, begin_call, mark_incomplete, and freeze on a retired rollout raise TokenCaptureRetiredError. It subclasses TokenCaptureFrozenError and so RuntimeError, so code that caught either still catches it.
• TokenCaptureStore.sweep_retired and TokenCaptureStore.sweep_stale are removed. Nothing in Gym or NeMo RL called either. A caller that swept abandoned captures by age now retires the rollouts it abandoned.

Issue

No tracking issue. This aligns the complete-record store's retirement with the capture ledger's and closes its gaps for masked, failed, and abandoned captures and for fences.

Validation

• New tests in tests/unit_tests/test_token_id_capture.py cover retiring live, frozen, and never-captured rollouts (records removed; later writes, begin_call, mark_incomplete, and freeze all raise TokenCaptureRetiredError; a repeat is a no-op), a late call after retire dropped by the sink without a capture failure or an incomplete marker, one directory fsync per retire batch and per drop, delete removing the fence so the ID can be reused, batch validation with bare-string rejection, every non-empty delete batch syncing, files removed before a later removal failed still being synced, a batch failing part way still syncing what it changed, a retire failing part way leaving the rollout fenced and a retry finishing it, an overlapping retire returning only after the fence is durable, a sync failure after an error not hiding that error, two directory syncs per retire batch and per drop, a retire that cannot write its fence changing nothing, and late calls without token ids or with a failed build not reported as failures. Existing tests that used sweep_retired or single-ID delete now use delete_now, and the sweep_stale tests are removed with it.
• New conformance checks unconditional_retirement and delete_clears_retirement pass for the file store and the in-memory reference backend, and reject a backend that lets a late mark_incomplete through after retire, rejects late writes with a plain or merely frozen error, or whose freeze fails after retire.
• The new tests fail when the retire step in retire_now, the fence removal in delete_now, the retired error type, or the per-batch fsync is reverted.
• Measured with the token-capture metrics from feat(token-capture): low-overhead metrics for capture operations, outcomes, digests, and caches #4218: 50 late calls to retired rollouts were counted as 50 capture_failed{write} and 50 marked_incomplete before the error-type change, and as 50 late_write_dropped after it. Retiring a rollout takes about 330 µs on a local disk, down from about 420 µs, and its system calls drop from three fsyncs to one.
• The combined tree (main with this PR, fix(token-capture): use striped lock files that are never deleted #3937, feat(token-capture): retire and delete external-staging capture ledgers #3938, feat(token-capture): low-overhead metrics for capture operations, outcomes, digests, and caches #4218, and fix(token-capture): bound the capture-ledger cache #4268): pytest tests/unit_tests passed.
• pytest on every token-capture, staging, trajectory-builder, rollout-collection, and model-server unit test file: passed.
• pytest responses_api_models/vllm_model/tests/test_token_capture_integration.py responses_api_models/vllm_model/tests/test_external_capture_streaming.py: passed.
• pre-commit run --files <changed files>: all hooks passed.

Rollout evidence

N/A: the delivery path still retires through the same drop, and the new operations run only when a caller invokes them, so rollouts behave as before.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[macandro96]

Both cases below are race conditions for very unlikely cases but fixes seem cheap

[github-actions]

🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/38026783272