lib/llm/src/preprocessor.rs (1)

585-594: Optional: collapse the two tokenizing branches.

The else if has_backend_instance_id { ... } and final else { ... } branches have identical result expressions — only the tracing::warn! differs. They can be flattened with a guard log, reducing the risk of future drift:

♻️ Suggested simplification

-                            } else if has_backend_instance_id {
-                                tracing::warn!(
-                                    "backend_instance_id provided but no token_data; tokenizing prompt"
-                                );
-                                let encoding = self.encode_with_timing(&prompt, tracker)?;
-                                (encoding.token_ids().to_vec(), false)
-                            } else {
-                                let encoding = self.encode_with_timing(&prompt, tracker)?;
-                                (encoding.token_ids().to_vec(), false)
-                            };
+                            } else {
+                                if has_backend_instance_id {
+                                    tracing::warn!(
+                                        "backend_instance_id provided but no token_data; tokenizing prompt"
+                                    );
+                                }
+                                let encoding = self.encode_with_timing(&prompt, tracker)?;
+                                (encoding.token_ids().to_vec(), false)
+                            };

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@lib/llm/src/preprocessor.rs` around lines 585 - 594, The two branches under
the else/else if both call self.encode_with_timing(&prompt, tracker)? and return
(encoding.token_ids().to_vec(), false), differing only by a tracing::warn! when
has_backend_instance_id is true; refactor by removing the duplicate branches: in
the outer else, if has_backend_instance_id { tracing::warn!(...); } then call
let encoding = self.encode_with_timing(&prompt, tracker)? and return
(encoding.token_ids().to_vec(), false) — this collapses the branches while
preserving the warning behavior and using the existing encode_with_timing,
prompt, tracker, and has_backend_instance_id symbols.

lib/llm/src/protocols/openai/chat_completions/delta.rs (1)

466-486: Nit: use std::mem::take to release the accumulator vector on emit.

self.accumulated_completion_token_ids.clone() retains the backing allocation on the generator even though no more chunks will follow finish_reason. Swapping in std::mem::take releases it promptly:

♻️ Proposed change

-        let completion_token_ids =
-            if self.options.enable_completion_token_ids && finish_reason.is_some() {
-                Some(self.accumulated_completion_token_ids.clone())
-            } else {
-                None
-            };
+        let completion_token_ids =
+            if self.options.enable_completion_token_ids && finish_reason.is_some() {
+                Some(std::mem::take(&mut self.accumulated_completion_token_ids))
+            } else {
+                None
+            };

Also, consider a unit test in this module covering: opt-in via extra_fields: ["completion_token_ids"], mid-stream chunks carry no completion_token_ids, and the final chunk's nvext.completion_token_ids equals the concatenation of per-chunk token_ids (hard invariant called out in nvext.rs).

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@lib/llm/src/protocols/openai/chat_completions/delta.rs` around lines 466 -
486, The code currently clones self.accumulated_completion_token_ids into
completion_token_ids causing the accumulator's allocation to be retained;
replace the clone with std::mem::take(&mut
self.accumulated_completion_token_ids) when building completion_token_ids (only
when self.options.enable_completion_token_ids && finish_reason.is_some()) so the
generator releases the vector's backing buffer immediately; ensure the rest of
the branch (where completion_token_ids is moved into NvExtResponse) still
compiles and preserves behavior of NvExtResponse::completion_token_ids and
related logic involving token_ids, worker_id_info, timing_info, and
routed_experts; optionally add a unit test that opts in via extra_fields:
["completion_token_ids"] confirming mid-stream chunks lack completion_token_ids
and the final nvext.completion_token_ids equals the concatenation of per-chunk
token_ids.

components/src/dynamo/vllm/worker_factory.py (1)

338-353: Extract duplicate engine-route registration to a helper function.

Both _create_decode_worker (lines 338–353) and _create_prefill_worker (lines 586–601) register the same 7 RL routes (pause_generation, resume_generation, flush_cache, update_weights_from_path, get_weight_version, load_lora_adapter, unload_lora_adapter) with an identical log message. This duplication invites drift—e.g., adding a route to one path and forgetting the other. Since both DecodeWorkerHandler and PrefillWorkerHandler inherit from BaseWorkerHandler and all methods are defined there, consider extracting:

♻️ Proposed refactor

_RL_ROUTES: tuple[str, ...] = (
    "pause_generation",
    "resume_generation",
    "flush_cache",
    "update_weights_from_path",
    "get_weight_version",
    "load_lora_adapter",
    "unload_lora_adapter",
)

def _register_lifecycle_and_rl_routes(runtime: DistributedRuntime, handler) -> None:
    runtime.register_engine_route("sleep", handler.sleep)
    runtime.register_engine_route("wake_up", handler.wake_up)
    runtime.register_engine_route("scale_elastic_ep", handler.scale_elastic_ep)
    for name in _RL_ROUTES:
        runtime.register_engine_route(name, getattr(handler, name))
    logger.info(
        "Registered engine routes: sleep, wake_up, scale_elastic_ep, %s",
        ", ".join(_RL_ROUTES),
    )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@components/src/dynamo/vllm/worker_factory.py` around lines 338 - 353, Extract
the duplicated engine-route registration in _create_decode_worker and
_create_prefill_worker into a helper function (e.g.,
_register_lifecycle_and_rl_routes) that takes the DistributedRuntime and a
handler instance; define a single _RL_ROUTES tuple containing the seven RL route
names and use a loop to call runtime.register_engine_route(name,
getattr(handler, name)) for each, also registering "sleep", "wake_up", and
"scale_elastic_ep" there, and replace the duplicate
runtime.register_engine_route calls and logger.info in both functions with a
single call to the new helper; reference BaseWorkerHandler (handler methods),
runtime.register_engine_route, _create_decode_worker, and _create_prefill_worker
when making the change.

lib/llm/src/tokenizers/fastokens.rs (1)

49-63: Document (and optionally test) the HF fallback for add_special_tokens=true.

fastokens doesn't support special-token-aware encoding, so this path transparently falls back to HuggingFaceTokenizer, which means callers lose the fastokens speedup whenever they request special tokens. That's a reasonable choice, but the struct-level doc (Lines 20–22) still claims "fast BPE encoding via fastokens, decoding via HuggingFace" which is now incomplete. A one-liner here plus an extension of test_fast_matches_hf_encoding with add_special_tokens=true would make the behavior explicit and guard against regressions (e.g., mismatched special-token IDs between the two backends loaded from the same tokenizer.json).

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@lib/llm/src/tokenizers/fastokens.rs` around lines 49 - 63, Update the
struct-level documentation for the fastokens tokenizer to state that
encode_with_special_tokens falls back to the HuggingFace tokenizer when
add_special_tokens=true (i.e., special-token-aware encoding is performed by
hf_decoder, so callers lose fastokens speedup in that case), and add/extend the
test_fast_matches_hf_encoding test to call encode_with_special_tokens with
add_special_tokens=true and assert the HF fallback produces matching token
ids/special-token ids compared to the fast path where applicable; locate the
behavior in the encode_with_special_tokens method (uses self.hf_decoder and
self.fast_encoder) and ensure the test compares outputs from
hf_decoder.encode_with_special_tokens(...) and fast_encoder.encode(...) wrapped
as Encoding::Sp (and guards for any differences in special-token id handling).

lib/llm/src/tokenizers.rs (1)

65-71: Silent default for encode_with_special_tokens is error-prone.

The default implementation silently ignores _add_special_tokens and delegates to encode. If a new Encoder backend is added later and the author forgets to override this method, callers requesting add_special_tokens = true will quietly get encoding without special tokens — a correctness bug with no compile-time or runtime signal. All current backends override this, so a stricter contract is free today.

♻️ Two safer alternatives

Either make it a required trait method (no default), or have the default fail loudly when the flag is set:

-        fn encode_with_special_tokens(
-            &self,
-            input: &str,
-            _add_special_tokens: bool,
-        ) -> Result<Encoding> {
-            self.encode(input)
-        }
+        fn encode_with_special_tokens(
+            &self,
+            input: &str,
+            add_special_tokens: bool,
+        ) -> Result<Encoding> {
+            if add_special_tokens {
+                return Err(Error::msg(
+                    "encode_with_special_tokens(add_special_tokens=true) not supported by this backend",
+                ));
+            }
+            self.encode(input)
+        }

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@lib/llm/src/tokenizers.rs` around lines 65 - 71, The trait's default impl of
encode_with_special_tokens silently ignores the _add_special_tokens flag; remove
the default implementation from the trait so encode_with_special_tokens becomes
a required method (i.e., declare fn encode_with_special_tokens(&self, input:
&str, add_special_tokens: bool) -> Result<Encoding> without a body) and then
update each Encoder implementation that currently relies on the default to
implement encode_with_special_tokens explicitly (search for
encode_with_special_tokens and implementations of the Encoder trait to add the
proper handling).

lib/llm/src/http/service/service_v2.rs (1)

546-550: Promote DYN_ENABLE_RL to an env_llm constant.

The rest of this file references environment variable names through the env_llm module (e.g., env_llm::DYN_ENABLE_STREAMING_TOOL_DISPATCH, env_llm::DYN_HTTP_GRACEFUL_SHUTDOWN_TIMEOUT_SECS). The hardcoded "DYN_ENABLE_RL" string here — plus the DYN_RL_WORKER_SYSTEM_URLS mentioned in the doc comment at Lines 249–251 — breaks that convention and makes these env vars harder to discover/grep and prone to typos.

♻️ Suggested change

Add the constant alongside existing ones in dynamo_runtime::config::environment_names::llm and use it here:

-        if config.enable_rl || env_is_truthy("DYN_ENABLE_RL") {
+        if config.enable_rl || env_is_truthy(env_llm::DYN_ENABLE_RL) {

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@lib/llm/src/http/service/service_v2.rs` around lines 546 - 550, Add a new
constant DYN_ENABLE_RL to the env_llm module
(dynamo_runtime::config::environment_names::llm) and replace the hardcoded
string usages with that constant; specifically change the condition that now
reads env_is_truthy("DYN_ENABLE_RL") to use
env_is_truthy(env_llm::DYN_ENABLE_RL) in the RL admin routes block alongside
config.enable_rl, and update any doc comments that mention the literal
DYN_RL_WORKER_SYSTEM_URLS to reference the env_llm constant name instead so all
environment names are discoverable and consistent.

components/src/dynamo/rl_admin/service.py (2)

151-168: Swallowed except Exception in /ready should log.

The two except Exception: ...False blocks hide the actual failure reason, so an operator watching a not-ready service has no signal as to whether the frontend is down, DNS is broken, or TLS handshakes are failing. Log at debug level so the reason is retrievable without changing the response shape.

♻️ Proposed fix

     try:
         frontend_ok = (await client.get(f"{DYNAMO_FRONTEND_URL}/v1/models")).status_code == 200
-    except Exception:
+    except httpx.HTTPError as e:
+        logger.debug("ready: frontend probe failed: %s", e)
         frontend_ok = False
     try:
         system_ok = (await client.get(f"{DYNAMO_SYSTEM_URL}/health")).status_code == 200
-    except Exception:
+    except httpx.HTTPError as e:
+        logger.debug("ready: system probe failed: %s", e)
         system_ok = False

As per coding guidelines: "catch only specific exceptions you can handle, and if you catch broadly (except Exception) you must log and re-raise."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@components/src/dynamo/rl_admin/service.py` around lines 151 - 168, The
ready() health check currently swallows exceptions when probing
DYNAMO_FRONTEND_URL and DYNAMO_SYSTEM_URL; update the two try/except blocks that
call _get_http_client().get(...) to catch the exception, log the full exception
details at debug (include context like which URL failed and the exception)
before setting frontend_ok/system_ok = False, and do not change the JSONResponse
shape or status codes; use the existing logger (or create one if none) to emit
the debug-level messages so operators can inspect failures without altering
behavior.

---

54-59: Use lifespan context manager instead of deprecated @app.on_event("shutdown").

The on_event decorator is deprecated in FastAPI (with deprecation warnings present). While still functional, it is slated for eventual removal. Since this is a new file, implement using the lifespan async context manager from the start to avoid a future migration.

Suggested refactor

from contextlib import asynccontextmanager

`@asynccontextmanager`
async def lifespan(app: FastAPI):
    yield
    global _http_client
    if _http_client is not None:
        await _http_client.aclose()
        _http_client = None

app = FastAPI(title="Dynamo RL Admin", lifespan=lifespan)

Add from contextlib import asynccontextmanager at the top of the file with the other imports (line 24–31).

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@components/src/dynamo/rl_admin/service.py` around lines 54 - 59, Replace the
deprecated FastAPI shutdown event handler by adding an asynccontextmanager-based
lifespan and passing it into the FastAPI constructor: import asynccontextmanager
from contextlib, define async def lifespan(app: FastAPI) decorated with
`@asynccontextmanager` that yields and then performs the same cleanup of the
global _http_client (close and set to None), and create app =
FastAPI(title="Dynamo RL Admin", lifespan=lifespan) instead of relying on the
_shutdown function; remove the `@app.on_event`("shutdown") _shutdown definition
afterward.

components/src/dynamo/vllm/handlers.py (2)

705-881: Lock-map growth: new RL LoRA handlers don't prune _lora_load_locks like their URI counterparts.

load_lora (lines 1202–1210) and unload_lora (lines 1323–1330) both wrap the critical section in a try/finally that removes the per-LoRA lock from _lora_load_locks once the adapter is no longer tracked. The new load_lora_adapter / unload_lora_adapter skip that cleanup, so over a long RL run every distinct lora_name ever loaded (or ever attempted on a failed load) leaves a permanent entry in the lock dict. It's not a functional bug, but for a training loop that hot-swaps adapters repeatedly it's worth mirroring the existing pattern for consistency.

♻️ Suggested pattern (apply to both new handlers)

             lock = self._get_lora_lock(lora_name)
             async with lock:
-                ...
+                try:
+                    ...
+                finally:
+                    with self._lora_load_locks_guard:
+                        if (
+                            lora_name not in self.loaded_loras
+                            and self._lora_load_locks.get(lora_name) is lock
+                        ):
+                            self._lora_load_locks.pop(lora_name, None)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@components/src/dynamo/vllm/handlers.py` around lines 705 - 881, The handlers
load_lora_adapter and unload_lora_adapter leak entries in the _lora_load_locks
map because they don't remove the per-lora lock after the critical section;
update both functions to mirror the existing pattern in load_lora/unload_lora by
wrapping the async with lock block in a try/finally and in the finally remove
the lock from self._lora_load_locks (use the same key lora_name) when the
adapter is no longer tracked (i.e., after successful unload or on failure) so
_get_lora_lock and loaded_loras remain consistent and the lock-map doesn't grow
unbounded.

---

829-881: unload_lora_adapter state ordering is fine, but note the asymmetric error semantics.

del self.loaded_loras[lora_name] on line 853 happens before unregister_model, so if unregister fails the adapter is gone from both the engine and local tracking yet the caller sees status=error. That matches the idempotent contract in the docstring (a retry hits the lora is None no-op branch and returns ok), but it differs from the older unload_lora which rolls the engine state back on unregister failure (lines 1278–1300). Worth a one-line comment noting the intentional divergence so a future maintainer doesn't "fix" it by copying the old rollback.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@components/src/dynamo/vllm/handlers.py` around lines 829 - 881, The
unload_lora_adapter method currently deletes the local tracking entry (del
self.loaded_loras[lora_name]) before calling unregister_model, which means if
unregister_model fails the adapter is already removed from engine and local
state but the method returns status=error; add a one-line comment immediately
above the del self.loaded_loras[lora_name] (or above the unregister_model call)
explaining that this ordering is intentional and documents the asymmetric error
semantics (idempotent no-op on retry) and that it intentionally differs from the
older unload_lora rollback behavior to avoid re-registering the adapter on
unregister failure.

Fix all unresolved CodeRabbit comments on this PR:

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

lib/llm/src/tokenizers.rs (1)

64-72: Default encode_with_special_tokens silently drops the flag — consider making it explicit.

The default impl ignores _add_special_tokens and just calls encode. Today every concrete backend (hf, fastokens, tiktoken) overrides it so this never runs, but if a future Encoder impl forgets to override it, callers will get the wrong token stream with no error. Either delete the default (force every impl to implement it) or make the divergence loud (e.g. debug-log once, or unimplemented!-style default returning an Err). Non-blocking for this PR.

♻️ One way to remove the footgun

     pub trait Encoder: Send + Sync {
         fn encode(&self, input: &str) -> Result<Encoding>;
         fn encode_batch(&self, inputs: &[&str]) -> Result<Vec<Encoding>>;
-
-        fn encode_with_special_tokens(
-            &self,
-            input: &str,
-            _add_special_tokens: bool,
-        ) -> Result<Encoding> {
-            self.encode(input)
-        }
+
+        /// Each backend MUST implement this; honoring `add_special_tokens`
+        /// is required for consistent behavior between `/v1/tokenize` and the
+        /// request preprocessor.
+        fn encode_with_special_tokens(
+            &self,
+            input: &str,
+            add_special_tokens: bool,
+        ) -> Result<Encoding>;
     }

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@lib/llm/src/tokenizers.rs` around lines 64 - 72, The default
encode_with_special_tokens currently ignores the _add_special_tokens flag and
just calls encode, which can silently produce incorrect token streams; change
the default implementation of encode_with_special_tokens to return an explicit
Err (or use unimplemented!()) with a clear message so implementors of the
Encoder must override it instead of silently dropping the flag — update the
method body for encode_with_special_tokens (referencing the
encode_with_special_tokens and encode symbols) to return an error describing
that encode_with_special_tokens must be implemented by backend encoders.

lib/llm/src/http/service/openai.rs (2)

1984-1994: This helper duplicates ErrorMessage::bad_request.

ErrorMessage::bad_request is defined at line 211 with identical behaviour. The free bad_request here just forwards to String and reconstructs the same ErrorResponse. Use the existing associated function and drop this duplicate to keep a single source of truth for 400 responses.

♻️ Proposed refactor

-fn bad_request<T: Into<String>>(message: T) -> ErrorResponse {
-    let code = StatusCode::BAD_REQUEST;
-    (
-        code,
-        Json(ErrorMessage {
-            message: message.into(),
-            error_type: map_error_code_to_error_type(code),
-            code: code.as_u16(),
-        }),
-    )
-}
-
 fn resolve_tokenizer_model_name(
     state: &Arc<service_v2::State>,
     requested_model: Option<&str>,
 ) -> Result<String, ErrorResponse> {
     if let Some(model) = requested_model {
         if state.manager().has_model_any(model) {
             return Ok(model.to_string());
         }
         return Err(ErrorMessage::model_not_found());
     }
     let served_models = state.manager().model_display_names();
     if served_models.len() == 1 {
         return Ok(served_models.into_iter().next().unwrap());
     }
-    Err(bad_request(
+    Err(ErrorMessage::bad_request(
         "Model must be specified when more than one model is served.",
     ))
 }

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@lib/llm/src/http/service/openai.rs` around lines 1984 - 1994, The standalone
function bad_request<T: Into<String>>(...) duplicates ErrorMessage::bad_request;
remove the free function and update any call sites to use the associated
constructor ErrorMessage::bad_request (or ErrorMessage::bad_request(message)) so
there is a single source of truth for 400 responses; ensure the return type
remains ErrorResponse and preserve mapping using map_error_code_to_error_type
where used.

---

897-958: Factor the RL nvext setup out of the two handlers.

This block and handler_chat_completions_tokens (lines 2307-2354) perform nearly identical work: take tokens → nvext.token_data, push "token_ids"/"completion_token_ids" into nvext.extra_fields, force logprobs = Some(true), and inject the "(token-in mode)" placeholder when messages is empty. The literal placeholder string is also duplicated, which is easy to drift on. Extract a single helper (e.g. fn rl_promote_token_fields(req: &mut NvCreateChatCompletionRequest, want_token_ids: bool) -> Result<(), ErrorResponse>) and call it from both handlers — this removes the risk of the two paths diverging (e.g. one path adds a new extra_field, the other forgets).

One minor concern on the placeholder itself: because validate_chat_completion_required_fields later rejects empty messages (line 1545), the placeholder is load-bearing, yet its content is an English magic string that will leak into downstream chat-template rendering if the template doesn't honour nvext.token_data. Please at least centralise the placeholder as a const next to the helper so it's greppable.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@lib/llm/src/http/service/openai.rs` around lines 897 - 958, Extract the
duplicated RL nvext setup into a single helper (e.g. fn
rl_promote_token_fields(req: &mut NvCreateChatCompletionRequest, want_token_ids:
bool) -> Result<(), ErrorResponse>) and call it from both the current handler
block and handler_chat_completions_tokens; the helper should: take
request.tokens and set nvext.token_data when present, ensure
request.inner.messages contains the same placeholder when empty, merge
"token_ids" and "completion_token_ids" into nvext.extra_fields when
want_token_ids is true, and set request.inner.logprobs = Some(true) if not
already set; centralize the "(token-in mode)" placeholder as a const next to the
helper and return/propagate an ErrorResponse on failure so callers can handle
validation consistently.

Fix all unresolved CodeRabbit comments on this PR:

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