feat(rl): RL integration layer for prime-rl - #9382
Conversation
WalkthroughThis PR introduces a complete RL (admin control plane) subsystem for distributed worker orchestration. It adds a new Rust library ( ChangesRL Admin Control Plane Implementation
🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ 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. Comment |
There was a problem hiding this comment.
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 winReplace internal Linear ticket IDs in source comments.
The new TODOs reference
DIS-1661andDIS-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 winAvoid 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
⛔ Files ignored due to path filters (2)
Cargo.lockis excluded by!**/*.locklib/bindings/python/Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (15)
Cargo.tomlcomponents/src/dynamo/vllm/handlers.pycomponents/src/dynamo/vllm/worker_factory.pylib/llm/Cargo.tomllib/llm/src/entrypoint/input/http.rslib/llm/src/http/service/openai.rslib/llm/src/http/service/service_v2.rslib/rl/Cargo.tomllib/rl/src/lib.rstests/rl/make_lora.pytests/rl/nccl_broadcaster.pytests/rl/smoke_test.shtests/rl/smoke_test_lora.shtests/rl/smoke_test_nccl.shtests/rl/smoke_test_no_extension.sh
| 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)} |
There was a problem hiding this comment.
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.
| 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, | ||
| } |
There was a problem hiding this comment.
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.
| 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, | ||
| } | ||
|
|
There was a problem hiding this comment.
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.
| 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, | ||
| } |
There was a problem hiding this comment.
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).
| 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); | ||
| } |
There was a problem hiding this comment.
🧩 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:
- 1: https://doc.rust-lang.org/stable/std/time/struct.Duration.html
- 2: https://doc.rust-lang.org/core/time/struct.Duration.html
- 3: https://doc.servo.org/time/struct.Duration.html
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.
| 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.
| // 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 { |
There was a problem hiding this comment.
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.
| fn call_options(&self, default_timeout: Duration) -> CallOptions { | ||
| let timeout = self | ||
| .timeout_secs | ||
| .map(Duration::from_secs_f64) | ||
| .unwrap_or(default_timeout); |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, locate the file and check its existence
find . -name "lib.rs" -path "*/rl/src/*" | head -20Repository: 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
fiRepository: 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 3Repository: 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:
- 1: https://doc.rust-lang.org/stable/std/time/struct.Duration.html
- 2: Tracking Issue for
Duration::try_from_secs_{f32, f64}rust-lang/rust#83400
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.
| 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.") |
There was a problem hiding this comment.
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.
| 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}" |
There was a problem hiding this comment.
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.
| # 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}" |
There was a problem hiding this comment.
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.
| # 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.
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)
…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>)
…_ENABLE_RL_ENDPOINTS
…mWeightUpdateWorker
…_from_distributed
….prompt_logprobs (plan A2)
…nk alias (plan A3/A4/A6/A9)
…lt + skip_special_tokens
…tion_config defaults in TITO mode
…(closes 17x Mismatch KL gap)
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
/v1/chat/completions/v1/rl/engineThis PR implements the RL admin path. The generation path (
dynamo_chat_nvext) is already on main.Architecture
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/rlExactly one path with two methods:
POST/v1/rl/engineinstance_id)GET/v1/rl/engineThat'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 methodRequest 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 }methodkwargs{})instance_idcomponentstimeout_secsResponse (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— describeReturns 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/describeendpoint.HTTP response codes
error_typeno_workersmembership_changedinstance_idpresent but instance not foundunknown_instancefanout_failedPer-worker errors are reported as
status="error"entries inside the 200FanoutReport, 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 asmethod.Control plane
liveness_probepause_generationkwargs:mode(keep|wait|abort, defaultkeep),clear_cache(bool, defaultfalse).resume_generationflush_cacheabort_requestWeight management
engine_rpcupdate_weights_from_diskkwargs:model_path,weight_version,engine_rpc.reload_weights(vLLM built-in)update_weights_from_distributedinit_weights_update_groupfirst.update_weights_from_tensorinit_weights_update_groupkwargs:master_address,master_port,rank_offset,world_size,timeout.destroy_weights_update_groupget_weight_versionLoRA (registered only when
--enable-lora)load_lora_adapteradd_lorafails.unload_lora_adapterengine_rpcfield — worker-extension dispatchvLLM's
--worker-extension-clsinjects custom methods into per-GPU Worker processes, reachable viacollective_rpc. Theengine_rpcfield on weight-update kwargs selects the target:FileSystemWeightUpdateWorkerNCCLWeightUpdateWorkerupdate_weights_from_diskreload_weightsupdate_weights_from_pathupdate_weights_from_distributedupdate_weights_from_pathinit_weights_update_groupinit_broadcasterBoth worker extensions are unchanged — they already expose the right
collective_rpctargets. The caller declares whichengine_rpcvalue to use; Dynamo does not probe the extension.Fan-out safety properties (unconditional)
instance_idis present and the target has vanished, returns409 unknown_instancerather than silently routing to a different instance.409 membership_changedif a worker joined / left mid-call. Prevents NCCL group corruption.Both are baked into the dispatcher — not configurable.
Environment variables
DYN_ENABLE_RLtrue, worker registersdyn://<ns>.<comp>.rlrequest-plane endpointDYN_ENABLE_RL_ENDPOINTStrue, frontend mounts/v1/rl/engineDYN_NAMESPACEdynamoDYN_RL_PORT8002DYN_RL_COMPONENTWhere the reviewer should start
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)components/src/dynamo/vllm/worker_factory.py—register_engine_routes()(lines 700–739): the canonical registered-method list, gated onlora_enabledcomponents/src/dynamo/vllm/handlers.py— all RL handler methods onBaseWorkerHandler. Note the pause guard onupdate_weights_from_*and atomic LoRA hot-swap / unload rollback.lib/llm/src/http/service/service_v2.rs— RL router mount + listener lifecycle (RL listener spawned only after main bind succeeds; hard error whenDYN_ENABLE_RL_ENDPOINTS=truebutruntimeis unset).lib/llm/src/http/service/openai.rs::rl_router— wiresdynamo-rlcrate into the HTTP service.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)Errinstead oftracing::warnwhenDYN_ENABLE_RL_ENDPOINTS=truebutHttpServiceConfig.runtimeis unset — prevents a misconfigured service from appearing healthy with no RL routes mountedPython (
handlers.py)_pausedguard onupdate_weights_from_diskandupdate_weights_from_distributed— weight updates while generation is active risk mixing weight versions across in-flight requestsloaded_loraspop is delayed untiladd_lorasucceeds; rolls back the old adapter onadd_lorafailureunload_lora_adapterrolls back (re-adds adapter to engine + restoresloaded_loras) whenunregister_modelfails, keeping engine and discovery in syncPython (
worker_factory.py)load_lora_adapter/unload_lora_adapterengine routes gated onlora_enabledso workers started without--enable-lorano longer advertise LoRA admin methods they cannot serveTests (
tests/rl/)nccl_broadcaster.py:--portis nowrequired=True(no hardcoded 29501 default) — prevents collisions between concurrent smoke runssmoke_test.sh,smoke_test_lora.sh,smoke_test_nccl.sh: requirePRIME_RL_SRCfrom the caller; removed author-local/home/biswaranjanpdefaultsRelated Issues
Supersedes #9131