Add router_url support for elastic inference pool - #2049
Conversation
When configured, inference requests route through a vllm-router while admin operations (weight updates, LoRA loading) still go directly to discovered pods via elastic DNS discovery.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1a4a886b8a
ℹ️ 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".
| if self.router_url: | ||
| urls = [self.router_url] |
There was a problem hiding this comment.
Keep router traffic gated on ready servers
When router_url is set, clients becomes non-empty even when ready_urls is empty or a newly discovered pod is still syncing its adapter. Scheduler._select_least_loaded_client() and the validation/teacher paths in src/prime_rl/orchestrator/orchestrator.py use a non-empty clients list as the signal that inference is ready, so after a scale-down or during a weight update they will keep sending requests to the router instead of waiting for ServerState.status == "ready". That bypasses the ready_urls filter that previously enforced health/LoRA readiness, so rollouts can be routed to backends with the wrong weights or to a router with no ready backend at all.
Useful? React with 👍 / 👎.
Only expose the router client when there are actually ready backends, preventing requests from being sent to a router with no healthy pods during weight updates or scale-down.
mikasenghaas
left a comment
There was a problem hiding this comment.
hm can we avoid prime-rl local and hosted rl diverging even more here? eg. would it make sense to deploy the router for local setups as well? maybe this would also remove the need for static vs. elastic inference pools as discovery could happen in the shared router component?
I only consider this as temporary fix for a bigger problem in hosted-rl. We still need elastic as a control plane. The router URL serves as data plane. |
Reconciles the raw multimodal offload work with main's iter_trainable_branches dedup, mm_kwargs packing, MXFP8, and seq_lens contract: - trajectories.py: main's iter_trainable_branches() loop is authoritative; this branch's _validate_image_spans + mm_refs build spliced into it. - trainer/model.py: kept the ForwardPolicy generalization over main's image_grid_thw string check in forward() — Kimi K2.5 and Qwen VL need opposite position_ids behavior, which the key-presence check can't express. Restored the numpy import main's _routed_experts_row_size needs (auto-merge dropped it; caught by tests, not conflicts). - trainer/batch.py: took main's mm_kwargs bin-packing machinery, with a guard that raw-ref (mm_refs) samples never pack — with text or each other: their placeholder offsets are sample-relative and _materialize_bin carries only the first sample's refs, so packing would drop or misalign images. Rewrote main's packing test to assert this branch's contract. - trainer/rl/train.py: kept both mm_forward_policy threading (ours) and the [model.vlm] guard for multimodal samples (main's). - Bumped deps/renderers (e64cc58) and deps/verifiers (2b1627d03) to their own main-reconciled companion-PR heads and relocked. The verifiers bump crosses the interleaving-agents refactor (#2049), which renamed textarena's seat: migrated env.agent -> env.player in the two wordle rl.toml configs. Validation: uv lock --check clean; tests/unit minus GPU-only tests/unit/train: 393 passed; tests/unit/train collects (160 tests) without import errors.

Summary
router_urlfield toClientConfigfor routing inference requests through a vllm-routerNote
Medium Risk
Changes how elastic inference clients are constructed by optionally routing all inference traffic through a single
router_url, which can affect availability/traffic flow if misconfigured. Admin operations still target individual pods, but the new branching logic could change behavior in multi-server readiness edge cases.Overview
Adds an optional
router_urltoClientConfigto support sending inference requests through a vLLM router while keeping elastic DNS discovery for admin operations.Updates
ElasticInferencePoolto accept/propagaterouter_urlfrom config and, when configured and at least one backend isready, exposes the router as the sole client URL instead of directly using discovered pod/v1endpoints.Written by Cursor Bugbot for commit 4783368. This will update automatically on new commits. Configure here.