From 8f32f8bb2bfc1308f81c1da50a21dbf0f2acce24 Mon Sep 17 00:00:00 2001 From: Ali Roshan Ghias Date: Thu, 10 Sep 2026 13:27:06 -0700 Subject: [PATCH 1/2] fix(capture): fingerprint reasoning in the served response dialect Signed-off-by: Ali Roshan Ghias --- responses_api_models/vllm_model/app.py | 14 +- .../tests/test_token_capture_integration.py | 135 +++++++++++++++++- 2 files changed, 146 insertions(+), 3 deletions(-) diff --git a/responses_api_models/vllm_model/app.py b/responses_api_models/vllm_model/app.py index 749272d88ac..1f996b4524d 100644 --- a/responses_api_models/vllm_model/app.py +++ b/responses_api_models/vllm_model/app.py @@ -1011,7 +1011,9 @@ async def chat_completions( ) if self._external_capture_enabled: - await self._finalize_external_capture(chat_completion_dict) + await self._finalize_external_capture( + chat_completion_dict, as_responses=request.url.path in {"/v1/responses", "/v1/messages"} + ) if self.config.return_token_id_information: message_dict = choice_dict["message"] @@ -1079,7 +1081,7 @@ async def chat_completions( return NeMoGymChatCompletion.model_validate(chat_completion_dict) - async def _finalize_external_capture(self, payload: Dict[str, Any]) -> None: + async def _finalize_external_capture(self, payload: Dict[str, Any], *, as_responses: bool = False) -> None: """Validate and record a response staged by the inference worker. The worker returns commit coordinates only after ``StagingSink.stage`` succeeds. @@ -1130,6 +1132,14 @@ async def _finalize_external_capture(self, payload: Dict[str, Any]) -> None: raise ValueError(f"served response for {coords.model_call_id} carries no envelope id") child_staging_chain = list(context.parent_staging_chain) + [str(coords.staging_key)] response_items, _ = strip_token_fields(response_to_output_items(payload)) + if as_responses: + # Match the representation served by responses(), including separate + # reasoning items. Chat callers instead echo the wrapped assistant text. + response_items = [ + item.model_dump() + for message in response_items + for item in self._converter.postprocess_assistant_message_dict(message) + ] # Compute one fingerprint for the response items. # Compute another for the request and response items together. # If either input cannot be fingerprinted, store no fingerprints and continue recording the call. diff --git a/responses_api_models/vllm_model/tests/test_token_capture_integration.py b/responses_api_models/vllm_model/tests/test_token_capture_integration.py index bfdd04ea520..2f727d88e17 100644 --- a/responses_api_models/vllm_model/tests/test_token_capture_integration.py +++ b/responses_api_models/vllm_model/tests/test_token_capture_integration.py @@ -6,8 +6,10 @@ from typing import Any from unittest.mock import AsyncMock, MagicMock +import pytest from fastapi.testclient import TestClient +from nemo_gym.base_responses_api_model import _CaptureMiddleware from nemo_gym.openai_utils import NeMoGymAsyncOpenAI from nemo_gym.server_utils import ServerClient from nemo_gym.token_id_capture import ( @@ -25,10 +27,13 @@ set_token_sink, trajectories_from_source, ) +from nemo_gym.token_id_capture.sink import NG_CAPTURE_FIELD, NG_COMMIT_COORDS_FIELD +from nemo_gym.token_id_capture.staging.capture import RolloutTokenCapture +from nemo_gym.token_id_capture.staging.records import CaptureAdmission, StageResult from responses_api_models.vllm_model.app import VLLMModel, VLLMModelConfig -def _model(client: NeMoGymAsyncOpenAI) -> VLLMModel: +def _model(client: NeMoGymAsyncOpenAI, **config_overrides: Any) -> VLLMModel: config = VLLMModelConfig( host="0.0.0.0", port=8080, @@ -42,11 +47,139 @@ def _model(client: NeMoGymAsyncOpenAI) -> VLLMModel: uses_interleaved_reasoning=False, supply_prefix_token_ids=True, ) + config = config.model_copy(update=config_overrides) model = VLLMModel(config=config, server_client=MagicMock(spec=ServerClient, global_config_dict={})) model._clients = [client] return model +@pytest.mark.parametrize("endpoint", ["responses", "chat/completions", "messages"]) +@pytest.mark.parametrize("reasoning_field", [None, "reasoning", "reasoning_content", "inline"]) +@pytest.mark.parametrize("tool_call", [False, True]) +def test_external_staging_resolves_served_reasoning_history(tmp_path, endpoint, reasoning_field, tool_call) -> None: + """Echoing served output must preserve capture ancestry across API conversion.""" + records = {} + + class Sink: + def stage(self, record): + key = f"{record.rollout_id}/{record.model_call_id}" + records[key] = record + return StageResult(ok=True, staging_key=key) + + capture = RolloutTokenCapture(sink=Sink(), weight_version_fn=lambda: 7) + admissions = [] + + async def complete(**body): + admission = CaptureAdmission.model_validate(body[NG_CAPTURE_FIELD]) if NG_CAPTURE_FIELD in body else None + admissions.append(admission) + index = len(admissions) + payload = _completion([], [], f"answer {index}") + message = payload["choices"][0]["message"] + if reasoning_field == "inline": + message["content"] = f"reasoning {index}" + message["content"] + elif reasoning_field: + message[reasoning_field] = f"reasoning {index}" + if tool_call: + message["tool_calls"] = [ + {"id": f"tool-{index}", "type": "function", "function": {"name": "lookup", "arguments": "{}"}} + ] + if admission is not None: + prefix = [token for key in admission.staging_chain for token in records[key].token_ids_delta] + active = capture.begin_call(admission, prefix_token_ids=prefix) + coords = capture.complete_call( + active, + prompt_token_ids=prefix + [10 * index], + generated_token_ids=[100 + index], + generated_logprobs=[-0.1], + ) + payload[NG_COMMIT_COORDS_FIELD] = coords.model_dump(mode="json") + return payload + + backend = MagicMock(spec=NeMoGymAsyncOpenAI) + backend.create_chat_completion = AsyncMock(side_effect=complete) + model = _model( + backend, + return_token_id_information=False, + supply_prefix_token_ids=False, + uses_reasoning_parser=reasoning_field is not None, + ) + model._external_capture_enabled = True + ledger = FileLineageStore(tmp_path) + history = [{"role": "user", "content": "question"}] + app = _CaptureMiddleware( + model.setup_webserver(), + store=None, + model_server_name="vllm_model", + lineage_store=ledger, + external_staging=True, + token_capture_enabled=True, + ) + try: + with TestClient(app) as client: + for index in range(1, 4): + body = {"input" if endpoint == "responses" else "messages": history} + if endpoint == "messages": + body.update(model="dummy_model", max_tokens=16) + response = client.post( + f"/ng-rollout/reasoning-rollout/training-token-capture/v1/{endpoint}", json=body + ) + assert response.status_code == 200, response.text + manifest = asyncio.run(ledger.manifest("reasoning-rollout")) + assert not manifest["failures"], manifest + assert len(manifest["records"]) == index + assert admissions[-1] is not None + assert admissions[-1].parent_call_id == (None if index == 1 else admissions[-2].model_call_id) + assert admissions[-1].prev_len == 2 * (index - 1) + assert NG_COMMIT_COORDS_FIELD not in response.text + assert "prompt_token_ids" not in response.text + served = response.json() + if endpoint == "responses": + output = served["output"] + elif endpoint == "messages": + output = [{"role": "assistant", "content": served["content"]}] + else: + output = [ + {key: value for key, value in served["choices"][0]["message"].items() if value is not None} + ] + history.extend(output) + if tool_call: + if endpoint == "responses": + history.append( + {"type": "function_call_output", "call_id": f"tool-{index}", "output": "result"} + ) + elif endpoint == "messages": + history.append( + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": f"tool-{index}", "content": "result"} + ], + } + ) + else: + history.append({"role": "tool", "tool_call_id": f"tool-{index}", "content": "result"}) + history.append({"role": "user", "content": "continue"}) + + manifest = asyncio.run(ledger.manifest("reasoning-rollout")) + assert len(manifest["records"]) == 3 + assert not manifest["failures"] + # Parent verification must still reject edits to the last served answer. + last_answer = history[-3 if tool_call else -2] + if last_answer.get("type") == "function_call": + last_answer["arguments"] = '{"changed": true}' + else: + last_answer["content"] = "different answer" + with TestClient(app) as client: + response = client.post(f"/ng-rollout/reasoning-rollout/training-token-capture/v1/{endpoint}", json=body) + assert response.status_code == 200, response.text + assert admissions[-1] is None + manifest = asyncio.run(ledger.manifest("reasoning-rollout")) + assert len(manifest["records"]) == 3 + assert manifest["failures"] + finally: + asyncio.run(ledger.close()) + + def _completion(prompt: list[int], generation: list[int], content: str) -> dict[str, Any]: return { "id": f"completion-{content}", From b5644eb03306141c7b87a5e2b2eed8c830585a1e Mon Sep 17 00:00:00 2001 From: Ali Roshan Ghias Date: Fri, 11 Sep 2026 09:50:25 -0700 Subject: [PATCH 2/2] fix(capture): finalize lineage after served API conversion Adapt the post-conversion finalization design from Ananth Subramaniam in NVIDIA-NeMo/Gym#3050, commit 4c67a8677ee79e38bd8177a13e22acfb51a2e89b. Retain the worker acknowledgement in request-scoped context and finalize after conversion and serialization. Add refusal, failure, identity and single-commit HTTP regressions. Co-authored-by: Ananth Subramaniam Signed-off-by: Ali Roshan Ghias --- nemo_gym/base_responses_api_model.py | 17 +++- nemo_gym/token_id_capture/records.py | 17 ++++ nemo_gym/token_id_capture/sink.py | 2 + responses_api_models/vllm_model/app.py | 36 +++---- .../tests/test_token_capture_integration.py | 94 ++++++++++++++++++- 5 files changed, 142 insertions(+), 24 deletions(-) diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 864c7c83fef..86173a61808 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -194,6 +194,9 @@ class BaseResponsesAPIModel(BaseServer): class SimpleResponsesAPIModel(BaseResponsesAPIModel, SimpleServer): + async def _finalize_served_response(self, response: Any) -> None: + """Finalize capture after conversion to the response returned to the client.""" + def setup_webserver(self) -> FastAPI: app = FastAPI() @@ -255,7 +258,10 @@ async def responses_dispatch(self, request: Request, body: dict = Body()): _reject_external_capture_streaming(body) if not body.get("stream"): params = _validate_responses_params(body) - return _orjson_dispatch_response(await self._invoke_responses(request, params)) + response = await self._invoke_responses(request, params) + dispatched = _orjson_dispatch_response(response) + await self._finalize_served_response(response) + return dispatched cleaned, ns_map = sanitize_streaming_responses_body(body) try: @@ -300,7 +306,10 @@ async def chat_completions_dispatch(self, request: Request, body: dict = Body()) _reject_external_capture_streaming(body) if body.get("stream") is not True: params = _validate_chat_params(body) - return _orjson_dispatch_response(await self._invoke_chat_completions(request, params)) + response = await self._invoke_chat_completions(request, params) + dispatched = _orjson_dispatch_response(response) + await self._finalize_served_response(response) + return dispatched cleaned, include_usage = sanitize_streaming_chat_body(body) params = _validate_chat_params(cleaned) @@ -354,7 +363,9 @@ async def messages(self, request: Request, body: dict = Body()): _ANTHROPIC_CONVERTER.anthropic_response_to_sse(anthropic_response), media_type="text/event-stream", ) - return _orjson_dispatch_response(anthropic_response) + dispatched = _orjson_dispatch_response(anthropic_response) + await self._finalize_served_response(anthropic_response) + return dispatched async def _invoke_responses( self, request: Request, params: NeMoGymResponseCreateParamsNonStreaming diff --git a/nemo_gym/token_id_capture/records.py b/nemo_gym/token_id_capture/records.py index 30bbd83a38c..09dd072a57f 100644 --- a/nemo_gym/token_id_capture/records.py +++ b/nemo_gym/token_id_capture/records.py @@ -266,11 +266,28 @@ def response_to_output_items(payload: dict) -> list[dict]: Responses payloads already carry ``output``. Chat payloads carry ``choices[*].message``. + Anthropic Messages payloads carry top-level assistant content. Wrap each assistant message as a Responses ``message`` item. """ output = payload.get("output") if isinstance(output, list) and output: return [item for item in output if isinstance(item, dict)] + if payload.get("type") == "message" and payload.get("role") == "assistant": + reasoning_items: list[dict] = [] + message_content: list[Any] = [] + for block in payload.get("content") or []: + if not isinstance(block, dict) or block.get("type") not in {"thinking", "redacted_thinking"}: + message_content.append(block) + continue + reasoning_item: dict[str, Any] = {"type": "reasoning", "summary": []} + if block.get("type") == "thinking" and isinstance(block.get("thinking"), str): + reasoning_item["summary"] = [{"type": "summary_text", "text": block["thinking"]}] + elif block.get("type") == "redacted_thinking" and block.get("data") is not None: + reasoning_item["encrypted_content"] = block["data"] + reasoning_items.append(reasoning_item) + if message_content: + reasoning_items.append({"type": "message", "role": "assistant", "content": message_content}) + return reasoning_items items: list[dict] = [] for choice in payload.get("choices") or []: message = (choice or {}).get("message") or {} diff --git a/nemo_gym/token_id_capture/sink.py b/nemo_gym/token_id_capture/sink.py index 1e41b27a115..1cd7fa3c721 100644 --- a/nemo_gym/token_id_capture/sink.py +++ b/nemo_gym/token_id_capture/sink.py @@ -105,6 +105,8 @@ class CaptureContext: # ``resolve_parent`` so the commit hook can publish the ledger row with # the exact representation the next request will echo. request_items: list[dict] | None = None + # Retain the worker acknowledgement privately until API conversion finishes. + external_commit_coords: dict[str, Any] | None = None @property def parent_call_id(self) -> str | None: diff --git a/responses_api_models/vllm_model/app.py b/responses_api_models/vllm_model/app.py index 1f996b4524d..e692aca5f5d 100644 --- a/responses_api_models/vllm_model/app.py +++ b/responses_api_models/vllm_model/app.py @@ -1011,9 +1011,7 @@ async def chat_completions( ) if self._external_capture_enabled: - await self._finalize_external_capture( - chat_completion_dict, as_responses=request.url.path in {"/v1/responses", "/v1/messages"} - ) + self._prepare_external_capture(chat_completion_dict) if self.config.return_token_id_information: message_dict = choice_dict["message"] @@ -1081,13 +1079,26 @@ async def chat_completions( return NeMoGymChatCompletion.model_validate(chat_completion_dict) - async def _finalize_external_capture(self, payload: Dict[str, Any], *, as_responses: bool = False) -> None: + def _prepare_external_capture(self, payload: Dict[str, Any]) -> None: + """Strip transport fields and retain acknowledgement until API conversion finishes.""" + context = current_capture_context() + if context is None or not context.external_staging: + return + context.external_commit_coords = payload.pop(NG_COMMIT_COORDS_FIELD, None) + self._strip_capture_transport_fields(payload) + + async def _finalize_served_response(self, response: Any) -> None: + """Publish lineage using the final Chat, Responses, or Messages representation.""" + context = current_capture_context() + if context is not None and context.external_staging: + await self._finalize_external_capture(_jsonable(response)) + + async def _finalize_external_capture(self, payload: Dict[str, Any]) -> None: """Validate and record a response staged by the inference worker. The worker returns commit coordinates only after ``StagingSink.stage`` succeeds. This method validates those coordinates against the active call. - It then records the call in the lineage store. - Finally, it removes token data and commit coordinates from the served response. + It records fingerprints from the final API response in the lineage store. """ context = current_capture_context() if context is None or not context.external_staging or context.lineage_store is None: @@ -1095,11 +1106,10 @@ async def _finalize_external_capture(self, payload: Dict[str, Any], *, as_respon ledger = context.lineage_store if not isinstance(ledger, CaptureLedger): raise ValueError("external staging requires a CaptureLedger on the capture context") - coords_payload = payload.pop(NG_COMMIT_COORDS_FIELD, None) + coords_payload = context.external_commit_coords admission = context.capture_admission if admission is None: # UNRESOLVED — the ledger already carries this call's poison row. - self._strip_capture_transport_fields(payload) return try: if coords_payload is None: @@ -1132,14 +1142,6 @@ async def _finalize_external_capture(self, payload: Dict[str, Any], *, as_respon raise ValueError(f"served response for {coords.model_call_id} carries no envelope id") child_staging_chain = list(context.parent_staging_chain) + [str(coords.staging_key)] response_items, _ = strip_token_fields(response_to_output_items(payload)) - if as_responses: - # Match the representation served by responses(), including separate - # reasoning items. Chat callers instead echo the wrapped assistant text. - response_items = [ - item.model_dump() - for message in response_items - for item in self._converter.postprocess_assistant_message_dict(message) - ] # Compute one fingerprint for the response items. # Compute another for the request and response items together. # If either input cannot be fingerprinted, store no fingerprints and continue recording the call. @@ -1207,8 +1209,6 @@ async def _finalize_external_capture(self, payload: Dict[str, Any], *, as_respon context.rollout_id, context.model_call_id, ) - finally: - self._strip_capture_transport_fields(payload) @staticmethod def _strip_capture_transport_fields(payload: Dict[str, Any]) -> None: diff --git a/responses_api_models/vllm_model/tests/test_token_capture_integration.py b/responses_api_models/vllm_model/tests/test_token_capture_integration.py index 2f727d88e17..4d6f7492b1e 100644 --- a/responses_api_models/vllm_model/tests/test_token_capture_integration.py +++ b/responses_api_models/vllm_model/tests/test_token_capture_integration.py @@ -9,6 +9,7 @@ import pytest from fastapi.testclient import TestClient +from nemo_gym.anthropic_converter import AnthropicConverter from nemo_gym.base_responses_api_model import _CaptureMiddleware from nemo_gym.openai_utils import NeMoGymAsyncOpenAI from nemo_gym.server_utils import ServerClient @@ -27,6 +28,7 @@ set_token_sink, trajectories_from_source, ) +from nemo_gym.token_id_capture.records import UNCOMMITTED_CALL_REASON from nemo_gym.token_id_capture.sink import NG_CAPTURE_FIELD, NG_COMMIT_COORDS_FIELD from nemo_gym.token_id_capture.staging.capture import RolloutTokenCapture from nemo_gym.token_id_capture.staging.records import CaptureAdmission, StageResult @@ -53,8 +55,15 @@ def _model(client: NeMoGymAsyncOpenAI, **config_overrides: Any) -> VLLMModel: return model -@pytest.mark.parametrize("endpoint", ["responses", "chat/completions", "messages"]) -@pytest.mark.parametrize("reasoning_field", [None, "reasoning", "reasoning_content", "inline"]) +@pytest.mark.parametrize( + "endpoint,reasoning_field", + [ + (endpoint, field) + for endpoint in ("responses", "chat/completions", "messages") + for field in (None, "reasoning", "reasoning_content", "inline") + ] + + [("responses", "refusal"), ("messages", "refusal")], +) @pytest.mark.parametrize("tool_call", [False, True]) def test_external_staging_resolves_served_reasoning_history(tmp_path, endpoint, reasoning_field, tool_call) -> None: """Echoing served output must preserve capture ancestry across API conversion.""" @@ -75,7 +84,9 @@ async def complete(**body): index = len(admissions) payload = _completion([], [], f"answer {index}") message = payload["choices"][0]["message"] - if reasoning_field == "inline": + if reasoning_field == "refusal": + message.update(content=None, refusal=f"refusal {index}") + elif reasoning_field == "inline": message["content"] = f"reasoning {index}" + message["content"] elif reasoning_field: message[reasoning_field] = f"reasoning {index}" @@ -105,6 +116,7 @@ async def complete(**body): ) model._external_capture_enabled = True ledger = FileLineageStore(tmp_path) + ledger.record = AsyncMock(wraps=ledger.record) history = [{"role": "user", "content": "question"}] app = _CaptureMiddleware( model.setup_webserver(), @@ -127,12 +139,15 @@ async def complete(**body): manifest = asyncio.run(ledger.manifest("reasoning-rollout")) assert not manifest["failures"], manifest assert len(manifest["records"]) == index + # Messages delegates through Responses and Chat, but commits only once. + assert ledger.record.await_count == index assert admissions[-1] is not None assert admissions[-1].parent_call_id == (None if index == 1 else admissions[-2].model_call_id) assert admissions[-1].prev_len == 2 * (index - 1) assert NG_COMMIT_COORDS_FIELD not in response.text assert "prompt_token_ids" not in response.text served = response.json() + assert manifest["records"][-1]["response_id"] == served["id"] if endpoint == "responses": output = served["output"] elif endpoint == "messages": @@ -176,6 +191,79 @@ async def complete(**body): manifest = asyncio.run(ledger.manifest("reasoning-rollout")) assert len(manifest["records"]) == 3 assert manifest["failures"] + assert ledger.record.await_count == 3 + finally: + asyncio.run(ledger.close()) + + +@pytest.mark.parametrize( + "endpoint,failure", + [("responses", "conversion"), ("messages", "conversion")] + + [(endpoint, "serialization") for endpoint in ("responses", "chat/completions", "messages")], +) +def test_external_staging_conversion_failure_does_not_commit(tmp_path, monkeypatch, endpoint, failure) -> None: + """Worker staging can succeed while the final API conversion fails.""" + staged = [] + + class Sink: + def stage(self, record): + staged.append(record) + return StageResult(ok=True, staging_key=f"{record.rollout_id}/{record.model_call_id}") + + capture = RolloutTokenCapture(sink=Sink(), weight_version_fn=lambda: 7) + + async def complete(**body): + admission = CaptureAdmission.model_validate(body[NG_CAPTURE_FIELD]) + coords = capture.complete_call( + capture.begin_call(admission, prefix_token_ids=[]), + prompt_token_ids=[10], + generated_token_ids=[100], + generated_logprobs=[-0.1], + ) + payload = _completion([], [], "answer") + payload[NG_COMMIT_COORDS_FIELD] = coords.model_dump(mode="json") + return payload + + backend = MagicMock(spec=NeMoGymAsyncOpenAI) + backend.create_chat_completion = AsyncMock(side_effect=complete) + model = _model(backend, return_token_id_information=False, supply_prefix_token_ids=False) + model._external_capture_enabled = True + + def fail_conversion(*args, **kwargs): + raise ValueError("injected final conversion failure") + + if failure == "serialization": + monkeypatch.setattr("nemo_gym.base_responses_api_model._orjson_dispatch_response", fail_conversion) + elif endpoint == "responses": + monkeypatch.setattr(type(model._converter), "chat_completion_to_response", fail_conversion) + else: + monkeypatch.setattr(AnthropicConverter, "responses_to_anthropic_response", fail_conversion) + ledger = FileLineageStore(tmp_path) + ledger.record = AsyncMock(wraps=ledger.record) + app = _CaptureMiddleware( + model.setup_webserver(), + store=None, + model_server_name="vllm_model", + lineage_store=ledger, + external_staging=True, + token_capture_enabled=True, + ) + body = {"input" if endpoint == "responses" else "messages": [{"role": "user", "content": "question"}]} + if endpoint == "messages": + body.update(model="dummy_model", max_tokens=16) + try: + with TestClient(app, raise_server_exceptions=False) as client: + response = client.post( + f"/ng-rollout/conversion-failure/training-token-capture/v1/{endpoint}", + json=body, + ) + assert response.status_code == 500 + assert len(staged) == 1, "failure must occur after worker staging" + ledger.record.assert_not_awaited() + manifest = asyncio.run(ledger.manifest("conversion-failure")) + assert not manifest["records"] + assert any(row["reason"] == UNCOMMITTED_CALL_REASON for row in manifest["failures"]) + assert NG_COMMIT_COORDS_FIELD not in response.text finally: asyncio.run(ledger.close())