Skip to content

feat(rl): RL integration layer for prime-rl - #9382

Closed
biswapanda wants to merge 30 commits into
mainfrom
rl-sdk-1
Closed

feat(rl): RL integration layer for prime-rl#9382
biswapanda wants to merge 30 commits into
mainfrom
rl-sdk-1

Conversation

@biswapanda

@biswapanda biswapanda commented May 11, 2026

Copy link
Copy Markdown
Contributor

Overview

RL admin control plane for Dynamo inference workers — pause / resume / weight-update / LoRA fan-out — to support prime-rl, Slime, and other RL training loops. Supersedes #9131.

One HTTP base URL. One worker request-plane endpoint. No K8s API access in the orchestrator. Customers plug in their own vLLM --worker-extension-cls (e.g. FileSystemWeightUpdateWorker, NCCLWeightUpdateWorker) and expose them over the same surface — no Dynamo changes required.

Two-path system

Property Generation RL admin
Frequency Thousands per rollout Once or twice per step
Fan-out None — frontend load-balances All workers
Transport OpenAI HTTP Fan-out proxy → request plane
URL /v1/chat/completions /v1/rl/engine

This PR implements the RL admin path. The generation path (dynamo_chat_nvext) is already on main.

Architecture

prime-rl orchestrator
  │  POST /v1/rl/engine  (one HTTP call)
  ▼
Dynamo frontend (lib/rl + lib/llm)
  │  discovers live workers (etcd / k8s — already configured)
  │  fans out via request plane (NATS / TCP)
  ├──► worker-0  →  registered handler
  ├──► worker-1  →  registered handler
  └──► worker-N  →  registered handler
  │  aggregates responses
  ▼
FanoutReport  →  prime-rl

The orchestrator is a plain HTTP client. It does not import dynamo._core, does not talk to NATS or etcd, and does not need to know the deployment topology.

Routes mounted under /v1/rl

Exactly one path with two methods:

Method Path Purpose
POST /v1/rl/engine Invoke a registered worker method — fan-out (all workers) or direct (one instance_id)
GET /v1/rl/engine Describe live workers + their registered method lists

That's the entire admin surface. There is no /v1/rl/pause, /v1/rl/resume, /v1/rl/update_weights, /v1/rl/state, /v1/rl/health, etc. — typed routes from #9131 were intentionally dropped in favor of the generic {method, kwargs} dispatch.

POST /v1/rl/engine — invoke a method

Request body:

{
  "method":       "update_weights_from_disk",
  "kwargs":       {"model_path": "/ckpt/step_42",
                   "weight_version": "step_42",
                   "engine_rpc": "update_weights_from_path"},
  "timeout_secs": 180,
  "instance_id":  null,
  "components":   null
}
Field Required Meaning
method yes Name of a registered worker route
kwargs no (default {}) Forwarded verbatim to the handler
instance_id no If present, call only this worker (strict-direct); if absent, fan-out to all
components no Restrict fan-out to workers with these component names
timeout_secs no Per-worker call timeout; overrides the policy default

Response (200):

{
  "epoch": 12345678901234,
  "workers": [
    {"instance_id": 9001, "component": "vllm", "status": "ok",
     "response": {"status": "ok", "version": "step_42"}},
    {"instance_id": 9002, "component": "vllm", "status": "ok",
     "response": {"status": "ok", "version": "step_42"}}
  ]
}

GET /v1/rl/engine — describe

Returns the live worker set and the method names each worker registered at startup.

Response (200):

{
  "epoch": 12345678901234,
  "workers": [
    {"instance_id": 9001, "component": "vllm",
     "registered_methods": ["liveness_probe", "pause_generation",
                            "resume_generation", "update_weights_from_disk",
                            "init_weights_update_group", "..."]}
  ]
}

Replaces a separate /v1/rl/describe endpoint.

HTTP response codes

Situation Code error_type
Fan-out completed (any per-worker outcome) 200
Zero live workers in the matched set 503 no_workers
Membership changed mid-fan-out 409 membership_changed
instance_id present but instance not found 409 unknown_instance
Malformed body 400
Runtime / NATS failure 500 fanout_failed

Per-worker errors are reported as status="error" entries inside the 200 FanoutReport, not via HTTP status.

Registered methods (canonical set)

Defined in components/src/dynamo/vllm/worker_factory.py::register_engine_routes. This is the source of truth for what callers can pass as method.

Control plane
Method Purpose
liveness_probe Engine event-loop probe. IPC round-trip is the liveness signal; wedged engines fail at the frontend's short timeout.
pause_generation Pause before weight update. kwargs: mode (keep|wait|abort, default keep), clear_cache (bool, default false).
resume_generation Resume after weight update.
flush_cache Invalidate prefix/KV cache.
abort_request Abort an in-flight request.
Weight management
Method Purpose Default engine_rpc
update_weights_from_disk Load weights from a shared filesystem checkpoint. kwargs: model_path, weight_version, engine_rpc. reload_weights (vLLM built-in)
update_weights_from_distributed Receive weights via NCCL / distributed transport. Requires init_weights_update_group first. (no default — extension required)
update_weights_from_tensor In-process tensor transfer (not implemented).
init_weights_update_group Set up the distributed weight-update communicator. kwargs: master_address, master_port, rank_offset, world_size, timeout. (no default — extension required)
destroy_weights_update_group Tear down the communicator.
get_weight_version Return the current weight version tag.
LoRA (registered only when --enable-lora)
Method Purpose
load_lora_adapter Load or hot-swap a LoRA adapter from a filesystem path. Atomic: rolls back to the old adapter if add_lora fails.
unload_lora_adapter Unload a LoRA adapter. Atomic: re-adds adapter if MDC unregister fails.

engine_rpc field — worker-extension dispatch

vLLM's --worker-extension-cls injects custom methods into per-GPU Worker processes, reachable via collective_rpc. The engine_rpc field on weight-update kwargs selects the target:

Route Default (no extension) FileSystemWeightUpdateWorker NCCLWeightUpdateWorker
update_weights_from_disk reload_weights update_weights_from_path
update_weights_from_distributed n/a update_weights_from_path
init_weights_update_group n/a init_broadcaster

Both worker extensions are unchanged — they already expose the right collective_rpc targets. The caller declares which engine_rpc value to use; Dynamo does not probe the extension.

Fan-out safety properties (unconditional)

  • Strict-direct: when instance_id is present and the target has vanished, returns 409 unknown_instance rather than silently routing to a different instance.
  • Abort on membership change: snapshots the membership epoch before and after fan-out; aborts with 409 membership_changed if a worker joined / left mid-call. Prevents NCCL group corruption.

Both are baked into the dispatcher — not configurable.

Environment variables

Env var Default Effect
DYN_ENABLE_RL unset When true, worker registers dyn://<ns>.<comp>.rl request-plane endpoint
DYN_ENABLE_RL_ENDPOINTS unset When true, frontend mounts /v1/rl/engine
DYN_NAMESPACE dynamo etcd namespace; both worker and frontend resolve via this var
DYN_RL_PORT 8002 Dedicated RL listener port; sits behind a different NetworkPolicy than OpenAI ingress
DYN_RL_COMPONENT unset Default component filter for the frontend's fan-out scope

Where the reviewer should start

  1. lib/rl/src/lib.rs — core fan-out logic (RlClient::engine_call, RlClient::engine_call_one, MembershipSnapshot, call_worker, strict-direct + abort-on-membership-change)
  2. components/src/dynamo/vllm/worker_factory.pyregister_engine_routes() (lines 700–739): the canonical registered-method list, gated on lora_enabled
  3. components/src/dynamo/vllm/handlers.py — all RL handler methods on BaseWorkerHandler. Note the pause guard on update_weights_from_* and atomic LoRA hot-swap / unload rollback.
  4. lib/llm/src/http/service/service_v2.rs — RL router mount + listener lifecycle (RL listener spawned only after main bind succeeds; hard error when DYN_ENABLE_RL_ENDPOINTS=true but runtime is unset).
  5. lib/llm/src/http/service/openai.rs::rl_router — wires dynamo-rl crate into the HTTP service.
  6. tests/rl/smoke_test.sh — E2E flow (pause → update_weights → get_weight_version → resume → inference check) to understand the API contract.

CodeRabbit review fixes (commit 87c89c834c)

Rust (service_v2.rs)

  • Return hard Err instead of tracing::warn when DYN_ENABLE_RL_ENDPOINTS=true but HttpServiceConfig.runtime is unset — prevents a misconfigured service from appearing healthy with no RL routes mounted
  • RL admin listener is now spawned after the main TCP/TLS bind succeeds; closes the race where the admin port is live while the main port is still in-flight or has already failed

Python (handlers.py)

  • _paused guard on update_weights_from_disk and update_weights_from_distributed — weight updates while generation is active risk mixing weight versions across in-flight requests
  • LoRA hot-swap is now transactional: loaded_loras pop is delayed until add_lora succeeds; rolls back the old adapter on add_lora failure
  • unload_lora_adapter rolls back (re-adds adapter to engine + restores loaded_loras) when unregister_model fails, keeping engine and discovery in sync

Python (worker_factory.py)

  • load_lora_adapter / unload_lora_adapter engine routes gated on lora_enabled so workers started without --enable-lora no longer advertise LoRA admin methods they cannot serve

Tests (tests/rl/)

  • nccl_broadcaster.py: --port is now required=True (no hardcoded 29501 default) — prevents collisions between concurrent smoke runs
  • smoke_test.sh, smoke_test_lora.sh, smoke_test_nccl.sh: require PRIME_RL_SRC from the caller; removed author-local /home/biswaranjanp defaults

Related Issues

Supersedes #9131

@biswapanda
biswapanda requested review from a team as code owners May 11, 2026 16:49
@biswapanda
biswapanda requested a review from a team May 11, 2026 16:49
@github-actions github-actions Bot added feat backend::vllm Relates to the vllm backend frontend `python -m dynamo.frontend` and `dynamo-run in=http|text|grpc` labels May 11, 2026
@biswapanda biswapanda self-assigned this May 11, 2026
@biswapanda biswapanda changed the title feat(RL): Dynamo RL integration layer to support Prime-RL feat(rl): RL admin control plane — pause/resume/update_weights fan-out for prime-rl May 11, 2026
@coderabbitai

coderabbitai Bot commented May 11, 2026

Copy link
Copy Markdown
Contributor

Walkthrough

This PR introduces a complete RL (admin control plane) subsystem for distributed worker orchestration. It adds a new Rust library (dynamo-rl) that discovers live workers and dispatches concurrent fan-out requests, integrates 14 new handler methods into vLLM workers for weight updates and LoRA management, wires the RL router into the HTTP service, and provides comprehensive smoke tests for weight/LoRA/NCCL update scenarios.

Changes

RL Admin Control Plane Implementation

Layer / File(s) Summary
Workspace & Crate Registration
Cargo.toml
Adds lib/rl to workspace members and declares dynamo-rl as a workspace dependency at version 1.2.0.
RL Crate Manifest
lib/rl/Cargo.toml
Defines the dynamo-rl Rust crate with metadata, dependency-direction constraint (must not depend on dynamo-llm), and core async/web dependencies (axum, tokio, serde, dynamo-runtime).
RL Client & Fan-out Logic
lib/rl/src/lib.rs
Implements RlClient for discovering workers via DistributedRuntime, snapshotting membership with epochs, and dispatching concurrent fan-out requests; includes error handling for membership changes, HTTP error mapping, and Axum router mounting POST /v1/rl/engine (call one/broadcast) and GET /v1/rl/engine (describe methods).
Worker Handler RL Endpoint Methods
components/src/dynamo/vllm/handlers.py
Adds 14 new handler methods to BaseWorkerHandler: rl_dispatch (async generator), generation control (pause_generation, resume_generation, liveness_probe), request/cache management (abort_request, flush_cache), weight tracking (get_weight_version), weight updates (update_weights_from_disk, update_weights_from_distributed, update_weights_from_tensor), group lifecycle (init_weights_update_group, destroy_weights_update_group), and LoRA management (load_lora_adapter with hot-swap and discovery registration, unload_lora_adapter with rollback).
Worker Factory RL Endpoint Wiring
components/src/dynamo/vllm/worker_factory.py
Wires RL endpoint creation and serving into both decode and prefill workers; expands register_engine_routes to populate handler._rl_routes dispatch table and register RL routes on the runtime.
HTTP Service RL Router Integration
lib/llm/Cargo.toml, lib/llm/src/entrypoint/input/http.rs, lib/llm/src/http/service/openai.rs, lib/llm/src/http/service/service_v2.rs
Adds dynamo-rl dependency, wires DistributedRuntime into HttpService for discovery/fan-out, implements rl_router factory, and adds RL configuration (enable_rl, rl_port, runtime) and conditional TCP listener spawning on host:rl_port during startup.
Test Utilities & Smoke Tests
tests/rl/make_lora.py, tests/rl/nccl_broadcaster.py, tests/rl/smoke_test.sh, tests/rl/smoke_test_lora.sh, tests/rl/smoke_test_nccl.sh, tests/rl/smoke_test_no_extension.sh
Provides LoRA adapter builder, NCCL broadcast sender, and four comprehensive integration tests validating weight updates, LoRA lifecycle, NCCL distributed updates, and baseline (no-extension) worker behavior.

🎯 4 (Complex) | ⏱️ ~60 minutes

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Docstring Coverage ✅ Passed Docstring coverage is 84.38% which is sufficient. The required threshold is 80.00%.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Description check ✅ Passed The PR description is comprehensive and well-structured, providing clear context, architecture details, and implementation guidance for reviewers.
Title check ✅ Passed The PR title clearly describes the main change: adding an RL (Prime-RL) integration layer to the Dynamo system with new RL admin endpoints, request dispatch, and supporting infrastructure across multiple components.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.


Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Actionable comments posted: 12

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
components/src/dynamo/vllm/handlers.py (1)

2595-2597: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Replace internal Linear ticket IDs in source comments.

The new TODOs reference DIS-1661 and DIS-1664, which the repo forbids in non-Markdown source. Please swap these to GitHub issue references instead.

As per coding guidelines, Source code must not carry references to internal Linear tickets (DIS-XXXX, DYN-XXXX, or <PROJECT>-NNNN).

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@components/src/dynamo/vllm/handlers.py` around lines 2595 - 2597, Replace the
internal Linear IDs in the two TODO comments — change "DIS-1661" and "DIS-1664"
to GitHub issue references (e.g., use an inline issue number like "#1661" /
"#1664" or a full GitHub issue URL) so the source comment no longer contains
"DIS-XXXX" tokens; update the two TODO lines shown (the video/audio
re-downloaded TODO and the mixed image+video TODO) to use the chosen GitHub
issue format.
🧹 Nitpick comments (1)
tests/rl/smoke_test_nccl.sh (1)

52-57: ⚡ Quick win

Avoid fixed ports in the NCCL smoke test.

Using fixed defaults for the frontend, NATS, NCCL rendezvous, and system ports makes this flaky under parallel CI runs or on any dev box that already has one of those ports in use. Please allocate free ports up front and thread them through every subprocess.

Based on learnings: Flag hard-coded portability-reducing constants in shell/scripts across the repository (e.g. static ports, hard-coded temporary file names).

Also applies to: 136-136

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/rl/smoke_test_nccl.sh` around lines 52 - 57, The smoke test currently
hardcodes ports via the variables HTTP_PORT, NATS_PORT, and NCCL_PORT (and
NCCL_HOST), which causes flakes when parallel CI or local processes use those
ports; instead, detect and allocate free ports at test startup (e.g., ephemeral
port allocation via a small helper or `python -c "import socket;
s=socket.socket(); s.bind(('',0)); print(s.getsockname()[1]); s.close()"`) and
assign them to HTTP_PORT, NATS_PORT and NCCL_PORT, then propagate those
variables into every subprocess or command invocation that uses them (ensure
MODEL and PRIME_RL_SRC remain as-is), updating any places that referenced the
fixed constants so the dynamically allocated ports are threaded through the
entire test run.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@components/src/dynamo/vllm/handlers.py`:
- Around line 1019-1073: Both update_weights_from_disk and
update_weights_from_distributed must reject weight updates unless the worker is
paused: add a guard at the start of each method (in update_weights_from_disk and
update_weights_from_distributed) that checks self._paused and if False
immediately returns an error (e.g. {"status":"error","message":"Worker must be
paused before updating weights"}) instead of calling
self.engine_client.collective_rpc; this enforces failing fast until
pause_generation() has been called and set self._paused.
- Around line 1258-1277: The unregister_model failure branch must roll back the
removal so the LoRA adapter and discovery state are restored: after catching the
Exception in the unregister_model call (inside the generate_endpoint != None
branch), re-insert the adapter entry back into loaded_loras (using the same
key/value shape used when it was removed) and call the same restoration logic
used by unload_lora() to reload the LoRA into the engine (or otherwise
re-register the model into the engine/runtime) so the MDC remains consistent;
include lora_name and lora_id in the restored entry and ensure any in-memory and
engine state mirrors the pre-removal state before returning the error response.
- Around line 1148-1185: The hot-swap sequence in the handler is not atomic: you
remove the old LoRA from self.loaded_loras before ensuring add_lora and
reset_prefix_cache succeed, which can leave discovery advertising lora_name
while the engine is inconsistent; fix by making the hot-swap transactional—when
is_hot_swap is True, (a) capture old_id and old_path from self.loaded_loras but
do not pop it yet, (b) call await self.engine_client.remove_lora(old_id) then
await self.engine_client.add_lora(LoRARequest(...)) and only after both succeed
update self.loaded_loras[lora_name] = LoRAInfo(...), (c) if add_lora or
reset_prefix_cache fails, attempt a rollback: try to re-add the old LoRA (await
self.engine_client.add_lora(...) with old_id/old_path), restore
self.loaded_loras[lora_name] to the old LoRAInfo, and log any rollback errors;
ensure all exceptions around remove_lora, add_lora, and reset_prefix_cache are
caught and return consistent error messages that include lora_name and lora_id.

In `@components/src/dynamo/vllm/worker_factory.py`:
- Around line 709-726: The rl_routes dictionary unconditionally registers LoRA
RPCs (handler.load_lora_adapter and handler.unload_lora_adapter) even when LoRA
support is disabled; change the construction of rl_routes so those two entries
are only added when the same LoRA capability flag used for the dedicated
/load_lora and /unload_lora endpoints is true (i.e., check the worker/handler
capability or config variable that gates LoRA support and conditionally insert
"load_lora_adapter" and "unload_lora_adapter" into rl_routes rather than
registering them unconditionally).

In `@lib/llm/src/http/service/openai.rs`:
- Around line 2660-2664: The code reading DYN_RL_DEFAULT_TIMEOUT_SECS and
calling std::time::Duration::from_secs_f64 can panic for negative/NaN/inf
values; in the block that sets config.default_request_timeout, validate the
parsed secs (from timeout.parse::<f64>()) is finite and >= 0.0 (e.g.,
secs.is_finite() && secs >= 0.0) before calling Duration::from_secs_f64; if
validation fails, skip setting config.default_request_timeout (or log/warn) so
startup falls back gracefully. Ensure you reference the existing env var parsing
and assign to config.default_request_timeout and use Duration::from_secs_f64
only after the checks.

In `@lib/llm/src/http/service/service_v2.rs`:
- Around line 321-349: You currently spawn the RL admin listener (using
rl_router, rl_addr, rl_cancel via cancel_token.child_token() and tokio::spawn +
axum::serve) before the primary server completes bind/TLS setup; change the
startup order so the RL listener is only spawned after the main HTTP/HTTPS bind
and TLS setup succeed (i.e., move the tokio::spawn block out of the early path
and into the success branch that confirms the main server is "known-good" or
gate it behind the same startup barrier/confirmation used for the main server),
or alternatively register the RL service with the same listener startup sequence
so it won't be independently live if main startup fails; keep the same
cancel/graceful shutdown handling (rl_cancel/with_graceful_shutdown) when you
move it.
- Around line 635-669: The code currently downgrades a requested RL admin plane
to a warning when config.runtime is None; instead make this a build-time error:
in the match arm where config.runtime is None (the branch under if
config.enable_rl || env_is_truthy("DYN_ENABLE_RL_ENDPOINTS")), replace the
tracing::warn and returning (router, None) with an early Err return (or
propagate an error) that clearly states RL endpoints were requested via
enable_rl/DYN_ENABLE_RL_ENDPOINTS but HttpServiceConfigBuilder.runtime is
missing; use the crate's existing error type (e.g. anyhow::anyhow! or the
function's Result error) so the caller fails to build rather than silently
omitting the rl_router (super::openai::rl_router).

In `@lib/rl/src/lib.rs`:
- Around line 446-449: The post-fanout epoch check rebuilds `after_opts` with
`CallOptions::new(timeout)` and thus loses the original `opts.components` filter
used for the first `snapshot`, causing false `MembershipChanged` errors; change
the construction of `after_opts` so it preserves the original component filter
(e.g. clone or derive from the original `opts` and only update the timeout)
before calling `self.snapshot(&after_opts).await?`, ensuring `opts.components`
is applied to both snapshots.
- Around line 579-583: In call_options, validate the user-supplied timeout_secs
before converting: check that the f64 is finite and > 0.0 (e.g. if let Some(s) =
self.timeout_secs { if s.is_finite() && s > 0.0 { Duration::from_secs_f64(s) }
else { /* use default_timeout or return a 4xx error */ } } else {
default_timeout }), or use
Duration::try_from_secs_f64(self.timeout_secs.unwrap_or_default()) and handle
the Err by falling back to default_timeout (or propagate a 4xx error) so
negative/NaN/inf values never panic; update call_options to use this safe
conversion logic referencing timeout_secs, Duration::from_secs_f64 /
Duration::try_from_secs_f64, and default_timeout.

In `@tests/rl/nccl_broadcaster.py`:
- Around line 94-97: The port argument in tests/rl/nccl_broadcaster.py is
hardcoded to 29501 via ap.add_argument("--port", type=int, default=29501);
remove the literal default and ensure the port is dynamically allocated or
required by the caller — either make the argument required
(ap.add_argument("--port", type=int, required=True)) so callers must supply a
port, or import and call allocate_port()/allocate_ports() from
tests.utils.port_utils (or rely on the dynamo_dynamic_ports fixture) to provide
a non-conflicting default.

In `@tests/rl/smoke_test_lora.sh`:
- Around line 37-40: Remove the developer-local default for PRIME_RL_SRC and
make its value explicit: stop setting PRIME_RL_SRC="${PRIME_RL_SRC:-/home/…}" in
the script and either derive it from the repository layout (e.g., resolve
relative to the script location using dirname "$0" or git rev-parse
--show-toplevel) or require the caller to pass PRIME_RL_SRC (fail with a clear
error if unset). Update the script to reference the PRIME_RL_SRC variable (no
hard-coded path) and add a brief usage/error message that exits if PRIME_RL_SRC
is not provided.

In `@tests/rl/smoke_test.sh`:
- Around line 36-39: The script currently defaults PRIME_RL_SRC to an
author-local absolute path; replace that with a portable resolution or make it
required: compute the repo root relative to the script (using the script's
directory) and set PRIME_RL_SRC to the repository's src directory only if not
provided, or else exit with a clear error asking the caller to set PRIME_RL_SRC;
update the assignment of PRIME_RL_SRC and keep the existing fallback for
PYTHON/SMOKE_PYTHON unchanged. Ensure you modify the variables PRIME_RL_SRC and
PYTHON in tests/rl/smoke_test.sh and add a short error message when PRIME_RL_SRC
remains unset.

---

Outside diff comments:
In `@components/src/dynamo/vllm/handlers.py`:
- Around line 2595-2597: Replace the internal Linear IDs in the two TODO
comments — change "DIS-1661" and "DIS-1664" to GitHub issue references (e.g.,
use an inline issue number like "#1661" / "#1664" or a full GitHub issue URL) so
the source comment no longer contains "DIS-XXXX" tokens; update the two TODO
lines shown (the video/audio re-downloaded TODO and the mixed image+video TODO)
to use the chosen GitHub issue format.

---

Nitpick comments:
In `@tests/rl/smoke_test_nccl.sh`:
- Around line 52-57: The smoke test currently hardcodes ports via the variables
HTTP_PORT, NATS_PORT, and NCCL_PORT (and NCCL_HOST), which causes flakes when
parallel CI or local processes use those ports; instead, detect and allocate
free ports at test startup (e.g., ephemeral port allocation via a small helper
or `python -c "import socket; s=socket.socket(); s.bind(('',0));
print(s.getsockname()[1]); s.close()"`) and assign them to HTTP_PORT, NATS_PORT
and NCCL_PORT, then propagate those variables into every subprocess or command
invocation that uses them (ensure MODEL and PRIME_RL_SRC remain as-is), updating
any places that referenced the fixed constants so the dynamically allocated
ports are threaded through the entire test run.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 27728e92-49d8-4bc1-9226-c29e03491177

📥 Commits

Reviewing files that changed from the base of the PR and between 5d9e5f6 and 9e6ac56.

⛔ Files ignored due to path filters (2)
  • Cargo.lock is excluded by !**/*.lock
  • lib/bindings/python/Cargo.lock is excluded by !**/*.lock
📒 Files selected for processing (15)
  • Cargo.toml
  • components/src/dynamo/vllm/handlers.py
  • components/src/dynamo/vllm/worker_factory.py
  • lib/llm/Cargo.toml
  • lib/llm/src/entrypoint/input/http.rs
  • lib/llm/src/http/service/openai.rs
  • lib/llm/src/http/service/service_v2.rs
  • lib/rl/Cargo.toml
  • lib/rl/src/lib.rs
  • tests/rl/make_lora.py
  • tests/rl/nccl_broadcaster.py
  • tests/rl/smoke_test.sh
  • tests/rl/smoke_test_lora.sh
  • tests/rl/smoke_test_nccl.sh
  • tests/rl/smoke_test_no_extension.sh

Comment on lines +1019 to +1073
async def update_weights_from_disk(self, body: dict) -> dict:
"""Load weights from a shared filesystem checkpoint.

Body:
model_path str path to the safetensors checkpoint directory
weight_version str version tag to record (default "unknown")
engine_rpc str collective_rpc target (default "reload_weights")

When engine_rpc is "reload_weights" (vLLM built-in), the kwargs key is
"weights_path". For the FileSystemWeightUpdateWorker extension the target
is "update_weights_from_path" and the kwargs key is "weight_path".
"""
body = body or {}
path = body.get("model_path")
if not path:
return {"status": "error", "message": "Missing 'model_path' in body"}
version = body.get("weight_version", "unknown")
rpc = body.get("engine_rpc", "reload_weights")
kwargs = {"weights_path": path} if rpc == "reload_weights" else {"weight_path": path}
try:
await self.engine_client.collective_rpc(rpc, kwargs=kwargs)
self._weight_version = version
logger.info(f"[RL] Weights loaded from {path} (version={version}, rpc={rpc})")
return {"status": "ok", "version": version}
except EngineDeadError as e:
self._shutdown_on_engine_dead(e)
except Exception as e:
logger.error(f"[RL] update_weights_from_disk failed: {e}")
return {"status": "error", "message": str(e)}

async def update_weights_from_distributed(self, body: dict) -> dict:
"""Receive weights via a distributed transport (e.g. NCCL).

Requires init_weights_update_group to have been called first.

Body:
weight_version str version tag to record (default "unknown")
engine_rpc str collective_rpc target (default "update_weights_from_path")
<any other key> forwarded to the rpc as kwargs (e.g. weight_dir for
NCCLWeightUpdateWorker.update_weights_from_path)
"""
body = body or {}
version = body.get("weight_version", "unknown")
rpc = body.get("engine_rpc", "update_weights_from_path")
rpc_kwargs = {k: v for k, v in body.items() if k not in ("engine_rpc", "weight_version")}
try:
await self.engine_client.collective_rpc(rpc, kwargs=rpc_kwargs)
self._weight_version = version
logger.info(f"[RL] Weights received via distributed (version={version}, rpc={rpc})")
return {"status": "ok", "version": version}
except EngineDeadError as e:
self._shutdown_on_engine_dead(e)
except Exception as e:
logger.error(f"[RL] update_weights_from_distributed failed: {e}")
return {"status": "error", "message": str(e)}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Reject RL weight updates until the worker is paused.

_paused is tracked but never enforced here. Both update paths can reload weights while generation is still active, which can mix versions across in-flight requests. Add a guard before the collective_rpc() call and fail fast until pause_generation() has succeeded.

🧰 Tools
🪛 Ruff (0.15.12)

[warning] 1045-1045: Do not catch blind exception: Exception

(BLE001)


[warning] 1071-1071: Do not catch blind exception: Exception

(BLE001)

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@components/src/dynamo/vllm/handlers.py` around lines 1019 - 1073, Both
update_weights_from_disk and update_weights_from_distributed must reject weight
updates unless the worker is paused: add a guard at the start of each method (in
update_weights_from_disk and update_weights_from_distributed) that checks
self._paused and if False immediately returns an error (e.g.
{"status":"error","message":"Worker must be paused before updating weights"})
instead of calling self.engine_client.collective_rpc; this enforces failing fast
until pause_generation() has been called and set self._paused.

Comment thread components/src/dynamo/vllm/handlers.py Outdated
Comment on lines +1148 to +1185
if is_hot_swap:
old_id = self.loaded_loras[lora_name].id
try:
await self.engine_client.remove_lora(old_id)
self.loaded_loras.pop(lora_name, None)
except Exception as e:
logger.error(
f"[RL] remove_lora({lora_name}, id={old_id}) failed during hot-swap: {e}"
)
return {
"status": "error",
"message": f"Failed to remove existing LoRA '{lora_name}' before hot-swap: {e}",
"lora_name": lora_name,
}

await self.engine_client.add_lora(
LoRARequest(
lora_name=lora_name,
lora_int_id=lora_id,
lora_path=lora_path,
)
)
self.loaded_loras[lora_name] = LoRAInfo(id=lora_id, path=lora_path)

if is_hot_swap:
try:
await self.engine_client.reset_prefix_cache()
except Exception as e:
logger.error(f"[RL] reset_prefix_cache after LoRA swap failed: {e}")
return {
"status": "error",
"message": (
f"LoRA '{lora_name}' loaded but prefix cache reset failed; "
"worker is not safe to serve until next successful swap."
),
"lora_name": lora_name,
"lora_id": lora_id,
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Make LoRA hot-swap atomic or roll it back on every failure path.

Once the old adapter is removed from loaded_loras, any later add_lora() or reset_prefix_cache() failure leaves discovery still advertising lora_name while the engine is either missing that LoRA or serving with stale prefix-cache entries. Requests routed to model=<lora_name> can then run against the wrong weights.

🧰 Tools
🪛 Ruff (0.15.12)

[warning] 1153-1153: Do not catch blind exception: Exception

(BLE001)


[warning] 1175-1175: Do not catch blind exception: Exception

(BLE001)

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@components/src/dynamo/vllm/handlers.py` around lines 1148 - 1185, The
hot-swap sequence in the handler is not atomic: you remove the old LoRA from
self.loaded_loras before ensuring add_lora and reset_prefix_cache succeed, which
can leave discovery advertising lora_name while the engine is inconsistent; fix
by making the hot-swap transactional—when is_hot_swap is True, (a) capture
old_id and old_path from self.loaded_loras but do not pop it yet, (b) call await
self.engine_client.remove_lora(old_id) then await
self.engine_client.add_lora(LoRARequest(...)) and only after both succeed update
self.loaded_loras[lora_name] = LoRAInfo(...), (c) if add_lora or
reset_prefix_cache fails, attempt a rollback: try to re-add the old LoRA (await
self.engine_client.add_lora(...) with old_id/old_path), restore
self.loaded_loras[lora_name] to the old LoRAInfo, and log any rollback errors;
ensure all exceptions around remove_lora, add_lora, and reset_prefix_cache are
caught and return consistent error messages that include lora_name and lora_id.

Comment on lines +1258 to +1277
if self.generate_endpoint is not None:
try:
await unregister_model(
endpoint=self.generate_endpoint,
lora_name=lora_name,
)
except Exception as e:
logger.error(
f"[RL] Failed to unregister LoRA '{lora_name}' MDC after engine removal: {e}"
)
return {
"status": "error",
"message": (
f"LoRA '{lora_name}' removed from engine but discovery "
f"unregister failed; frontend may still route here: {e}"
),
"lora_name": lora_name,
"lora_id": lora_id,
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Restore the adapter if discovery unregister fails.

This path removes the LoRA from the engine and loaded_loras before calling unregister_model(). If discovery unregister fails, the MDC stays published and the frontend can keep routing model=<lora_name> here even though the worker now falls back to base-model behavior. The legacy unload_lora() implementation below already rolls this case back; the RL path should do the same.

🧰 Tools
🪛 Ruff (0.15.12)

[warning] 1264-1264: Do not catch blind exception: Exception

(BLE001)

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@components/src/dynamo/vllm/handlers.py` around lines 1258 - 1277, The
unregister_model failure branch must roll back the removal so the LoRA adapter
and discovery state are restored: after catching the Exception in the
unregister_model call (inside the generate_endpoint != None branch), re-insert
the adapter entry back into loaded_loras (using the same key/value shape used
when it was removed) and call the same restoration logic used by unload_lora()
to reload the LoRA into the engine (or otherwise re-register the model into the
engine/runtime) so the MDC remains consistent; include lora_name and lora_id in
the restored entry and ensure any in-memory and engine state mirrors the
pre-removal state before returning the error response.

Comment on lines +709 to +726
rl_routes: dict = {
# Control plane
"liveness_probe": handler.liveness_probe,
"pause_generation": handler.pause_generation,
"resume_generation": handler.resume_generation,
"flush_cache": handler.flush_cache,
"abort_request": handler.abort_request,
# Weight management
"update_weights_from_disk": handler.update_weights_from_disk,
"update_weights_from_distributed": handler.update_weights_from_distributed,
"update_weights_from_tensor": handler.update_weights_from_tensor,
"init_weights_update_group": handler.init_weights_update_group,
"destroy_weights_update_group": handler.destroy_weights_update_group,
"get_weight_version": handler.get_weight_version,
# LoRA
"load_lora_adapter": handler.load_lora_adapter,
"unload_lora_adapter": handler.unload_lora_adapter,
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Only advertise RL LoRA methods when LoRA support is enabled.

These entries are registered unconditionally, unlike the dedicated /load_lora and /unload_lora endpoints above. That means GET /v1/rl/engine will claim load_lora_adapter / unload_lora_adapter support on workers started without LoRA enabled, and the first POST will fail at runtime. Gate these route registrations on the same capability check.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@components/src/dynamo/vllm/worker_factory.py` around lines 709 - 726, The
rl_routes dictionary unconditionally registers LoRA RPCs
(handler.load_lora_adapter and handler.unload_lora_adapter) even when LoRA
support is disabled; change the construction of rl_routes so those two entries
are only added when the same LoRA capability flag used for the dedicated
/load_lora and /unload_lora endpoints is true (i.e., check the worker/handler
capability or config variable that gates LoRA support and conditionally insert
"load_lora_adapter" and "unload_lora_adapter" into rl_routes rather than
registering them unconditionally).

Comment on lines +2660 to +2664
if let Ok(timeout) = std::env::var("DYN_RL_DEFAULT_TIMEOUT_SECS")
&& let Ok(secs) = timeout.parse::<f64>()
{
config.default_request_timeout = std::time::Duration::from_secs_f64(secs);
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🌐 Web query:

Does Rust's std::time::Duration::from_secs_f64 panic on negative, NaN, or infinite inputs?

💡 Result:

Yes. std::time::Duration::from_secs_f64 panics when secs is negative, NaN, or infinite. Rust’s official docs state: “This constructor will panic if secs is negative, overflows Duration or not finite.” [1][2]. Since NaN and ±∞ are not finite, they fall under “not finite,” so they also panic. [1][2] Note: If you need non-panicking behavior, use Duration::checked_from_secs_f64/try_from_secs_f64 (fallible) or Duration::saturating_seconds_f64 (saturating; NaN becomes 0 seconds). [3]

Citations:


Guard the RL timeout environment variable before converting it to Duration.

std::time::Duration::from_secs_f64 panics on negative, NaN, or infinite inputs. A malformed DYN_RL_DEFAULT_TIMEOUT_SECS would crash the HTTP service at startup instead of falling back gracefully.

Suggested fix
-    if let Ok(timeout) = std::env::var("DYN_RL_DEFAULT_TIMEOUT_SECS")
-        && let Ok(secs) = timeout.parse::<f64>()
-    {
-        config.default_request_timeout = std::time::Duration::from_secs_f64(secs);
-    }
+    if let Ok(timeout) = std::env::var("DYN_RL_DEFAULT_TIMEOUT_SECS") {
+        match timeout.parse::<f64>() {
+            Ok(secs) if secs.is_finite() && secs >= 0.0 => {
+                config.default_request_timeout = std::time::Duration::from_secs_f64(secs);
+            }
+            _ => {
+                tracing::warn!(
+                    value = %timeout,
+                    "Ignoring invalid DYN_RL_DEFAULT_TIMEOUT_SECS"
+                );
+            }
+        }
+    }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if let Ok(timeout) = std::env::var("DYN_RL_DEFAULT_TIMEOUT_SECS")
&& let Ok(secs) = timeout.parse::<f64>()
{
config.default_request_timeout = std::time::Duration::from_secs_f64(secs);
}
if let Ok(timeout) = std::env::var("DYN_RL_DEFAULT_TIMEOUT_SECS") {
match timeout.parse::<f64>() {
Ok(secs) if secs.is_finite() && secs >= 0.0 => {
config.default_request_timeout = std::time::Duration::from_secs_f64(secs);
}
_ => {
tracing::warn!(
value = %timeout,
"Ignoring invalid DYN_RL_DEFAULT_TIMEOUT_SECS"
);
}
}
}
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/llm/src/http/service/openai.rs` around lines 2660 - 2664, The code
reading DYN_RL_DEFAULT_TIMEOUT_SECS and calling
std::time::Duration::from_secs_f64 can panic for negative/NaN/inf values; in the
block that sets config.default_request_timeout, validate the parsed secs (from
timeout.parse::<f64>()) is finite and >= 0.0 (e.g., secs.is_finite() && secs >=
0.0) before calling Duration::from_secs_f64; if validation fails, skip setting
config.default_request_timeout (or log/warn) so startup falls back gracefully.
Ensure you reference the existing env var parsing and assign to
config.default_request_timeout and use Duration::from_secs_f64 only after the
checks.

Comment thread lib/rl/src/lib.rs
Comment on lines +446 to +449
// Abort-on-membership-change: check epoch after fan-out completes.
let after_opts = CallOptions::new(timeout);
let after = self.snapshot(&after_opts).await?;
if after.epoch != snapshot.epoch {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Keep the original component filter in the post-fanout epoch check.

The first snapshot honors opts.components, but the second one is rebuilt with CallOptions::new(timeout) and no component filter. A worker joining or leaving some other RL-enabled component in the same namespace will now turn an otherwise stable scoped call into a false MembershipChanged 409.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/rl/src/lib.rs` around lines 446 - 449, The post-fanout epoch check
rebuilds `after_opts` with `CallOptions::new(timeout)` and thus loses the
original `opts.components` filter used for the first `snapshot`, causing false
`MembershipChanged` errors; change the construction of `after_opts` so it
preserves the original component filter (e.g. clone or derive from the original
`opts` and only update the timeout) before calling
`self.snapshot(&after_opts).await?`, ensuring `opts.components` is applied to
both snapshots.

Comment thread lib/rl/src/lib.rs
Comment on lines +579 to +583
fn call_options(&self, default_timeout: Duration) -> CallOptions {
let timeout = self
.timeout_secs
.map(Duration::from_secs_f64)
.unwrap_or(default_timeout);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

# First, locate the file and check its existence
find . -name "lib.rs" -path "*/rl/src/*" | head -20

Repository: ai-dynamo/dynamo

Length of output: 80


🏁 Script executed:

# Read the code around lines 579-583 to understand the context
if [ -f "lib/rl/src/lib.rs" ]; then
  sed -n '570,590p' lib/rl/src/lib.rs
fi

Repository: ai-dynamo/dynamo

Length of output: 830


🏁 Script executed:

# Search for timeout_secs usage and related code
rg -n "timeout_secs" lib/rl/src/lib.rs -B 3 -A 3

Repository: ai-dynamo/dynamo

Length of output: 1014


🌐 Web query:

Duration::from_secs_f64 panic behavior Rust

💡 Result:

Rust’s std::time::Duration::from_secs_f64 panics when the input seconds is negative, not finite (e.g., NaN or ±∞), or would overflow the Duration representation [1]. Practical implication: if you want to avoid panics when converting from f64, use the fallible APIs that do not panic (Duration::try_from_secs_f64), which return a Result [2].

Citations:


Validate timeout_secs before converting it to Duration.

Duration::from_secs_f64() panics on negative, NaN, or infinite inputs. Since timeout_secs originates from user-supplied HTTP requests, an unvalidated request can crash this handler instead of returning a 4xx error. Validate that the value is positive and finite before calling Duration::from_secs_f64(), or use Duration::try_from_secs_f64() to handle errors gracefully.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/rl/src/lib.rs` around lines 579 - 583, In call_options, validate the
user-supplied timeout_secs before converting: check that the f64 is finite and >
0.0 (e.g. if let Some(s) = self.timeout_secs { if s.is_finite() && s > 0.0 {
Duration::from_secs_f64(s) } else { /* use default_timeout or return a 4xx error
*/ } } else { default_timeout }), or use
Duration::try_from_secs_f64(self.timeout_secs.unwrap_or_default()) and handle
the Err by falling back to default_timeout (or propagate a 4xx error) so
negative/NaN/inf values never panic; update call_options to use this safe
conversion logic referencing timeout_secs, Duration::from_secs_f64 /
Duration::try_from_secs_f64, and default_timeout.

Comment on lines +94 to +97
ap.add_argument("--host", default="127.0.0.1")
ap.add_argument("--port", type=int, default=29501)
ap.add_argument("--world-size", type=int, default=2,
help="Total ranks (trainer + inference workers). Default 2 = 1 trainer + 1 worker.")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Make the NCCL rendezvous port explicit or dynamically allocated.

29501 is a fixed test port under tests/, so parallel smoke runs can collide on the store. Please require the caller to pass a dynamically allocated port instead of defaulting to a literal here.

As per coding guidelines, Never hardcode literal port numbers in test code; always use dynamo_dynamic_ports fixture or allocate_port()/allocate_ports() from tests.utils.port_utils.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/rl/nccl_broadcaster.py` around lines 94 - 97, The port argument in
tests/rl/nccl_broadcaster.py is hardcoded to 29501 via ap.add_argument("--port",
type=int, default=29501); remove the literal default and ensure the port is
dynamically allocated or required by the caller — either make the argument
required (ap.add_argument("--port", type=int, required=True)) so callers must
supply a port, or import and call allocate_port()/allocate_ports() from
tests.utils.port_utils (or rely on the dynamo_dynamic_ports fixture) to provide
a non-conflicting default.

Comment thread tests/rl/smoke_test_lora.sh Outdated
Comment on lines +37 to +40
MODEL="${1:-Qwen/Qwen3-0.6B}"
HTTP_PORT="${DYN_HTTP_PORT:-8000}"
NATS_PORT="${NATS_PORT:-4222}"
PRIME_RL_SRC="${PRIME_RL_SRC:-/home/biswaranjanp/dev/rl/prime-rl/src}"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Remove the developer-local PRIME_RL_SRC default here too.

This script is only runnable as-is on one workstation. Either derive the location from the repo layout or make PRIME_RL_SRC a required input so failures are explicit.

Based on learnings: Flag hard-coded portability-reducing constants in shell/scripts across the repository.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/rl/smoke_test_lora.sh` around lines 37 - 40, Remove the developer-local
default for PRIME_RL_SRC and make its value explicit: stop setting
PRIME_RL_SRC="${PRIME_RL_SRC:-/home/…}" in the script and either derive it from
the repository layout (e.g., resolve relative to the script location using
dirname "$0" or git rev-parse --show-toplevel) or require the caller to pass
PRIME_RL_SRC (fail with a clear error if unset). Update the script to reference
the PRIME_RL_SRC variable (no hard-coded path) and add a brief usage/error
message that exits if PRIME_RL_SRC is not provided.

Comment thread tests/rl/smoke_test.sh
Comment on lines +36 to +39
# prime_rl source must be on PYTHONPATH so the spawned vLLM worker subprocess
# can import prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker.
PRIME_RL_SRC="${PRIME_RL_SRC:-/home/biswaranjanp/dev/rl/prime-rl/src}"
PYTHON="${SMOKE_PYTHON:-python}"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Don't default PRIME_RL_SRC to an author-local path.

This makes the smoke test fail everywhere except the original dev machine unless the caller already knows to override it. Resolve the path from the repo layout or require PRIME_RL_SRC explicitly.

Suggested fix
-PRIME_RL_SRC="${PRIME_RL_SRC:-/home/biswaranjanp/dev/rl/prime-rl/src}"
+: "${PRIME_RL_SRC:?Set PRIME_RL_SRC to the prime-rl src directory before running this smoke test}"

Based on learnings: Flag hard-coded portability-reducing constants in shell/scripts across the repository.

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
# prime_rl source must be on PYTHONPATH so the spawned vLLM worker subprocess
# can import prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker.
PRIME_RL_SRC="${PRIME_RL_SRC:-/home/biswaranjanp/dev/rl/prime-rl/src}"
PYTHON="${SMOKE_PYTHON:-python}"
# prime_rl source must be on PYTHONPATH so the spawned vLLM worker subprocess
# can import prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker.
: "${PRIME_RL_SRC:?Set PRIME_RL_SRC to the prime-rl src directory before running this smoke test}"
PYTHON="${SMOKE_PYTHON:-python}"
🧰 Tools
🪛 Shellcheck (0.11.0)

[warning] 39-39: PYTHON appears unused. Verify use (or export if used externally).

(SC2034)

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/rl/smoke_test.sh` around lines 36 - 39, The script currently defaults
PRIME_RL_SRC to an author-local absolute path; replace that with a portable
resolution or make it required: compute the repo root relative to the script
(using the script's directory) and set PRIME_RL_SRC to the repository's src
directory only if not provided, or else exit with a clear error asking the
caller to set PRIME_RL_SRC; update the assignment of PRIME_RL_SRC and keep the
existing fallback for PYTHON/SMOKE_PYTHON unchanged. Ensure you modify the
variables PRIME_RL_SRC and PYTHON in tests/rl/smoke_test.sh and add a short
error message when PRIME_RL_SRC remains unset.

@biswapanda biswapanda changed the title feat(rl): RL admin control plane — pause/resume/update_weights fan-out for prime-rl feat(rl): RL integration layer for prime-rl May 11, 2026
biswapanda added a commit that referenced this pull request May 11, 2026
service_v2.rs:
- Fail fast (return Err) instead of warn when DYN_ENABLE_RL_ENDPOINTS=true
  but DistributedRuntime is not set; prevents a misconfigured service from
  appearing healthy with no RL routes mounted
- Move RL admin listener spawn to after main TCP/TLS bind succeeds so the
  RL port is never live while the inference port is still in-flight or has
  already failed (closes CodeRabbit :349 and :669 comments)

handlers.py:
- Add _paused guard to update_weights_from_disk and update_weights_from_distributed:
  weight updates while generation is active risk mixing weight versions across
  in-flight requests (closes :1073)
- Make LoRA hot-swap transactional: delay popping loaded_loras until add_lora
  succeeds; attempt rollback (re-add old adapter) on add_lora failure so engine
  and discovery stay in sync (closes :1185)
- Restore adapter on unregister_model failure in unload_lora_adapter: re-add
  to engine and restore loaded_loras so MDC and engine remain consistent;
  return retriable error instead of leaving state split (closes :1277)

worker_factory.py:
- Gate load_lora_adapter / unload_lora_adapter engine-route registration on
  lora_enabled so workers started without --enable-lora do not advertise
  LoRA admin methods they cannot serve (closes :726)

tests/:
- nccl_broadcaster.py: make --port required (no hardcoded 29501 default) to
  prevent port collisions between concurrent smoke runs (closes :97)
- smoke_test.sh, smoke_test_lora.sh, smoke_test_nccl.sh: require PRIME_RL_SRC
  to be set by the caller; remove author-local /home/biswaranjanp default
  (closes :39 and :40)
biswapanda added 13 commits May 12, 2026 21:54
…outes

Add RL admin handler methods to BaseWorkerHandler (handlers.py):
- liveness_probe, pause_generation, resume_generation, flush_cache, abort_request
- update_weights_from_disk (engine_rpc dispatch), update_weights_from_distributed,
  update_weights_from_tensor (stub), init_weights_update_group, destroy_weights_update_group
- get_weight_version, load_lora_adapter (with hot-swap + MDC), unload_lora_adapter

Add _shutdown_on_engine_dead helper and _paused / _weight_version state.

Register all RL routes in WorkerFactory.register_engine_routes (worker_factory.py).
These routes are reachable via POST /v1/rl/engine {"method": "<name>", "kwargs": {...}}
once the Stage 2 frontend surface lands.
- BaseWorkerHandler._rl_routes dict populated by register_engine_routes
- rl_dispatch async generator: receives {method, kwargs} from the frontend
  fan-out (NATS/TCP), dispatches to _rl_routes; __describe__ returns method list
- worker_factory: serve rl_dispatch on dyn://<ns>.<comp>.rl endpoint for both
  decode and prefill workers
- register_engine_routes: dual-registers RL routes in _rl_routes (request plane)
  and runtime.register_engine_route (system status server /engine/<name>)
@biswapanda biswapanda changed the title feat(rl): RL integration layer for prime-rl [DRAFT] feat(rl): RL integration layer for prime-rl May 14, 2026
@rmccorm4
rmccorm4 marked this pull request as draft May 14, 2026 05:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

backend::vllm Relates to the vllm backend feat frontend `python -m dynamo.frontend` and `dynamo-run in=http|text|grpc` size/XXL

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant