feat: dynamo prime-rl integration#3062
Conversation
Signed-off-by: Biswa Panda <biswa.panda@gmail.com>
16db6e6 to
ae95617
Compare
Signed-off-by: Biswa Panda <biswa.panda@gmail.com>
| """Narrow compatibility surface for vLLM's token-in/token-out API move. | ||
|
|
||
| vLLM moved these types from ``entrypoints.serve.disagg`` to | ||
| ``entrypoints.scale_out.token_in_token_out`` without a compatibility alias. | ||
| Keep the version branch in one module so Prime's serving and client code do not | ||
| grow parallel implementations. | ||
| """ |
There was a problem hiding this comment.
I didn't follow why we need this compat layers, is this related to a breaking change in upstream vllm that we used for dynamo ?
| Everything else (request/response schema, sampling params, error handling) | ||
| delegates to upstream so we track future vLLM changes for free. | ||
| """ | ||
| """Small Prime extensions to vLLM's canonical token-in/token-out handler.""" |
There was a problem hiding this comment.
same, I didn't follow here. Maybe would be cleaner to do a seperate for this vllm upstream so that dynamo change are clear
| component = raw_worker.get("component") | ||
| instance_id = raw_worker.get("instance_id") | ||
| admin_base_url = raw_worker.get("admin_base_url") | ||
| system_url = raw_worker.get("system_url") | ||
| world_size = raw_worker.get("world_size") | ||
| system_routes = raw_worker.get("system_routes", []) | ||
| if not isinstance(component, str) or not component: | ||
| raise ValueError("Dynamo RL worker is missing component identity") | ||
| if not isinstance(instance_id, int) or isinstance(instance_id, bool) or instance_id < 0: | ||
| raise ValueError("Dynamo RL worker has an invalid instance_id") | ||
| if not isinstance(admin_base_url, str) or not admin_base_url: | ||
| raise ValueError("Dynamo RL worker is missing admin_base_url") | ||
| if system_url is not None and (not isinstance(system_url, str) or not system_url): | ||
| raise ValueError("Dynamo RL worker has an invalid system_url") | ||
| if not isinstance(world_size, int) or isinstance(world_size, bool) or world_size <= 0: | ||
| raise ValueError("Dynamo RL worker has an invalid world_size") | ||
| if not isinstance(system_routes, list) or any(not isinstance(route, str) or not route for route in system_routes): | ||
| raise ValueError("Dynamo RL worker has an invalid system_routes list") |
There was a problem hiding this comment.
can we use a pydantic object here instead of doing validation/parsing by hand ?
| ) | ||
|
|
||
|
|
||
| def _validate_dynamo_snapshot(workers: tuple[DiscoveredDynamoWorker, ...]) -> None: |
There was a problem hiding this comment.
same here if we could leverage pydantic
Signed-off-by: Biswa Panda <biswa.panda@gmail.com>
Signed-off-by: Biswa Panda <biswa.panda@gmail.com>
Signed-off-by: Biswa Panda <biswa.panda@gmail.com>
Signed-off-by: Biswa Panda <biswa.panda@gmail.com>
Signed-off-by: Biswa Panda <biswa.panda@gmail.com>
Signed-off-by: Biswa Panda <biswa.panda@gmail.com>
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 2 potential issues.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Want reviews to match your repository better? Bugbot Learning can learn team-specific rules from PR activity. A team admin can enable Learning in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 5a8827c. Configure here.

Summary
This PR lets Prime-RL use an externally deployed Dynamo inference stack while Prime continues to own the trainer and orchestrator.
/v1/rl/workersAPI instead of hard-coding per-worker admin URLs.world_sizevalues produce cumulative NCCL rank offsets; their sum must match an explicitly configured inference world size before communicator creation.Companion PRs:
Published, cross-repository tested heads:
17f67eb0f00c3440c711c46392ddb97b08297c41b0b51b2783346a6ff6566fcaecc7d35fcd2ee80b2d2bfbe3765cfb5c05567ad8b2338564faedf5e3Minimal user-facing configuration
Prime launches no local inference process for this shape, so
deployment.num_infer_gpusmay be zero and no[inference]block is required.admin_base_urlremains supported for ordinary native-vLLM deployments and debugging, but is not required for Dynamo.inference_world_sizeremains explicit because Dynamo protocol v1 reports registered workers, not an atomic topology-complete epoch. Prime uses the configured total both to size the trainer's NCCL group and to reject a valid-looking partial discovery snapshot. For example, engine sizes[2, 2, 2, 2]produce rank offsets[0, 2, 4, 6]and ranks0..7; a six-rank snapshot cannot initialize an undersized communicator when eight are expected.Discovery and control lifecycle
When
dynamo_discovery_urlis configured, Prime:<dynamo_discovery_url>/v1/rl/workerswithinwait_for_ready_timeout.sum(world_size) == inference_world_size.A record represents one engine endpoint, not one GPU:
{ "protocol_version": 1, "namespace": "training", "workers": [ { "component": "backend", "instance_id": 10, "model": "Qwen/Qwen3-0.6B", "admin_base_url": "http://worker:8120", "system_url": "http://worker:8181", "world_size": 2, "system_routes": ["update/load_lora"], "error": null } ] }For NCCL updates, Prime logs every local-worker to global-inference-rank mapping and requires completion evidence from all ranks. The GLM fix in this head correctly combines
data_parallel_index, worker rank, and engine rank offset, including managed-DP engines whose local worker rank restarts at zero for each replica.Why the vLLM compatibility files are included
These files are not a second Dynamo protocol. The companion vLLM branch moved the native token-in/token-out Python package from
vllm.entrypoints.serve.disaggtovllm.entrypoints.scale_out.token_in_token_out, and renamed serving attributes while retaining legacy import aliases. Prime must run against both the existing supported vLLM tree and the reviewed companion tree.The compatibility layer is intentionally narrow:
compat.pyresolves the old and new native package locations.PrimeRlServingTokensdelegates validation, sampling, DP routing, multimodal preprocessing, streaming, usage, logprobs, and cancellation to vLLM.kv_transfer_paramsand compacting routed-expert data.This replaces the previous large copied serving handler with a small adapter. It can be reviewed separately as a vLLM-compatibility concern, but it is required for this three-repository stack to execute without
VLLM_PLUGINSor a parallelopenengine.v1schema.Examples
examples/dynamo/qwen3_06b_math/examples/dynamo/qwen3_30b_Thinking/examples/dynamo/glm52_fp8_r2e/The GLM example documents a separately managed multi-node DGD serving
zai-org/GLM-5.2-FP8, a 16-GPU Prime trainer, two discovered engine cohorts totaling 16 inference ranks, R2E multi-turn session affinity, NCCL transport requirements, readiness checks, and the required post-update rollout gate. It intentionally keeps DGD and Helm lifecycle outside Prime's chart.Review fixes in the current head
inference_world_size=1from the NCCL default using Pydantic field-set state.Validation
Focused review suite at the published Prime head:
The broader selected suite also passed before this final fence (
187 passed); the 107-test exact-head rerun covers the modified config, discovery, client, and orchestrator setup surfaces.The focused suite covers discovery/Pydantic validation, retries, duplicate and incomplete snapshots, config fencing, heterogeneous engine sizes, NCCL rank mapping, weight-update evidence, and existing static/native-vLLM behavior.
Exact-head GPU validation used Prime
17f67eb0f, Dynamob0b51b278, and vLLM2d2bfbe376as source overlays plus sidecar binaryb0b51b278for the two Qwen gates. The GLM runs used the immutable image lineage described after the list:agents.x-k8s.iosandbox: reward 1.0, 35 turns,agent_completed, PASS.The GLM image was built from Prime
abcb05c34c0bdacfcfda3e0d4223f23ac5f88bd2, Dynamo6c688b197512e86e3b5b957f2c9b54d9b223bba3, and vLLM2d2bfbe3765cfb5c05567ad8b2338564faedf5e3. The Prime and Dynamo image commits are ancestors of the current published heads; the current Prime discovery and weight-update fences were validated by both exact-head Qwen aggregate and disaggregate gates.The validated runtime image was rebuilt with DIND and pushed as:
Scope
openengine.v1orprime_rl.engine.v2protocol is introduced.VLLM_PLUGINSenvironment variable is required.rank_startconfiguration is introduced.Note
High Risk
Touches distributed NCCL weight broadcast, discovery-driven admin endpoints, and inference rank mapping across heterogeneous Dynamo topologies—misconfiguration or partial snapshots can hang or mis-apply updates.
Overview
Adds external Dynamo inference so Prime can train against a separately deployed Dynamo stack (frontend + vLLM engines) without launching local
[inference].Discovery and config:
dynamo_discovery_urlon the orchestrator client polls/v1/rl/workers, validates protocol v1 snapshots (per-engineadmin_base_url,world_size, optional LoRA system routes), and builds admin clients with cumulative NCCL rank offsets.weight_broadcast.inference_world_sizeis required for Dynamo and must match the sum of discovered engine sizes; RL config allows NCCL with zero local infer GPUs when that total is set explicitly. Full-weight Dynamo training requires NCCL broadcast; filesystem broadcast remains for LoRA.Control path:
DynamoInferencePoolwires pause/resume, NCCL init, and weight updates through vLLM/collective_rpc(and Dynamo/v1/loraswhen advertised). NCCL rank mapping is centralized inglobal_inference_rankwith optionalengine_world_sizefor managed-DP topologies; checkpoint reloads consume per-broadcast state dict chunks and fail on empty transfers.vLLM serving: A narrow
compatlayer tracks vLLM’s token-in/token-out package move;PrimeRlServingTokensshrinks to KV handoff forwarding plus compact routed-expert encoding, withbuild_prime_serving_tokensreplacing unsafe__dict__copying.Docs/examples: New
examples/dynamo/*recipes (Qwen smoke, 30B disagg, GLM R2E) document deployment overlays and trainer/orchestrator TOML shapes.Reviewed by Cursor Bugbot for commit 17f67eb. Bugbot is set up for automated code reviews on this repo. Configure here.