Skip to content

[TRTLLM-14604][fix] add auth for RL endpoints - #17434

Open
hchings wants to merge 1 commit into
NVIDIA:mainfrom
hchings:nvbug_6507663
Open

[TRTLLM-14604][fix] add auth for RL endpoints#17434
hchings wants to merge 1 commit into
NVIDIA:mainfrom
hchings:nvbug_6507663

Conversation

@hchings

@hchings hchings commented Aug 8, 2026

Copy link
Copy Markdown
Collaborator

Dev Engineer Review

  • Added optional bearer-token authentication for RL control endpoints.
  • Added configuration validation for rl_control_api_key and AsyncLLM.
  • Registered /release_memory, /resume_memory, and /update_weights only when enabled.
  • Added constant-time token comparison and HTTP 401 responses for unauthorized requests.
  • Changed MemoryUpdateRequest.tags to require at least one item.
  • Review focus: verify compatibility with existing OpenAIServer callers and confirm secure handling of missing or invalid bearer tokens.

QA Engineer Review

  • Added tests:
    • test_rl_control_routes_require_key
    • test_rl_control_routes_require_async_llm
    • test_rl_control_routes_disabled_by_default
    • test_rl_control_routes_require_auth
    • test_rl_control_routes_accept_valid_token
    • test_release_memory_requires_tags
  • The tests cover configuration validation, route registration, authentication failures and success, RPC dispatch, and request validation.
  • No matching entries were identified in tests/integration/test_lists/.
  • Verdict: insufficient. The test functions need CI or manual QA test-list coverage.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
@hchings hchings self-assigned this Aug 8, 2026
@hchings
hchings requested a review from a team as a code owner August 8, 2026 00:37
@hchings hchings changed the title [][fix] add auth for RL endpoints [TRTLLM-14604][fix] add auth for RL endpoints Aug 8, 2026
@coderabbitai

coderabbitai Bot commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Walkthrough

Changes

RL control endpoints

Layer / File(s) Summary
Memory update request contract
tensorrt_llm/serve/openai_protocol.py
MemoryUpdateRequest.tags no longer has a default value and must contain at least one item.
Server configuration and protected routes
tensorrt_llm/serve/openai_server.py
OpenAIServer validates RL endpoint configuration, conditionally registers the control routes, and requires a valid bearer token before dispatching RPCs.
Configuration and endpoint validation
tests/unittest/llmapi/test_rl_control_auth.py
Tests cover configuration errors, disabled routes, authentication failures, successful RPC dispatch, and missing release tags.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Client
  participant FastAPI
  participant OpenAIServer
  participant AsyncLLM
  Client->>FastAPI: Send bearer-authenticated RL request
  FastAPI->>OpenAIServer: Validate token and request
  OpenAIServer->>AsyncLLM: Dispatch collective RPC
  AsyncLLM-->>OpenAIServer: Return RPC result
  OpenAIServer-->>Client: Return HTTP response
Loading

Suggested reviewers: bowenfu

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Description check ⚠️ Warning The description contains the template but leaves the Description and Test Coverage sections empty and does not complete the PR checklist. Add a concise issue and solution summary, list the relevant tests, and complete the checklist items that apply to this pull request.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly identifies the ticket, fix type, and main change: authentication for RL endpoints.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
✨ Finishing Touches 💡 1
🛠️ Fix failing CI checks 💡
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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

@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: 1

🧹 Nitpick comments (1)
tests/unittest/llmapi/test_rl_control_auth.py (1)

70-118: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Add invalid-credential and empty-tags tests.

  • Test Authorization: Bearer wrong and {"tags": []}. Assert 401 or 422, respectively, and no collective_rpc dispatch.
  • Add unittest/llmapi/test_rl_control_auth.py to tests/integration/test_lists/test-db/l0_cpu.yml. No current test-db/ or qa/ list includes this file.
  • Run pytest tests/unittest/.

Coverage remains insufficient for invalid credentials and the min_length=1 constraint.

🤖 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/unittest/llmapi/test_rl_control_auth.py` around lines 70 - 118, Extend
test_rl_control_routes_require_auth with an invalid-token case using
Authorization: Bearer wrong, asserting 401 and no collective_rpc dispatch, and
add an empty-tags case for release_memory asserting 422 with no dispatch to
cover the min_length=1 constraint. Register
tests/unittest/llmapi/test_rl_control_auth.py in
tests/integration/test_lists/test-db/l0_cpu.yml, then run pytest
tests/unittest/.

Sources: Coding guidelines, Path instructions

🤖 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 `@tensorrt_llm/serve/openai_protocol.py`:
- Line 1382: Update the tags field declaration to use the modern list[str]
annotation instead of List[str], and add a clear description argument to Field
so the generated RL endpoint request schema documents the field’s purpose.

---

Nitpick comments:
In `@tests/unittest/llmapi/test_rl_control_auth.py`:
- Around line 70-118: Extend test_rl_control_routes_require_auth with an
invalid-token case using Authorization: Bearer wrong, asserting 401 and no
collective_rpc dispatch, and add an empty-tags case for release_memory asserting
422 with no dispatch to cover the min_length=1 constraint. Register
tests/unittest/llmapi/test_rl_control_auth.py in
tests/integration/test_lists/test-db/l0_cpu.yml, then run pytest
tests/unittest/.
🪄 Autofix

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: Enterprise

Run ID: 978eaf53-4115-4bb4-8bc9-4841c4ed1593

📥 Commits

Reviewing files that changed from the base of the PR and between bcc0327 and c74e7f6.

📒 Files selected for processing (3)
  • tensorrt_llm/serve/openai_protocol.py
  • tensorrt_llm/serve/openai_server.py
  • tests/unittest/llmapi/test_rl_control_auth.py


class MemoryUpdateRequest(OpenAIBaseModel):
tags: List[str] = Field(default=["model", "kv_cache"])
tags: List[str] = Field(min_length=1)

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.

📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win

Document and modernize the Pydantic field.

Use list[str] and add a description to Field. This improves the generated request schema for RL endpoint clients.

Proposed fix
-    tags: List[str] = Field(min_length=1)
+    tags: list[str] = Field(
+        min_length=1,
+        description="Memory resource tags to release or resume.",
+    )
📝 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
tags: List[str] = Field(min_length=1)
tags: list[str] = Field(
min_length=1,
description="Memory resource tags to release or resume.",
)
🤖 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 `@tensorrt_llm/serve/openai_protocol.py` at line 1382, Update the tags field
declaration to use the modern list[str] annotation instead of List[str], and add
a clear description argument to Field so the generated RL endpoint request
schema documents the field’s purpose.

Source: Coding guidelines

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant