Skip to content

fix(capture): preserve token capture after synthetic completions - #3670

Open
pthombre wants to merge 1 commit into
mainfrom
pranav/vllm_token_limit_fix
Open

pthombre wants to merge 1 commit into
mainfrom
pranav/vllm_token_limit_fix

Conversation

@pthombre

Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes external token capture when Gym serves an empty synthetic completion after a vLLM context/token-limit error or the sequential-reasoning guard.

The served-response finalization introduced in #3301 correctly commits fingerprints after API conversion, but it also runs for responses that did not come from a normal worker completion. Two existing paths hit this case:

  • A backend HTTP 400 for context length or max_tokens is converted into an empty completion with finish_reason="length".
  • With sequential_reasoning_allowed=false, a continuation that hits the reasoning guard returns an empty completion with finish_reason="content_filter".

Neither path reaches the worker-response preparation hook, so neither has worker commit coordinates. Finalization currently records worker_response_missing_commit_coordinates, which can poison the entire rollout and discard valid tokens captured in earlier turns. For example, a successful tool-use turn followed by a context-limit completion loses its already-staged training row.

This PR adds a request-scoped external_worker_response_seen flag and sets it only when a normal worker completion reaches prepare_response(). If no worker completion was received, the shared finalizer returns without committing or poisoning the call. Existing middleware still records request_finished_without_staged_coordinates; a consumer selecting the preceding committed terminal can retain the valid token chain.

The fix lives in the shared handler introduced by #2823 and covers both vLLM and Megatron worker capture. A real worker completion missing its acknowledgement still fails with worker_response_missing_commit_coordinates; malformed acknowledgements and explicit worker capture failures retain their existing behavior. Sending a request alone does not set the flag, because the backend may return a context-limit error.

Adds 44 regression cases covering both backends, JSON and buffered SSE, Chat Completions/Responses/Messages/compaction routes, overflow before or after a committed call, and reasoning-only continuations. The tests verify lineage, preserved tokens/masks/logprobs, request isolation, and transport-field stripping. A synthetic response explicitly selected as the terminal remains unattributable; an initial overflow creates no trainable row.

No separate issue is needed for this focused correction to the existing capture lifecycle; the reproduction and affected behavior are documented here.

Validation

  • Slurm job 3966743: 1,055 tests passed, zero failures/errors/skips, using container-native Python 3.13.14 in the requested NeMo RL nightly image. Both changed test files were also overlaid on the unmodified base: all 44 new regression cases failed with the missing-coordinates behavior before the production fix.
  • Ruff lint, formatting, and changed-file pre-commit passed in that job. Targeted coverage exercised every added production line and both outcomes of the new finalizer guard.
  • pre-commit run --all-files passed before publishing, and git diff --check passed.
  • CI classification is full because shared core code changes. The full core/sandbox coverage gate and eight-shard server suite were not run locally; the focused capture/model-server suite and real rollouts were used for this patch. Targeted coverage does not establish the repository-wide coverage threshold.

Container used for both Slurm jobs:

/lustre/fsw/portfolios/nemotron/projects/nemotron_sw_post/users/pthombre/enroot-images/nvcr.io+nvidian+nemo-rl+nightly.squashfs
Validation commands and local artifacts

Job submissions from the fix worktree:

sbatch cache/worker-bypass-validation/native-validation.sbatch
sbatch cache/worker-bypass-validation/smoke.sbatch

The native validation runner invokes pytest through a bootstrap that selects the current checkout and asserts the imported Gym/model-server code comes from that checkout. Its passing test command uses this selection and configuration:

python cache/worker-bypass-validation/pytest_bootstrap.py \
  -o addopts= -o log_cli=false --import-mode=importlib -q --tb=short \
  tests/unit_tests/test_external_capture_handlers.py \
  tests/unit_tests/test_base_responses_api_model.py \
  tests/unit_tests/test_chat_completions_streaming.py \
  tests/unit_tests/test_responses_api_model_streaming.py \
  tests/unit_tests/test_token_id_capture.py \
  tests/unit_tests/test_token_capture*.py \
  responses_api_models/vllm_model/tests \
  responses_api_models/vllm_model_with_compaction/tests \
  --cov=nemo_gym.token_id_capture.sink \
  --cov=nemo_gym.token_id_capture.external_capture \
  --cov-report=term-missing --cov-fail-under=0 \
  --cov-config=cache/worker-bypass-validation/coverage.ini \
  --cov-report=xml:cache/worker-bypass-validation/coverage.xml \
  --junitxml=cache/worker-bypass-validation/regressions.xml

The temporary baseline checkout runs the two changed test files with -k synthetic_completion. The GPU smoke job runs python cache/worker-bypass-validation/real_worker_smoke.py.

All-files hooks were invoked with:

PRE_COMMIT_HOME="$PWD/cache/worker-bypass-validation/pr-pre-commit" \
UV_CACHE_DIR="$PWD/cache/worker-bypass-validation/uv-cache" \
uv tool run --python /cm/local/apps/python3/bin/python3 --from pre-commit \
  pre-commit run --all-files
git diff --check

Validation scripts, logs, JUnit XML, coverage, and rollout evidence are local artifacts, not committed source. They are retained under:

/scratch/fsw/portfolios/nemotron/projects/nemotron_sw_post/users/pthombre/GymPR/Gym/.worktrees/capture-worker-bypass/cache/worker-bypass-validation/

Key artifacts: results.json, baseline.xml, regressions.xml, pre-commit-all-files.log, and real-rollout/evidence.json.

Rollouts

Slurm GPU job 3966328 completed successfully using Qwen/Qwen3-0.6B, vLLM 0.25.1, and the real NeMo RL HTTP worker with local durable staging, served through Gym Responses SSE:

  • A two-turn add(2, 2) tool rollout returned 4, passed the smoke verifier (reward 1.0), and produced two chained committed calls: 234 total tokens, including 27 training tokens, with finite logprobs.
  • A subsequent real backend context-overflow error left those committed records and the reconstructed training row unchanged, adding only the uncommitted-call record.
  • A separate reasoning-only generation followed by the reasoning guard retained all 16 generated tokens, again adding only the uncommitted-call record.

This is a vLLM real-model smoke test; Megatron behavior is covered by the parameterized handler and route tests, not a live Megatron deployment.

Compatibility and benchmark impact

No public API, configuration, ledger/staging schema, or migration changes. Synthetic API responses keep their existing finish reasons. Valid earlier generations can remain trainable when the rollout ends in a synthetic completion, eliminating this source of dropped training rows. Missing acknowledgements from actual worker completions still fail closed. Benchmark scoring logic is unchanged.

Documentation: N/A; this restores existing capture behavior without adding user configuration or a public interface.

Checklist

  • I have read the contributing guidelines.
  • The change is focused; no unrelated edits.
  • Tests added and the focused suite passes.
  • Documentation assessed: N/A, as explained above.
  • pre-commit run --all-files passes.
  • No new source files; existing SPDX headers retained.
  • All commits have DCO sign-off.

Track normal worker completions per request so token-limit and reasoning-guard responses remain uncommitted without poisoning earlier captured turns. Preserve acknowledgement validation for actual worker responses across vLLM and Megatron.

Add 44 regression cases across handlers, API routes, and JSON/SSE. Validation: 1,055 focused tests and real vLLM rollouts pass in Slurm; all-files pre-commit passes.

Signed-off-by: Pranav Thombre <pthombre@nvidia.com>
@pthombre pthombre added bug Something isn't working area:training Training framework integrations and training-data interfaces labels Sep 23, 2026
@copy-pr-bot

copy-pr-bot Bot commented Sep 23, 2026

Copy link
Copy Markdown

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.

@pthombre

Copy link
Copy Markdown
Contributor Author

/ok to test 3c8f7ba

@pthombre
pthombre marked this pull request as ready for review September 24, 2026 00:40
@pthombre
pthombre requested a review from ananthsub September 24, 2026 00:41
@yaoyu-33 yaoyu-33 added complexity:low Localized change in one scope with a small, straightforward review surface needs-review PR is ready for code review and waiting on a reviewer labels Sep 24, 2026
@github-actions github-actions Bot added the sla:review-overdue Review response is over the one-business-day SLA label Sep 25, 2026
# it to an earlier generated turn.
attribution = resolve_terminal(manifest.records, served, declared_response_id=served["id"])
assert not attribution.attributed
assert "declared_terminal_not_captured" in attribution.reason

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This PR decides whether a completion is synthetic by checking whether prepare_response() ran for the request.
That makes the prepare_response() call in VLLMModel.chat_completions() part of the fail-closed guarantee.

Suppose a later change skips that call for a real worker completion, for example by calling it only when ng_commit_coords is present.
A worker response that lost its acknowledgement would then be recorded only as request_finished_without_staged_coordinates.
A training framework that treats that reason as non-fatal could then train on a chain with a missing turn.

Could you add a test that sends a request through the model server and removes ng_commit_coords from the fake worker's response?
It should check that the call still records worker_response_missing_commit_coordinates.
The test below covers both backends, with and without streaming.
It passes on this branch and fails with the change described above.
It also needs WORKER_MISSING_COMMIT_COORDS_REASON added to the nemo_gym.token_id_capture.staging.records import at the top of the file.

Suggested change
assert "declared_terminal_not_captured" in attribution.reason
assert "declared_terminal_not_captured" in attribution.reason
@pytest.mark.parametrize("backend", ["vllm_worker", "megatron_worker"])
@pytest.mark.parametrize("stream", [False, True])
async def test_worker_completion_without_coordinates_still_fails_closed(make_harness, monkeypatch, backend, stream):
# A real worker completion that is missing its acknowledgement must still poison the call.
# It must not be treated like a synthetic completion, which is left uncommitted.
h = make_harness("responses", backend=backend)
original = h.worker.create_chat_completion
async def drop_coords(**body):
payload = await original(**body)
payload.pop("ng_commit_coords")
return payload
monkeypatch.setattr(h.worker, "create_chat_completion", drop_coords)
messages = await _request(h.app, _path("responses"), _body("responses", stream))
assert messages[0]["status"] == 200
assert h.finalize.await_count == 1
manifest = RolloutManifest.model_validate(await h.ledger.manifest("r1"))
assert manifest.records == []
assert [failure.reason for failure in manifest.failures] == [
WORKER_MISSING_COMMIT_COORDS_REASON,
UNCOMMITTED_CALL_REASON,
]

Comment on lines +142 to +146
if not context.external_worker_response_seen:
# Guard and context-overflow completions have no worker acknowledgement.
# Leave the call uncommitted for the middleware to record, without
# treating a synthetic completion as a lost worker acknowledgement.
return

@ananthsub ananthsub Sep 28, 2026 •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After this change, finalize_response() does nothing unless prepare_response() ran earlier in the same request.

This should be documented in the ExternalCaptureHandler protocol prepare_response() docstring. A new model server that uses this handler but skips prepare_response() would have its real completions left uncommitted.
They would no longer fail with worker_response_missing_commit_coordinates.

The code comment above the early return also says the synthetic completions have no worker acknowledgement. That is true, but it is also true of the case this code must still reject. A synthetic completion is built by the model server itself: the reasoning guard never calls the worker, and on context overflow the worker returns an HTTP 400 error instead of a completion.

In a lost-acknowledgement case, the worker does return a completion, but ng_commit_coords is missing from it. The new flag tells these apart by checking whether any worker response arrived, so the comment should describe that.

@github-actions github-actions Bot removed the sla:review-overdue Review response is over the one-business-day SLA label Sep 28, 2026

This branch was successfully deployed

1 active deployment
public — 3c8f7ba5 Deployed Sep 24, 2026 by copy-pr-bot[bot] via release / finalize / notify #3634
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:training Training framework integrations and training-data interfaces bug Something isn't working complexity:low Localized change in one scope with a small, straightforward review surface needs-review PR is ready for code review and waiting on a reviewer

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants