Repository navigation
perf(openai-utils): discriminate response item unions - #3217
Merged
bxyu-nvidia merged 1 commit intoSep 10, 2026
Merged
bxyu-nvidia merged 1 commit into
bxyu-nvidia merged 1 commit into
Conversation
Route response input and output items through callable Pydantic discriminators instead of trying every union member. Shape-based tags resolve shared message and function-call-output literals while preserving permissive typeless input handling. Require token metadata to remain atomic on both input and output items so an incomplete payload cannot select a base model and lose fields. Add regression coverage for concrete routing, training variants, and dump/revalidation. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
|
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. |
Contributor
Author
|
/ok to test 1560587 |
bxyu-nvidia
approved these changes
Sep 10, 2026
jkyi-nvidia
pushed a commit
that referenced
this pull request
Sep 11, 2026
## Summary Restores the callable discriminated response item unions from #2889 on current `main` after #3183 reverted the original merge. `NeMoGymResponseInputItem` and `NeMoGymResponseOutputItem` are large untagged Pydantic unions. Validating each response item tries union members sequentially, so validation cost grows with both union size and trajectory length. The callable discriminators select one concrete model using item type and stable shape fields while preserving permissive typeless input behavior. Complete token metadata selects the corresponding training model; partial metadata is rejected instead of being silently discarded by a base model. This version is compatible with the refusal models added by #3165. Replay normalization uses `exclude_unset=True`, which can omit a refusal content part’s default `type="refusal"`. The content discriminator now recognizes that serialized shape by its `refusal` field, keeping refusal-only assistant messages routed to `NeMoGymResponseOutputMessage`. ## Pydantic validation microbenchmark The original #2889 measurements compared mixed message, reasoning, function, MCP, and training payloads against its then-current base: - 80 items: input union 22.19 -> 1.47 microseconds/item (15.1x); output union 22.70 -> 1.57 (14.4x); full request model 22.12 -> 1.60 (13.9x) - 320 items: input union 22.40 -> 1.48 microseconds/item (15.2x); output union 23.53 -> 1.58 (14.8x); full request model 22.24 -> 1.58 (14.1x) ## Test plan - [x] `uv run pytest tests/unit_tests/test_openai_utils.py tests/unit_tests/test_responses_api_model_streaming.py tests/unit_tests/test_responses_converter.py -q` (501 passed) - [x] `uv run pytest responses_api_models/inference_provider/tests/test_app.py -q` (23 passed) - [x] `uv run pytest tests/unit_tests/ -q` (4477 passed, 104 skipped, 20 subtests passed) - [x] `uv run pre-commit run --files nemo_gym/openai_utils.py tests/unit_tests/test_openai_utils.py` Related: #2889, #3165, #3183 Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com> Signed-off-by: Junkeun Yi <jkyi@nvidia.com>
5 tasks done
This branch was previously deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Restores the callable discriminated response item unions from #2889 on current
mainafter #3183 reverted the original merge.NeMoGymResponseInputItemandNeMoGymResponseOutputItemare large untagged Pydantic unions. Validating each response item tries union members sequentially, so validation cost grows with both union size and trajectory length. The callable discriminators select one concrete model using item type and stable shape fields while preserving permissive typeless input behavior. Complete token metadata selects the corresponding training model; partial metadata is rejected instead of being silently discarded by a base model.This version is compatible with the refusal models added by #3165. Replay normalization uses
exclude_unset=True, which can omit a refusal content part’s defaulttype="refusal". The content discriminator now recognizes that serialized shape by itsrefusalfield, keeping refusal-only assistant messages routed toNeMoGymResponseOutputMessage.Pydantic validation microbenchmark
The original #2889 measurements compared mixed message, reasoning, function, MCP, and training payloads against its then-current base:
Test plan
uv run pytest tests/unit_tests/test_openai_utils.py tests/unit_tests/test_responses_api_model_streaming.py tests/unit_tests/test_responses_converter.py -q(501 passed)uv run pytest responses_api_models/inference_provider/tests/test_app.py -q(23 passed)uv run pytest tests/unit_tests/ -q(4477 passed, 104 skipped, 20 subtests passed)uv run pre-commit run --files nemo_gym/openai_utils.py tests/unit_tests/test_openai_utils.pyRelated: #2889, #3165, #3183