Conversation
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>
|
/ok to test 3c8f7ba |
| # 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 |
There was a problem hiding this comment.
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.
| 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, | |
| ] |
| 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 |
There was a problem hiding this comment.
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.
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:
max_tokensis converted into an empty completion withfinish_reason="length".sequential_reasoning_allowed=false, a continuation that hits the reasoning guard returns an empty completion withfinish_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_seenflag and sets it only when a normal worker completion reachesprepare_response(). If no worker completion was received, the shared finalizer returns without committing or poisoning the call. Existing middleware still recordsrequest_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
pre-commit run --all-filespassed before publishing, andgit diff --checkpassed.Container used for both Slurm jobs:
Validation commands and local artifacts
Job submissions from the fix worktree:
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.xmlThe temporary baseline checkout runs the two changed test files with
-k synthetic_completion. The GPU smoke job runspython cache/worker-bypass-validation/real_worker_smoke.py.All-files hooks were invoked with:
Validation scripts, logs, JUnit XML, coverage, and rollout evidence are local artifacts, not committed source. They are retained under:
Key artifacts:
results.json,baseline.xml,regressions.xml,pre-commit-all-files.log, andreal-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:
add(2, 2)tool rollout returned4, passed the smoke verifier (reward 1.0), and produced two chained committed calls: 234 total tokens, including 27 training tokens, with finite logprobs.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
pre-commit run --all-filespasses.