[None][chore] Align LlmArgs with some Pydantic best practices - #11158
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #34275 [ run ] triggered by Bot. Commit: |
|
PR_Github #34275 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #34290 [ run ] triggered by Bot. Commit: |
|
PR_Github #34290 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #34349 [ run ] triggered by Bot. Commit: |
|
PR_Github #34349 [ run ] completed with state
|
e8e43e0 to
d1fd5a4
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #34371 [ run ] triggered by Bot. Commit: |
|
PR_Github #34371 [ run ] completed with state
|
4a9a133 to
7017bb8
Compare
|
/bot run --disable-fail-fast |
|
@coderabbitai full review |
✅ Actions performedFull review triggered. |
|
PR_Github #34438 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis PR refactors the codebase to adopt Pydantic-based configuration models using StrictBaseModel with strict validation. Core changes include replacing legacy patterns (BaseModel, manual init, from_kwargs, to_dict) with modern Pydantic practices (Field descriptors, model_dump, discriminated unions), converting CpType from IntEnum to StrEnum, and modernizing configuration serialization and construction throughout the LLM API and model layers. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
examples/models/core/mllama/convert_checkpoint.py (1)
1-10:⚠️ Potential issue | 🟡 MinorMissing NVIDIA copyright header.
Per coding guidelines, all TensorRT-LLM source files (including
.pyfiles) should contain an NVIDIA copyright header with the year of the latest meaningful modification. This file is missing the required header.Proposed fix
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import jsontensorrt_llm/_torch/speculative/save_hidden_state.py (2)
1-1:⚠️ Potential issue | 🟠 MajorAdd the NVIDIA copyright header.
This source file is missing the required NVIDIA header with the year of latest meaningful modification. Please add the standard header used elsewhere in the repo (use 2026 for the year).
1-7:⚠️ Potential issue | 🟡 MinorSwitch to namespace import for
_utilsto follow coding guidelines.The import should maintain the module namespace. Please change to
import tensorrt_llm._utils as trt_llm_utilsand update the two call sites (lines 35 and 66).Suggested change
-from tensorrt_llm._utils import local_mpi_rank +import tensorrt_llm._utils- if local_mpi_rank() == 0: + if tensorrt_llm._utils.local_mpi_rank() == 0:- if local_mpi_rank() == 0: + if tensorrt_llm._utils.local_mpi_rank() == 0:
🤖 Fix all issues with AI agents
In `@tensorrt_llm/builder.py`:
- Around line 25-33: The PR imports Pydantic's Field and StrictBaseModel from
llmapi.utils directly which breaks the module-namespace guideline; change the
imports to import the modules instead (import pydantic and import
tensorrt_llm.llmapi.utils as llmapi_utils or similar) and update all usages of
Field and StrictBaseModel in the file to use the module-qualified names (e.g.,
pydantic.Field and llmapi_utils.StrictBaseModel) so references in
functions/classes like any dataclass/model definitions consistently use the
module namespace.
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 949-953: The validator method validate_speculative_model currently
annotates its return type as -> None but returns self; update the annotation to
-> 'EagleDecodingConfig' (matching other validators) so the signature reflects
the actual return value; locate the `@model_validator`(mode="after") def
validate_speculative_model method and change its return type annotation
accordingly.
- Around line 654-663: The type annotation for draft_len_schedule uses Python
3.9+ syntax (dict[int, int]); update it to use typing.Dict[int, int] for Python
3.8 compatibility and add the corresponding import from typing if missing.
Locate the Field declaration for draft_len_schedule in llm_args.py (the
draft_len_schedule: Optional[dict[int, int]] = Field(...) line) and change the
annotation to Optional[Dict[int, int]] and ensure Dict is imported, keeping the
rest of the Field parameters and the validation text (including references to
max_draft_len) unchanged.
In `@tensorrt_llm/llmapi/utils.py`:
- Around line 35-45: The file is missing the required NVIDIA copyright header;
add the standard NVIDIA copyright header (including the year of the latest
meaningful modification) at the very top of the source file before any imports
or definitions so it applies to the whole module; ensure the header appears
above the StrictBaseModel class (which inherits from BaseModel) and other
top-level symbols so the file complies with TensorRT‑LLM source file guidelines.
In `@tensorrt_llm/mapping.py`:
- Around line 70-79: The code currently treats non-str cp_config["cp_type"]
values (like ints/IntEnum) as valid and later causes silent misconfiguration;
update the handling in the cp_config block so that if "cp_type" exists it must
be either a str or an instance of CpType: if it's a str, normalize with
.upper(), validate against CpType.__members__ and convert to
CpType(cp_type_str); if it's already a CpType leave it unchanged; otherwise
raise a ValueError rejecting the invalid type. Ensure you reference and adjust
the existing symbols cp_config, "cp_type", cp_type_str, and CpType and preserve
the final assignment cp_type = cp_config.get("cp_type", CpType.ULYSSES).
🧹 Nitpick comments (10)
CODING_GUIDELINES.md (1)
440-442: Consider mentioningmodel_validate()for dict-to-model conversion.The guidance to avoid
from_dict()is good, but the alternativeMyModel(**kwargs)doesn't provide the same validation behavior for nested models as Pydantic'smodel_validate(). For converting dictionaries (especially those with nested structures),model_validate()is often the better choice.Suggested enhancement
- Avoid defining `from_dict()` / `from_kwargs()` methods - prefer constructing the class directly from arguments. - - Good: `MyModel(**kwargs)` + - Good: `MyModel(**kwargs)` or `MyModel.model_validate(kwargs)` (the latter handles nested model coercion) - Bad: `MyModel.from_dict(kwargs)`, `MyModel.from_kwargs(kwargs)`tests/unittest/trt/quantization/test_quant.py (1)
1-1: Update the copyright year.The copyright header shows 2022-2024, but meaningful modifications are being made in 2026. As per coding guidelines, source files should contain the year of latest meaningful modification.
Suggested fix
-# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.tensorrt_llm/lora_helper.py (1)
18-127: Use namespacedpydantic.Fieldimports.This keeps module namespace per repo convention.
♻️ Suggested change
-from pydantic import Field +import pydantic @@ -class LoraConfig(StrictBaseModel): - lora_dir: List[str] = Field( +class LoraConfig(StrictBaseModel): + lora_dir: List[str] = pydantic.Field( @@ - lora_ckpt_source: Literal["hf", "nemo"] = Field( + lora_ckpt_source: Literal["hf", "nemo"] = pydantic.Field( @@ - max_lora_rank: int = Field( + max_lora_rank: int = pydantic.Field( @@ - lora_target_modules: List[str] = Field( + lora_target_modules: List[str] = pydantic.Field( @@ - trtllm_modules_to_hf_modules: Dict[str, str] = Field( + trtllm_modules_to_hf_modules: Dict[str, str] = pydantic.Field( @@ - max_loras: Optional[int] = Field( + max_loras: Optional[int] = pydantic.Field( @@ - max_cpu_loras: Optional[int] = Field( + max_cpu_loras: Optional[int] = pydantic.Field( @@ - swap_gate_up_proj_lora_b_weight: bool = Field( + swap_gate_up_proj_lora_b_weight: bool = pydantic.Field(As per coding guidelines, Always maintain the namespace when importing Python modules, even if only one class or function from a module is used.
tensorrt_llm/mapping.py (1)
18-33: Use namespacedstrenum.StrEnumimport.Keeps module namespace consistent with repo conventions.
♻️ Suggested change
-from strenum import StrEnum +import strenum @@ -class CpType(StrEnum): +class CpType(strenum.StrEnum):As per coding guidelines, Always maintain the namespace when importing Python modules, even if only one class or function from a module is used.
tests/unittest/llmapi/test_llm_args.py (3)
896-896: Unused loop variablefield_name.The loop variable
field_nameis not used within the loop body. Consider renaming it to_field_nameto indicate it's intentionally unused.♻️ Suggested fix
# Visit all field types - for field_name, field_info in model_cls.model_fields.items(): + for _field_name, field_info in model_cls.model_fields.items(): annotation = field_info.annotation
928-951: Consider annotating mutable class attributes withClassVar.The
_COMPATIBILITY_EXEMPT_FIELDSdictionary is a mutable class attribute that should be annotated withtyping.ClassVarto indicate it's shared across instances and not an instance attribute.♻️ Suggested fix
+from typing import ClassVar class TestPydanticBestPractices: ... # Fields exempt from Pydantic compatibility checks due to typing limitations or other edge cases. - _COMPATIBILITY_EXEMPT_FIELDS: dict[type, list[str]] = { + _COMPATIBILITY_EXEMPT_FIELDS: ClassVar[dict[type, list[str]]] = {Apply the same pattern to
_FORBIDDEN_METHODSand_FORBIDDEN_METHODS_EXEMPT_CLASSES.
1054-1056: Remove extraneousfprefix from strings without placeholders.Several f-strings in this test class don't contain any placeholders. Remove the
fprefix for cleaner code.♻️ Suggested fix
if violations: pytest.fail( - f"The following fields are missing descriptions:\n" + + "The following fields are missing descriptions:\n" + "\n".join(violations) + "\n\nPlease add a description to each by using Field(description=\"...\")." )Apply the same fix to lines 1077, 1118, 1141, and 1168.
tensorrt_llm/plugin/plugin.py (1)
1-1: Copyright year may need updating.The copyright header shows "2022-2024" but the current year is 2026. As per coding guidelines, TensorRT-LLM source files should contain the year of latest meaningful modification.
tensorrt_llm/llmapi/llm_args.py (2)
280-286: Side effect in query method may cause unexpected behavior.
needs_separate_short_long_cuda_graphs()mutatesself.seq_len_thresholdas a side effect. This is unusual for a method that reads like a predicate query. Consider initializingseq_len_thresholdin amodel_validatorinstead, or document this behavior clearly.♻️ Suggested refactor using model_validator
+ `@model_validator`(mode='after') + def _set_seq_len_threshold(self) -> 'DeepSeekSparseAttentionConfig': + if self.skip_indexer_for_short_seqs: + self.seq_len_threshold = self.index_topk + return self + def needs_separate_short_long_cuda_graphs(self) -> bool: """ Whether to capture separate CUDA graphs for short and long sequences. Use seq_len_threshold to determine the threshold for separating short and long sequences. """ - self.seq_len_threshold = self.index_topk return self.skip_indexer_for_short_seqs
3144-3148: Consider usingValueErrorinstead ofassertfor validation.Using
assertfor configuration validation can be disabled withpython -O. For consistent error handling with other validators in this file, consider usingraise ValueError(...).♻️ Suggested refactor
if cp_tokens_per_block is not None: kv_tokens_per_block = self.kv_cache_config.tokens_per_block - assert cp_tokens_per_block == kv_tokens_per_block, ( - f"When HELIX parallelism is active, cp_config.tokens_per_block ({cp_tokens_per_block}) " - f"must match kv_cache_config.tokens_per_block ({kv_tokens_per_block})." - ) + if cp_tokens_per_block != kv_tokens_per_block: + raise ValueError( + f"When HELIX parallelism is active, cp_config.tokens_per_block ({cp_tokens_per_block}) " + f"must match kv_cache_config.tokens_per_block ({kv_tokens_per_block})." + )
|
/bot kill |
|
PR_Github #34447 [ kill ] triggered by Bot. Commit: |
|
PR_Github #34447 [ kill ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #34452 [ run ] triggered by Bot. Commit: |
|
PR_Github #34452 [ run ] completed with state
|
f88596d to
a31bc23
Compare
|
/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #36311 [ run ] triggered by Bot. Commit: |
|
PR_Github #36311 [ run ] completed with state
|
d8a056e to
a86f080
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #36358 [ run ] triggered by Bot. Commit: |
|
PR_Github #36358 [ run ] completed with state
|
a86f080 to
2f308c7
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #36370 [ run ] triggered by Bot. Commit: |
|
PR_Github #36370 [ run ] completed with state
|
Signed-off-by: Anish Shanbhag <ashanbhag@nvidia.com>
Signed-off-by: Anish Shanbhag <ashanbhag@nvidia.com>
Signed-off-by: Anish Shanbhag <ashanbhag@nvidia.com>
Signed-off-by: Anish Shanbhag <ashanbhag@nvidia.com>
Signed-off-by: Anish Shanbhag <ashanbhag@nvidia.com>
2f308c7 to
06a773c
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #36491 [ run ] triggered by Bot. Commit: |
|
PR_Github #36491 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #36546 [ run ] triggered by Bot. Commit: |
|
PR_Github #36546 [ run ] completed with state |
…#11158) Signed-off-by: Anish Shanbhag <ashanbhag@nvidia.com>
11 participants
Motivation
LlmArgsis the main user-facing interface for configuring TRTLLM, and it uses Pydantic for validating custom user configuration, including YAML files passed with--config/--extra_llm_api_options. It's important thatLlmArgsfollows Pydantic best practices so that (a) it's easier for users to understand how to configure the LLM API, and (b) we reduce the surface area for subtle bugs related to configuration parsing / validation.This PR is also a preparatory change for #8331 which will significantly simplify the parsing of
LlmArgsfrom config YAML files. The changes here would enable a variety of custom parsing logic (e.g.update_llm_args_with_extra_dict) to be completely removed.Description
This PR aims to align
LlmArgswith some Pydantic best practices that have been added toCODING_GUIDELINES.md.The diff is unfortunately pretty large due to the interdependent nature of the changes here. To help with this, I recommend that reviewers read the changes in this order:
CODING_GUIDELINES.md. Some highlights include:model_validator/model_dumpetc. instead of using custom methods likevalidate()orto_dict()test_llm_args.py::TestPydanticBestPractices, which enforce a subset of the guidelines above.LlmArgsinto compliance with the above guidelines. Some highlights include:QuantConfigto a Pydantic modelPositiveInt/NonNegativeInt/ etc. instead of custom validatorsThe intention here is a pure refactoring PR that contains minimal to no behavior changes or user-facing breaking changes.
I've also added PR comments to many of the changes here to explain the motivation behind them.
At the time of writing this, all but 3 single GPU tests are passing (https://prod.blsm.nvidia.com/sw-tensorrt-top-1/blue/organizations/jenkins/LLM%2Fmain%2FL0_Test-x86_64-Single-GPU/detail/L0_Test-x86_64-Single-GPU/16435/pipeline) and ARM multi-GPU tests are 100% passing (https://prod.blsm.nvidia.com/sw-tensorrt-top-1/blue/organizations/jenkins/LLM%2Fmain%2FL0_Test-SBSA-Multi-GPU/detail/L0_Test-SBSA-Multi-GPU/6875/pipeline). Currently verifying fixes for the remaining tests.
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)
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
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]to print this help message.See details below for each supported subcommand.
Details
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--disable-fail-fast(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--test-backend "pytorch, cpp"(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.--only-multi-gpu-test(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test(OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.--post-merge(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-listparameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip testing for latest commit on pull request.
--comment "Reason for skipping build/test"is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipelineReuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.