[https://nvbugs/6226016][fix] Avoid trusting request-controlled router tokenizers#16697
[https://nvbugs/6226016][fix] Avoid trusting request-controlled router tokenizers#16697yibinl-nvidia wants to merge 4 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
Walkthrough
ChangesTokenizer routing
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/disaggregated/test_router.py (1)
1787-1805: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExpand
BlockHashMixin._get_tokenizercoverage.test_block_hash_mixin_routes_through_transformers_tokenizeronly exercisesTransformersTokenizer.from_pretrained; addload_custom_tokenizercases for default trust, explicittokenizer_dir, andtokenizer_dir=""so the remote-code/path behavior stays covered.Coverage summary: modified
test_block_hash_mixin_routes_through_transformers_tokenizer; no tests added or removed. Listed intests/integration/test_lists/test-db/l0_a10.ymlandtests/integration/test_lists/test-db/l0_h100.yml; no matchingtests/integration/test_lists/qa/entry found. Verdict: needs follow-up.🤖 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/disaggregated/test_router.py` around lines 1787 - 1805, Expand test_block_hash_mixin_routes_through_transformers_tokenizer to also patch and exercise load_custom_tokenizer through BlockHashMixin._get_tokenizer. Cover default trust behavior, an explicit tokenizer_dir, and tokenizer_dir=""; assert each call receives the expected path and trust_remote_code value, and verify the returned and cached tokenizer objects.Source: 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/router_utils.py`:
- Line 209: Update the tokenizer-path handling around trust_remote_code and
model_path so empty tokenizer_dir overrides are treated consistently: reject
empty overrides or derive trust_remote_code from the same normalized-path
predicate used by model_path. Ensure remote-code trust is enabled only when the
effective tokenizer path is non-empty.
---
Nitpick comments:
In `@tests/unittest/disaggregated/test_router.py`:
- Around line 1787-1805: Expand
test_block_hash_mixin_routes_through_transformers_tokenizer to also patch and
exercise load_custom_tokenizer through BlockHashMixin._get_tokenizer. Cover
default trust behavior, an explicit tokenizer_dir, and tokenizer_dir=""; assert
each call receives the expected path and trust_remote_code value, and verify the
returned and cached tokenizer objects.
🪄 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: Enterprise
Run ID: 815bceca-27c9-4940-b2d2-6307a0c8b547
📒 Files selected for processing (2)
tensorrt_llm/serve/router_utils.pytests/unittest/disaggregated/test_router.py
|
/bot run |
|
PR_Github #60834 [ run ] triggered by Bot. Commit: |
|
PR_Github #60834 [ run ] completed with state
|
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #60985 [ run ] triggered by Bot. Commit: |
|
PR_Github #60985 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61016 [ run ] triggered by Bot. Commit: |
|
PR_Github #61016 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61098 [ run ] triggered by Bot. Commit: |
|
PR_Github #61098 [ run ] completed with state
|
c00e7c9 to
c3cfdf1
Compare
|
/bot run |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tests/unittest/disaggregated/test_router.py`:
- Around line 1787-1805: Add a custom-tokenizer branch to
test_block_hash_mixin_routes_through_transformers_tokenizer that mocks
load_custom_tokenizer and verifies it receives trust_remote_code=False for
request-controlled model paths and True when _init_block_hashing sets
tokenizer_dir, asserting the returned and cached tokenizer values. Register the
expanded test in the appropriate tests/integration/test_lists/test-db or
tests/integration/test_lists/qa test list.
🪄 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: Enterprise
Run ID: a6989a1f-e79d-464a-b80f-721b0787c21a
📒 Files selected for processing (2)
tensorrt_llm/serve/router_utils.pytests/unittest/disaggregated/test_router.py
|
PR_Github #61108 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #61108 [ run ] completed with state
|
|
PR_Github #61138 [ run ] triggered by Bot. Commit: |
| def _get_tokenizer(self, model: str): | ||
| if model not in self._tokenizers: | ||
| model_path = self._tokenizer_dir or model | ||
| trust_remote_code = self._tokenizer_dir is not None |
There was a problem hiding this comment.
Please add some comments here.
9f3f385 to
e9db3e5
Compare
|
/bot run |
|
PR_Github #61335 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #61335 [ run ] completed with state
|
|
PR_Github #61339 [ run ] triggered by Bot. Commit: |
|
PR_Github #61339 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61393 [ run ] triggered by Bot. Commit: |
|
PR_Github #61393 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61423 [ run ] triggered by Bot. Commit: |
|
PR_Github #61423 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61446 [ run ] triggered by Bot. Commit: |
|
PR_Github #61446 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #61448 [ run ] triggered by Bot. Commit: |
|
/bot kill |
|
/bot run --disable-fail-fast |
4 participants
Dev Engineer Review
BlockHashMixin._get_tokenizerby settingtrust_remote_code = (self._tokenizer_dir is not None), aligning tokenizer loading behavior with whether the tokenizer directory is operator-configured vs request-controlled.trust_remote_codeconsistently to both loading paths:load_custom_tokenizer(..., trust_remote_code=trust_remote_code)when using an operator-providedtokenizer_dir.TransformersTokenizer.from_pretrained(..., trust_remote_code=trust_remote_code)for the standard HF loading path.test_block_hash_mixin_routes_through_transformers_tokenizerto verify:tokenizer_dir):TransformersTokenizer.from_pretrained(..., trust_remote_code=False).tokenizer_dirmode:TransformersTokenizer.from_pretrained(..., trust_remote_code=True)._Probe._tokenizers[model]are the underlying raw HF tokenizer object (.tokenizer), not theTransformersTokenizerwrapper, in both modes.QA Engineer Review
test_block_hash_mixin_routes_through_transformers_tokenizer(tests/unittest/disaggregated/test_router.py)trust_remote_codehandling and tokenizer caching/identity (raw HF tokenizer object).tests/integration/test_lists/test-db/via inclusion ofunittest/disaggregated/test_router.pyin:l0_h100.ymll0_a10.ymlDescription
Harden kv-cache-aware router tokenizer loading for the request-controlled
modelpath.Previously, when
tokenizer_dirwas unset,_get_tokenizer()used the request'smodelvalue as the tokenizer source while still allowing trusted remote code. That meant a client-provided model reference could point at a malicious Hugging Face tokenizer repo and cause custom tokenizer Python code to be loaded by the router.This change keeps online standard tokenizer loading available, but only enables
trust_remote_codewhen the tokenizer source comes from an operator-configuredtokenizer_dir. Request-derived model paths now load withtrust_remote_code=Falsefor both custom and Transformers tokenizer paths, while explicitly configured tokenizer directories preserve the existing custom-tokenizer behavior.Test Coverage
python -m pytest tests/unittest/disaggregated/test_router.py -k block_hash_mixin_routes_through_transformers_tokenizerPR 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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.