[TRTLLM-14831][chore] Consolidate package bootstrap and relocate restricted deserialization - #17254
[TRTLLM-14831][chore] Consolidate package bootstrap and relocate restricted deserialization#17254YihuiLu512 wants to merge 2 commits into
Conversation
…ricted deserialization Reduce the package root's mechanics to a single private bootstrap module, and move the restricted unpickler into the package that consumes it. Each change is described below together with the reasoning it rests on and what it changes for callers. 1) tensorrt_llm/_bootstrap.py (new), tensorrt_llm/__init__.py _common.py's _init() and the three setup functions that lived in __init__.py -- DLL directory, Python-library preload, vendored-triton_kernels precedence -- merge into one private bootstrap module. __init__.py drops from 176 to 103 lines and is now only the package's compatibility surface: the intentional re-exports and __all__. The two existing phases keep their order, which is the constraint the file is arranged around: _prepare_environment() runs before `import torch`, _init() runs after the package's own imports, and its early-return guard is unchanged. That ordering forces the one change here that is not a pure move. _bootstrap.py is imported ahead of torch, so its module scope is restricted to the standard library; a module-level `import torch` or `from .bindings import MpiComm` would load Torch and the compiled extension before phase 1 had prepared for them. Phase 2's four imports -- torch, tensorrt_llm.bindings, ._utils and .logger -- therefore move into _init()'s body. This defers statements, not any module's first import: _init() is the last statement of __init__.py, and all four are already imported earlier on the `import tensorrt_llm._torch.models` path. Tracing module completion order with -X importtime over `import tensorrt_llm` confirms it: 2393 modules on both sides, and the whole difference is +tensorrt_llm._bootstrap at the front and -tensorrt_llm._common where it used to sit. _common.py is absent from legacy-files.txt while __init__.py is on it, so the three relocated functions land in ruff's scope instead of yapf's and are reformatted accordingly. That is formatting only and confined to those three functions; _init() was already ruff-clean and moves unchanged. 2) tensorrt_llm/_common.py becomes a shim The old path keeps a definition-free shim that re-exports _init and warns on import. The module is private, and every source of evidence that can be checked inside this repository is clean -- but one cannot be checked here, namely consumers outside the repository, and only the module's owner can answer for those. That question was never asked, so removing the path was never certifiable. Ten lines against a breakage nobody here can observe is the same asymmetry that keeps a shim on every public path in this migration, so the shim goes in rather than the question going to an owner. The export set is recorded rather than computed: both names in the old module are underscored, so a private module has no computable public surface, and the shim re-exports the one symbol this repository itself consumed. _inited is deliberately left out -- it is mutable module state, and a re-export would bind the value once and then diverge from the guard it appears to mirror. What the shim preserves is one function object, not a copy: `_common._init is _bootstrap._init is tensorrt_llm._init`. A copy would give the guard two _inited flags, and a call through the old path would run phase 2 -- custom op registration and MPI init -- a second time, which it is not designed for; calling _common._init() again is verified to be a no-op. `import tensorrt_llm` does not pull the shim in, since __init__.py now imports _bootstrap, so the DeprecationWarning still means "you named the old path" rather than firing for every user. 3) tensorrt_llm/serialization.py moves to tensorrt_llm/llmapi/serialization.py The restricted unpickler moves into the package that consumes it, and both consumers at HEAD are rewritten to the canonical path: llmapi/rlhf_utils.py, the only production consumer, and tests/unittest/llmapi/test_serialization.py. A stale reference in a model_engine.py comment is retargeted as well. Nothing in the tree goes through the shim -- in-repo references, lint lists, ownership rules and derived configs are all rewritten as if the old path had disappeared. The file gains the NVIDIA copyright header it never had, and the four allowlist module strings travel with it unchanged; later tasks of this migration update their contents at the new location. Ownership moves with the code. That is the one effective-CODEOWNERS change in this commit and it is intended: serialization.py inherited the package-root default and now falls under llmapi/'s rule, whose team owns its only consumer and its only test. No file loses an owner. The legacy-files.txt entry likewise follows the implementation to the new path, so the module stays in yapf's scope and the entry count is unchanged at 765. 4) tensorrt_llm/serialization.py becomes a shim The old path keeps a definition-free shim over all seven public names, warning on import. serialization has no leading underscore, so it is a public module path and keeps a shim regardless of what the evidence says; what the measurement earns is a shorter deprecation window, and it is what fixes the export set -- with no evidence source naming a symbol, the supported surface falls back to the module's own public surface: BASE_EXAMPLE_CLASSES, Unpickler, dump, dumps, load, loads and register_approved_class. A shim, and only a shim, preserves `from tensorrt_llm import serialization`. That form works today only because serialization is a submodule of the package; __init__.py never imports it, so the top-level re-export that rescues other modules' symbols cannot help here -- what is imported is the module itself, and that goes through the import finder. Because the shim re-exports objects rather than copies, isinstance and pickle keep working: every name satisfies `tensorrt_llm.serialization.X is tensorrt_llm.llmapi.serialization.X`, seven of seven, nothing is defined at the old path and no __module__ is rewritten. The two module objects are necessarily distinct -- a re-export shim shares symbols, not identity -- so per-symbol identity is the assertion that holds, and it is the property pickle and isinstance actually depend on. The security boundary forwards too, not just the names: a payload dumped through either path loads through the other, and an unapproved class is still rejected when loaded through the shim. A shim that forwarded only the success path would quietly widen the boundary while every round-trip test passed. 5) .pre-commit-config.yaml, pyproject.toml, ruff-legacy.toml Regenerated from legacy-files.txt, so no derived config still names the retired path; check-configs passes and the ruff baseline needs no key renames. Both shims are new files and therefore not on the legacy list, which puts them in ruff's scope; their formatting is already what ruff produces, so the hooks leave them alone. Left alone deliberately version.py stays at the package root: setup.py, lock-file generation and the packaging metadata read that exact source path. _deprecation.py is confirmed absent at HEAD with no remaining reference. Nothing outside the two deprecation warnings changes for callers: the top-level __all__ is byte-identical and hashes identically at runtime, the api-stability reference diff is empty, the golden manifest's sha256 is unchanged, no test moves so node IDs, test lists and waives.txt are untouched, and the import graph keeps its five module-level cycles with no new cross-package edge and neither shim joining one. Against a build of the base commit with this change overlaid, test_serialization.py passes 6 before and after and pre-migration pickles replay 10 of 10. The eleventh class in that compatibility matrix, GreedyDecodingParams, does not exist at this commit -- the name appears nowhere in the repository except as a string in the allowlist. Left as-is and recorded for the task that updates it. Signed-off-by: Yihui Lu <269394165+YihuiLu512@users.noreply.github.com>
|
/bot run |
|
PR_Github #63761 [ run ] triggered by Bot. Commit: |
WalkthroughThe change moves package initialization into a two-phase bootstrap module and relocates restricted pickle serialization under ChangesPackage bootstrap initialization
Serialization module relocation
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant PythonImport
participant tensorrt_llm.__init__
participant _bootstrap
participant torch
participant NativeRuntime
PythonImport->>tensorrt_llm.__init__: import package
tensorrt_llm.__init__->>_bootstrap: _prepare_environment()
_bootstrap->>_bootstrap: configure DLLs, Python library, and vendored triton_kernels
tensorrt_llm.__init__->>torch: import torch
tensorrt_llm.__init__->>_bootstrap: _init()
_bootstrap->>NativeRuntime: load custom libraries and initialize MPI
🚥 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: 4
🧹 Nitpick comments (8)
tensorrt_llm/_bootstrap.py (5)
163-174: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate
_print_stacksand move the period read above the closure.
_print_stackshas no return annotation. It also readsprint_stacks_periodfrom the enclosing scope, and that name is bound at line 171, after the definition. The current order works because the thread starts at line 174. Read the environment variable before the definition, or pass the period as an argument, so the closure does not depend on statement order.♻️ Proposed reorder
- def _print_stacks(): + print_stacks_period = int(os.getenv("TRTLLM_PRINT_STACKS_PERIOD", "-1")) + + def _print_stacks() -> None: counter = 0 while True: time.sleep(print_stacks_period) counter += 1 logger.error(f"Printing stacks {counter} times") print_all_stacks() - print_stacks_period = int(os.getenv("TRTLLM_PRINT_STACKS_PERIOD", "-1")) if print_stacks_period > 0:As per coding guidelines: "Annotate every function, use
Nonefor non-returning functions".🤖 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/_bootstrap.py` around lines 163 - 174, Update the stack-printing setup around `_print_stacks` by annotating the function with a `None` return type and reading `TRTLLM_PRINT_STACKS_PERIOD` before the closure is defined, or pass the period explicitly as an argument. Preserve the existing thread-start behavior and positive-period guard.Source: Coding guidelines
120-134: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winType
log_levelprecisely.
objectaccepts any value and gives no useful checking. The value is forwarded tologger.set_level(), which indexesseverity_mapwith it, so an invalid value fails inside the logger instead of at the call site. Use the severity type thatlogger.set_levelaccepts, for examplestr | None.♻️ Proposed signature
-def _init(log_level: object = None) -> None: +def _init(log_level: str | None = None) -> None:As per coding guidelines: "avoid
Anyand unnecessary type ignores, ... useLiteral,overload,TypeVar, orProtocolwhen appropriate".🤖 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/_bootstrap.py` around lines 120 - 134, Update the _init function’s log_level annotation from object to the precise severity type accepted by logger.set_level, such as str | None, while preserving the existing optional forwarding behavior and logger.set_level call.Source: Coding guidelines
48-49: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd return annotations to the phase-1 helpers.
_add_trt_llm_dll_directory,_preload_python_liband_setup_vendored_triton_kernelshave no return annotation._prepare_environmentand_initin the same file already use-> None. Annotate these three the same way for consistency.♻️ Proposed annotations
-def _add_trt_llm_dll_directory(): +def _add_trt_llm_dll_directory() -> None:-def _preload_python_lib(): +def _preload_python_lib() -> None:-def _setup_vendored_triton_kernels(): +def _setup_vendored_triton_kernels() -> None:As per coding guidelines: "Annotate every function, use
Nonefor non-returning functions".Also applies to: 56-57, 81-82
🤖 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/_bootstrap.py` around lines 48 - 49, Add -> None return annotations to the phase-1 helper functions _add_trt_llm_dll_directory, _preload_python_lib, and _setup_vendored_triton_kernels, matching the existing annotations on _prepare_environment and _init.Source: Coding guidelines
103-110: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winRestore
sys.pathwithtry/finally.If
import triton_kernelsraises, line 110 never runs. The package root then stays at the front ofsys.pathfor the rest of the process, and it shadows later imports of any top-level module that also exists under the repository root. The priorsys.modulesdeletion also leaves no cachedtriton_kernels, so the state after the failure is worse than before the call.♻️ Proposed fix
should_add_to_path = str(root) not in sys.path if should_add_to_path: sys.path.insert(0, str(root)) - import triton_kernels # noqa: F401 - - if should_add_to_path: - sys.path.remove(str(root)) + try: + import triton_kernels # noqa: F401 + finally: + if should_add_to_path: + sys.path.remove(str(root))🤖 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/_bootstrap.py` around lines 103 - 110, Wrap the temporary path modification and import in a try/finally block around the existing triton_kernels import, ensuring the inserted root is removed whenever should_add_to_path is true even if import triton_kernels raises. Preserve the current path ordering and avoid changing the import behavior.
149-159: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winNarrow exception types and chain the original exception.
The exception handler at these lines has two issues:
raise ImportError(str(e) + msg)drops the original traceback. Usefrom eto chain exceptions.except Exceptioncatches failures fromtorch.classes.load_library(), from importing._torch.custom_ops, and from calling_register_fake(). All three failure modes report the same error message, but they have different causes. Per the coding guidelines, catch the narrowest possible exception type.Separate the
torch.classes.load_library()call into its own try-except block.torch.classes.load_library()raisesOSErrorwhen the shared library fails to load (it usesctypes.CDLLinternally). The import and_register_fake()call can raiseImportError. Addfrom eto both handlers:Proposed fix
try: torch.classes.load_library(ft_decoder_lib) + except OSError as e: + msg = ( + "\nFATAL: Decoding operators failed to load. This may be caused by an incompatibility " + "between PyTorch and TensorRT-LLM. Please rebuild and install TensorRT-LLM." + ) + raise ImportError(str(e) + msg) from e + + try: from ._torch.custom_ops import _register_fake _register_fake() - except Exception as e: + except ImportError as e: msg = ( "\nFATAL: Decoding operators failed to load. This may be caused by an incompatibility " "between PyTorch and TensorRT-LLM. Please rebuild and install TensorRT-LLM." ) - raise ImportError(str(e) + msg) + raise ImportError(str(e) + msg) from e🤖 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/_bootstrap.py` around lines 149 - 159, Split the initialization flow around torch.classes.load_library, ._torch.custom_ops import, and _register_fake into separate handlers: catch OSError only for the shared-library load and ImportError for the import/registration path. Preserve the existing fatal message, and chain the original exception with from e in both ImportError raises.Sources: Coding guidelines, Linters/SAST tools
tensorrt_llm/_common.py (1)
15-28: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReduce the dependence on unresolvable internal references.
The docstring cites "Epic §0.4 and §6.1" and "Epic decision D4 (b)". A reader with only the repository cannot resolve those sections. State the two constraints inline, for example: the shim must not define new objects, and it must not rewrite
__module__, because both break object identity for pre-migration pickles. Keep the ticket IDs as pointers.🤖 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/_common.py` around lines 15 - 28, Update the module docstring in the compatibility shim to replace the unresolved Epic section and decision references with inline explanations that the shim must not define new objects or rewrite __module__, because either breaks object identity for pre-migration pickles. Preserve the existing ticket IDs as repository pointers and retain the explicit-list/no-import-* constraint.tests/unittest/llmapi/test_serialization.py (1)
3-3: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a regression test for the legacy
tensorrt_llm.serializationshim.Extend the registered
tests/unittest/llmapi/test_serialization.pytest. Assert identity for all seven re-exported names and assert that importing the shim emitsDeprecationWarning. Existing tests cover canonical serialization only; compatibility coverage is insufficient.🤖 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_serialization.py` at line 3, Add a regression test in the existing serialization test module for importing the legacy tensorrt_llm.serialization shim. Verify that all seven re-exported names are identical to their canonical tensorrt_llm.llmapi.serialization counterparts, and capture the import to assert it emits DeprecationWarning while preserving the existing canonical serialization tests.Source: Path instructions
tensorrt_llm/llmapi/serialization.py (1)
16-21: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDeclare the canonical serializer exports.
This module has no
__all__. A wildcard import also exportsio,pickle,re, andpartial. Declare the intended serializer names and verify the API-stability fixture after the change.Proposed export declaration
+__all__ = ( + "BASE_EXAMPLE_CLASSES", + "register_approved_class", + "Unpickler", + "dump", + "dumps", + "load", + "loads", +)As per coding guidelines, “keep
__all__updated for public interfaces.”🤖 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/llmapi/serialization.py` around lines 16 - 21, Add an __all__ declaration in serialization.py listing only the module’s intended public serializer symbols, excluding imported helpers such as io, pickle, re, and partial. Verify the API-stability fixture passes with the canonical export set.Source: Coding guidelines
🤖 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/_common.py`:
- Line 1: Restore the copyright header year range at the file header to
2022-2026, preserving the original start year for this modified file.
In `@tensorrt_llm/llmapi/serialization.py`:
- Line 150: Update the serialization registry handling around the initializer
and the load()/loads() methods to use None as the omitted-registry sentinel,
resolving it to the shared baseline/registered approved-class mapping while
preserving an explicitly supplied empty mapping as default-deny. Ensure
register_approved_class() updates the registry used by default loads, and add
regression coverage verifying a registered class can round-trip through dumps()
and default loads()/loads().
- Around line 164-167: Update the approved-class resolution in the serializer’s
find_class method to replace broad regex matching such as ^torch.* with an exact
allowlist of permitted (module, name) pairs. Ensure serialization.loads can
resolve only explicitly approved torch symbols, rejecting torch.hub.load and all
other unlisted callables before deserialization proceeds, while preserving
approved non-torch module handling as appropriate.
- Line 127: Add type annotations to all six functions and methods lacking them:
_register_class (parameters dict and obj, return type), register_approved_class
(parameter obj, return type), Unpickler.__init__ (all parameters and return
type), Unpickler.find_class (parameters module and name, return type), load (all
parameters including file and return type), and loads (all parameters including
s and return type). Use dict[str, list[str]] for allowlist dictionaries,
list[str] | None for pattern lists, BinaryIO for file-like input streams, and
Any for pickle return values. Ensure Python 3.10+ syntax is used throughout to
comply with the module's strict mypy checking requirement.
---
Nitpick comments:
In `@tensorrt_llm/_bootstrap.py`:
- Around line 163-174: Update the stack-printing setup around `_print_stacks` by
annotating the function with a `None` return type and reading
`TRTLLM_PRINT_STACKS_PERIOD` before the closure is defined, or pass the period
explicitly as an argument. Preserve the existing thread-start behavior and
positive-period guard.
- Around line 120-134: Update the _init function’s log_level annotation from
object to the precise severity type accepted by logger.set_level, such as str |
None, while preserving the existing optional forwarding behavior and
logger.set_level call.
- Around line 48-49: Add -> None return annotations to the phase-1 helper
functions _add_trt_llm_dll_directory, _preload_python_lib, and
_setup_vendored_triton_kernels, matching the existing annotations on
_prepare_environment and _init.
- Around line 103-110: Wrap the temporary path modification and import in a
try/finally block around the existing triton_kernels import, ensuring the
inserted root is removed whenever should_add_to_path is true even if import
triton_kernels raises. Preserve the current path ordering and avoid changing the
import behavior.
- Around line 149-159: Split the initialization flow around
torch.classes.load_library, ._torch.custom_ops import, and _register_fake into
separate handlers: catch OSError only for the shared-library load and
ImportError for the import/registration path. Preserve the existing fatal
message, and chain the original exception with from e in both ImportError
raises.
In `@tensorrt_llm/_common.py`:
- Around line 15-28: Update the module docstring in the compatibility shim to
replace the unresolved Epic section and decision references with inline
explanations that the shim must not define new objects or rewrite __module__,
because either breaks object identity for pre-migration pickles. Preserve the
existing ticket IDs as repository pointers and retain the
explicit-list/no-import-* constraint.
In `@tensorrt_llm/llmapi/serialization.py`:
- Around line 16-21: Add an __all__ declaration in serialization.py listing only
the module’s intended public serializer symbols, excluding imported helpers such
as io, pickle, re, and partial. Verify the API-stability fixture passes with the
canonical export set.
In `@tests/unittest/llmapi/test_serialization.py`:
- Line 3: Add a regression test in the existing serialization test module for
importing the legacy tensorrt_llm.serialization shim. Verify that all seven
re-exported names are identical to their canonical
tensorrt_llm.llmapi.serialization counterparts, and capture the import to assert
it emits DeprecationWarning while preserving the existing canonical
serialization tests.
🪄 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: f9601d36-7f07-4fe0-8431-716416ab78a8
📒 Files selected for processing (12)
.pre-commit-config.yamllegacy-files.txtpyproject.tomlruff-legacy.tomltensorrt_llm/__init__.pytensorrt_llm/_bootstrap.pytensorrt_llm/_common.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/llmapi/rlhf_utils.pytensorrt_llm/llmapi/serialization.pytensorrt_llm/serialization.pytests/unittest/llmapi/test_serialization.py
| } | ||
|
|
||
|
|
||
| def _register_class(dict, obj): |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
# First, locate and examine the serialization.py file
fd -t f "serialization.py" tensorrt_llm/llmapi/Repository: NVIDIA/TensorRT-LLM
Length of output: 194
🏁 Script executed:
# Get an outline of the file to understand its structure
ast-grep outline tensorrt_llm/llmapi/serialization.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 550
🏁 Script executed:
# Read the specific line ranges mentioned in the review comment
sed -n '120,150p' tensorrt_llm/llmapi/serialization.py | cat -n
sed -n '155,160p' tensorrt_llm/llmapi/serialization.py | cat -n
sed -n '185,215p' tensorrt_llm/llmapi/serialization.py | cat -nRepository: NVIDIA/TensorRT-LLM
Length of output: 2687
🏁 Script executed:
# Check for mypy configuration in the repository
fd -t f "pyproject.toml|setup.cfg|mypy.ini|.mypy.ini" | head -10Repository: NVIDIA/TensorRT-LLM
Length of output: 617
🏁 Script executed:
# Check if CODING_GUIDELINES.md exists and contains annotation requirements
find . -name "CODING_GUIDELINES.md" -type fRepository: NVIDIA/TensorRT-LLM
Length of output: 180
🏁 Script executed:
# Check mypy configuration in pyproject.toml
sed -n '/\[tool.mypy\]/,/^\[/p' pyproject.toml | head -50Repository: NVIDIA/TensorRT-LLM
Length of output: 550
🏁 Script executed:
# Read the entire serialization.py file to assess annotation status
cat -n tensorrt_llm/llmapi/serialization.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 11911
🏁 Script executed:
# Check CODING_GUIDELINES.md for annotation requirements
grep -A 5 -B 2 "Annotate every function" CODING_GUIDELINES.mdRepository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
# Read CODING_GUIDELINES.md to find annotation guidance
cat CODING_GUIDELINES.md | grep -A 5 -B 2 "Annotate\|type annotation"Repository: NVIDIA/TensorRT-LLM
Length of output: 706
🏁 Script executed:
# Check if serialization.py is included in mypy overrides (could be excluded)
sed -n '/\[\[tool.mypy.overrides\]\]/,/^\[\[/p' pyproject.tomlRepository: NVIDIA/TensorRT-LLM
Length of output: 1215
🏁 Script executed:
# Verify that serialization.py is not in any mypy exclude patterns
grep -r "serialization" pyproject.tomlRepository: NVIDIA/TensorRT-LLM
Length of output: 252
🏁 Script executed:
# Check if there's any type checking in place currently by checking imports in serialization.py
head -25 tensorrt_llm/llmapi/serialization.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1346
🏁 Script executed:
# Get the full mypy configuration for llmapi module if any
grep -A 20 "llmapi" pyproject.tomlRepository: NVIDIA/TensorRT-LLM
Length of output: 8581
Add required type annotations to all public and module-level functions in serialization.py.
Functions _register_class(), register_approved_class(), load(), and loads() have no parameter or return type annotations. The Unpickler.__init__() and find_class() methods also lack annotations. Provide precise Python 3.10+ types for all parameters and return values: use dict[str, list[str]] for the allowlist dictionaries, list[str] | None for pattern lists, BinaryIO for file-like input streams, and Any for pickle return values. The module is subject to global strict mypy checking and must comply with the coding guideline requirement to "annotate every function."
Functions requiring annotation:
- Line 127:
_register_class(dict, obj) - Line 137:
register_approved_class(obj) - Lines 143–148:
Unpickler.__init__() - Line 155:
Unpickler.find_class(module, name) - Lines 185–193:
load(file, ...) - Lines 204–213:
loads(s, ...)
🧰 Tools
🪛 Ruff (0.16.0)
[error] 127-127: Function argument dict is shadowing a Python builtin
(A002)
🤖 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/llmapi/serialization.py` at line 127, Add type annotations to
all six functions and methods lacking them: _register_class (parameters dict and
obj, return type), register_approved_class (parameter obj, return type),
Unpickler.__init__ (all parameters and return type), Unpickler.find_class
(parameters module and name, return type), load (all parameters including file
and return type), and loads (all parameters including s and return type). Use
dict[str, list[str]] for allowlist dictionaries, list[str] | None for pattern
lists, BinaryIO for file-like input streams, and Any for pickle return values.
Ensure Python 3.10+ syntax is used throughout to comply with the module's strict
mypy checking requirement.
Source: Coding guidelines
|
PR_Github #63761 [ run ] completed with state
|
|
/bot help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse 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. |
|
/bot -h run |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse 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. |
|
/bot run --disable-fail-fast |
|
PR_Github #63872 [ run ] triggered by Bot. Commit: |
|
PR_Github #63872 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63888 [ run ] triggered by Bot. Commit: |
|
PR_Github #63888 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63965 [ run ] triggered by Bot. Commit: |
|
PR_Github #63965 [ run ] completed with state
|
|
/bot run |
|
PR_Github #64051 [ run ] triggered by Bot. Commit: |
|
PR_Github #64051 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64166 [ run ] triggered by Bot. Commit: |
|
PR_Github #64166 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64351 [ run ] triggered by Bot. Commit: |
|
PR_Github #64351 [ run ] completed with state |
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
tensorrt_llm/llmapi/serialization.py (1)
137-155: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the public restricted-serialization API.
register_approved_class,Unpickler,load, andloadshave no docstrings. Add Google-style docstrings that describe the allowlist parameters and the default-deny behavior. Document thatapproved_module_patternsauthorizes all names in a matched module.As per coding guidelines, use docstrings rather than comments for externally usable interfaces, and use Google-style docstrings for classes and functions.
Also applies to: 185-224
🤖 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/llmapi/serialization.py` around lines 137 - 155, Add Google-style docstrings to the public APIs register_approved_class, Unpickler, load, and loads, covering their allowlist parameters, default-deny behavior, and the security boundary enforced by Unpickler.find_class. Explicitly state that approved_module_patterns authorizes every name in a matched module, and replace interface comments with docstrings where applicable.Source: Coding guidelines
tensorrt_llm/_bootstrap.py (1)
48-56: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd
Nonereturn annotations to procedures.
_add_trt_llm_dll_directory,_preload_python_lib,_setup_vendored_triton_kernels, and_print_stacksdo not return values. Add-> Noneto each definition.As per coding guidelines, “Annotate every function, use
Nonefor procedures.”Also applies to: 81-113, 163-170
🤖 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/_bootstrap.py` around lines 48 - 56, Add -> None return annotations to the procedure definitions _add_trt_llm_dll_directory, _preload_python_lib, _setup_vendored_triton_kernels, and _print_stacks, preserving their existing implementations.Source: Coding guidelines
🤖 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/_bootstrap.py`:
- Around line 120-125: Update _init so _inited is not set until Torch import,
custom-op loading, fake registration, and MPI initialization all complete
successfully; ensure a failed setup does not cause later calls to return as
initialized, either by leaving _inited false or preserving and re-raising the
initialization failure.
- Around line 103-110: Update the import setup around triton_kernels so the
vendored root is always moved to index zero, even when it already exists later
in sys.path. Save its original position, remove the existing entry before
insertion, and wrap the import in a finally block that restores the prior
sys.path state.
- Around line 149-159: In the library-loading block, keep only
torch.classes.load_library(ft_decoder_lib) inside the try, catch OSError and
RuntimeError, and raise ImportError with the existing context using raise ...
from exc to preserve the cause. Move the _register_fake import and invocation
outside the try/except so registration failures propagate independently.
In `@tensorrt_llm/_common.py`:
- Around line 24-25: Correct the removal-schedule text by completing the
sentence about removal being tracked by ticket T25 (TRTLLM-14855), then state
the at-least-one-release compatibility window in a separate sentence. Do not add
definitions or otherwise rewrite the surrounding content.
---
Nitpick comments:
In `@tensorrt_llm/_bootstrap.py`:
- Around line 48-56: Add -> None return annotations to the procedure definitions
_add_trt_llm_dll_directory, _preload_python_lib, _setup_vendored_triton_kernels,
and _print_stacks, preserving their existing implementations.
In `@tensorrt_llm/llmapi/serialization.py`:
- Around line 137-155: Add Google-style docstrings to the public APIs
register_approved_class, Unpickler, load, and loads, covering their allowlist
parameters, default-deny behavior, and the security boundary enforced by
Unpickler.find_class. Explicitly state that approved_module_patterns authorizes
every name in a matched module, and replace interface comments with docstrings
where applicable.
🪄 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: 416854da-c26e-4fd7-84f5-5e50763fad18
📒 Files selected for processing (12)
.pre-commit-config.yamllegacy-files.txtpyproject.tomlruff-legacy.tomltensorrt_llm/__init__.pytensorrt_llm/_bootstrap.pytensorrt_llm/_common.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/llmapi/rlhf_utils.pytensorrt_llm/llmapi/serialization.pytensorrt_llm/serialization.pytests/unittest/llmapi/test_serialization.py
🚧 Files skipped from review as they are similar to previous changes (8)
- tensorrt_llm/llmapi/rlhf_utils.py
- legacy-files.txt
- tests/unittest/llmapi/test_serialization.py
- .pre-commit-config.yaml
- tensorrt_llm/init.py
- ruff-legacy.toml
- pyproject.toml
- tensorrt_llm/serialization.py
| should_add_to_path = str(root) not in sys.path | ||
| if should_add_to_path: | ||
| sys.path.insert(0, str(root)) | ||
|
|
||
| import triton_kernels # noqa: F401 | ||
|
|
||
| if should_add_to_path: | ||
| sys.path.remove(str(root)) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Place the vendored root first on every import.
If root already occurs later in sys.path, Line 103 skips insertion. Line 107 can then import an external triton_kernels before the vendored package. Remove the existing entry, insert it at index zero, and restore its prior position in a finally block.
Proposed fix
- should_add_to_path = str(root) not in sys.path
- if should_add_to_path:
- sys.path.insert(0, str(root))
-
- import triton_kernels # noqa: F401
-
- if should_add_to_path:
- sys.path.remove(str(root))
+ root_str = str(root)
+ previous_index = sys.path.index(root_str) if root_str in sys.path else None
+ if previous_index is not None:
+ del sys.path[previous_index]
+ sys.path.insert(0, root_str)
+ try:
+ import triton_kernels # noqa: F401
+ finally:
+ sys.path.pop(0)
+ if previous_index is not None:
+ sys.path.insert(previous_index, root_str)📝 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.
| should_add_to_path = str(root) not in sys.path | |
| if should_add_to_path: | |
| sys.path.insert(0, str(root)) | |
| import triton_kernels # noqa: F401 | |
| if should_add_to_path: | |
| sys.path.remove(str(root)) | |
| root_str = str(root) | |
| previous_index = sys.path.index(root_str) if root_str in sys.path else None | |
| if previous_index is not None: | |
| del sys.path[previous_index] | |
| sys.path.insert(0, root_str) | |
| try: | |
| import triton_kernels # noqa: F401 | |
| finally: | |
| sys.path.pop(0) | |
| if previous_index is not None: | |
| sys.path.insert(previous_index, root_str) |
🤖 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/_bootstrap.py` around lines 103 - 110, Update the import setup
around triton_kernels so the vendored root is always moved to index zero, even
when it already exists later in sys.path. Save its original position, remove the
existing entry before insertion, and wrap the import in a finally block that
restores the prior sys.path state.
| def _init(log_level: object = None) -> None: | ||
| """Phase 2: custom-op registration and MPI initialization, after imports.""" | ||
| global _inited | ||
| if _inited: | ||
| return | ||
| _inited = True |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Do not mark initialization complete before fallible setup.
Line 125 runs before the Torch import, custom-op load, fake registration, and MPI initialization. If any step fails, later _init() calls return successfully although initialization did not complete. Set _inited only after successful initialization, or retain and re-raise a saved initialization failure.
🤖 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/_bootstrap.py` around lines 120 - 125, Update _init so _inited
is not set until Torch import, custom-op loading, fake registration, and MPI
initialization all complete successfully; ensure a failed setup does not cause
later calls to return as initialized, either by leaving _inited false or
preserving and re-raising the initialization failure.
| try: | ||
| torch.classes.load_library(ft_decoder_lib) | ||
| from ._torch.custom_ops import _register_fake | ||
|
|
||
| _register_fake() | ||
| except Exception as e: | ||
| msg = ( | ||
| "\nFATAL: Decoding operators failed to load. This may be caused by an incompatibility " | ||
| "between PyTorch and TensorRT-LLM. Please rebuild and install TensorRT-LLM." | ||
| ) | ||
| raise ImportError(str(e) + msg) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
For the PyTorch version declared by this project, what exception types can torch.classes.load_library() raise?
💡 Result:
The torch.classes.load_library function in PyTorch does not explicitly define a custom set of exception types that it raises. Instead, it relies on underlying mechanisms—specifically torch.ops.load_library and the Python ctypes module—to handle the loading of shared libraries [1][2]. Consequently, the primary exception type raised by torch.classes.load_library is OSError [3][2]. This exception occurs when the operating system encounters an error during the library loading process, such as when the specified file path cannot be found (e.g., "cannot open shared object file: No such file or directory") or when there are issues with the library's symbols (e.g., "undefined symbol" errors) [3][2]. Because torch.classes.load_library internally uses ctypes.CDLL to load the library, it may also propagate other exceptions standard to the Python runtime depending on the environment and the nature of the failure [2]. Users should be prepared to catch OSError and potentially other standard Python exceptions when calling this function [3][2].
Citations:
- 1: https://github.com/pytorch/pytorch/blob/main/torch/_classes.py
- 2: https://discuss.pytorch.org/t/extending-torchscript-with-custom-c-classes-undefined-symbol-error/81948
- 3: https://discuss.pytorch.org/t/libth-transformer-so-cannot-open-shared-object-file-no-such-file-or-directory/185256
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- candidate file ---'
sed -n '120,175p' tensorrt_llm/_bootstrap.py
printf '%s\n' '--- PyTorch version declarations ---'
rg -n --glob '!*lock*' --glob '!*.md' 'torch(==|>=|<=|~=|[[:space:]])|pytorch' pyproject.toml setup.py setup.cfg requirements* tensorrt_llm 2>/dev/null | head -120
printf '%s\n' '--- load_library call sites ---'
rg -n 'load_library|_register_fake' .Repository: NVIDIA/TensorRT-LLM
Length of output: 16750
🌐 Web query:
PyTorch torch.classes.load_library implementation and documented exception behavior for the PyTorch version used by TensorRT-LLM
💡 Result:
The torch.classes.load_library function is a utility in PyTorch used to dynamically load C++ shared libraries into the Python process, enabling the registration of custom C++ classes and operators with the PyTorch JIT/TorchScript runtime [1][2]. Implementation Details The function is defined in torch/_classes.py as a wrapper around the internal C++ operator torch.ops.load_library [1]. When called, it invokes this underlying C++ routine to load the specified shared library (.so on Linux,.pyd on Windows) [1][3]. This library typically contains static global initialization code that registers custom classes with the PyTorch runtime [1]. Once successfully loaded, the library path is added to the torch.classes.loaded_libraries set, which allows for tracking of loaded dependencies [1]. Exception Behavior The function does not have complex, custom-defined exception behavior explicitly documented in its Python wrapper. Instead, it relies on the underlying C++ library loading mechanism (typically using platform-specific APIs like dlopen on Linux or LoadLibrary on Windows). Consequently, if the library cannot be found, if it fails to link, or if the initialization code within the library raises an error, the operation will raise a standard Python RuntimeError or OSError propagated from the underlying C++ loading process [1][3]. Context for TensorRT-LLM TensorRT-LLM is built to be compatible with specific versions of PyTorch, generally ranging between 2.9.x and 2.11.x (as of early 2026, depending on the specific release and build environment) [4][5]. Because TensorRT-LLM relies on these standard PyTorch versions, the behavior of torch.classes.load_library remains consistent with the upstream PyTorch implementation [1][6][7]. Users should ensure their environment matches the PyTorch version expected by their specific TensorRT-LLM wheel, as mismatches can lead to runtime environment inconsistencies [6][7]. Top results: [1][3][6][7][4][5]
Citations:
- 1: https://github.com/pytorch/pytorch/blob/main/torch/_classes.py
- 2: https://docs.pytorch.org/tutorials/advanced/custom_classes.md
- 3: https://discuss.pytorch.org/t/torch-classes-has-no-attribute-load-library/70151
- 4: https://github.com/NVIDIA/TensorRT-LLM/blob/fc2bf2790d6e02e7028f0b51cf967f41d7df45f3/requirements.txt
- 5: https://github.com/NVIDIA/TensorRT-LLM/blob/15d06c0923b63ac1781784d5f59e1747bb47d5f1/requirements.txt
- 6: https://nvidia.github.io/TensorRT-LLM/installation/linux.html
- 7: https://nvidia.github.io/TensorRT-LLM/installation/installation-guide.html
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- dependency files ---'
git ls-files | rg '(^|/)(pyproject\.toml|setup\.py|setup\.cfg|requirements[^/]*|environment[^/]*)$' | head -80
printf '%s\n' '--- declared torch constraints ---'
for f in $(git ls-files | rg '(^|/)(pyproject\.toml|setup\.py|setup\.cfg|requirements[^/]*|environment[^/]*)$' | head -80); do
rg -n -i 'torch|pytorch' "$f" || true
done
printf '%s\n' '--- bootstrap context ---'
cat -n tensorrt_llm/_bootstrap.py | sed -n '125,170p'
printf '%s\n' '--- related definitions and callers ---'
rg -n -C 4 '_register_fake|load_library|_bootstrap' tensorrt_llmRepository: NVIDIA/TensorRT-LLM
Length of output: 34337
🌐 Web query:
site:github.com/pytorch/pytorch torch/_classes.py "def load_library" "torch.classes"
💡 Result:
The torch.classes module is a specialized namespace in PyTorch used for accessing custom C++ classes that have been registered with the PyTorch runtime [1]. The file torch/_classes.py defines the Classes class, which acts as this module and handles the dynamic discovery and exposure of these custom classes [1]. Key functions and concepts include: 1. torch.classes.load_library: This function is used to load a shared library (e.g., a.so or.dll file) into the current process [1]. Loading the library triggers the execution of global initialization code, which typically registers custom C++ classes using the torch::class API [1][2]. Once loaded, the library path is added to torch.classes.loaded_libraries [1]. 2. Dynamic Namespace Access: The torch.classes module uses a custom getattr implementation [1]. When you access an attribute (e.g., torch.classes.my_namespace), it dynamically returns a ClassNamespace object [1]. This allows PyTorch to handle class registrations that may occur at runtime after the initial module import [3]. 3. Relation to C++ Registration: Developers register C++ classes by instantiating torch::class in C++ code [2]. By specifying a namespace and class name (e.g., 'foo' and 'Bar'), the class becomes accessible in Python via torch.classes.foo.Bar [2]. While both torch.ops and torch.classes provide mechanisms for loading shared libraries to extend PyTorch, torch.ops is primarily used for custom operators, whereas torch.classes is used for custom C++ classes [1][4]. In many cases, torch.classes.load_library internally calls torch.ops.load_library to perform the actual dynamic loading [1].
Citations:
- 1: https://github.com/pytorch/pytorch/blob/main/torch/_classes.py
- 2: https://github.com/pytorch/pytorch/blob/main/torch/custom_class.h
- 3: torch.classes module doesn't set all module attributes pytorch/pytorch#158871
- 4: https://github.com/pytorch/pytorch/blob/main/torch/_ops.py
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- all tracked dependency/config files mentioning torch ---'
git grep -n -i -E 'torch|pytorch' -- \
'pyproject.toml' 'setup.py' 'setup.cfg' 'requirements*.txt' \
'*requirements*.txt' 'environment*.yml' 'environment*.yaml' \
'docker/**' '.github/**' 2>/dev/null | head -200 || true
printf '%s\n' '--- exact bootstrap block ---'
nl -ba tensorrt_llm/_bootstrap.py | sed -n '140,165p'
printf '%s\n' '--- exception handling and registration definitions ---'
git grep -n -C 5 -E 'load_library|_register_fake|register_fake' -- tensorrt_llm torch 2>/dev/null | head -240 || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 44544
🏁 Script executed:
#!/bin/bash
set -eu
python3 - <<'PY'
from urllib.request import Request, urlopen
versions = ("v2.11.0", "v2.12.0", "v2.13.0a0")
paths = ("torch/_classes.py", "torch/_ops.py")
for version in versions:
print(f"=== {version} ===")
for path in paths:
url = f"https://raw.githubusercontent.com/pytorch/pytorch/{version}/{path}"
try:
request = Request(url, headers={"User-Agent": "CodeRabbit-review-check"})
text = urlopen(request, timeout=15).read().decode("utf-8")
except Exception as exc:
print(f"{path}: FETCH_ERROR {type(exc).__name__}: {exc}")
continue
lines = text.splitlines()
matches = [
index for index, line in enumerate(lines)
if "def load_library" in line or "_dlopen" in line
]
print(f"--- {path} ---")
for index in matches:
start = max(0, index - 3)
end = min(len(lines), index + 28)
print("\n".join(f"{number + 1}: {lines[number]}" for number in range(start, end)))
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 1251
🌐 Web query:
PyTorch v2.11.0 torch.classes.load_library raises OSError RuntimeError source _ops.py _classes.py
💡 Result:
When torch.classes.load_library raises an OSError or RuntimeError during the loading of a shared library, it is almost always caused by an issue within the library itself or the environment in which it is being loaded, rather than a bug in PyTorch's _classes.py or _ops.py internal logic [1][2]. Because torch.classes.load_library internally invokes torch.ops.load_library, which uses ctypes.CDLL(path) to load the object [2], these errors usually indicate one of the following root causes: 1. Missing Shared Dependencies: The most common cause for an OSError is that the shared library cannot be loaded because its own dependencies (such as CUDA libraries or other system files) are not found or are incompatible [3]. Ensure that your library's dependencies are correctly installed and that your LD_LIBRARY_PATH (on Linux) or system PATH (on Windows) includes the necessary directories [4][3]. 2. ABI Incompatibility: A frequent source of RuntimeError or symbol-related errors is a mismatch between the C++ ABI (Application Binary Interface) used to compile your custom library and the ABI used to build your installed version of PyTorch (often related to the _GLIBCXX_USE_CXX11_ABI flag) [5][6]. If your library uses a different ABI than the PyTorch binary, it will fail to load or resolve symbols [5][6]. You may need to recompile your library with -D_GLIBCXX_USE_CXX11_ABI=0 or 1 to match your PyTorch installation [5]. 3. Incorrect File Path: Ensure that the path provided to load_library is accurate and that the file actually exists [4]. Even if the file exists, the dynamic linker might fail if it is the wrong architecture or format for your system [4]. 4. Version Mismatch: If you are using dependent packages like torchaudio or torchvision, ensure they are strictly compatible with your installed version of PyTorch [6]. Mismatched versions of these packages frequently cause "undefined symbol" errors because they may be trying to link against incompatible internal PyTorch APIs [6]. To debug, you can inspect the specific exception message associated with the OSError, as it often provides details about which file is missing or which symbol failed to resolve [5][4][3]. Verification of your environment by confirming versions of torch and its associated libraries is a recommended first step [6].
Citations:
- 1: https://github.com/pytorch/pytorch/blob/main/torch/_classes.py
- 2: https://github.com/pytorch/pytorch/blob/main/torch/_ops.py
- 3: OSError: libtorch_cuda.so: cannot open shared object file: No such file or directory pytorch/pytorch#111983
- 4: https://discuss.pytorch.org/t/libth-transformer-so-cannot-open-shared-object-file-no-such-file-or-directory/185256
- 5: https://discuss.pytorch.org/t/extending-torchscript-with-custom-c-classes-undefined-symbol-error/81948
- 6: https://discuss.pytorch.org/t/undefined-symbol-znk-nameb5cxx11ev/204738
Narrow the exception boundary and preserve the original cause.
Keep only torch.classes.load_library(ft_decoder_lib) inside the try and handle its loading failures (OSError and RuntimeError). Move the import and _register_fake() outside the boundary so registration errors are not reported as library-load failures. Use raise ImportError(...) from exc.
🧰 Tools
🪛 Ruff (0.16.1)
[warning] 154-154: Do not catch blind exception: Exception
(BLE001)
[warning] 159-159: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
🤖 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/_bootstrap.py` around lines 149 - 159, In the library-loading
block, keep only torch.classes.load_library(ft_decoder_lib) inside the try,
catch OSError and RuntimeError, and raise ImportError with the existing context
using raise ... from exc to preserve the cause. Move the _register_fake import
and invocation outside the try/except so registration failures propagate
independently.
Sources: Coding guidelines, Linters/SAST tools
| Removal is tracked by the removal ticket T25 (TRTLLM-14855) delivers; the compatibility window spans at least | ||
| one release (Epic decision D4 (b)). Do not add definitions, do not rewrite |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Correct the removal-schedule sentence.
Line 24 ends with “delivers”, so the sentence is incomplete. State the ticket and compatibility window as separate sentences.
Proposed fix
-Removal is tracked by the removal ticket T25 (TRTLLM-14855) delivers; the compatibility window spans at least
+Removal is tracked by ticket T25 (TRTLLM-14855). The compatibility window spans at least📝 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.
| Removal is tracked by the removal ticket T25 (TRTLLM-14855) delivers; the compatibility window spans at least | |
| one release (Epic decision D4 (b)). Do not add definitions, do not rewrite | |
| Removal is tracked by ticket T25 (TRTLLM-14855). The compatibility window spans at least | |
| one release (Epic decision D4 (b)). Do not add definitions, do not rewrite |
🤖 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/_common.py` around lines 24 - 25, Correct the removal-schedule
text by completing the sentence about removal being tracked by ticket T25
(TRTLLM-14855), then state the at-least-one-release compatibility window in a
separate sentence. Do not add definitions or otherwise rewrite the surrounding
content.
|
PR_Github #64370 [ run ] triggered by Bot. Commit: |
8 participants
Reduce the package root's mechanics to a single private bootstrap module, and move the restricted unpickler into the package that consumes it. Each change is described below together with the reasoning it rests on and what it changes for callers.
_common.py's _init() and the three setup functions that lived in init.py -- DLL directory, Python-library preload, vendored-triton_kernels precedence -- merge into one private bootstrap module. init.py drops from 176 to 103 lines and is now only the package's compatibility surface: the intentional re-exports and all.
The two existing phases keep their order, which is the constraint the file is arranged around: _prepare_environment() runs before
import torch, _init() runs after the package's own imports, and its early-return guard is unchanged.That ordering forces the one change here that is not a pure move. _bootstrap.py is imported ahead of torch, so its module scope is restricted to the standard library; a module-level
import torchorfrom .bindings import MpiCommwould load Torch and the compiled extension before phase 1 had prepared for them. Phase 2's four imports -- torch, tensorrt_llm.bindings, ._utils and .logger -- therefore move into _init()'s body. This defers statements, not any module's first import: _init() is the last statement of init.py, and all four are already imported earlier on theimport tensorrt_llm._torch.modelspath. Tracing module completion order with -X importtime overimport tensorrt_llmconfirms it: 2393 modules on both sides, and the whole difference is +tensorrt_llm._bootstrap at the front and -tensorrt_llm._common where it used to sit._common.py is absent from legacy-files.txt while init.py is on it, so the three relocated functions land in ruff's scope instead of yapf's and are reformatted accordingly. That is formatting only and confined to those three functions; _init() was already ruff-clean and moves unchanged.
The old path keeps a definition-free shim that re-exports _init and warns on import. The module is private, and every source of evidence that can be checked inside this repository is clean -- but one cannot be checked here, namely consumers outside the repository, and only the module's owner can answer for those. That question was never asked, so removing the path was never certifiable. Ten lines against a breakage nobody here can observe is the same asymmetry that keeps a shim on every public path in this migration, so the shim goes in rather than the question going to an owner.
The export set is recorded rather than computed: both names in the old module are underscored, so a private module has no computable public surface, and the shim re-exports the one symbol this repository itself consumed. _inited is deliberately left out -- it is mutable module state, and a re-export would bind the value once and then diverge from the guard it appears to mirror.
What the shim preserves is one function object, not a copy:
_common._init is _bootstrap._init is tensorrt_llm._init. A copy would give the guard two _inited flags, and a call through the old path would run phase 2 -- custom op registration and MPI init -- a second time, which it is not designed for; calling _common._init() again is verified to be a no-op.import tensorrt_llmdoes not pull the shim in, since init.py now imports _bootstrap, so the DeprecationWarning still means "you named the old path" rather than firing for every user.The restricted unpickler moves into the package that consumes it, and both consumers at HEAD are rewritten to the canonical path: llmapi/rlhf_utils.py, the only production consumer, and tests/unittest/llmapi/test_serialization.py. A stale reference in a model_engine.py comment is retargeted as well. Nothing in the tree goes through the shim -- in-repo references, lint lists, ownership rules and derived configs are all rewritten as if the old path had disappeared. The file gains the NVIDIA copyright header it never had, and the four allowlist module strings travel with it unchanged; later tasks of this migration update their contents at the new location.
Ownership moves with the code. That is the one effective-CODEOWNERS change in this commit and it is intended: serialization.py inherited the package-root default and now falls under llmapi/'s rule, whose team owns its only consumer and its only test. No file loses an owner. The legacy-files.txt entry likewise follows the implementation to the new path, so the module stays in yapf's scope and the entry count is unchanged at 765.
The old path keeps a definition-free shim over all seven public names, warning on import. serialization has no leading underscore, so it is a public module path and keeps a shim regardless of what the evidence says; what the measurement earns is a shorter deprecation window, and it is what fixes the export set -- with no evidence source naming a symbol, the supported surface falls back to the module's own public surface: BASE_EXAMPLE_CLASSES, Unpickler, dump, dumps, load, loads and register_approved_class.
A shim, and only a shim, preserves
from tensorrt_llm import serialization. That form works today only because serialization is a submodule of the package; init.py never imports it, so the top-level re-export that rescues other modules' symbols cannot help here -- what is imported is the module itself, and that goes through the import finder.Because the shim re-exports objects rather than copies, isinstance and pickle keep working: every name satisfies
tensorrt_llm.serialization.X is tensorrt_llm.llmapi.serialization.X, seven of seven, nothing is defined at the old path and no module is rewritten. The two module objects are necessarily distinct -- a re-export shim shares symbols, not identity -- so per-symbol identity is the assertion that holds, and it is the property pickle and isinstance actually depend on. The security boundary forwards too, not just the names: a payload dumped through either path loads through the other, and an unapproved class is still rejected when loaded through the shim. A shim that forwarded only the success path would quietly widen the boundary while every round-trip test passed.Regenerated from legacy-files.txt, so no derived config still names the retired path; check-configs passes and the ruff baseline needs no key renames. Both shims are new files and therefore not on the legacy list, which puts them in ruff's scope; their formatting is already what ruff produces, so the hooks leave them alone.
Left alone deliberately
version.py stays at the package root: setup.py, lock-file generation and the packaging metadata read that exact source path. _deprecation.py is confirmed absent at HEAD with no remaining reference.
Nothing outside the two deprecation warnings changes for callers: the top-level all is byte-identical and hashes identically at runtime, the api-stability reference diff is empty, the golden manifest's sha256 is unchanged, no test moves so node IDs, test lists and waives.txt are untouched, and the import graph keeps its five module-level cycles with no new cross-package edge and neither shim joining one. Against a build of the base commit with this change overlaid, test_serialization.py passes 6 before and after and pre-migration pickles replay 10 of 10. The eleventh class in that compatibility matrix, GreedyDecodingParams, does not exist at this commit -- the name appears nowhere in the repository except as a string in the allowlist. Left as-is and recorded for the task that updates it.
Dev Engineer Review
tensorrt_llm/_bootstrap.py._init().tensorrt_llm/__init__.pyas the package compatibility surface.tensorrt_llm/_common.pyinto a deprecation shim.tensorrt_llm/llmapi/serialization.py.version.pyremains unchanged.QA Engineer Review
tests/unittest/llmapi/test_serialization.py.test_serialization_allowed_classtest_serialization_disallowed_classtest_serialization_basic_objecttest_serialization_complex_object_allowed_classtest_serialization_complex_object_partially_allowed_classtest_serialization_complex_object_disallowed_classtest-db/,qa/, orwaives.txtchanges were reported.tests/integration/test_lists/is unavailable.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-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.