[ananthsub]

Branch ananthsub/token-capture-ledger-retire, based on main, with no dependencies. #3939 sits on it in the partial-rollout checkpoint stack, but doesn't need it: the independent form of #3939 (branch ananthsub/token-capture-store-retire-independent) carries identical validate_rollout_ids and RolloutRemoval and applies to main alone. Related and independent: #3937 (striped locks) and #4268 (bounded ledger cache, split out of this PR).

Action for NeMo RL

Required: retire every rollout attempt's ledger once NeMo RL is done with it. Without this, the model server's lineage directory grows by one ledger per attempt for the whole run, and nothing in Gym can clean it up on NeMo RL's behalf.

Recommended first: keep the capture directory on each node's local disk. Today token_capture.capture_dir defaults to <logger.log_dir>/gym_token_capture, which is on Lustre, so every captured model call does several metadata round trips and an fsync against shared storage. Nothing outside the model server reads these files: the workers of one model server run on the actor's node, the manifest is fetched over HTTP, the reassembler works from receipts and TransferQueue, and Gym's partial-rollout checkpoints move ledger rows through the checkpoint itself. Setting capture_dir to a node-local path, ideally one that is new for each launch (for example, one that includes the Slurm job ID), removes that traffic. It also means fences and leftover ledgers disappear when the job's local storage is wiped, which bounds the open problem described below to a single launch.

The steps below assume node-local directories, so each request has to reach the actor whose node holds the ledger:

1. Retire finished attempts from the actor that ran them. In NemoGym._postprocess_receipt_mode (nemo_rl/environments/nemo_gym.py), once the manifest has been fetched and the receipt assembled (or assembly has failed), send the attempt's rollout ID to POST /training-token-capture/control/rollouts/retire through the same _control helper. That actor is the one whose model server holds the ledger, so no extra bookkeeping is needed. There is no need to wait for the receipt to be persisted, because after a restore every attempt that was not sealed is retried under a fresh attempt ID (prepare_incomplete_retry), so its old manifest is never read again. One request per rollout is fine to start with, since it matches the manifest fetches the actor already makes.
2. Record which actor ran each prompt group. The rollout manager picks the actor in _dispatch with shard_set.pick_handle(...) and already computes its instance_label, but only uses it for timers and error messages; the receipt, the recovery ledger, and the reassembly request don't keep it. Store it on the group's record in the recovery ledger. Every sibling of a prompt group is sent to the same actor, so one label per group is enough.
3. Retire abandoned attempts through that actor. Drive this from the recovery ledger rather than from each failure path: when it moves an attempt to ABANDONED or FAILED, or replaces it with a fresh attempt, send the attempt's gate rollout ID to the actor recorded for its group. If that actor has died, skip it, because its ledgers went with its node. Abandoned agents may still be running, and the fence discards their later writes.
4. Never retire an attempt a restorable checkpoint may continue, or one whose manifest is still needed. With today's rollout checkpoints that excludes nothing, because restored attempts get fresh IDs. It becomes a real constraint once partial-rollout checkpointing continues an attempt under its original ID.
5. Retry a failed request with the same batch: both calls are idempotent, and a request takes at most 4096 IDs (RolloutControlClient.retire splits larger lists). A rollout reported as absent that you know made calls means the request reached an actor on a different node.
6. Unchanged: clearing staged token rows from the TransferQueue stays NeMo RL's job, and it must keep rows a restorable checkpoint still references.

If the capture directory stays on shared storage instead, steps 2 and 3 become simpler, because every actor's model server then sees every ledger and any live actor can retire any attempt. The cost is the per-call Lustre traffic described above, plus the cross-node flock requirement from #3937.

Related issue: re-sent rows reuse their rollout ID. When a row fails with a transport error, _dispatch sends it again to the same actor under the same _ng_rollout_id. If the first run is still going, both runs write into the same ledger and the manifest can contain extra records. This is the duplicate-execution problem that idempotent /run addresses; until then, give a re-sent row a fresh rollout ID, or retire and then delete the old one before reusing it.

Open problem: removing fences during a run. Each retired attempt leaves one empty <rollout_id>.lineage.retired file, and this PR gives no safe point during a run to delete it: a finished attempt can still write because a client retry can run it twice, and an abandoned attempt's agent may still be running. A launch-scoped node-local directory bounds the growth to one launch. Within a launch, idempotent /run would let NeMo RL delete finished attempts instead of retiring them, leaving fences only for the rare abandoned ones. Until one of these is in place, use delete only at the end of a run or after every Gym server sharing the directory has restarted.

Once #4218 is in, check adoption with gym.token_capture.cache.size{cache=ledger_rollouts}, which should track in-flight rollouts instead of growing, and compare ledger.record.in_thread with the capture directory on Lustre and on local disk to measure the saving.

What changed and why

With token_id_capture.external_staging: true, the policy model server keeps a token-free capture ledger for every rollout: one <rollout_id>.lineage.jsonl file in the lineage store's directory, plus a cached copy of the rows in each model-server worker. Nothing ever removes either. The training framework reads a rollout's manifest when the rollout ends, and from then on the ledger is dead weight, but Gym has no way to know that. A long training run therefore leaves one ledger file per rollout attempt, and each worker caches the rows of every rollout it served. Retiring a rollout drops its cached rows in the worker that serves the retire request; the other workers keep theirs until #4268's bound evicts them, so that bound is what caps every worker's cache.

The complete-record token store already has the two operations this needs, so the ledger gets the same two, with the same names and meaning:

Operation	Complete-record token store (existing)	Capture ledger (this PR)	Use it for
Retire: remove the data, keep a fence so late writes are discarded	TokenSource.drop leaves a retired fence	CaptureLedger.retire(rollout_ids) leaves an empty <rollout_id>.lineage.retired fence	Every rollout the framework is done with: finished (receipt durable) or abandoned.
Delete: remove the data and any fence	TokenCaptureStore.delete	CaptureLedger.delete(rollout_ids)	Retired rollouts once nothing of them can still write, to remove their fences. A rollout ID about to be reused.

• retire validates the batch, writes all the fences and syncs the directory, then deletes the ledgers and their cached rows. A later record or record_failure for a retired rollout is dropped with one warning. Writes look for the fence only when the rollout has no ledger rows, so a normal captured call costs the same file-system operations as on main. Finished rollouts need the fence too: a client retry can run a rollout twice, and the second execution writes after the first returned. A stress run showed this (see Validation).
• delete removes the ledger, the fence if any, and the cached rows.
• Both are batched and safe to repeat, validate the whole batch before changing anything, reject a bare string (a string is a sequence of one-character IDs), and report which rollouts had a ledger (removed) and which did not (absent). Their result is typed as the RolloutRemovalPayload TypedDict.
• Retries write nothing, but every non-empty batch syncs. Retire creates only missing fences, and syncs them before removing any ledger and before returning, even when it wrote none: an overlapping or crashed retire may have created them without syncing yet. Delete syncs every non-empty batch for the same reason, so a caller reusing the rollout ID never sees the old fence or ledger come back after a crash. Retire re-checks each fence under the rollout's lock before removing its ledger, so a rollout ID that a delete cleared and a new attempt reused in between keeps its new ledger.
• A retired rollout reads as retired, not as empty. Its manifest raises RolloutRetiredError, which the manifest route returns as HTTP 410 and RolloutControlClient.manifest raises again. has_rows raises it too.
• Calls that start after retire are not captured. Admission checks for the fence before granting a capture, so a late call of an abandoned attempt, or the first call of a duplicate run, stages nothing. A call that was already in flight may still stage tokens before its row is discarded; the warning then names its staging key so the framework can remove them.
• Routes: POST /training-token-capture/control/rollouts/retire and …/rollouts/delete, next to the manifest route and behind the same bearer token, with at most 4096 IDs per request. RolloutControlClient.retire() and .delete() validate and deduplicate the whole list, then split it into batches, so an invalid ID fails before any batch is sent. RolloutControlClient.manifest() validates its rollout ID and rejects a manifest for a different rollout.
• Docs: a "Retire capture ledgers" section in model-server/model-call-capture.mdx.

Neither operation touches staged token data. The framework that staged it owns it, and must keep any staged row that a checkpoint it may restore still refers to.

How it works

Where retire and delete sit

flowchart LR
  subgraph F["Training framework"]
    D["Rollout finished:<br/>read manifest, receipt durable"]
    A["Rollout abandoned:<br/>gave up or replaced"]
    R["Retire batch"]
    X["Delete batch<br/>(nothing can still write)"]
    S["Clear staged token rows<br/>not referenced by any<br/>restorable checkpoint"]
  end
  subgraph G["Policy model server that served the rollout"]
    CR["Control routes:<br/>GET …/manifest<br/>POST …/rollouts/retire<br/>POST …/rollouts/delete"]
    L[("Lineage store directory:<br/>&lt;id&gt;.lineage.jsonl<br/>&lt;id&gt;.lineage.retired")]
  end
  T[("Framework token store:<br/>staged token rows")]
  D --> R
  A --> R
  R -. "later" .-> X
  D -- "GET manifest" --> CR
  X -- "POST delete" --> CR
  R -- "POST retire" --> CR
  CR --> L
  D --> S --> T
  A --> S
  classDef new fill:#fde68a,stroke:#b45309,stroke-width:2px,color:#1f2937
  class X,R,CR new

Loading

One rollout's ledger

flowchart TD
  A["Absent<br/>no ledger, no fence"] -- "first record or record_failure" --> L["Live<br/>ledger file exists"]
  L -- "record, record_failure, manifest" --> L
  L -- "retire (finished or abandoned)" --> R["Retired<br/>no ledger, fence file exists"]
  L -- "delete (reset before reuse)" --> A
  A -- "retire" --> R
  R -- "late record or record_failure (discarded)" --> R
  R -- "retire again (no-op)" --> R
  R -- "delete (nothing can still write)" --> A

Loading

Retiring a rollout while another execution keeps calling

sequenceDiagram
  participant F as Training framework
  participant A as Model-server worker A
  participant B as Model-server worker B
  participant D as Lineage store directory
  F->>A: POST /rollouts/retire {rollout_ids: [r]}
  A->>D: under r's lock: create r.lineage.retired
  A->>D: fsync directory
  A->>D: under r's lock: delete r.lineage.jsonl
  A->>D: fsync directory
  A-->>F: {removed: [r], absent: []}
  Note over B: a late call to r finishes (abandoned agent, or a duplicate execution)
  B->>D: under r's lock: no r.lineage.jsonl, and r.lineage.retired exists
  B->>D: discard the row
  Note over B: one warning, and the model call still succeeds

Loading

A write reads the rollout's ledger under its lock and, only when the ledger has no rows, checks for the fence. That is enough: retire makes every fence in the batch durable before it deletes any ledger, so a rollout whose ledger still has rows either was not retired yet, or its retire is still running and will delete the ledger next. If retire fails between the two steps, the caller retries it, and the retry deletes the ledger. A late row therefore never starts a new ledger after retire, whichever worker handles it.

What the caller must do

The caller is the training framework that consumes external-staging ledgers, today NeMo RL. rollout_collection never uses external staging and does not call these routes.

Topic	Expectation
Which ID	The exact rollout ID the model server recorded under, the one used for GET …/rollouts/{rollout_id}/manifest. For NeMo RL that is the attempt-scoped ID passed as _ng_rollout_id. IDs are never reused without a delete first.
Which server	A model server whose lineage directory holds the rollout's ledger. With a node-local directory, that is the server that ran the rollout, or another server on the same node configured with the same directory; with a shared directory, any server configured with it. A model server only sees the ledgers in its own lineage directory and reports any other rollout as absent, so absent for a rollout you know recorded calls means the request went to a server with a different directory.
Finished rollout	retire it once the manifest has been read and the receipt built from it, as long as the framework will never need that manifest again. If a crash before the receipt is saved makes the framework rerun the rollout under the same ID, wait until the receipt is saved; if it reruns under a fresh ID, as NeMo RL does, retire right away.
Abandoned rollout	retire it once the framework has decided not to use the attempt: a dispatch failure, a retry, a stale abort, or an attempt superseded after a restart. Drive this from the one place that records attempt outcomes (in NeMo RL, the recovery ledger) rather than from each failure path.
Fences	One empty file per retired rollout. delete retired rollouts once nothing of those attempts can still write: at the end of the run, or after every Gym server sharing the directory has restarted. There is no safe point during a run yet; see the open problem under "Action for NeMo RL".
Never remove	An attempt the framework may still continue or whose manifest it still needs. With partial-rollout checkpointing, that includes a restored attempt that has not been re-dispatched yet.
Input	Always a list of rollout IDs. A bare string is rejected with TypeError.
Batching	Up to 4096 IDs per request (MAX_LEDGER_BATCH). The client splits longer lists. Batch per training step or per group rather than per rollout.
Retries	Both operations are idempotent. After a timeout or an error, resend the same batch.
Errors	401: missing or wrong bearer token. 422: malformed body (empty list, more than 4096 IDs, unknown field). 400: an ID that is not a valid rollout ID; nothing in the batch was changed.
After removal	Don't read the rollout's manifest again; it returns an empty ledger.
Staged token rows	Not touched. The framework clears them itself, and must keep every row named by a ledger it still holds or by the retained staging keys of any checkpoint it may restore.

Compatibility

retire and delete are now part of the CaptureLedger protocol. A custom ledger configured through token_id_capture.lineage_store without them now fails the model server's startup check, which lists the required methods. Gym's FileLineageStore, the ledger NeMo RL configures, implements both.

Nothing changes until a caller sends a retire or delete request: model calls, manifests, and the complete-record token store behave as before.

Without #3937, every rollout also leaves a <rollout_id>.tokens.lock file that ledger operations lock and that delete cannot safely remove, because another process may be waiting on it. #3937 replaces those with a fixed set of shared lock files.

Issue

No tracking issue. This closes the cleanup gap that external staging has had since the capture ledger was introduced, using the complete-record token store's operations; the change and its caller contract are described here.

Validation

• New tests in tests/unit_tests/test_token_capture_ledger.py cover retire, repeated retires, writes discarded after retirement, delete removing the fence so an ID can be reused, delete of a live ledger, retirement seen by every worker and by a fresh process, no ledger or fence files left after delete, whole-batch validation and bare-string rejection for both operations, the fence being checked only before a ledger is created, no fence touches or directory syncs on retried or no-op batches, a fence synced before its ledger is deleted, the staging key in the late-row warning, a retired rollout's manifest and has_rows raising RolloutRetiredError, calls on retired rollouts not admitted for capture, and removals leaving no cached rows behind, mostly against both the file-backed and in-memory ledgers. The new tests/unit_tests/test_token_capture_control_routes.py covers both routes' results, the bearer token, malformed batches, invalid IDs, client batching and errors, and, against the real routes and ledger, the client rejecting a bare string, an invalid ID in a later batch, and IDs repeated across batches, plus the manifest's HTTP 410 and its typed error.
• The new tests fail when the fence check on writes, the fence removal in delete, or the bare-string guard is disabled.
• pytest on every token-capture, staging, trajectory-builder, rollout-collection, and model-server unit test file: passed.
• A process-level stress run (real Gym servers, a fake inference worker using Gym's staging code, up to 16,000 concurrent rollouts on one workstation) found retire and delete correct under load, and showed why finished rollouts need the fence: with a quarter of 4,000 rollouts sent twice, as a client retry does, deleting finished rollouts let 345 duplicates recreate their ledger; retiring them let none.
• Measured with the token-capture metrics from feat(token-capture): low-overhead metrics for capture operations, outcomes, digests, and caches #4218 on 20,000 rollouts of four calls, 1,024 in flight, retired in batches of 256: no ledger bytes and no cached ledger rows are left (main: 72 MB of ledgers and 20,000 cached rollouts), and ledger.record takes the same time as on main (7.2 ms against 7.0 ms on average, same p99 histogram bucket). A captured call makes the same system calls as on main. Retiring costs about 65 µs per rollout and deleting about 41 µs, on a local disk.
• 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 complete-record captures #3939, 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 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 for this PR: no caller sends retire or delete requests yet, so rollouts behave exactly as before. The real check comes with the NeMo RL change that calls them: each lineage directory should stay roughly the size of its in-flight rollouts plus the fences not yet deleted.

[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.

[github-actions]

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