address comments, pre-commit

Signed-off-by: Erin Ho <erinh@nvidia.com>

Author: hchings

nemo_rl/data/llm_message_utils.py

@@ -377,7 +377,8 @@ def batched_message_log_to_flat_message(
         try:
             _validate_tensor_consistency(tensors)
         except RuntimeError as e:
-            raise RuntimeError(f"[key={key!r}] {e}") from e
+            e.add_note(f"[key={key!r}]")
+            raise
 
         # Create zero tensors for None values
         filled_values: list[Tensor] = [

nemo_rl/environments/nemo_gym.py

@@ -345,7 +345,10 @@ def _postprocess_nemo_gym_to_nemo_rl_result(
 
             # Note that NeMo-Gym will only return token ids on "assistant" messages and not other message types.
             # Also skip if generation_token_ids is present but empty, e.g. all-EOS generation stripped to [] — torch.tensor([]) defaults to float32 and breaks batch dtype consistency.
-            if "generation_token_ids" not in output_item_dict or not output_item_dict["generation_token_ids"]:
+            if (
+                "generation_token_ids" not in output_item_dict
+                or not output_item_dict["generation_token_ids"]
+            ):
                 continue
 
             assert (

nemo_rl/models/generation/openai_server_utils.py

@@ -120,4 +120,6 @@ def replace_prefix_tokens(
         f"Template repr (detokenized): {repr(tokenizer.decode(template_token_ids))}"
     )
 
-    return model_prefix_token_ids[:model_cut_end] + template_token_ids[template_cut_start:]
+    return (
+        model_prefix_token_ids[:model_cut_end] + template_token_ids[template_cut_start:]
+    )

nemo_rl/models/generation/trtllm/trtllm_http_server.py

@@ -24,18 +24,19 @@
 import threading
 import time
 import uuid
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any
 
 from nemo_rl.models.generation.openai_server_utils import (
     replace_prefix_tokens,
 )
 
+if TYPE_CHECKING:
+    from fastapi import FastAPI
+
 logger = logging.getLogger(__name__)
 
 _TOOL_BOT = "<tool_call>\n"
-_TOOL_CALL_RE = re.compile(
-    r"<tool_call>\n(.*?)\n</tool_call>", re.DOTALL
-)
+_TOOL_CALL_RE = re.compile(r"<tool_call>\n(.*?)\n</tool_call>", re.DOTALL)
 
 
 def create_app(
@@ -59,6 +60,7 @@ def create_app(
     # so model output starts inside the reasoning block without a leading <think> tag.
     # parse() is stateless — safe for concurrent requests.
     from tensorrt_llm.llmapi.reasoning_parser import DeepSeekR1Parser
+
     _reasoning_parser = DeepSeekR1Parser(reasoning_at_start=True)
 
     # Stop tokens TRT-LLM appends to gen_token_ids but NOT reproduced by apply_chat_template.
@@ -87,7 +89,10 @@ async def chat_completions(request: Request):
         logprobs_requested = body.get("logprobs", False)
 
         prompt_token_ids = _build_prompt_token_ids(
-            messages, tokenizer, tools=tools, default_template_kwargs=_default_template_kwargs
+            messages,
+            tokenizer,
+            tools=tools,
+            default_template_kwargs=_default_template_kwargs,
         )
 
         # Prefix splice: empty required_prefix_ids (no prior assistant turn) → returns template unchanged.
@@ -102,14 +107,18 @@ async def chat_completions(request: Request):
             template_token_ids=prompt_token_ids,
         )
 
-        max_tokens_requested = body.get("max_tokens") or body.get("max_completion_tokens") or max_seq_len
+        max_tokens_requested = (
+            body.get("max_tokens") or body.get("max_completion_tokens") or max_seq_len
+        )
         remaining_ctx = max(0, max_seq_len - len(adj_prompt))
 
         # Return HTTP 400 on context exhaustion.
         if remaining_ctx == 0:
             return JSONResponse(
                 status_code=400,
-                content={"error": f"context length exceeded: prompt ({len(adj_prompt)} tokens) exhausted context window ({max_seq_len})"},
+                content={
+                    "error": f"context length exceeded: prompt ({len(adj_prompt)} tokens) exhausted context window ({max_seq_len})"
+                },
             )
 
         max_tokens = min(int(max_tokens_requested), remaining_ctx)
@@ -131,7 +140,11 @@ async def chat_completions(request: Request):
             )
         except Exception as e:
             err = str(e)
-            if "prompt length" in err or "max_num_tokens" in err or "max_seq_len" in err:
+            if (
+                "prompt length" in err
+                or "max_num_tokens" in err
+                or "max_seq_len" in err
+            ):
                 return JSONResponse(
                     status_code=400,
                     content={"error": f"context length exceeded: {err}"},
@@ -144,13 +157,14 @@ async def chat_completions(request: Request):
 
         gen_logprobs: list[float] = []
         if gen.logprobs:
-            for lp in gen.logprobs:
+            # TRT-LLM returns floats in simple format and token-indexed dicts otherwise.
+            for token_id, lp in zip(gen_token_ids, gen.logprobs, strict=True):
                 if isinstance(lp, (int, float)):
                     gen_logprobs.append(float(lp))
                 elif isinstance(lp, dict):
-                    gen_logprobs.append(float(next(iter(lp.values())).logprob))
+                    gen_logprobs.append(float(lp[token_id].logprob))
                 else:
-                    gen_logprobs.append(float(lp))
+                    raise TypeError(f"Unsupported TRT-LLM logprob type: {type(lp)}")
 
         # Strip trailing stop tokens TRT-LLM appends — apply_chat_template doesn't reproduce
         # <|endoftext|>, so they'd break seen_token_ids contiguity. Trim logprobs in lockstep.
@@ -239,6 +253,7 @@ async def chat_completions(request: Request):
 #  Tool-call parsing — mirrors Qwen3ToolParser.detect_and_parse() exactly
 # ---------------------------------------------------------------------------
 
+
 def _parse_tool_calls(text: str) -> list[dict[str, Any]]:
     """Extract <tool_call> blocks; arg serialisation is byte-identical to Qwen3ToolParser."""
     results: list[dict[str, Any]] = []
@@ -253,14 +268,16 @@ def _parse_tool_calls(text: str) -> list[dict[str, Any]]:
             obj.get("parameters") or obj.get("arguments", {}),
             ensure_ascii=False,
         )
-        results.append({
-            "id": f"call_{uuid.uuid4().hex[:8]}",
-            "type": "function",
-            "function": {
-                "name": func_name,
-                "arguments": arguments,
-            },
-        })
+        results.append(
+            {
+                "id": f"call_{uuid.uuid4().hex[:8]}",
+                "type": "function",
+                "function": {
+                    "name": func_name,
+                    "arguments": arguments,
+                },
+            }
+        )
     return results
 
 
@@ -276,6 +293,7 @@ def _strip_tool_call_tags(text: str) -> str:
 #  Prompt construction
 # ---------------------------------------------------------------------------
 
+
 def _to_int_ids(enc: Any) -> list[int]:
     """Coerce apply_chat_template output to flat list[int] (handles transformers v5 BatchEncoding)."""
     if hasattr(enc, "input_ids"):  # v5 BatchEncoding
@@ -329,17 +347,23 @@ def _compute_splice_inputs(
     required_prefix_ids: list[int] = []
     for _m in reversed(messages):
         if _m.get("role") == "assistant" and "prompt_token_ids" in _m:
-            required_prefix_ids = (
-                list(_m["prompt_token_ids"]) + list(_m.get("generation_token_ids") or [])
+            required_prefix_ids = list(_m["prompt_token_ids"]) + list(
+                _m.get("generation_token_ids") or []
             )
             break
 
     _clean = [_strip_token_fields(m) for m in messages]
     _last_asst_idx = next(
-        (i for i in reversed(range(len(_clean))) if _clean[i].get("role") == "assistant"),
+        (
+            i
+            for i in reversed(range(len(_clean)))
+            if _clean[i].get("role") == "assistant"
+        ),
         None,
     )
-    _msgs_to_last_asst = _clean[:_last_asst_idx + 1] if _last_asst_idx is not None else _clean
+    _msgs_to_last_asst = (
+        _clean[: _last_asst_idx + 1] if _last_asst_idx is not None else _clean
+    )
     _prefix_tkw: dict[str, Any] = {
         **default_template_kwargs,
         "tokenize": True,
@@ -357,6 +381,7 @@ def _compute_splice_inputs(
 #  Server lifecycle
 # ---------------------------------------------------------------------------
 
+
 def start_server(
     llm: Any,
     tokenizer: Any,
@@ -368,6 +393,7 @@ def start_server(
 ) -> "tuple[threading.Thread, str, Any]":
     """Start the HTTP server in a daemon thread and return (thread, base_url, server)."""
     import uvicorn
+
     from nemo_rl.distributed.virtual_cluster import (
         _get_free_port_local,
         _get_node_ip_local,

nemo_rl/models/generation/trtllm/trtllm_worker_async.py

@@ -271,7 +271,8 @@ def start_http_server(self, port: int = 0) -> str:
         from nemo_rl.models.generation.trtllm.trtllm_http_server import start_server
 
         tokenizer = AutoTokenizer.from_pretrained(
-            self.model_name, trust_remote_code=True,
+            self.model_name,
+            trust_remote_code=True,
         )
         self._http_thread, self._http_base_url, self._http_server = start_server(
             llm=self.llm,

tests/unit/models/generation/test_vllm_generation.py

@@ -40,7 +40,9 @@
 from nemo_rl.models.generation.vllm.vllm_worker import (
     _resolve_enable_prefix_caching,
 )
-from nemo_rl.models.generation.vllm.vllm_worker_async import VllmAsyncGenerationWorkerImpl
+from nemo_rl.models.generation.vllm.vllm_worker_async import (
+    VllmAsyncGenerationWorkerImpl,
+)
 from nemo_rl.models.policy import LoRAConfig, PolicyConfig
 from nemo_rl.models.policy.lm_policy import Policy