feat(rl): stamp weight versions on generation responses#672
Conversation
c6ce7d9 to
6bea365
Compare
Stamp every generation response with the current weight version in meta_info["weight_version"], so RL trainers can detect which policy version produced each sample and identify staleness. - Add `weight_version` field to ServerArgs (default "default") - Stamp weight_version into meta_info in output_processor - Add `weight_version` field to all update_weights request inputs - Auto-increment version in WeightTransferManager.finish_update() when no explicit version is provided by the trainer - Add GET /get_weight_version and POST /update_weight_version endpoints - Add GET /model_info endpoint returning model path + weight version - Proxy new endpoints through control_server Follows the same design pattern as SGLang: the trainer optionally passes a weight_version with each update request; on success the version is applied to server_args and stamped into all subsequent generation responses. Signed-off-by: Hejian Sang <sanghj0923@gmail.com>
- Fix backward-compat regression in /finish_weight_update: replace
Content-Type header truthiness check with tolerant body parsing
(empty body → {}) so bodyless POSTs through the sidecar proxy still
work and don't wedge the lifecycle state machine.
- Remove auto-increment counter: it could reuse version strings after
explicit versions were set, defeating staleness detection. Now
weight_version is only updated when the trainer explicitly supplies
one (matching SGLang's semantics).
- Defer version application from mid-chunk to finish_update: avoids
stamping responses with a new version while the model holds
partially-updated weights. Chunks record the pending version;
finish_update applies it atomically after all weights are loaded.
- Extract _stamp_weight_version helper in sglang_compat_http to
eliminate the 3-way copy-paste across update handlers.
- Initialize _version_set_this_update and _pending_weight_version in
__init__ to prevent AttributeError on fresh managers.
Signed-off-by: Hejian Sang <sanghj0923@gmail.com>
Signed-off-by: Hejian Sang <sanghj0923@gmail.com>
6bea365 to
3bb827d
Compare
HJSang
changed the title
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3bb827df5a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # response's meta_info so RL trainers know which policy version produced each | ||
| # sample. Updated atomically after a successful weight push when the trainer | ||
| # supplies a new version string. | ||
| weight_version: str = "default" |
There was a problem hiding this comment.
This new field is documented as --weight-version, but ServerArgs.add_cli_args never registers that flag (the RL args only add --weight-transfer-config and --rl-control-port). Because ts serve derives engine-recognized flags from this parser, a deployment that tries to set the initial policy version on startup will not populate this field (and may pass the unknown flag to the gateway instead), so initial generation metadata is incorrectly stamped as default until an HTTP update arrives.
Useful? React with 👍 / 👎.
| @app.post("/update_weight_version") | ||
| async def update_weight_version(request: Request): | ||
| return await _proxy_to_rl_control(request) |
There was a problem hiding this comment.
The in-engine SGLang router now defines GET /model_info and the docs advertise it, but the public sidecar only proxies /get_weight_version and /update_weight_version here. In normal ts serve deployments clients call control_server rather than the private rl_control_port, so GET /model_info returns 404 instead of exposing the model/version pair.
Useful? React with 👍 / 👎.
1 participant
Summary
meta_info["weight_version"]so RL trainers can identify the policy version that produced each sample.GET /get_weight_version,POST /update_weight_version, andGET /model_info, and document the control-plane semantics.Behavior
ServerArgs.weight_versionis the single source of truth and defaults to"default".finish_weight_update; a version on the finish request takes precedence.finish_weight_updaterequests retain backward-compatible behavior.Test plan
main.runtime-1gpumatrix: PR-relevant tests passed. Two FlashInfernumTokensPerPage=1088model failures were reproduced unchanged onmain.pre-commit run --all-files.The branch is rebased on
origin/mainat61183a2.