Skip to content

feat: dynamo prime-rl integration#3062

Open
biswapanda wants to merge 26 commits into
PrimeIntellect-ai:mainfrom
biswapanda:feat/dyn-pi-sidecar-v2
Open

feat: dynamo prime-rl integration#3062
biswapanda wants to merge 26 commits into
PrimeIntellect-ai:mainfrom
biswapanda:feat/dyn-pi-sidecar-v2

Conversation

@biswapanda

@biswapanda biswapanda commented Jul 16, 2026

Copy link
Copy Markdown

Summary

This PR lets Prime-RL use an externally deployed Dynamo inference stack while Prime continues to own the trainer and orchestrator.

  • Rollout requests use the normal Dynamo OpenAI-compatible frontend.
  • Prime discovers engine control endpoints from Dynamo's opt-in /v1/rl/workers API instead of hard-coding per-worker admin URLs.
  • One validated discovery snapshot is converted into the existing Prime vLLM control clients.
  • Per-engine world_size values produce cumulative NCCL rank offsets; their sum must match an explicitly configured inference world size before communicator creation.
  • Pause, resume, full-weight updates, and filesystem updates use vLLM's existing HTTP control surface. LoRA may use the discovered Dynamo system route when every worker advertises it.
  • Existing non-Dynamo static/native-vLLM and elastic inference paths remain supported.
  • No Helm, DGD, or Kubernetes chart code is added to Prime.

Companion PRs:

Published, cross-repository tested heads:

Repository Commit
Prime-RL 17f67eb0f00c3440c711c46392ddb97b08297c41
Dynamo b0b51b2783346a6ff6566fcaecc7d35fcd2ee80b
vLLM 2d2bfbe3765cfb5c05567ad8b2338564faedf5e3

Minimal user-facing configuration

[orchestrator.model.client]
base_url = ["http://dynamo-frontend:8000/v1"]
dynamo_discovery_url = "http://dynamo-frontend:8001"

[weight_broadcast]
type = "nccl"
host = "trainer-rank-zero"
inference_world_size = 8

Prime launches no local inference process for this shape, so deployment.num_infer_gpus may be zero and no [inference] block is required. admin_base_url remains supported for ordinary native-vLLM deployments and debugging, but is not required for Dynamo.

inference_world_size remains 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 ranks 0..7; a six-rank snapshot cannot initialize an undersized communicator when eight are expected.

Discovery and control lifecycle

When dynamo_discovery_url is configured, Prime:

  1. Polls <dynamo_discovery_url>/v1/rl/workers within wait_for_ready_timeout.
  2. Parses protocol-v1 records with Pydantic models.
  3. Selects the configured model and rejects empty, errored, malformed, duplicate, or incomplete snapshots.
  4. Requires positive per-engine world sizes, unique engine identities/control URLs, and sum(world_size) == inference_world_size.
  5. Builds the existing vLLM admin clients and cumulative rank offsets.
  6. Waits for direct engine readiness and Dynamo frontend model publication.
  7. Treats the accepted snapshot as immutable for the lifetime of the communicator.

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.disagg to vllm.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.py resolves the old and new native package locations.
  • PrimeRlServingTokens delegates validation, sampling, DP routing, multimodal preprocessing, streaming, usage, logprobs, and cancellation to vLLM.
  • Prime keeps only its actual behavior: forwarding P/D kv_transfer_params and compacting routed-expert data.
  • Response reconstruction preserves canonical request metadata and future response/choice fields.

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_PLUGINS or a parallel openengine.v1 schema.

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

  • Replace hand-written Dynamo discovery parsing with strict Pydantic protocol models.
  • Distinguish explicitly configured inference_world_size=1 from the NCCL default using Pydantic field-set state.
  • Preserve the non-Dynamo local-GPU guard.
  • Remove the unused routed-expert capture class.
  • Harden duplicate endpoint/identity validation and reject partial snapshots for both NCCL and filesystem/LoRA modes.
  • Correct managed-DP/NCCL rank calculation and add completion evidence for every rank.
  • Tighten GLM readiness and post-update rollout checks.
  • Reject Dynamo full-weight updates until NCCL initialization has completed; filesystem broadcast remains the LoRA-only path.
  • Require an explicit expected inference world size for every Dynamo discovery pool so auxiliary pools cannot accept partial snapshots.

Validation

Focused review suite at the published Prime head:

107 passed, 83 deselected

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, Dynamo b0b51b278, and vLLM 2d2bfbe376 as source overlays plus sidecar binary b0b51b278 for the two Qwen gates. The GLM runs used the immutable image lineage described after the list:

  • Qwen3-0.6B disaggregated 1-prefill/1-decode: three trainer and rollout steps, v0/v1 full-weight updates on both engines, post-v1 generation, PASS.
  • Qwen3-0.6B aggregated DP2: DP index 0→global rank 0 and DP index 1→global rank 1, three trainer and rollout steps, v0/v1 updates on both ranks, post-v1 generation, PASS.
  • GLM-5.2 FP8 Math DGD: five steps, 18/18 generation requests, all inference receivers updated, no worker restarts or fatal logs, PASS.
  • GLM-5.2 FP8 standalone R2E against the live Dynamo frontend using a local agents.x-k8s.io sandbox: reward 1.0, 35 turns, agent_completed, PASS.

The GLM image was built from Prime abcb05c34c0bdacfcfda3e0d4223f23ac5f88bd2, Dynamo 6c688b197512e86e3b5b957f2c9b54d9b223bba3, and vLLM 2d2bfbe3765cfb5c05567ad8b2338564faedf5e3. 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:

nvcr.io/nvidian/dynamo-dev/biswa@sha256:47e6e8301f3bfdfebd80b2d364b7532341816610e6e46a5d77b4ba39abef54ca

Scope

  • No openengine.v1 or prime_rl.engine.v2 protocol is introduced.
  • No VLLM_PLUGINS environment variable is required.
  • No manual rank_start configuration is introduced.
  • Active NCCL communicators are not elastic; autoscaling, worker replacement, topology epochs, and communicator reconstruction remain follow-up work.
  • This PR contains no Helm, DGD, or Kubernetes manifests.

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_url on the orchestrator client polls /v1/rl/workers, validates protocol v1 snapshots (per-engine admin_base_url, world_size, optional LoRA system routes), and builds admin clients with cumulative NCCL rank offsets. weight_broadcast.inference_world_size is 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: DynamoInferencePool wires pause/resume, NCCL init, and weight updates through vLLM /collective_rpc (and Dynamo /v1/loras when advertised). NCCL rank mapping is centralized in global_inference_rank with optional engine_world_size for managed-DP topologies; checkpoint reloads consume per-broadcast state dict chunks and fail on empty transfers.

vLLM serving: A narrow compat layer tracks vLLM’s token-in/token-out package move; PrimeRlServingTokens shrinks to KV handoff forwarding plus compact routed-expert encoding, with build_prime_serving_tokens replacing 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.

@biswapanda biswapanda changed the title feat: discover Dynamo inference workers feat: dynamo integration Jul 16, 2026
@biswapanda biswapanda changed the title feat: dynamo integration feat: dynamo prime-rl integration Jul 16, 2026
Comment thread src/prime_rl/utils/client.py Outdated
Comment thread src/prime_rl/utils/client.py Outdated
Comment thread src/prime_rl/utils/client.py Outdated
@biswapanda
biswapanda force-pushed the feat/dyn-pi-sidecar-v2 branch from 16db6e6 to ae95617 Compare July 17, 2026 19:56
Comment thread src/prime_rl/orchestrator/utils.py Outdated
Comment thread src/prime_rl/utils/client.py Outdated
Comment thread src/prime_rl/utils/client.py Outdated
Comment thread src/prime_rl/utils/client.py Outdated
Comment thread src/prime_rl/orchestrator/utils.py
Comment thread packages/prime-rl-configs/src/prime_rl/configs/rl.py Outdated
Comment thread src/prime_rl/inference/vllm/routed_experts.py Outdated
Comment on lines +1 to +7
"""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.
"""

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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."""

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same, I didn't follow here. Maybe would be cleaner to do a seperate for this vllm upstream so that dynamo change are clear

Comment thread src/prime_rl/utils/dynamo.py Outdated
Comment on lines +56 to +73
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")

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we use a pydantic object here instead of doing validation/parsing by hand ?

)


def _validate_dynamo_snapshot(workers: tuple[DiscoveredDynamoWorker, ...]) -> None:

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same here if we could leverage pydantic

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes and found 2 potential issues.

Fix All in Cursor

❌ 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.

Comment thread src/prime_rl/utils/dynamo.py
Comment thread src/prime_rl/utils/dynamo.py
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants