feat(vllm-model): supply the previous call's exact training tokens

[ananthsub]

Supplies a request-time resolved parent's exact cumulative tokens to compatible vLLM generation backends, verifies generation-time proof, and adds delta storage in the initial schema-v1 contract for long multi-call rollouts.

Prefix supply and compact publication

sequenceDiagram
    participant H as Harness
    participant M as Gym vLLM model server
    participant L as LineageStore
    participant V as vLLM backend
    participant S as TokenSink

    H->>M: continuation request
    M->>L: resolve(request items as received)
    alt RESOLVED
        L-->>M: parent call and exact cumulative tokens
        M->>V: generation request with required_prefix_token_ids
        V-->>M: response and generation-time prompt_token_ids
        M->>M: verify observed prompt starts with requested prefix
        M->>M: retain full cumulative digest and replace prompt with parent-relative suffix when enabled
        M->>S: put(TokenEntry, prefix_requested=true, prefix_supplied=true, prompt_is_delta=optional)
    else ROOT or UNRESOLVED
        L-->>M: no proven parent tokens
        M->>V: ordinary generation request
        V-->>M: response
        M->>S: put(full-prompt TokenEntry, prefix_requested=false)
    end

Loading

Prefix intent and proof are separate persisted facts. Setting required_prefix_token_ids is not proof that the engine used them. prefix_supplied becomes true only after the generation response demonstrates that the served prompt extended the exact requested tokens.

Summary

• Consumes only a typed RESOLVED lineage result computed before request conversion. ROOT, UNRESOLVED, and lookup failures leave the ordinary generation request unchanged.
• Requests the verified parent's exact cumulative token sequence, avoiding a lossy second chat-template rendering at the continuation boundary.
• Reads generation-time prompt proof from both supported message-bundle and top-level response shapes, using the same precedence as capture extraction.
• Persists prefix_requested separately from prefix_supplied and keeps concurrency-safe eligible, requested, and successful supply diagnostics.
• Rejects supply at startup for completions-API and responses-native paths that cannot satisfy the extension, and requires returned token information for proof.
• Keeps TokenEntry at initial schema version 1; prefix evidence and parent-relative prompt deltas are part of that unreleased contract. Roots and unresolved calls remain full-prompt anchors.
• Keeps cum_len and digest defined over the full cumulative sequence even when storage contains only a suffix. The builder and incremental resolver walk parent links and verify each reconstructed sequence before use.
• Walks delta chains iteratively with a depth bound and masks a broken, cyclic, missing, or digest-invalid chain instead of guessing. The in-memory resolver refuses delta entries because it lacks a durable parent log.
• Includes an integration simulation that drives local file storage and framework-provided sink/source/resolver protocols through parent resolution, prefix proof, capture, freeze, and trajectory reconstruction.

Builds on the lineage resolver merged in #2180. Final integration documentation and adapter coverage are in #2349.

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

[ananthsub]

/claude review

[claude]

SHIP — no correctness or async-safety concerns.

Reviewed the full diff plus the PR-head app.py, sink.py, records.py, the resolve_parent call sites in base_responses_api_model.py, and the lineage/client APIs. This is a tight, well-tested opt-in feature (prefix supply for chained training trajectories).

What I checked and why it's safe:

• Fail-closed verification. _verify_generation_prefix requires generation-time prompt_token_ids proving the engine applied the prefix and raises RuntimeError (propagates to a 500, fails the run) if the backend ignores or omits it. A misconfigured backend fails loud instead of silently recording a re-rendered prompt against a spliced generation — correct for training-data integrity.
• Conservative supply. Only fires on a unique verified parent resolved once before dispatch; otherwise the request is forwarded untouched (cold cache, never a wrong prefix). Fork/ambiguous/fingerprint-miss cases are covered by tests.
• Tokenize path. Correctly skipped only when supply verified, substituting the proven prompt ids; TestPrefixSupplyReachesTokenize guards the both-calls-must-agree hazard.
• Config guard. supply_prefix_token_ids ⇒ return_token_id_information enforced by a model_validator, which the verification path relies on.
• Schema v3. Backward-compatible — prefix_supplied defaults False, and _refuse_a_newer_record still refuses genuinely-newer records loudly.

Two NOTEs (author's call, both fail safe, not merge-blocking):

• completions-API path. _apply_prefix_supply runs inside _preprocess_chat_completion_create_params, which the use_completions_api=true path never invokes (per the comment at app.py:1140). Enabling both flags together makes supply silently no-op (records via tokenize, still correct — just no chaining benefit and no accounting signal). Given the effort spent on auditability, a one-line config warning or validator note when both are set would prevent an inconclusive experiment.
• _prefix_supply_counts accounting. In-process counter mutated without a lock under concurrent requests; the += reads/writes have no intervening await so it's effectively fine under the GIL, and it's diagnostic-only. No correctness impact.

Test coverage is real behavioral assertions, not pass-throughs. LGTM.

[ananthsub]

/claude review

[claude]

Review incomplete: the PR revision changed during analysis. The base SHA captured at the start of this review (7c505f01bfaa163f65cd276d815666f059469620) no longer matches the current base (4df7a80816817fd97cbcdde44ae5ad18833c0092), so the diff I analyzed may no longer reflect the head under review. No inline findings are being posted. Please re-run the review against the updated revision.

[ananthsub]

/claude review

[ananthsub]

/claude review

[claude]

SHIP WITH CARE — one reliability risk worth resolving before this runs on deep rollouts.

This adds delta (parent-relative suffix) token storage and vLLM prefix supply, both fail-closed by design: prefix supply requires generation-time prompt_token_ids as proof and declines on any ambiguity; delta reconstruction re-verifies the full sequence against the stored digest and masks on a broken chain. Digest checks, the RESOLVED-parent invariant on delta records, the schema bump to v5 with the min floor held at 3, and the two-worker integration test are all solid.

One finding (inline, RISK):

• _materialize_delta_prompts.cum_of (builder.py:213) is recursive over the delta parent chain. With delta_records=true every RESOLVED continuation is a delta, so recursion depth = chain length; a long sequential rollout can raise RecursionError, which is not in the caught tuple at consumer.py:126 and escapes the to_thread build outside the try at consumer.py:261 — failing the batch instead of masking the sample. Both path_to (builder.py:358) and FileLineageStore._materialize (lineage.py:619) already solve this iteratively with a depth cap; cum_of should match.

No async-client, ray.get, or public-API compatibility concerns. Config/records changes keep new fields defaulted off and YAML-driven.

[claude]

RISK — cum_of recurses once per delta ancestor, so recursion depth equals the delta-chain length. With delta_records=true, every RESOLVED continuation in a rollout is a delta, so a long multi-call agent rollout (hundreds+ of sequential model calls) produces a chain deep enough to hit Python's recursion limit (~1000).

What breaks: RecursionError (a RuntimeError) raised here. _assemble (consumer.py:126) only catches (AssertionError, ValueError, KeyError, IndexError, TypeError), and trajectories_from_source runs _assemble via asyncio.to_thread outside its try/except (consumer.py:261-270). So RecursionError escapes the "mask instead of fail" contract and fails the whole rollout / training batch — the exact outcome consumer.py:115 warns against.

Blast radius: any run with delta_records enabled and deep sequential rollouts; silent until a long trajectory trips it, then it takes down the batch rather than masking one sample.

This is a known foot-gun in this very file: path_to (line 358) was made iterative precisely because "agent rollouts can exceed the recursion limit," and the sibling delta-walk in FileLineageStore._materialize (lineage.py:619-636) is iterative with an explicit depth > 10_000 cap. cum_of should follow the same iterative pattern with a depth guard.

[ananthsub]

Fixed in 57ef86c. Delta materialization now walks parent links iteratively with cycle detection and the same 10,000-link depth bound used by the lineage store. A 1,100-turn delta-chain regression test verifies this path no longer depends on Python recursion.

[github-actions]

🌿 Preview your docs: https://nvidia-preview-ananthsub-tokidcap-supplying.docs.buildwithfern.com/nemo/gym

Here are the markdown pages you've updated:

• nemo/gym/tutorials/training-tutorials/external-agent-harnesses