feat(rollout-gateway): Add ThreadedGatewayServer. Change integration tests. - #93
Merged
Merged
Conversation
Linbo-Liu
force-pushed
the
integrate-rollout-gateway
branch
from
July 28, 2026 17:04
b3dfa88 to
156fde4
Compare
Linbo-Liu
force-pushed
the
integrate-rollout-gateway
branch
from
July 28, 2026 21:55
156fde4 to
0d4c049
Compare
lyzustc
approved these changes
Jul 29, 2026
luyuzhe111
added a commit
that referenced
this pull request
Aug 4, 2026
…emas, auto-detected from the chat template (#97) * feat(rollout-gateway): schema-based derender via tokenizer.parse_response, auto-detected from the chat template The default derender's XML tool-call regex silently missed Qwen3's JSON <tool_call> format: zero matches, ill_formed=False hardcoded, the agent saw a plain-text turn and stopped — the validated GSM8K run was single-turn/no-tool without any signal. Fix the default instead of adding config: - response_schemas.py: response-schema dicts vendored from huggingface/trl (Apache-2.0, baseline 7073af94; see NOTICE) for qwen3 (JSON <tool_call>: Qwen2.5/Qwen3/Instruct-2507/VL), qwen3_5 (XML <function=>: Qwen3.5/3.6, Nemotron-3), llama3, glm4moe, gptoss — plus a sha256(chat_template) -> schema table for auto-detection. Byte-exact matching: fine-tunes inherit the template and keep resolving; a revised template stops matching loudly rather than misparsing silently. - HfTemplateRenderer: when the tokenizer's template hash matches, derender the whole output in one tokenizer.parse_response pass (reasoning + text + tool calls, ordered by the schema's anchored regex, so a literal <tool_call> inside a think block cannot corrupt the split). Parse failures and structurally invalid tool calls degrade in place to raw text with ill_formed=True — flagged, never a 500, never partial extraction. - No implicit tool parser anymore: with no matched schema and no injected tool_parser, a tools-bearing parse raises instead of guessing with the one-format regex (parsing.parse_tool_uses remains available by explicit injection). Injected stage parsers disable detection entirely — explicit wins (slime's SGLang detectors path is unchanged). - Drop the two no-op vLLM parser flags from the GSM8K script: they configure vLLM's chat layer, which the token-in/token-out path never touches. - gateway extra: transformers floor 4.44 -> 5.0 (tokenizer.parse_response and the schema engine shipped in 5.0.0); declare jmespath (schema transforms). The verl integration needs zero changes: gateway_host's bare HfTemplateRenderer(tokenizer) construction now does the right thing for every covered model family. * feat(verl): pass the capture session key to the agent via _rollout.api_key The agent previously derived its gateway session key from the ACR runtime session id (context.session_id). Make the trainer supply it explicitly instead: AgentCoreAgentLoop passes api_key=sid through RolloutClient's _rollout overrides, and the math agent reads payload['_rollout'].get('api_key') — the key is trainer-supplied configuration like base_url/model_id, not something the agent infers from its runtime. This also removes rl_app.py's dead assignment left by the #91/#93 merge (the payload-derived api_key was immediately overwritten by context.session_id). Backward compatible both ways: the loop's sid still doubles as the ACR runtimeSessionId, so an already-deployed image reading context.session_id produces the same key; and 'EMPTY' still covers local runs and evaluation endpoints that ignore the api key. * refactor(rollout-gateway): request bare token ids from apply_chat_template Pass return_dict=False instead of accepting the BatchEncoding default and unwrapping input_ids after the fact. The bundled attention mask is a padding artifact the training backend builds itself when batching rows (verl's own apply_chat_template wrapper defaults return_dict=False for the same reason). The stub tokenizer now mirrors the real API's dict-by-default behavior, so dropping the kwarg fails tests on behavior, not just the kwargs echo. * refactor(rollout-gateway): drop the llama3 response schema Remove LLAMA3_SCHEMA, its registry entry, and the Llama 3.1/3.2 template hashes: Llama 3.x is not a fine-tuning target worth carrying a vendored schema for. Also drops the only hub-gated repos from template verification, so hash checks need no HF token. * test(rollout-gateway): cut tautological renderer tests, test schemas directly test_render.py: drop test_tool_parser_ill_formed_propagates (asserted a hard-coded lambda return, already covered by the stage-sequence test) and test_recognized_chat_template_switches_to_schema_derender (duplicated the schema-path coverage and leaked a fake hash into the global registry). test_response_schemas.py: rewrite without the sentinel-template/hash- registration machinery — tests assign the schema on the renderer directly. Down to what only this repo can break: one golden-path test per vendored schema dict (a transcription error in the hand-copied regexes would parse wrong silently) and the _parse_with_schema degradation contract (parse failures degrade to raw text flagged ill_formed, never an exception or a silently dropped call). Adversarial regex cases are TRL's own test surface; detection is covered live in test_template_hashes_live.py. * test(rollout-gateway): verify vendored template hashes against the hub Detection is byte-exact and the offline unit tests bypass it, so a mistranscribed registry hash — or a vendor revising a chat template upstream — would silently drop real models into the no-schema rejection. Fetch each registered hash's actual template from the hub (a few KB per repo, never weights; one repo per hash, e.g. both Qwen3.5 think/nothink variants) and assert it resolves to the expected schema. Always-on: runs in the default suite and CI. * style(rollout-gateway): hoist test imports to module top
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
2 participants
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.
feat(rollout-gateway): Add ThreadedGatewayServer.
Issue #, if available:
A sync trainer blocks its thread waiting for episode results, but the gateway must keep answering the agent's LLM calls the whole time. Served on the trainer's thread, the block would freeze the event loop and deadlock the episode. ThreadedGatewayServer runs gateway.app on a daemon thread with its own long-lived event loop, while the session API (create_session / finish_session / drop_session) stays safe to call from the trainer thread. Async trainers don't need it — they mount gateway.app on their own loop.
Description of changes:
Adds
ThreadedGatewayServer (rollout_gateway/server.py), the deployment shape for serving a RolloutGateway from synchronous trainers (slime, verl), and switches the integration tests to exercise the gateway through it.By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.