[None][chore] load models lazily - #17281
Open
QiJune wants to merge 5 commits into
Open
Conversation
get_steady_clock_now_in_seconds was defined in serve/responses_utils, so the executor imported the serve package (openai types, openai_harmony, transformers processors) at module level just for a one-line clock wrapper. Move it to tensorrt_llm._utils next to the other bindings helpers; responses_utils re-exports it for its serve-side importers. Signed-off-by: junq <22017000+QiJune@users.noreply.github.com>
_torch/models/__init__.py eagerly imported every modeling_* module, so any process importing tensorrt_llm executed ~50 model files' class bodies and registration decorators at startup. Replace this with: - _arch_index.py: static tables mapping architecture name and public class name to the providing module (generated from the @register_auto_model decorators and the previous eager import list); - PEP 562 __getattr__ on the package: attribute access imports just the providing module; bare modeling_* submodule access keeps working; - ensure_model_registered() in modeling_utils: architecture-based resolution (AutoModelForCausalLM, get_model_architecture, model_loader preference lookup) imports the providing module on demand so its decorators run; - commands/utils checks the static index instead of importing the zoo; - py_executor matches Llama4 by class name and model_loader builds its post-transform profile registry on first use, so neither imports a model-zoo module at module level; - drop the shadowed legacy MistralForCausalLM registration in modeling_llama (modeling_mistral always won under eager import order; under lazy loading the duplicate would be order-dependent). Signed-off-by: junq <22017000+QiJune@users.noreply.github.com>
'import tensorrt_llm' eagerly pulled the whole public surface (both model zoos, quantization, runtime, tools, llmapi, visual_gen), which executed ~99% of product modules in every process at startup. Keep the environment setup, _common._init(), logger and version banner eager; resolve every public name on first attribute access instead, with a plain-submodule fallback so previously-reachable module attributes keep working. A TYPE_CHECKING block preserves the original imports for static tooling. Also maps KvCacheConfig, which __all__ listed but the eager chain never actually imported. Signed-off-by: junq <22017000+QiJune@users.noreply.github.com>
QiJune
requested review from
JunyiXu-nv,
Tabrizian,
brnguyen2,
reasonsolo,
schetlur-nv and
zhaoyangwang-nvidia
August 5, 2026 04:04
Contributor
WalkthroughChangesThe package now loads public modules and symbols lazily. Static model indexes support on-demand provider imports and registration. Executor and serving paths defer model-related imports. Shared steady-clock conversion moved to Lazy loading and registration
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant ModelLoader
participant MODEL_ARCH_TO_MODULE
participant ProviderModule
participant MODEL_CLASS_MAPPING
ModelLoader->>MODEL_ARCH_TO_MODULE: resolve checkpoint architecture
MODEL_ARCH_TO_MODULE->>ProviderModule: import provider module
ProviderModule->>MODEL_CLASS_MAPPING: register model class
ModelLoader->>MODEL_CLASS_MAPPING: resolve registered class
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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 (1)
tensorrt_llm/_torch/models/modeling_auto.py (1)
36-40: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winAlign EAGLE3 architecture naming. When
architectures[0]isDeepseekV3ForCausalLM, this code generatesEAGLE3DeepseekV3ForCausalLM. NeitherMODEL_ARCH_TO_MODULEnorregister_auto_modeldefines that exact key; the registered key isEagle3DeepSeekV3ForCausalLM.ensure_model_registeredtherefore skips the lazy import, and model resolution returnsNone. Align the rewrite, registration, and index entries, and add a regression test.🤖 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/_torch/models/modeling_auto.py` around lines 36 - 40, Update the EAGLE3 architecture rewrite in the model-architecture resolution flow around ensure_model_registered so DeepseekV3ForCausalLM produces the exact registered key Eagle3DeepSeekV3ForCausalLM, using consistent spelling and casing across MODEL_ARCH_TO_MODULE and register_auto_model. Add a regression test covering this architecture and verifying successful model resolution.
🧹 Nitpick comments (3)
tensorrt_llm/_torch/models/__init__.py (1)
111-112: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueInclude the lazy names in
__dir__.
__dir__reports__all__and the already-resolved globals. It omits anyMODEL_CLASS_TO_MODULEkey that is not in__all__, so those names are invisible todir()and to tab completion until first access. The top-level package uses the union with its lazy table; match that here.♻️ Proposed fix
-def __dir__(): - return sorted(set(__all__) | set(globals())) +def __dir__() -> list[str]: + return sorted(set(__all__) | set(globals()) | set(MODEL_CLASS_TO_MODULE))🤖 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/_torch/models/__init__.py` around lines 111 - 112, Update __dir__ to include the keys from MODEL_CLASS_TO_MODULE alongside __all__ and existing globals, matching the top-level package’s lazy-name discovery behavior so unresolved model names appear in dir() and tab completion.tensorrt_llm/__init__.py (2)
195-199: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winDo not convert unrelated import failures into
AttributeError.The fallback catches every
ModuleNotFoundError. A submodule that exists but imports a missing optional dependency also raisesModuleNotFoundError. The user then seesmodule 'tensorrt_llm' has no attribute 'X'instead of the real missing dependency. Check that the failed module is the requested one before converting.♻️ Proposed fix
try: return importlib.import_module(f'.{name}', __name__) - except ModuleNotFoundError: - raise AttributeError( - f"module {__name__!r} has no attribute {name!r}") from None + except ModuleNotFoundError as e: + if e.name != f'{__name__}.{name}': + raise # a real missing dependency inside the submodule + raise AttributeError( + f"module {__name__!r} has no attribute {name!r}") from None🤖 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/__init__.py` around lines 195 - 199, Update the import fallback around importlib.import_module in __getattr__ so it converts ModuleNotFoundError to AttributeError only when the missing module is the requested submodule; re-raise unrelated dependency import failures unchanged.
185-185: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd type annotations to the module hooks.
The coding guidelines require an annotation on every function. Use
def __getattr__(name: str) -> object:anddef __dir__() -> list[str]:.As per coding guidelines: "Annotate every function, use
Nonefor procedures".Also applies to: 202-202
🤖 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/__init__.py` at line 185, Update the module hooks __getattr__ and __dir__ with the required annotations: use a str parameter and object return type for __getattr__, and a list[str] return type for __dir__. Keep their existing behavior unchanged.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/_torch/models/__init__.py`:
- Around line 97-108: The __getattr__ fallback in __init__ must also resolve
direct submodule attributes such as checkpoints and hf_parameter_utils, not only
names beginning with modeling_. Attempt importing the requested submodule for
other names as well, convert only genuine module-not-found misses to
AttributeError, and preserve the existing modeling_multimodal_utils coverage and
model-class resolution behavior.
In `@tensorrt_llm/_torch/models/_arch_index.py`:
- Line 93: Remove the "SomeVLModel" mapping from MODEL_ARCH_TO_MODULE, leaving
the surrounding architecture registrations unchanged.
- Around line 18-99: Add a CI consistency check for MODEL_ARCH_TO_MODULE that
discovers every architecture registered by `@register_auto_model`, including
qualified and attribute-based decorator forms, and compares those registrations
against the mapping in both directions. Fail the check when an architecture is
missing from either set, so ensure_model_registered can import and register
every decorated model.
In `@tensorrt_llm/_torch/models/modeling_utils.py`:
- Around line 883-887: Record caught ImportError instances in a module-level
dictionary keyed by architecture within the lazy import path, while retaining
the existing warning. Update get_model_architecture to append the stored import
failure to its unknown-architecture error when the architecture is indexed but
its import failed. Keep non-ImportError exceptions propagating unchanged.
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Line 4884: Update the comment near the token-range checks to replace the
Unicode hyphen in “token‐range” with an ASCII hyphen, yielding “token-range” and
satisfying Ruff RUF003.
---
Outside diff comments:
In `@tensorrt_llm/_torch/models/modeling_auto.py`:
- Around line 36-40: Update the EAGLE3 architecture rewrite in the
model-architecture resolution flow around ensure_model_registered so
DeepseekV3ForCausalLM produces the exact registered key
Eagle3DeepSeekV3ForCausalLM, using consistent spelling and casing across
MODEL_ARCH_TO_MODULE and register_auto_model. Add a regression test covering
this architecture and verifying successful model resolution.
---
Nitpick comments:
In `@tensorrt_llm/__init__.py`:
- Around line 195-199: Update the import fallback around importlib.import_module
in __getattr__ so it converts ModuleNotFoundError to AttributeError only when
the missing module is the requested submodule; re-raise unrelated dependency
import failures unchanged.
- Line 185: Update the module hooks __getattr__ and __dir__ with the required
annotations: use a str parameter and object return type for __getattr__, and a
list[str] return type for __dir__. Keep their existing behavior unchanged.
In `@tensorrt_llm/_torch/models/__init__.py`:
- Around line 111-112: Update __dir__ to include the keys from
MODEL_CLASS_TO_MODULE alongside __all__ and existing globals, matching the
top-level package’s lazy-name discovery behavior so unresolved model names
appear in dir() and tab completion.
🪄 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: 5305a36f-c8f2-44e1-bb21-0350d829c389
📒 Files selected for processing (12)
tensorrt_llm/__init__.pytensorrt_llm/_torch/models/__init__.pytensorrt_llm/_torch/models/_arch_index.pytensorrt_llm/_torch/models/modeling_auto.pytensorrt_llm/_torch/models/modeling_llama.pytensorrt_llm/_torch/models/modeling_utils.pytensorrt_llm/_torch/pyexecutor/model_loader.pytensorrt_llm/_torch/pyexecutor/perf_metrics_manager.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_utils.pytensorrt_llm/commands/utils.pytensorrt_llm/serve/responses_utils.py
Comment on lines
+97
to
+108
| def __getattr__(name: str): | ||
| module_name = MODEL_CLASS_TO_MODULE.get(name) | ||
| if module_name is None: | ||
| # Also resolve bare submodule access (models.modeling_llama) so callers | ||
| # that relied on the previously-eager submodule attributes keep working. | ||
| if name.startswith("modeling_"): | ||
| return importlib.import_module(f".{name}", __name__) | ||
| raise AttributeError(f"module {__name__!r} has no attribute {name!r}") | ||
| module = importlib.import_module(f".{module_name}", __name__) | ||
| attr = getattr(module, name) | ||
| globals()[name] = attr # cache: subsequent access skips __getattr__ | ||
| return attr |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Find attribute access on tensorrt_llm._torch.models submodules that are not modeling_*.
set -uo pipefail
# Direct attribute access via an imported package alias.
rg -nP -C3 --type=py '\bmodels\.(?!modeling_)[a-z_]+\b' -g '!**/models/**'
# Submodules that exist under the package.
fd -t d -d 1 . tensorrt_llm/_torch/models
fd -t f -d 1 -e py . tensorrt_llm/_torch/modelsRepository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
#!/bin/bash
set -u
printf '%s\n' '--- package file list ---'
git ls-files tensorrt_llm/_torch/models | sed -n '1,160p'
printf '%s\n' '--- __init__.py structure ---'
ast-grep outline tensorrt_llm/_torch/models/__init__.py --view expanded 2>/dev/null || true
printf '%s\n' '--- __init__.py relevant sections ---'
sed -n '1,135p' tensorrt_llm/_torch/models/__init__.py
printf '%s\n' '--- non-modeling attribute references outside the package ---'
rg -nP -C3 --type=py '\bmodels\.(?!modeling_)[a-z_]+\b' -g '!tensorrt_llm/_torch/models/**' . || true
printf '%s\n' '--- package submodule candidates ---'
python3 - <<'PY'
from pathlib import Path
root = Path("tensorrt_llm/_torch/models")
for path in sorted(root.iterdir()):
if path.name == "__pycache__":
continue
print(path)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -u
printf '%s\n' '--- package file list ---'
git ls-files tensorrt_llm/_torch/models | sed -n '1,160p'
printf '%s\n' '--- __init__.py structure ---'
ast-grep outline tensorrt_llm/_torch/models/__init__.py --view expanded 2>/dev/null || true
printf '%s\n' '--- __init__.py relevant sections ---'
sed -n '1,135p' tensorrt_llm/_torch/models/__init__.py
printf '%s\n' '--- non-modeling attribute references outside the package ---'
rg -nP -C3 --type=py '\bmodels\.(?!modeling_)[a-z_]+\b' -g '!tensorrt_llm/_torch/models/**' . || true
printf '%s\n' '--- package submodule candidates ---'
python3 - <<'PY'
from pathlib import Path
root = Path("tensorrt_llm/_torch/models")
for path in sorted(root.iterdir()):
if path.name != "__pycache__":
print(path)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -u
printf '%s\n' '--- focused initial log ---'
sed -n '1,220p' /tmp/coderabbit-shell-logs/shell-output-EFTFt9
printf '%s\n' '--- current package initializer ---'
cat -n tensorrt_llm/_torch/models/__init__.py
printf '%s\n' '--- parent initializer, if available ---'
git show HEAD^:tensorrt_llm/_torch/models/__init__.py 2>/dev/null | cat -n | sed -n '1,180p' || true
printf '%s\n' '--- focused diff for initializer ---'
git diff HEAD^ -- tensorrt_llm/_torch/models/__init__.py | sed -n '1,240p'
printf '%s\n' '--- immediate package children ---'
python3 - <<'PY'
from pathlib import Path
root = Path("tensorrt_llm/_torch/models")
for path in sorted(root.iterdir()):
if path.name == "__pycache__":
continue
if path.is_dir() and (path / "__init__.py").exists():
print(f"package:{path.name}")
elif path.is_file() and path.suffix == ".py":
print(f"module:{path.stem}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 20943
🏁 Script executed:
#!/bin/bash
set -u
python3 - <<'PY'
from __future__ import annotations
import ast
from pathlib import Path
ROOT = Path("tensorrt_llm/_torch/models")
PACKAGE = "tensorrt_llm._torch.models"
def relative_import_name(node: ast.ImportFrom, current: str) -> str | None:
if node.level == 0:
return node.module
parts = current.split(".")
base = parts[: len(parts) - node.level]
if node.module:
base.append(node.module)
return ".".join(base)
def imported_local_modules(path: Path, module_name: str) -> set[str]:
tree = ast.parse(path.read_text(), filename=str(path))
result: set[str] = set()
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for alias in node.names:
if alias.name == PACKAGE or alias.name.startswith(PACKAGE + "."):
result.add(alias.name)
elif isinstance(node, ast.ImportFrom):
name = relative_import_name(node, module_name)
if name and (name == PACKAGE or name.startswith(PACKAGE + ".")):
result.add(name)
elif name == PACKAGE:
for alias in node.names:
if alias.name != "*":
result.add(f"{PACKAGE}.{alias.name}")
return result
print("--- current initializer imports ---")
for name in sorted(imported_local_modules(ROOT / "__init__.py", PACKAGE)):
print(name)
print("--- transitive local imports from initializer dependencies ---")
module_paths: dict[str, Path] = {PACKAGE: ROOT / "__init__.py"}
for path in ROOT.rglob("*.py"):
relative = path.relative_to(ROOT).with_suffix("")
suffix = ".".join(relative.parts)
module_paths[f"{PACKAGE}.{suffix}"] = path
seen = {PACKAGE}
queue = [PACKAGE]
edges: dict[str, set[str]] = {}
while queue:
current = queue.pop()
path = module_paths.get(current)
if path is None:
continue
deps = imported_local_modules(path, current)
edges[current] = deps
for dep in deps:
if dep in module_paths and dep not in seen:
seen.add(dep)
queue.append(dep)
for current in sorted(edges):
for dep in sorted(edges[current]):
print(f"{current} -> {dep}")
print("--- direct child modules and __getattr__ coverage ---")
children: list[str] = []
for path in ROOT.iterdir():
if path.name == "__pycache__":
continue
if path.is_file() and path.suffix == ".py":
children.append(path.stem)
elif path.is_dir() and (path / "__init__.py").exists():
children.append(path.name)
init_text = (ROOT / "__init__.py").read_text()
for child in sorted(children):
if child.startswith("modeling_"):
coverage = "prefix fallback"
else:
coverage = "not covered by prefix fallback"
print(f"{child}: {coverage}")
print("--- package-alias attribute access outside the package ---")
hits: list[tuple[str, int, str]] = []
for path in Path(".").rglob("*.py"):
if ROOT in path.parents or path == ROOT / "__init__.py":
continue
try:
tree = ast.parse(path.read_text(), filename=str(path))
except (OSError, SyntaxError):
continue
aliases: set[str] = set()
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for alias in node.names:
if alias.name == PACKAGE:
aliases.add(alias.asname or alias.name.rsplit(".", 1)[-1])
elif isinstance(node, ast.ImportFrom):
if node.module == "tensorrt_llm._torch" and any(
alias.name == "models" for alias in node.names
):
for alias in node.names:
if alias.name == "models":
aliases.add(alias.asname or alias.name)
for node in ast.walk(tree):
if isinstance(node, ast.Attribute) and node.value.__class__ is ast.Name:
if node.value.id in aliases:
hits.append((str(path), node.lineno, node.attr))
for path, line, attr in sorted(hits):
print(f"{path}:{line}: models.{attr}")
print("--- git state and initializer history ---")
PY
git status --short
git log --oneline -8 -- tensorrt_llm/_torch/models/__init__.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 3803
🏁 Script executed:
#!/bin/bash
set -u
python3 - <<'PY'
from __future__ import annotations
import ast
from pathlib import Path
ROOT = Path("tensorrt_llm/_torch/models")
PACKAGE = "tensorrt_llm._torch.models"
def relative_import_name(node: ast.ImportFrom, current: str) -> str | None:
if node.level == 0:
return node.module
parts = current.split(".")
base = parts[: len(parts) - node.level]
if node.module:
base.append(node.module)
return ".".join(base)
def imported_local_modules(path: Path, module_name: str) -> set[str]:
tree = ast.parse(path.read_text(), filename=str(path))
result: set[str] = set()
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for alias in node.names:
if alias.name == PACKAGE or alias.name.startswith(PACKAGE + "."):
result.add(alias.name)
elif isinstance(node, ast.ImportFrom):
name = relative_import_name(node, module_name)
if name and (name == PACKAGE or name.startswith(PACKAGE + ".")):
result.add(name)
elif name == PACKAGE:
for alias in node.names:
if alias.name != "*":
result.add(f"{PACKAGE}.{alias.name}")
return result
print("--- current initializer imports ---")
for name in sorted(imported_local_modules(ROOT / "__init__.py", PACKAGE)):
print(name)
print("--- transitive local imports from initializer dependencies ---")
module_paths: dict[str, Path] = {PACKAGE: ROOT / "__init__.py"}
for path in ROOT.rglob("*.py"):
relative = path.relative_to(ROOT).with_suffix("")
suffix = ".".join(relative.parts)
module_paths[f"{PACKAGE}.{suffix}"] = path
seen = {PACKAGE}
queue = [PACKAGE]
edges: dict[str, set[str]] = {}
while queue:
current = queue.pop()
path = module_paths.get(current)
if path is None:
continue
deps = imported_local_modules(path, current)
edges[current] = deps
for dep in deps:
if dep in module_paths and dep not in seen:
seen.add(dep)
queue.append(dep)
for current in sorted(edges):
for dep in sorted(edges[current]):
print(f"{current} -> {dep}")
print("--- direct child modules and __getattr__ coverage ---")
children: list[str] = []
for path in ROOT.iterdir():
if path.name == "__pycache__":
continue
if path.is_file() and path.suffix == ".py":
children.append(path.stem)
elif path.is_dir() and (path / "__init__.py").exists():
children.append(path.name)
for child in sorted(children):
coverage = (
"prefix fallback"
if child.startswith("modeling_")
else "not covered by prefix fallback"
)
print(f"{child}: {coverage}")
print("--- package-alias attribute access outside the package ---")
hits: list[tuple[str, int, str]] = []
for path in Path(".").rglob("*.py"):
if ROOT in path.parents or path == ROOT / "__init__.py":
continue
try:
tree = ast.parse(path.read_text(), filename=str(path))
except (OSError, SyntaxError):
continue
aliases: set[str] = set()
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for alias in node.names:
if alias.name == PACKAGE:
aliases.add(alias.asname or alias.name.rsplit(".", 1)[-1])
elif isinstance(node, ast.ImportFrom):
if node.module == "tensorrt_llm._torch" and any(
alias.name == "models" for alias in node.names
):
for alias in node.names:
if alias.name == "models":
aliases.add(alias.asname or alias.name)
for node in ast.walk(tree):
if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name):
if node.value.id in aliases:
hits.append((str(path), node.lineno, node.attr))
for path, line, attr in sorted(hits):
print(f"{path}:{line}: models.{attr}")
print("--- git state and initializer history ---")
PY
git status --short
git log --oneline -8 -- tensorrt_llm/_torch/models/__init__.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 3803
🏁 Script executed:
#!/bin/bash
set -u
printf '%s\n' '--- old eager initializer ---'
git show 22833f245^:tensorrt_llm/_torch/models/__init__.py | cat -n | sed -n '1,220p'
printf '%s\n' '--- lazy initializer commit diff ---'
git diff 22833f245^ 22833f245 -- tensorrt_llm/_torch/models/__init__.py | sed -n '1,280p'
printf '%s\n' '--- corrected static import and compatibility verifier ---'
python3 - <<'PY'
from __future__ import annotations
import ast
from pathlib import Path
ROOT = Path("tensorrt_llm/_torch/models")
PACKAGE = "tensorrt_llm._torch.models"
def relative_import_name(node: ast.ImportFrom, current: str) -> str | None:
if node.level == 0:
return node.module
parts = current.split(".")
# level=1 means the current package; level=2 means its parent.
base = parts[: len(parts) - node.level + 1]
if node.module:
base.append(node.module)
return ".".join(base)
def imports_from(path: Path, module_name: str) -> set[str]:
tree = ast.parse(path.read_text(), filename=str(path))
result: set[str] = set()
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for alias in node.names:
if alias.name == PACKAGE or alias.name.startswith(PACKAGE + "."):
result.add(alias.name)
elif isinstance(node, ast.ImportFrom):
name = relative_import_name(node, module_name)
if name and (name == PACKAGE or name.startswith(PACKAGE + ".")):
if node.module is None:
for alias in node.names:
if alias.name != "*":
result.add(f"{name}.{alias.name}")
else:
result.add(name)
return result
module_paths: dict[str, Path] = {PACKAGE: ROOT / "__init__.py"}
for path in ROOT.rglob("*.py"):
relative = path.relative_to(ROOT).with_suffix("")
module_paths[f"{PACKAGE}.{".".join(relative.parts)}"] = path
print("--- direct imports from current initializer ---")
for name in sorted(imports_from(ROOT / "__init__.py", PACKAGE)):
print(name)
print("--- transitive imports reachable from current initializer ---")
seen = {PACKAGE}
queue = [PACKAGE]
edges: dict[str, set[str]] = {}
while queue:
current = queue.pop()
path = module_paths.get(current)
if path is None:
continue
deps = imports_from(path, current)
edges[current] = deps
for dep in deps:
if dep in module_paths and dep not in seen:
seen.add(dep)
queue.append(dep)
for current in sorted(edges):
for dep in sorted(edges[current]):
print(f"{current} -> {dep}")
print("--- old eager direct imports that are now absent ---")
old_text = Path("/tmp/old_models_init.py")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 15615
🏁 Script executed:
#!/bin/bash
set -u
printf '%s\n' '--- old eager initializer ---'
git show 22833f245^:tensorrt_llm/_torch/models/__init__.py | cat -n | sed -n '1,220p'
printf '%s\n' '--- lazy initializer commit diff ---'
git diff 22833f245^ 22833f245 -- tensorrt_llm/_torch/models/__init__.py | sed -n '1,280p'
printf '%s\n' '--- corrected static import and compatibility verifier ---'
python3 - <<'PY'
from __future__ import annotations
import ast
from pathlib import Path
ROOT = Path("tensorrt_llm/_torch/models")
PACKAGE = "tensorrt_llm._torch.models"
def relative_import_name(node: ast.ImportFrom, current: str) -> str | None:
if node.level == 0:
return node.module
parts = current.split(".")
base = parts[: len(parts) - node.level + 1]
if node.module:
base.append(node.module)
return ".".join(base)
def imports_from(path: Path, module_name: str) -> set[str]:
tree = ast.parse(path.read_text(), filename=str(path))
result: set[str] = set()
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for alias in node.names:
if alias.name == PACKAGE or alias.name.startswith(PACKAGE + "."):
result.add(alias.name)
elif isinstance(node, ast.ImportFrom):
name = relative_import_name(node, module_name)
if name and (name == PACKAGE or name.startswith(PACKAGE + ".")):
if node.module is None:
for alias in node.names:
if alias.name != "*":
result.add(f"{name}.{alias.name}")
else:
result.add(name)
return result
module_paths: dict[str, Path] = {PACKAGE: ROOT / "__init__.py"}
for path in ROOT.rglob("*.py"):
relative = path.relative_to(ROOT).with_suffix("")
suffix = ".".join(relative.parts)
module_paths[f"{PACKAGE}.{suffix}"] = path
print("--- direct imports from current initializer ---")
for name in sorted(imports_from(ROOT / "__init__.py", PACKAGE)):
print(name)
print("--- transitive imports reachable from current initializer ---")
seen = {PACKAGE}
queue = [PACKAGE]
edges: dict[str, set[str]] = {}
while queue:
current = queue.pop()
path = module_paths.get(current)
if path is None:
continue
deps = imports_from(path, current)
edges[current] = deps
for dep in deps:
if dep in module_paths and dep not in seen:
seen.add(dep)
queue.append(dep)
for current in sorted(edges):
for dep in sorted(edges[current]):
print(f"{current} -> {dep}")
print("--- direct child modules not covered by the current fallback ---")
for path in sorted(ROOT.iterdir()):
if path.name == "__pycache__":
continue
if path.is_file() and path.suffix == ".py":
child = path.stem
elif path.is_dir() and (path / "__init__.py").exists():
child = path.name
else:
continue
if not child.startswith("modeling_"):
print(child)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 17586
🏁 Script executed:
#!/bin/bash
set -u
python3 - <<'PY'
from __future__ import annotations
import ast
import subprocess
from pathlib import Path
ROOT = Path("tensorrt_llm/_torch/models")
PACKAGE = "tensorrt_llm._torch.models"
def module_name_for(path: Path) -> str:
relative = path.relative_to(ROOT)
if relative.name == "__init__.py":
parts = relative.parent.parts
else:
parts = relative.with_suffix("").parts
return PACKAGE + (("." + ".".join(parts)) if parts else "")
module_paths = {
module_name_for(path): path
for path in ROOT.rglob("*.py")
}
def package_of(module_name: str) -> str:
if module_name in module_paths and module_paths[module_name].name == "__init__.py":
return module_name
return module_name.rsplit(".", 1)[0]
def relative_base(node: ast.ImportFrom, current: str) -> str | None:
if node.level == 0:
return node.module
base = package_of(current)
for _ in range(node.level - 1):
base = base.rsplit(".", 1)[0]
if node.module:
return f"{base}.{node.module}"
return base
def imports_from(path: Path, current: str) -> set[str]:
tree = ast.parse(path.read_text(), filename=str(path))
result: set[str] = set()
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for alias in node.names:
if alias.name == PACKAGE or alias.name.startswith(PACKAGE + "."):
result.add(alias.name)
elif isinstance(node, ast.ImportFrom):
base = relative_base(node, current)
if not base or not (
base == PACKAGE or base.startswith(PACKAGE + ".")
):
continue
if node.module is not None:
result.add(base)
else:
for alias in node.names:
if alias.name != "*":
result.add(f"{base}.{alias.name}")
return result
def reachable(root_text: str, root_name: str) -> set[str]:
root_path = Path("/tmp/initializer_for_scan.py")
root_path.write_text(root_text)
# The temporary root is parsed only; repository files are not modified.
seen = {root_name}
queue = [root_name]
while queue:
current = queue.pop()
path = root_path if current == root_name else module_paths.get(current)
if path is None:
continue
deps = imports_from(path, current)
for dep in deps:
if dep in module_paths and dep not in seen:
seen.add(dep)
queue.append(dep)
root_path.unlink()
return seen
current_text = Path(ROOT / "__init__.py").read_text()
old_text = subprocess.check_output(
[
"git",
"show",
"22833f245^:tensorrt_llm/_torch/models/__init__.py",
],
text=True,
)
old_loaded = reachable(old_text, PACKAGE)
current_loaded = reachable(current_text, PACKAGE)
def direct_children(loaded: set[str]) -> set[str]:
children = set()
prefix = PACKAGE + "."
for name in loaded:
if name.startswith(prefix):
children.add(name[len(prefix):].split(".", 1)[0])
return children
old_children = direct_children(old_loaded)
current_children = direct_children(current_loaded)
print("--- direct child attributes loaded by old eager initializer ---")
for name in sorted(old_children):
print(name)
print("--- direct child attributes loaded by current initializer ---")
for name in sorted(current_children):
print(name)
print("--- old direct child attributes absent after current initialization ---")
for name in sorted(old_children - current_children):
print(name)
print("--- absent old attributes not covered by current modeling_* fallback ---")
for name in sorted(old_children - current_children):
if not name.startswith("modeling_"):
print(name)
PY
printf '%s\n' '--- modeling_utils checkpoint import ---'
rg -n -C3 'checkpoints|hf_parameter_utils|multimodal_encoder_graph|dspark' \
tensorrt_llm/_torch/models/modeling_utils.py \
tensorrt_llm/_torch/models/modeling_*.py \
tensorrt_llm/_torch/models/_arch_index.py | sed -n '1,220p'Repository: NVIDIA/TensorRT-LLM
Length of output: 20403
🏁 Script executed:
#!/bin/bash
set -u
python3 - <<'PY'
from __future__ import annotations
import ast
import subprocess
from pathlib import Path
ROOT = Path("tensorrt_llm/_torch/models")
PACKAGE = "tensorrt_llm._torch.models"
def module_name_for(path: Path) -> str:
relative = path.relative_to(ROOT)
if relative.name == "__init__.py":
parts = relative.parent.parts
else:
parts = relative.with_suffix("").parts
return PACKAGE + (("." + ".".join(parts)) if parts else "")
module_paths = {
module_name_for(path): path
for path in ROOT.rglob("*.py")
}
def package_of(module_name: str) -> str:
if module_name in module_paths and module_paths[module_name].name == "__init__.py":
return module_name
return module_name.rsplit(".", 1)[0]
def relative_base(node: ast.ImportFrom, current: str) -> str | None:
if node.level == 0:
return node.module
base = package_of(current)
for _ in range(node.level - 1):
base = base.rsplit(".", 1)[0]
if node.module:
return f"{base}.{node.module}"
return base
def imports_from_tree(tree: ast.AST, current: str) -> set[str]:
result: set[str] = set()
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for alias in node.names:
if alias.name == PACKAGE or alias.name.startswith(PACKAGE + "."):
result.add(alias.name)
elif isinstance(node, ast.ImportFrom):
base = relative_base(node, current)
if not base or not (
base == PACKAGE or base.startswith(PACKAGE + ".")
):
continue
if node.module is not None:
result.add(base)
else:
for alias in node.names:
if alias.name != "*":
result.add(f"{base}.{alias.name}")
return result
def reachable(root_tree: ast.AST, root_name: str) -> set[str]:
seen = {root_name}
queue = [(root_name, root_tree)]
while queue:
current, tree = queue.pop()
deps = imports_from_tree(tree, current)
for dep in deps:
if dep in module_paths and dep not in seen:
seen.add(dep)
queue.append((dep, ast.parse(module_paths[dep].read_text())))
return seen
current_tree = ast.parse((ROOT / "__init__.py").read_text())
old_text = subprocess.check_output(
[
"git",
"show",
"22833f245^:tensorrt_llm/_torch/models/__init__.py",
],
text=True,
)
old_tree = ast.parse(old_text)
old_loaded = reachable(old_tree, PACKAGE)
current_loaded = reachable(current_tree, PACKAGE)
def direct_children(loaded: set[str]) -> set[str]:
prefix = PACKAGE + "."
return {
name[len(prefix):].split(".", 1)[0]
for name in loaded
if name.startswith(prefix)
}
old_children = direct_children(old_loaded)
current_children = direct_children(current_loaded)
print("--- old direct child attributes ---")
for name in sorted(old_children):
print(name)
print("--- current direct child attributes ---")
for name in sorted(current_children):
print(name)
print("--- old attributes absent after current initialization ---")
for name in sorted(old_children - current_children):
print(name)
print("--- absent old attributes not covered by modeling_* fallback ---")
for name in sorted(old_children - current_children):
if not name.startswith("modeling_"):
print(name)
PY
printf '%s\n' '--- relevant local imports ---'
rg -n -C3 'checkpoints|hf_parameter_utils|multimodal_encoder_graph|dspark' \
tensorrt_llm/_torch/models/modeling_utils.py \
tensorrt_llm/_torch/models/modeling_*.py \
tensorrt_llm/_torch/models/_arch_index.py | sed -n '1,220p'Repository: NVIDIA/TensorRT-LLM
Length of output: 20814
🏁 Script executed:
#!/bin/bash
set -u
python3 - <<'PY'
from pathlib import Path
import re
root = Path("tensorrt_llm/_torch/models")
old_model_names = re.findall(
r"from \.((?:modeling_[A-Za-z0-9_]+)) import",
__import__("subprocess").check_output(
["git", "show", "22833f245^:tensorrt_llm/_torch/models/__init__.py"],
text=True,
),
)
patterns = [
re.compile(r"from\s+tensorrt_llm\._torch\.models\.([A-Za-z0-9_\.]+)\s+import"),
re.compile(r"import\s+tensorrt_llm\._torch\.models\.([A-Za-z0-9_\.]+)"),
re.compile(r"from\s+\.([A-Za-z0-9_]+)\s+import"),
]
print("--- non-modeling imports from old eager model modules ---")
seen = set()
for name in old_model_names:
path = root / f"{name}.py"
if not path.exists():
continue
for line_number, line in enumerate(path.read_text().splitlines(), 1):
for pattern in patterns:
match = pattern.search(line)
if not match:
continue
imported = match.group(1).split(".", 1)[0]
if not imported.startswith("modeling_") and imported not in seen:
seen.add(imported)
print(f"{path}:{line_number}: {imported}: {line.strip()}")
break
print("--- non-modeling imports from current initializer dependencies ---")
for name in ("modeling_auto", "modeling_utils"):
path = root / f"{name}.py"
for line_number, line in enumerate(path.read_text().splitlines(), 1):
for pattern in patterns:
match = pattern.search(line)
if not match:
continue
imported = match.group(1).split(".", 1)[0]
if not imported.startswith("modeling_"):
print(f"{path}:{line_number}: {imported}: {line.strip()}")
break
print("--- package children and fallback result ---")
for path in sorted(root.iterdir()):
if path.name == "__pycache__":
continue
if path.is_file() and path.suffix == ".py":
child = path.stem
elif path.is_dir() and (path / "__init__.py").exists():
child = path.name
else:
continue
print(f"{child}: {'handled' if child.startswith('modeling_') else 'not handled'}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 3047
🏁 Script executed:
#!/bin/bash
set -u
python3 - <<'PY'
from pathlib import Path
import re
import subprocess
root = Path("tensorrt_llm/_torch/models")
old_init = subprocess.check_output(
[
"git",
"show",
"22833f245^:tensorrt_llm/_torch/models/__init__.py",
],
text=True,
)
old_model_names = re.findall(
r"from \.((?:modeling_[A-Za-z0-9_]+)) import",
old_init,
)
patterns = [
re.compile(r"from\s+tensorrt_llm\._torch\.models\.([A-Za-z0-9_\.]+)\s+import"),
re.compile(r"import\s+tensorrt_llm\._torch\.models\.([A-Za-z0-9_\.]+)"),
re.compile(r"from\s+\.([A-Za-z0-9_]+)\s+import"),
]
print("--- non-modeling imports from old eager model modules ---")
seen = set()
for name in old_model_names:
path = root / f"{name}.py"
if not path.exists():
continue
for line_number, line in enumerate(path.read_text().splitlines(), 1):
for pattern in patterns:
match = pattern.search(line)
if not match:
continue
imported = match.group(1).split(".", 1)[0]
if not imported.startswith("modeling_") and imported not in seen:
seen.add(imported)
print(f"{path}:{line_number}: {imported}: {line.strip()}")
break
print("--- non-modeling imports from current initializer dependencies ---")
for name in ("modeling_auto", "modeling_utils"):
path = root / f"{name}.py"
for line_number, line in enumerate(path.read_text().splitlines(), 1):
for pattern in patterns:
match = pattern.search(line)
if not match:
continue
imported = match.group(1).split(".", 1)[0]
if not imported.startswith("modeling_"):
print(f"{path}:{line_number}: {imported}: {line.strip()}")
break
print("--- package children and fallback result ---")
for path in sorted(root.iterdir()):
if path.name == "__pycache__":
continue
if path.is_file() and path.suffix == ".py":
child = path.stem
elif path.is_dir() and (path / "__init__.py").exists():
child = path.name
else:
continue
print(f"{child}: {'handled' if child.startswith('modeling_') else 'not handled'}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 3047
🏁 Script executed:
#!/bin/bash
set -u
printf '%s\n' '--- modeling_utils checkpoint import context ---'
cat -n tensorrt_llm/_torch/models/modeling_utils.py | sed -n '1008,1035p'
printf '%s\n' '--- current initializer dependency imports ---'
rg -n '^from |^import ' \
tensorrt_llm/_torch/models/modeling_auto.py \
tensorrt_llm/_torch/models/modeling_utils.py |
rg 'models|^\s*(from \.|import )' | sed -n '1,180p'
printf '%s\n' '--- all references to direct helper submodule names ---'
rg -n \
'multimodal_encoder_graph|hf_parameter_utils|models\.checkpoints|models\.hf_parameter_utils|models\.multimodal_encoder_graph' \
tensorrt_llm/_torch tests scripts examples 2>/dev/null | sed -n '1,220p'Repository: NVIDIA/TensorRT-LLM
Length of output: 26201
Preserve non-modeling_ submodule attributes.
The previous eager initializer exposed models.checkpoints and models.hf_parameter_utils. The current initializer raises AttributeError until a model module that imports them is loaded. Extend the fallback to direct submodules and convert only genuine import misses to AttributeError. models.modeling_multimodal_utils is already covered.
🤖 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/_torch/models/__init__.py` around lines 97 - 108, The
__getattr__ fallback in __init__ must also resolve direct submodule attributes
such as checkpoints and hf_parameter_utils, not only names beginning with
modeling_. Attempt importing the requested submodule for other names as well,
convert only genuine module-not-found misses to AttributeError, and preserve the
existing modeling_multimodal_utils coverage and model-class resolution behavior.
Comment on lines
+18
to
+99
| MODEL_ARCH_TO_MODULE = { | ||
| "AfmoeForCausalLM": "modeling_afmoe", | ||
| "BartForConditionalGeneration": "modeling_bart", | ||
| "BertForSequenceClassification": "modeling_bert", | ||
| "CLIPVisionModel": "modeling_clip", | ||
| "Cohere2ForCausalLM": "modeling_cohere2", | ||
| "Cosmos3ForConditionalGeneration": "modeling_cosmos3", | ||
| "DeciLMForCausalLM": "modeling_nemotron_nas", | ||
| "DeepseekV32ForCausalLM": "modeling_deepseekv3", | ||
| "DeepseekV3ForCausalLM": "modeling_deepseekv3", | ||
| "DeepseekV4ForCausalLM": "modeling_deepseekv4", | ||
| "EAGLE3LlamaForCausalLM": "modeling_speculative", | ||
| "Eagle3DeepSeekV3ForCausalLM": "modeling_speculative", | ||
| "Exaone4ForCausalLM": "modeling_exaone4", | ||
| "Exaone4_5_ForConditionalGeneration": "modeling_exaone4_5", | ||
| "ExaoneMoEForCausalLM": "modeling_exaone_moe", | ||
| "Gemma3ForCausalLM": "modeling_gemma3", | ||
| "Gemma3ForConditionalGeneration": "modeling_gemma3vl", | ||
| "Gemma4AssistantForCausalLM": "modeling_gemma4", | ||
| "Gemma4ForCausalLM": "modeling_gemma4", | ||
| "Gemma4ForConditionalGeneration": "modeling_gemma4mm", | ||
| "Gemma4UnifiedForConditionalGeneration": "modeling_gemma4_unified", | ||
| "Glm4MoeForCausalLM": "modeling_glm", | ||
| "GlmMoeDsaForCausalLM": "modeling_deepseekv3", | ||
| "GptOssForCausalLM": "modeling_gpt_oss", | ||
| "HCXVisionForCausalLM": "modeling_hyperclovax", | ||
| "HCXVisionModel": "modeling_hyperclovax", | ||
| "HunYuanDenseV1ForCausalLM": "modeling_hunyuan_dense", | ||
| "HunYuanMoEV1ForCausalLM": "modeling_hunyuan_moe", | ||
| "KimiK25ForConditionalGeneration": "modeling_kimi_k25", | ||
| "LagunaForCausalLM": "modeling_laguna", | ||
| "Llama4ForConditionalGeneration": "modeling_llama", | ||
| "LlamaForCausalLM": "modeling_llama", | ||
| "LlavaLlamaModel": "modeling_vila", | ||
| "LlavaNextForConditionalGeneration": "modeling_llava_next", | ||
| "MBartForConditionalGeneration": "modeling_bart", | ||
| "MTPDraftModelForCausalLM": "modeling_speculative", | ||
| "MiniCPMV4_6ForConditionalGeneration": "modeling_minicpmv4_6", | ||
| "MiniMaxM2ForCausalLM": "modeling_minimaxm2", | ||
| "MiniMaxM3SparseForCausalLM": "modeling_minimaxm3", | ||
| "MiniMaxM3SparseForConditionalGeneration": "modeling_minimaxm3", | ||
| "Mistral3ForConditionalGeneration": "modeling_mistral", | ||
| "MistralForCausalLM": "modeling_mistral", | ||
| "MistralLarge3EagleForCausalLM": "modeling_speculative", | ||
| "MistralLarge3ForCausalLM": "modeling_mistral_large3", | ||
| "MixtralForCausalLM": "modeling_mixtral", | ||
| "MllamaForConditionalGeneration": "modeling_mllama", | ||
| "NemotronForCausalLM": "modeling_nemotron", | ||
| "NemotronHForCausalLM": "modeling_nemotron_h", | ||
| "NemotronHPuzzleForCausalLM": "modeling_nemotron_h", | ||
| "NemotronH_Nano_Omni_Reasoning_V3": "modeling_nemotron_nano", | ||
| "NemotronH_Nano_VL_V2": "modeling_nemotron_nano", | ||
| "Phi3ForCausalLM": "modeling_phi3", | ||
| "Phi4MMForCausalLM": "modeling_phi4mm", | ||
| "PixtralForConditionalGeneration": "modeling_mistral", | ||
| "PixtralVisionModel": "modeling_pixtral", | ||
| "Qwen2ForCausalLM": "modeling_qwen", | ||
| "Qwen2ForProcessRewardModel": "modeling_qwen", | ||
| "Qwen2ForRewardModel": "modeling_qwen", | ||
| "Qwen2MoeForCausalLM": "modeling_qwen_moe", | ||
| "Qwen2VLForConditionalGeneration": "modeling_qwen2vl", | ||
| "Qwen2_5_VLForConditionalGeneration": "modeling_qwen2vl", | ||
| "Qwen3ForCausalLM": "modeling_qwen3", | ||
| "Qwen3ForTextEmbedding": "modeling_qwen3", | ||
| "Qwen3MoeForCausalLM": "modeling_qwen3_moe", | ||
| "Qwen3NextForCausalLM": "modeling_qwen3_next", | ||
| "Qwen3VLForConditionalGeneration": "modeling_qwen3vl", | ||
| "Qwen3VLMoeForConditionalGeneration": "modeling_qwen3vl_moe", | ||
| "Qwen3_5ForCausalLM": "modeling_qwen3_5", | ||
| "Qwen3_5ForConditionalGeneration": "modeling_qwen3_5", | ||
| "Qwen3_5MoeForCausalLM": "modeling_qwen3_5", | ||
| "Qwen3_5MoeForConditionalGeneration": "modeling_qwen3_5", | ||
| "QwenImageBenchForConditionalGeneration": "modeling_qwen_image_bench", | ||
| "SeedOssForCausalLM": "modeling_seedoss", | ||
| "SiglipVisionModel": "modeling_siglip", | ||
| "SomeVLModel": "modeling_utils", | ||
| "Starcoder2ForCausalLM": "modeling_starcoder2", | ||
| "Step3p5ForCausalLM": "modeling_step3p7", | ||
| "Step3p7ForConditionalGeneration": "modeling_step3p7vl", | ||
| "T5ForConditionalGeneration": "modeling_t5", | ||
| "WhisperForConditionalGeneration": "modeling_whisper", | ||
| } |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Compare `@register_auto_model` architectures with MODEL_ARCH_TO_MODULE.
set -uo pipefail
python - <<'PY'
import pathlib, re, ast
root = pathlib.Path("tensorrt_llm/_torch/models")
index_src = (root / "_arch_index.py").read_text()
tree = ast.parse(index_src)
tables = {}
for node in tree.body:
if isinstance(node, ast.Assign) and isinstance(node.targets[0], ast.Name):
try:
tables[node.targets[0].id] = ast.literal_eval(node.value)
except ValueError:
pass
arch_index = tables.get("MODEL_ARCH_TO_MODULE", {})
decorated = {}
pat = re.compile(r'`@register_auto_model`\(\s*["\']([^"\']+)["\']')
for path in sorted(root.rglob("modeling_*.py")):
for arch in pat.findall(path.read_text()):
decorated.setdefault(arch, set()).add(path.stem)
missing = sorted(a for a in decorated if a not in arch_index)
extra = sorted(a for a in arch_index if a not in decorated)
wrong = sorted(
(a, arch_index[a], sorted(mods))
for a, mods in decorated.items()
if a in arch_index and arch_index[a] not in mods
)
print("registered but MISSING from MODEL_ARCH_TO_MODULE:", missing)
print("in MODEL_ARCH_TO_MODULE but NOT registered:", extra)
print("mapped to the WRONG module:", wrong)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 327
🏁 Script executed:
#!/bin/bash
set -euo pipefail
root="tensorrt_llm/_torch/models"
printf '%s\n' '--- index and registration references ---'
rg -n -C 4 'MODEL_ARCH_TO_MODULE|ensure_model_registered|register_auto_model|LlavaLlamaModel|PixtralVisionModel' "$root" \
-g '*.py'
printf '%s\n' '--- candidate module outlines ---'
for file in "$root"/_arch_index.py "$root"/modeling_vila.py "$root"/modeling_pixtral.py; do
if [ -f "$file" ]; then
echo "### $file"
ast-grep outline "$file"
fi
done
printf '%s\n' '--- index module existence ---'
python3 - <<'PY'
import ast
import pathlib
root = pathlib.Path("tensorrt_llm/_torch/models")
tree = ast.parse((root / "_arch_index.py").read_text())
table = next(
node.value
for node in tree.body
if isinstance(node, ast.Assign)
and any(isinstance(target, ast.Name) and target.id == "MODEL_ARCH_TO_MODULE"
for target in node.targets)
)
mapping = ast.literal_eval(table)
for architecture, module in mapping.items():
path = root / f"{module}.py"
print(f"{architecture}: {module}.py {'OK' if path.is_file() else 'MISSING'}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
root="tensorrt_llm/_torch/models"
printf '%s\n' '--- relevant lazy-loading code ---'
rg -n -C 8 'def ensure_model_registered|def register_auto_model|MODEL_ARCH_TO_MODULE|import_module' \
"$root/_arch_index.py" "$root/modeling_auto.py" "$root/modeling_utils.py"
printf '%s\n' '--- non-literal registration sites for indexed architectures ---'
rg -n -C 3 'LlavaLlamaModel|PixtralVisionModel|model_architecture|register_auto_model' \
"$root/modeling_vila.py" "$root/modeling_pixtral.py" "$root/modeling_mistral.py"
printf '%s\n' '--- AST comparison, including constant expressions ---'
python3 - <<'PY'
import ast
import pathlib
root = pathlib.Path("tensorrt_llm/_torch/models")
index_tree = ast.parse((root / "_arch_index.py").read_text())
index_node = next(
node.value for node in index_tree.body
if isinstance(node, ast.Assign)
and any(isinstance(target, ast.Name)
and target.id == "MODEL_ARCH_TO_MODULE"
for target in node.targets)
)
index = ast.literal_eval(index_node)
registrations = {}
for path in sorted(root.glob("modeling_*.py")):
tree = ast.parse(path.read_text())
constants = {
node.targets[0].id: ast.literal_eval(node.value)
for node in tree.body
if isinstance(node, ast.Assign)
and len(node.targets) == 1
and isinstance(node.targets[0], ast.Name)
and isinstance(node.value, ast.Constant)
and isinstance(node.value.value, str)
}
for node in ast.walk(tree):
if not isinstance(node, ast.Call):
continue
func = node.func
if not (isinstance(func, ast.Name) and func.id == "register_auto_model"):
continue
if len(node.args) != 1:
continue
arg = node.args[0]
value = None
if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
value = arg.value
elif (isinstance(arg, ast.Attribute)
and isinstance(arg.value, ast.Name)
and arg.value.id in constants
and arg.attr == "model_architecture"):
value = constants[arg.value.id]
if value is not None:
registrations.setdefault(value, set()).add(path.stem)
missing = sorted(set(registrations) - set(index))
extra = sorted(set(index) - set(registrations))
wrong = sorted(
(arch, index[arch], sorted(registrations[arch]))
for arch in sorted(set(index) & set(registrations))
if index[arch] not in registrations[arch]
)
print("registered but missing:", missing)
print("indexed but not found by AST scan:", extra)
print("mapped to a different module:", wrong)
for arch in ("LlavaLlamaModel", "PixtralVisionModel"):
print(f"{arch}: registered in {sorted(registrations.get(arch, ()))}, "
f"indexed as {index.get(arch)!r}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 13967
🏁 Script executed:
#!/bin/bash
set -euo pipefail
root="tensorrt_llm/_torch/models"
python3 - <<'PY'
import pathlib
import re
root = pathlib.Path("tensorrt_llm/_torch/models")
index = {}
entry = re.compile(r'^\s*"([^"]+)":\s*"([^"]+)",\s*$')
for line in (root / "_arch_index.py").read_text().splitlines():
match = entry.match(line)
if match:
index[match.group(1)] = match.group(2)
registrations = {}
unresolved = []
decorator = re.compile(
r'^\s*@(?:[\w.]+\.)?register_auto_model\(\s*([^)]*?)\s*\)\s*$'
)
literal = re.compile(r"""^["']([^"']+)["']$""")
attribute = re.compile(r'^([A-Za-z_]\w*)\.([A-Za-z_]\w*)$')
for path in sorted(root.glob("modeling_*.py")):
lines = path.read_text().splitlines()
constants = {}
for line in lines:
match = re.match(
r'^\s*(\w+)\s*=\s*["\']([^"\']+)["\']\s*$', line
)
if match:
constants[match.group(1)] = match.group(2)
match = re.match(
r'^\s*model_architecture\s*=\s*["\']([^"\']+)["\']\s*$', line
)
if match:
constants["model_architecture"] = match.group(1)
for line_number, line in enumerate(lines, 1):
match = decorator.match(line)
if not match:
continue
argument = match.group(1).strip()
literal_match = literal.match(argument)
if literal_match:
architecture = literal_match.group(1)
else:
attribute_match = attribute.match(argument)
architecture = (
constants.get(attribute_match.group(2))
if attribute_match and attribute_match.group(1) == "VilaConfig"
else None
)
if architecture is None:
unresolved.append(f"{path}:{line_number}: {argument}")
else:
registrations.setdefault(architecture, set()).add(path.stem)
missing = sorted(set(registrations) - set(index))
extra = sorted(set(index) - set(registrations))
wrong = sorted(
(architecture, index[architecture], sorted(registrations[architecture]))
for architecture in sorted(set(index) & set(registrations))
if index[architecture] not in registrations[architecture]
)
print("registered but missing:", missing)
print("indexed but not registered:", extra)
print("mapped to a different module:", wrong)
print("unresolved decorators:", unresolved)
print("LlavaLlamaModel:", sorted(registrations.get("LlavaLlamaModel", ())))
print("PixtralVisionModel:", sorted(registrations.get("PixtralVisionModel", ())))
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 747
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="tensorrt_llm/_torch/models/_arch_index.py"
sed -n '1,220p' "$file"
python3 - <<'PY'
import pathlib
import re
root = pathlib.Path("tensorrt_llm/_torch/models")
text = (root / "_arch_index.py").read_text()
match = re.search(
r"MODEL_ARCH_TO_MODULE\s*=\s*\{(?P<body>.*?)\n\}",
text,
re.S,
)
if not match:
raise SystemExit("MODEL_ARCH_TO_MODULE block not found")
index = dict(re.findall(
r'^\s*"([^"]+)":\s*"([^"]+)",\s*$',
match.group("body"),
re.M,
))
registrations = {}
unresolved = []
decorator = re.compile(
r'^\s*@(?:[\w.]+\.)?register_auto_model\(\s*([^)]*?)\s*\)\s*$'
)
literal = re.compile(r"""^["']([^"']+)["']$""")
attribute = re.compile(r'^([A-Za-z_]\w*)\.([A-Za-z_]\w*)$')
for path in sorted(root.glob("modeling_*.py")):
lines = path.read_text().splitlines()
class_attributes = {}
current_class = None
for line in lines:
class_match = re.match(r'^\s*class\s+(\w+)', line)
if class_match:
current_class = class_match.group(1)
attr_match = re.match(
r'^\s*model_architecture\s*=\s*["\']([^"\']+)["\']\s*$',
line,
)
if attr_match and current_class:
class_attributes[current_class] = attr_match.group(1)
for line_number, line in enumerate(lines, 1):
match = decorator.match(line)
if not match:
continue
argument = match.group(1).strip()
literal_match = literal.match(argument)
if literal_match:
architecture = literal_match.group(1)
else:
attribute_match = attribute.match(argument)
architecture = (
class_attributes.get(attribute_match.group(1))
if attribute_match else None
)
if architecture is None:
unresolved.append(f"{path}:{line_number}: {argument}")
else:
registrations.setdefault(architecture, set()).add(path.stem)
print("registered but missing:", sorted(set(registrations) - set(index)))
print("indexed but not registered:", sorted(set(index) - set(registrations)))
print("mapped to a different module:", sorted(
(architecture, index[architecture], sorted(registrations[architecture]))
for architecture in sorted(set(index) & set(registrations))
if index[architecture] not in registrations[architecture]
))
print("unresolved decorators:", unresolved)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 8682
Add an automated consistency check for MODEL_ARCH_TO_MODULE.
If an architecture is absent from the index, ensure_model_registered skips its import and runtime registration fails. Add a CI check that compares @register_auto_model registrations with MODEL_ARCH_TO_MODULE, including qualified and attribute-based decorators.
🤖 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/_torch/models/_arch_index.py` around lines 18 - 99, Add a CI
consistency check for MODEL_ARCH_TO_MODULE that discovers every architecture
registered by `@register_auto_model`, including qualified and attribute-based
decorator forms, and compares those registrations against the mapping in both
directions. Fail the check when an architecture is missing from either set, so
ensure_model_registered can import and register every decorated model.
Comment on lines
+883
to
+887
| try: | ||
| importlib.import_module(f"tensorrt_llm._torch.models.{module_name}") | ||
| except ImportError as e: | ||
| logger.warning(f"Lazy import of {module_name} for architecture " | ||
| f"{model_arch} failed: {e!r}") |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Preserve the import failure so callers can report it.
The architecture is present in the static index, so a failed import is a real defect, not an unknown architecture. Today the traceback is dropped into a warning, and get_model_architecture then raises Unknown model architecture: <arch>. That message hides the root cause (a missing optional dependency, an incompatible transformers version, or a broken modeling_* module). Under the previous eager import, the original ImportError reached the user.
Record the failure and let the caller include it. For example, keep a module-level dict[str, ImportError] of failed architectures and append the stored error to the Unknown model architecture message.
Also note that only ImportError is caught. A modeling_* module that raises any other exception at import time still propagates, so the two failure modes behave differently.
🤖 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/_torch/models/modeling_utils.py` around lines 883 - 887, Record
caught ImportError instances in a module-level dictionary keyed by architecture
within the lazy import path, while retaining the existing warning. Update
get_model_architecture to append the stored import failure to its
unknown-architecture error when the architecture is indexed but its import
failed. Keep non-ImportError exceptions propagating unchanged.
The OpenAI server and the serve/bench CLIs imported the visual_gen tree at module level, so every plain LLM serving process paid the import cost of the whole visual_gen package. Type-only uses move under TYPE_CHECKING, runtime uses import locally inside the VisualGen code paths, and the isinstance checks go through a sys.modules probe: if visual_gen was never imported, the generator cannot be a VisualGen instance. Signed-off-by: junq <22017000+QiJune@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py (1)
203-234: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAdd a focused model-lookup contract test.
Test coverage summary: No test function covers
_lookup_model_clsorget_registered_model_class. The test file is scheduled intests/integration/test_lists/test-db/l0_a10.yml. Add tests for architecture lookup and the no-architecture fallback.🤖 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/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py` around lines 203 - 234, Add focused tests in the precheck test module covering the model-lookup contract: verify _lookup_model_cls resolves the expected class for a declared architecture and that get_registered_model_class handles the no-architecture case with its documented fallback. Reuse the existing _INTERNAL_APIS setup and registered model symbols, and keep the tests limited to these two lookup behaviors.Sources: Coding guidelines, Path instructions
tensorrt_llm/serve/openai_video_routes.py (1)
1-2: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd the required NVIDIA copyright header to both modified Python files.
Both files begin with
from __future__ import annotationsand have no header before the first import. Add or update the NVIDIA header and set the year to 2026.
tensorrt_llm/serve/openai_video_routes.py#L1-L2: add the header before the future import.tensorrt_llm/serve/visual_gen_utils.py#L1-L2: add the header before the future import.As per coding guidelines, modified source files must contain the NVIDIA copyright header with the year of the latest meaningful modification.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/serve/openai_video_routes.py` around lines 1 - 2, Add the required NVIDIA copyright and SPDX license header with year 2026 before the future import in tensorrt_llm/serve/openai_video_routes.py at lines 1-2 and tensorrt_llm/serve/visual_gen_utils.py at lines 1-2; apply the same header format consistently in both files.Source: Coding guidelines
🧹 Nitpick comments (4)
tests/unittest/others/test_lazy_model_zoo.py (1)
117-143: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd return annotations to all functions.
Add
-> Noneto the ninetest_*functions. Add-> tuple[dict[str, set[str]], dict[str, set[str]]]to_decorated_registrations.Test coverage: nine test functions are present, with no matching entries in
tests/integration/test_lists/. No CBTS artifact is available; coverage 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/others/test_lazy_model_zoo.py` around lines 117 - 143, Add return annotations to every function in the test module: annotate all nine test_* functions with -> None, and annotate _decorated_registrations with -> tuple[dict[str, set[str]], dict[str, set[str]]].Source: Coding guidelines
tensorrt_llm/_torch/models/_arch_index.py (1)
11-15: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReplace the reference to "the lazy-import PR".
The regeneration instruction points to a pull request instead of a checked-in artifact. A maintainer who adds a model cannot find that script later. Name the script path, or state that entries are added by hand.
🤖 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/_torch/models/_arch_index.py` around lines 11 - 15, Update the regeneration comment in _arch_index.py to remove the reference to the lazy-import PR, and instead state that entries should be added by hand or identify the checked-in script path used to regenerate them.tensorrt_llm/_torch/models/modeling_utils.py (1)
882-889: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate the decorator signatures.
decorator(cls)has no parameter or return annotation, andregister_auto_modelreturns an unannotated callable. Add annotations, as the coding guidelines require.🔧 Proposed fix
-def register_auto_model(name: str): - - def decorator(cls): +def register_auto_model(name: str) -> Callable[[type], type]: + + def decorator(cls: type) -> type:As per coding guidelines: "Annotate every function, use
Nonefor procedures".🤖 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/_torch/models/modeling_utils.py` around lines 882 - 889, Update register_auto_model and its nested decorator function to include parameter and return type annotations, including the decorator’s cls parameter and register_auto_model’s callable return type. Ensure the decorator’s return annotation matches the class it returns, and annotate any procedure return as None where applicable.Source: Coding guidelines
tensorrt_llm/serve/openai_server.py (1)
116-124: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate
objin_is_visual_gen_instance.Add
obj: objector another precise type. The function already annotates its return value.As per coding guidelines, every function must have annotations.
Proposed fix
-def _is_visual_gen_instance(obj) -> bool: +def _is_visual_gen_instance(obj: object) -> bool:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/serve/openai_server.py` around lines 116 - 124, Annotate the obj parameter in _is_visual_gen_instance with object (or another precise compatible type), preserving the existing return annotation and visual-gen detection logic.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/_torch/models/modeling_utils.py`:
- Around line 869-871: The built-in module checks use an overly broad prefix
match, misclassifying sibling packages as built-in. Update
_is_builtin_model_class in tensorrt_llm/_torch/models/modeling_utils.py (lines
869-871) and the registrant_module check in tensorrt_llm/inputs/registry.py
(lines 764-766) to accept only the exact built-in package name or names
beginning with that package followed by a dot.
In `@tensorrt_llm/inputs/registry.py`:
- Around line 872-894: Update get_registered_image_model_types,
get_registered_video_model_types, and get_registered_audio_model_types to wrap
each generator expression in tuple(), matching get_registered_model_types.
Preserve the existing provider initialization and modality filtering while
ensuring each method returns a reusable, indexable tuple.
In `@tests/unittest/llmapi/test_config_database.py`:
- Line 27: Add tests/unittest/llmapi/test_config_database.py to the test entries
in llm_config_database.yml, preserving the existing list structure. No test
implementation changes are needed; verify the updated test list with pytest
tests/unittest/.
In `@tests/unittest/others/test_lazy_model_zoo.py`:
- Around line 154-161: Extend the index validation around MODEL_ARCH_TO_MODULE
to detect stale architecture entries: identify keys present in
MODEL_ARCH_TO_MODULE but absent from arch_truth, and assert that this set is
empty with a diagnostic listing the stale entries. Preserve the existing missing
and wrong-module checks.
---
Outside diff comments:
In `@tensorrt_llm/serve/openai_video_routes.py`:
- Around line 1-2: Add the required NVIDIA copyright and SPDX license header
with year 2026 before the future import in
tensorrt_llm/serve/openai_video_routes.py at lines 1-2 and
tensorrt_llm/serve/visual_gen_utils.py at lines 1-2; apply the same header
format consistently in both files.
In `@tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py`:
- Around line 203-234: Add focused tests in the precheck test module covering
the model-lookup contract: verify _lookup_model_cls resolves the expected class
for a declared architecture and that get_registered_model_class handles the
no-architecture case with its documented fallback. Reuse the existing
_INTERNAL_APIS setup and registered model symbols, and keep the tests limited to
these two lookup behaviors.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/_arch_index.py`:
- Around line 11-15: Update the regeneration comment in _arch_index.py to remove
the reference to the lazy-import PR, and instead state that entries should be
added by hand or identify the checked-in script path used to regenerate them.
In `@tensorrt_llm/_torch/models/modeling_utils.py`:
- Around line 882-889: Update register_auto_model and its nested decorator
function to include parameter and return type annotations, including the
decorator’s cls parameter and register_auto_model’s callable return type. Ensure
the decorator’s return annotation matches the class it returns, and annotate any
procedure return as None where applicable.
In `@tensorrt_llm/serve/openai_server.py`:
- Around line 116-124: Annotate the obj parameter in _is_visual_gen_instance
with object (or another precise compatible type), preserving the existing return
annotation and visual-gen detection logic.
In `@tests/unittest/others/test_lazy_model_zoo.py`:
- Around line 117-143: Add return annotations to every function in the test
module: annotate all nine test_* functions with -> None, and annotate
_decorated_registrations with -> tuple[dict[str, set[str]], dict[str,
set[str]]].
🪄 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: f97f4a63-e85b-46ec-a97d-69425f5083a5
📒 Files selected for processing (19)
tensorrt_llm/__init__.pytensorrt_llm/_torch/models/__init__.pytensorrt_llm/_torch/models/_arch_index.pytensorrt_llm/_torch/models/modeling_auto.pytensorrt_llm/_torch/models/modeling_mistral.pytensorrt_llm/_torch/models/modeling_utils.pytensorrt_llm/_torch/pyexecutor/model_loader.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/commands/serve.pytensorrt_llm/commands/utils.pytensorrt_llm/executor/proxy.pytensorrt_llm/executor/worker.pytensorrt_llm/inputs/registry.pytensorrt_llm/serve/openai_server.pytensorrt_llm/serve/openai_video_routes.pytensorrt_llm/serve/visual_gen_utils.pytests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.pytests/unittest/llmapi/test_config_database.pytests/unittest/others/test_lazy_model_zoo.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tensorrt_llm/_torch/pyexecutor/model_loader.py
- tensorrt_llm/_torch/models/init.py
- tensorrt_llm/init.py
Comment on lines
+869
to
+871
| def _is_builtin_model_class(cls) -> bool: | ||
| return getattr(cls, "__module__", | ||
| "").startswith("tensorrt_llm._torch.models") |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Built-in detection uses a non-boundary prefix match in both registries. Both sites classify a module as built-in with startswith("tensorrt_llm._torch.models"), which also matches sibling packages such as tensorrt_llm._torch.models_custom. An external registration from such a module is then treated as built-in and is dropped when the slot is already filled, which inverts the stated precedence rule.
tensorrt_llm/_torch/models/modeling_utils.py#L869-L871: in_is_builtin_model_class, accept the exact package name or a name starting with the package name plus a dot.tensorrt_llm/inputs/registry.py#L764-L766: apply the same boundary test toregistrant_moduleagainst_BUILTIN_MODULE_PREFIX.
📍 Affects 2 files
tensorrt_llm/_torch/models/modeling_utils.py#L869-L871(this comment)tensorrt_llm/inputs/registry.py#L764-L766
🤖 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/_torch/models/modeling_utils.py` around lines 869 - 871, The
built-in module checks use an overly broad prefix match, misclassifying sibling
packages as built-in. Update _is_builtin_model_class in
tensorrt_llm/_torch/models/modeling_utils.py (lines 869-871) and the
registrant_module check in tensorrt_llm/inputs/registry.py (lines 764-766) to
accept only the exact built-in package name or names beginning with that package
followed by a dot.
| from starlette.middleware.base import BaseHTTPMiddleware | ||
|
|
||
| from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPING | ||
| from tensorrt_llm._torch.models.modeling_utils import get_registered_model_class |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline tests/unittest/llmapi/test_config_database.py \
--items all --type function
rg -n -C 3 \
'get_registered_model_class|_get_default_values_for_config|model_defaults' \
tests/unittest/llmapi/test_config_database.py
(fd -t f -e yml -e yaml . tests/integration/test_lists || true) |
while IFS= read -r list_file; do
rg -n -C 2 'test_config_database|llmapi' "$list_file" || true
doneRepository: NVIDIA/TensorRT-LLM
Length of output: 28378
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- diff ---'
git diff -- tests/unittest/llmapi/test_config_database.py
printf '%s\n' '--- file header and helper ---'
sed -n '1,125p' tests/unittest/llmapi/test_config_database.py
printf '%s\n' '--- matching test-list entries with file names ---'
while IFS= read -r list_file; do
matches=$(rg -n '(^|/)(test_config_database\.py|test_config_database)(:| |$)' "$list_file" || true)
if [ -n "$matches" ]; then
printf '%s\n%s\n' "$list_file" "$matches"
fi
done < <(fd -t f -e yml -e yaml . tests/integration/test_lists)Repository: NVIDIA/TensorRT-LLM
Length of output: 4838
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- matching test-list entries ---'
find tests/integration/test_lists -type f \( -name '*.yml' -o -name '*.yaml' \) -print0 |
while IFS= read -r -d '' list_file; do
if rg -n '(^|/)(test_config_database\.py|test_config_database)(:| |$)' "$list_file"; then
printf 'FILE: %s\n' "$list_file"
fi
done
printf '%s\n' '--- helper call sites ---'
rg -n -C 4 '_get_default_values_for_config|test_database_yaml_config_does_not_set_default_leaves' \
tests/unittest/llmapi/test_config_database.py
printf '%s\n' '--- registry definitions and call sites ---'
rg -n -C 5 'def get_registered_model_class|MODEL.*REGISTR|register.*model|import.*model' \
tensorrt_llm/_torch/models tensorrt_llm/_torch/models/modeling_utils.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- custom model defaults ---'
rg -n -C 4 'def get_model_defaults|get_model_defaults\s*=|get_model_defaults\(' \
tensorrt_llm/_torch/models --glob '*.py'
printf '%s\n' '--- lookup architectures and models ---'
rg -n '^[[:space:]]*(model|arch|config_path):' examples/configs/database/lookup.yaml \
examples/configs/curated/lookup.yaml | head -160
printf '%s\n' '--- architecture index entries for custom-default classes ---'
rg -n -C 2 'MODEL_ARCH_TO_MODULE|MODEL_CLASS_TO_MODULE|Deepseek|Qwen|Llama|Whisper|Gemma|Phi|Mistral' \
tensorrt_llm/_torch/models/_arch_index.py | head -240
printf '%s\n' '--- test module imports of model providers ---'
rg -n '^import |^from ' tests/unittest/llmapi/test_config_database.py \
| rg 'model|torch|tensorrt_llm' || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 30942
🏁 Script executed:
python3 - <<'PY'
import ast
from pathlib import Path
test_path = Path("tests/unittest/llmapi/test_config_database.py")
source = test_path.read_text(encoding="utf-8")
tree = ast.parse(source)
functions = {
node.name: node
for node in ast.walk(tree)
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
}
helper = functions["_get_default_values_for_config"]
coverage_test = functions["test_database_yaml_config_does_not_set_default_leaves"]
helper_calls_resolver = any(
isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "get_registered_model_class"
for node in ast.walk(helper)
)
coverage_test_calls_helper = any(
isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "_get_default_values_for_config"
for node in ast.walk(coverage_test)
)
provider_imported_by_test = "modeling_deepseekv4" in source
lookup = Path("examples/configs/curated/lookup.yaml").read_text(encoding="utf-8")
provider = Path("tensorrt_llm/_torch/models/modeling_deepseekv4.py").read_text(
encoding="utf-8"
)
arch_index = Path("tensorrt_llm/_torch/models/_arch_index.py").read_text(
encoding="utf-8"
)
print(f"helper_calls_get_registered_model_class={helper_calls_resolver}")
print(f"parametrized_test_calls_helper={coverage_test_calls_helper}")
print(f"lookup_contains_DeepseekV4ForCausalLM={'DeepseekV4ForCausalLM' in lookup}")
print(f"provider_defines_get_model_defaults={'def get_model_defaults' in provider}")
print(
"arch_index_maps_DeepseekV4ForCausalLM="
f"{'\"DeepseekV4ForCausalLM\": \"modeling_deepseekv4\"' in arch_index}"
)
print(f"test_directly_imports_modeling_deepseekv4={provider_imported_by_test}")
assert helper_calls_resolver
assert coverage_test_calls_helper
assert "DeepseekV4ForCausalLM" in lookup
assert "def get_model_defaults" in provider
assert '"DeepseekV4ForCausalLM": "modeling_deepseekv4"' in arch_index
assert not provider_imported_by_test
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 260
🏁 Script executed:
python3 - <<'PY'
import ast
from pathlib import Path
test_path = Path("tests/unittest/llmapi/test_config_database.py")
source = test_path.read_text(encoding="utf-8")
tree = ast.parse(source)
functions = {
node.name: node
for node in ast.walk(tree)
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
}
helper = functions["_get_default_values_for_config"]
coverage_test = functions["test_database_yaml_config_does_not_set_default_leaves"]
helper_calls_resolver = any(
isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "get_registered_model_class"
for node in ast.walk(helper)
)
coverage_test_calls_helper = any(
isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "_get_default_values_for_config"
for node in ast.walk(coverage_test)
)
provider_imported_by_test = "modeling_deepseekv4" in source
lookup = Path("examples/configs/curated/lookup.yaml").read_text(encoding="utf-8")
provider = Path("tensorrt_llm/_torch/models/modeling_deepseekv4.py").read_text(
encoding="utf-8"
)
arch_index = Path("tensorrt_llm/_torch/models/_arch_index.py").read_text(
encoding="utf-8"
)
deepseek_arch_entry = '"DeepseekV4ForCausalLM": "modeling_deepseekv4"'
print(f"helper_calls_get_registered_model_class={helper_calls_resolver}")
print(f"parametrized_test_calls_helper={coverage_test_calls_helper}")
print(f"lookup_contains_DeepseekV4ForCausalLM={'DeepseekV4ForCausalLM' in lookup}")
print(f"provider_defines_get_model_defaults={'def get_model_defaults' in provider}")
print(f"arch_index_maps_DeepseekV4ForCausalLM={deepseek_arch_entry in arch_index}")
print(f"test_directly_imports_modeling_deepseekv4={provider_imported_by_test}")
assert helper_calls_resolver
assert coverage_test_calls_helper
assert "DeepseekV4ForCausalLM" in lookup
assert "def get_model_defaults" in provider
assert deepseek_arch_entry in arch_index
assert not provider_imported_by_test
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 413
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- test-db list files ---'
find tests/integration/test_lists -type f -path '*test-db*' \
\( -name '*.yml' -o -name '*.yaml' \) -print | sort
printf '%s\n' '--- qa list files ---'
find tests/integration/test_lists -type f -path '*qa*' \
\( -name '*.yml' -o -name '*.yaml' \) -print | sort
printf '%s\n' '--- list files containing the llmapi unittest group ---'
find tests/integration/test_lists -type f \
\( -name '*.yml' -o -name '*.yaml' \) -print0 |
while IFS= read -r -d '' list_file; do
if rg -q 'unittest/llmapi/' "$list_file"; then
printf '%s\n' "$list_file"
fi
done | sortRepository: NVIDIA/TensorRT-LLM
Length of output: 4861
Add test_config_database.py to the test list.
Test coverage summary: No test function changed. The existing parametrized test covers _get_default_values_for_config, including lazy loading for DeepseekV4ForCausalLM. Add unittest/llmapi/test_config_database.py to tests/integration/test_lists/qa/llm_config_database.yml. Run pytest tests/unittest/ before merge.
🤖 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_config_database.py` at line 27, Add
tests/unittest/llmapi/test_config_database.py to the test entries in
llm_config_database.yml, preserving the existing list structure. No test
implementation changes are needed; verify the updated test list with pytest
tests/unittest/.
Sources: Coding guidelines, Path instructions
Comment on lines
+154
to
+161
| missing = set(arch_truth) - set(MODEL_ARCH_TO_MODULE) | ||
| assert not missing, f"architectures missing from _arch_index: {missing}" | ||
| wrong = { | ||
| arch: (MODEL_ARCH_TO_MODULE[arch], arch_truth[arch]) | ||
| for arch in MODEL_ARCH_TO_MODULE | ||
| if arch in arch_truth and MODEL_ARCH_TO_MODULE[arch] not in arch_truth[arch] | ||
| } | ||
| assert not wrong, f"index points at the wrong module: {wrong}" |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Check stale model-architecture index entries.
The test rejects stale entries from MULTIMODAL_MODEL_TYPE_TO_MODULE, but it does not reject stale entries from MODEL_ARCH_TO_MODULE. A stale architecture entry can make lazy resolution import a provider that no longer registers that architecture.
Proposed test update
missing = set(arch_truth) - set(MODEL_ARCH_TO_MODULE)
assert not missing, f"architectures missing from _arch_index: {missing}"
+ stale = set(MODEL_ARCH_TO_MODULE) - set(arch_truth)
+ assert not stale, f"stale architectures in _arch_index: {stale}"
wrong = {📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| missing = set(arch_truth) - set(MODEL_ARCH_TO_MODULE) | |
| assert not missing, f"architectures missing from _arch_index: {missing}" | |
| wrong = { | |
| arch: (MODEL_ARCH_TO_MODULE[arch], arch_truth[arch]) | |
| for arch in MODEL_ARCH_TO_MODULE | |
| if arch in arch_truth and MODEL_ARCH_TO_MODULE[arch] not in arch_truth[arch] | |
| } | |
| assert not wrong, f"index points at the wrong module: {wrong}" | |
| missing = set(arch_truth) - set(MODEL_ARCH_TO_MODULE) | |
| assert not missing, f"architectures missing from _arch_index: {missing}" | |
| stale = set(MODEL_ARCH_TO_MODULE) - set(arch_truth) | |
| assert not stale, f"stale architectures in _arch_index: {stale}" | |
| wrong = { | |
| arch: (MODEL_ARCH_TO_MODULE[arch], arch_truth[arch]) | |
| for arch in MODEL_ARCH_TO_MODULE | |
| if arch in arch_truth and MODEL_ARCH_TO_MODULE[arch] not in arch_truth[arch] | |
| } | |
| assert not wrong, f"index points at the wrong module: {wrong}" |
🤖 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/others/test_lazy_model_zoo.py` around lines 154 - 161, Extend
the index validation around MODEL_ARCH_TO_MODULE to detect stale architecture
entries: identify keys present in MODEL_ARCH_TO_MODULE but absent from
arch_truth, and assert that this set is empty with a diagnostic listing the
stale entries. Preserve the existing missing and wrong-module checks.
1 task
…ority, multimodal registry Review follow-ups on the lazy-loading change, squashed. Resolver API. Architecture lookups collapse into a single entry point that imports the built-in provider on demand: get_registered_model_class and get_registered_vision_encoder (ensure_model_registered becomes the private _ensure_model_registered). Each resolver short-circuits on its own registry only, so an external registration satisfies the model-class lookup without importing the built-in module, while a vision-encoder or placeholder lookup with an empty slot still pulls the provider in. All consumers (modeling_auto, model_loader's transceiver preference, get_model_architecture, the cache-transceiver precheck, test_config_database) go through the resolvers instead of pairing a manual import with a raw MODEL_CLASS_MAPPING.get(). Registration priority. One rule, applied by every registry (model class, vision encoder, placeholder metadata): built-in registrations only fill empty slots and never overwrite, external registrations always win. This is equivalent to main's eager order -- built-ins ran first there, and no architecture is double-registered among built-ins -- and needs no registrant bookkeeping. register_auto_model records the declared architectures on the class itself so register_vision_encoder no longer scans the mapping by identity (an external override used to make the built-in provider's import raise). Multimodal registry. The placeholder registry resolves its provider module on demand through a new static MULTIMODAL_MODEL_TYPE_TO_MODULE index, so model_type-keyed queries (trtllm-bench dataset prep, quickstart) work in a fresh process that never loads a model; enumeration APIs import all indexed providers first. Error handling. Lazy-import failures only swallow "the requested module itself does not exist" (both the model zoo and the PEP 562 package fallbacks); a missing dependency inside an existing module propagates instead of surfacing as unknown architecture or a missing attribute. The Llama4 check in py_executor probes sys.modules and uses a real isinstance; the VisualGen endpoint tests patch the probe helper. Drops the SomeVLModel index entry picked up from a docstring example. Adds fresh-process tests pinning the lazy contracts: import stays thin, attribute access resolves and caches, the static index matches the registration decorators, external registrations survive built-in provider imports (including the Qwen3VL vision-encoder path), and placeholder lookups resolve without a loaded model. Signed-off-by: junq <22017000+QiJune@users.noreply.github.com>
Contributor
|
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. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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/_torch/pyexecutor/py_executor.py`:
- Line 4885: Update the comment near the token-range check to replace the
Unicode hyphen in “token‐range” with the ASCII “-”, resolving Ruff RUF003
without changing the surrounding logic.
In `@tests/unittest/others/test_lazy_model_zoo.py`:
- Around line 183-213: Update test_resolver_keeps_existing_registration to
execute its registry assertions and cleanup through _run_fresh, isolating
MODEL_CLASS_MAPPING and imported provider state from the surrounding test
process. Preserve the existing-registration assertion and unknown-architecture
check within that fresh execution, and remove the unconditional restoration that
overwrites prior global state.
- Around line 64-343: The new tests in
tests/unittest/others/test_lazy_model_zoo.py are not included in any test-list
configuration, leaving their coverage insufficient. Add
tests/unittest/others/test_lazy_model_zoo.py to the appropriate test-db or qa
test list and run pytest tests/unittest/;
tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py requires no
direct change because its existing l0_b200.yml entry already covers the
_create_server-only modification.
🪄 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: 4c5ae1b9-dfe1-46f7-89d0-86cbba5f8fee
📒 Files selected for processing (24)
tensorrt_llm/__init__.pytensorrt_llm/_torch/models/__init__.pytensorrt_llm/_torch/models/_arch_index.pytensorrt_llm/_torch/models/modeling_auto.pytensorrt_llm/_torch/models/modeling_llama.pytensorrt_llm/_torch/models/modeling_mistral.pytensorrt_llm/_torch/models/modeling_utils.pytensorrt_llm/_torch/pyexecutor/model_loader.pytensorrt_llm/_torch/pyexecutor/perf_metrics_manager.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_utils.pytensorrt_llm/commands/serve.pytensorrt_llm/commands/utils.pytensorrt_llm/executor/proxy.pytensorrt_llm/executor/worker.pytensorrt_llm/inputs/registry.pytensorrt_llm/serve/openai_server.pytensorrt_llm/serve/openai_video_routes.pytensorrt_llm/serve/responses_utils.pytensorrt_llm/serve/visual_gen_utils.pytests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.pytests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.pytests/unittest/llmapi/test_config_database.pytests/unittest/others/test_lazy_model_zoo.py
🚧 Files skipped from review as they are similar to previous changes (21)
- tests/unittest/llmapi/test_config_database.py
- tensorrt_llm/serve/openai_video_routes.py
- tensorrt_llm/executor/proxy.py
- tensorrt_llm/_torch/models/modeling_mistral.py
- tensorrt_llm/serve/responses_utils.py
- tensorrt_llm/executor/worker.py
- tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py
- tensorrt_llm/_torch/models/modeling_llama.py
- tensorrt_llm/serve/visual_gen_utils.py
- tensorrt_llm/init.py
- tensorrt_llm/_torch/models/modeling_auto.py
- tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py
- tensorrt_llm/_utils.py
- tensorrt_llm/commands/serve.py
- tensorrt_llm/_torch/models/init.py
- tensorrt_llm/_torch/models/_arch_index.py
- tensorrt_llm/serve/openai_server.py
- tensorrt_llm/commands/utils.py
- tensorrt_llm/_torch/pyexecutor/model_loader.py
- tensorrt_llm/_torch/models/modeling_utils.py
- tensorrt_llm/inputs/registry.py
| # Only skip token‐range checks for Llama4 when the request has multimodal data | ||
| if isinstance(self.model_engine.model, | ||
| Llama4ForConditionalGeneration): | ||
| # Only skip token‐range checks for Llama4 when the request has |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Replace the Unicode hyphen in token‐range.
Ruff RUF003 still flags Line 4885. Replace ‐ with ASCII -. This is the same finding reported in the previous review and remains in the current code.
Proposed fix
- # Only skip token‐range checks for Llama4 when the request has
+ # Only skip token-range checks for Llama4 when the request has🧰 Tools
🪛 Ruff (0.16.1)
[warning] 4885-4885: Comment contains ambiguous ‐ (HYPHEN). Did you mean - (HYPHEN-MINUS)?
(RUF003)
🤖 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/_torch/pyexecutor/py_executor.py` at line 4885, Update the
comment near the token-range check to replace the Unicode hyphen in
“token‐range” with the ASCII “-”, resolving Ruff RUF003 without changing the
surrounding logic.
Source: Linters/SAST tools
Comment on lines
+64
to
+343
| def test_import_does_not_load_zoo_or_visual_gen(): | ||
| _run_fresh( | ||
| textwrap.dedent("""\ | ||
| import sys | ||
| import tensorrt_llm | ||
|
|
||
| loaded = [ | ||
| m for m in sys.modules | ||
| if m.startswith("tensorrt_llm._torch.models.modeling_") | ||
| or m == "tensorrt_llm.visual_gen" | ||
| or m.startswith("tensorrt_llm.visual_gen.") | ||
| ] | ||
| assert not loaded, f"import tensorrt_llm eagerly loaded: {loaded}" | ||
| """) | ||
| ) | ||
|
|
||
|
|
||
| def test_lazy_attribute_access_resolves_and_caches(): | ||
| _run_fresh( | ||
| textwrap.dedent("""\ | ||
| import tensorrt_llm | ||
|
|
||
| sp = tensorrt_llm.SamplingParams | ||
| assert sp is tensorrt_llm.SamplingParams # cached in globals() | ||
| assert "SamplingParams" in vars(tensorrt_llm) | ||
| assert "SamplingParams" in dir(tensorrt_llm) | ||
|
|
||
| try: | ||
| tensorrt_llm.definitely_not_an_attribute | ||
| except AttributeError: | ||
| pass | ||
| else: | ||
| raise AssertionError("missing attribute did not raise") | ||
| """) | ||
| ) | ||
|
|
||
|
|
||
| def test_placeholder_registry_resolves_in_fresh_process(): | ||
| # trtllm-bench dataset prep queries the placeholder registry by | ||
| # model_type in a process that never loads a model; the registry must | ||
| # import the provider on demand. | ||
| _run_fresh( | ||
| textwrap.dedent("""\ | ||
| from tensorrt_llm.inputs.registry import MULTIMODAL_PLACEHOLDER_REGISTRY | ||
|
|
||
| assert MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid("llama4", "image") | ||
| assert MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder( | ||
| "llama4", "image") | ||
| assert "qwen2_vl" in MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_model_types() | ||
| """) | ||
| ) | ||
|
|
||
|
|
||
| def _decorated_registrations(): | ||
| """AST-scan the modeling files for the registrations the index mirrors.""" | ||
| arch_to_modules = {} | ||
| model_type_to_modules = {} | ||
| for path in sorted(_MODELS_DIR.glob("*.py")): | ||
| tree = ast.parse(path.read_text()) | ||
| for node in ast.walk(tree): | ||
| if not isinstance(node, ast.Call): | ||
| continue | ||
| func = node.func | ||
| name = getattr(func, "id", None) or getattr(func, "attr", None) | ||
| if name == "register_auto_model": | ||
| if node.args and isinstance(node.args[0], ast.Constant): | ||
| arch_to_modules.setdefault(node.args[0].value, set()).add(path.stem) | ||
| elif name in ("register_input_processor", "set_placeholder_metadata"): | ||
| model_type = None | ||
| if name == "register_input_processor": | ||
| if len(node.args) >= 2 and isinstance(node.args[1], ast.Constant): | ||
| model_type = node.args[1].value | ||
| elif node.args and isinstance(node.args[0], ast.Constant): | ||
| model_type = node.args[0].value | ||
| for kw in node.keywords: | ||
| if kw.arg == "model_type" and isinstance(kw.value, ast.Constant): | ||
| model_type = kw.value.value | ||
| if model_type is not None: | ||
| model_type_to_modules.setdefault(model_type, set()).add(path.stem) | ||
| return arch_to_modules, model_type_to_modules | ||
|
|
||
|
|
||
| def test_arch_index_matches_decorators(): | ||
| from tensorrt_llm._torch.models._arch_index import ( | ||
| MODEL_ARCH_TO_MODULE, | ||
| MULTIMODAL_MODEL_TYPE_TO_MODULE, | ||
| ) | ||
|
|
||
| arch_truth, model_type_truth = _decorated_registrations() | ||
|
|
||
| missing = set(arch_truth) - set(MODEL_ARCH_TO_MODULE) | ||
| assert not missing, f"architectures missing from _arch_index: {missing}" | ||
| wrong = { | ||
| arch: (MODEL_ARCH_TO_MODULE[arch], arch_truth[arch]) | ||
| for arch in MODEL_ARCH_TO_MODULE | ||
| if arch in arch_truth and MODEL_ARCH_TO_MODULE[arch] not in arch_truth[arch] | ||
| } | ||
| assert not wrong, f"index points at the wrong module: {wrong}" | ||
|
|
||
| missing = set(model_type_truth) - set(MULTIMODAL_MODEL_TYPE_TO_MODULE) | ||
| assert not missing, f"model types missing from _arch_index: {missing}" | ||
| stale = set(MULTIMODAL_MODEL_TYPE_TO_MODULE) - set(model_type_truth) | ||
| assert not stale, f"stale model types in _arch_index: {stale}" | ||
| wrong = { | ||
| mt: (MULTIMODAL_MODEL_TYPE_TO_MODULE[mt], model_type_truth[mt]) | ||
| for mt in MULTIMODAL_MODEL_TYPE_TO_MODULE | ||
| if MULTIMODAL_MODEL_TYPE_TO_MODULE[mt] not in model_type_truth[mt] | ||
| } | ||
| assert not wrong, f"index points at the wrong module: {wrong}" | ||
|
|
||
|
|
||
| def test_models_package_missing_submodule_is_attribute_error(): | ||
| # The PEP 562 fallback must translate only "no such submodule" into | ||
| # AttributeError (so hasattr works), same as the top-level package. | ||
| import tensorrt_llm._torch.models as torch_models | ||
|
|
||
| assert not hasattr(torch_models, "modeling_definitely_not_a_model") | ||
|
|
||
|
|
||
| def test_resolver_keeps_existing_registration(): | ||
| # A registration made by user code (--custom_module_dirs) must win over | ||
| # the built-in module, exactly as it does with the eager import order on | ||
| # main where the zoo loads first and custom code overrides it. | ||
| from tensorrt_llm._torch.models.modeling_utils import ( | ||
| MODEL_CLASS_MAPPING, | ||
| get_registered_model_class, | ||
| ) | ||
|
|
||
| arch = "MistralForCausalLM" | ||
|
|
||
| class _CustomStub: | ||
| pass | ||
|
|
||
| MODEL_CLASS_MAPPING[arch] = _CustomStub | ||
| try: | ||
| assert get_registered_model_class(arch) is _CustomStub, ( | ||
| "resolving an architecture overrode its existing registration" | ||
| ) | ||
| finally: | ||
| # Restore the real class rather than deleting the entry: the | ||
| # built-in decorator skips occupied slots, so if the provider was | ||
| # (or gets) imported while the stub held the slot, a bare delete | ||
| # would leave the architecture unresolvable for the rest of the | ||
| # process (module imports are cached). | ||
| from tensorrt_llm._torch.models.modeling_mistral import MistralForCausalLM | ||
|
|
||
| MODEL_CLASS_MAPPING[arch] = MistralForCausalLM | ||
|
|
||
| # Unknown architectures resolve to None, left to the caller's handling. | ||
| assert get_registered_model_class("DefinitelyNotARegisteredArch") is None | ||
|
|
||
|
|
||
| def test_builtin_decorator_does_not_override_external_registration(): | ||
| # The priority guarantee lives in the registry itself: a built-in module | ||
| # may run its decorators after an external registration (direct imports, | ||
| # sibling architectures from the same module), and must not clobber it. | ||
| from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPING, register_auto_model | ||
|
|
||
| arch = "LazyZooTestOnlyArch" | ||
| assert arch not in MODEL_CLASS_MAPPING | ||
|
|
||
| class _External: | ||
| pass | ||
|
|
||
| class _Builtin: | ||
| pass | ||
|
|
||
| _Builtin.__module__ = "tensorrt_llm._torch.models.modeling_fake" | ||
|
|
||
| try: | ||
| register_auto_model(arch)(_External) | ||
| register_auto_model(arch)(_Builtin) | ||
| assert MODEL_CLASS_MAPPING[arch] is _External, ( | ||
| "built-in decorator overrode an external registration" | ||
| ) | ||
|
|
||
| # The other direction stays last-wins: external code registering | ||
| # over a built-in is exactly the --custom_module_dirs use case. | ||
| del MODEL_CLASS_MAPPING[arch] | ||
| register_auto_model(arch)(_Builtin) | ||
| register_auto_model(arch)(_External) | ||
| assert MODEL_CLASS_MAPPING[arch] is _External | ||
| finally: | ||
| MODEL_CLASS_MAPPING.pop(arch, None) | ||
|
|
||
|
|
||
| def test_custom_registration_survives_direct_provider_import(): | ||
| # Regression for the full production path: a custom implementation is | ||
| # registered, then some other code path imports the built-in provider | ||
| # directly (e.g. model_loader's post-transform profile registry imports | ||
| # modeling_llama) -- the custom registration must survive the built-in | ||
| # module's decorators. Fresh process so modeling_llama is genuinely not | ||
| # imported yet when the custom registration happens. | ||
| _run_fresh( | ||
| textwrap.dedent("""\ | ||
| import importlib | ||
| from tensorrt_llm._torch.models.modeling_utils import ( | ||
| MODEL_CLASS_MAPPING, register_auto_model) | ||
|
|
||
| @register_auto_model("LlamaForCausalLM") | ||
| class CustomLlama: | ||
| pass | ||
|
|
||
| importlib.import_module( | ||
| "tensorrt_llm._torch.models.modeling_llama") | ||
| assert MODEL_CLASS_MAPPING["LlamaForCausalLM"] is CustomLlama, ( | ||
| "direct import of the built-in provider overrode the custom " | ||
| "registration") | ||
| """) | ||
| ) | ||
|
|
||
|
|
||
| def test_external_multimodal_override_keeps_provider_importable(): | ||
| # An external override of a multimodal architecture must not break the | ||
| # built-in provider's import: register_vision_encoder used to locate the | ||
| # freshly decorated class in MODEL_CLASS_MAPPING by identity and raise | ||
| # when the external registration had won the slot. The built-in vision | ||
| # encoder still fills the empty sibling slot, like the eager import | ||
| # order on main. | ||
| _run_fresh( | ||
| textwrap.dedent("""\ | ||
| import importlib | ||
| from tensorrt_llm._torch.models.modeling_utils import ( | ||
| MODEL_CLASS_MAPPING, MODEL_CLASS_VISION_ENCODER_MAPPING, | ||
| register_auto_model) | ||
|
|
||
| arch = "Qwen3VLForConditionalGeneration" | ||
|
|
||
| @register_auto_model(arch) | ||
| class CustomQwen3VL: | ||
| pass | ||
|
|
||
| importlib.import_module( | ||
| "tensorrt_llm._torch.models.modeling_qwen3vl") | ||
|
|
||
| assert MODEL_CLASS_MAPPING[arch] is CustomQwen3VL, ( | ||
| "built-in provider import overrode the external registration") | ||
| assert MODEL_CLASS_VISION_ENCODER_MAPPING.get(arch) is not None, ( | ||
| "built-in vision encoder did not fill the empty sibling slot") | ||
| """) | ||
| ) | ||
|
|
||
|
|
||
| def test_external_sibling_registrations_not_clobbered(): | ||
| # When the external implementation brings its own vision encoder and | ||
| # placeholder metadata, a later built-in import must not overwrite them. | ||
| _run_fresh( | ||
| textwrap.dedent("""\ | ||
| import importlib | ||
| from tensorrt_llm._torch.models.modeling_utils import ( | ||
| MODEL_CLASS_VISION_ENCODER_MAPPING, register_auto_model, | ||
| register_vision_encoder) | ||
| from tensorrt_llm.inputs.registry import ( | ||
| MULTIMODAL_PLACEHOLDER_REGISTRY, MultimodalPlaceholderMetadata) | ||
|
|
||
| arch = "Qwen3VLForConditionalGeneration" | ||
|
|
||
| class CustomEncoder: | ||
| pass | ||
|
|
||
| @register_vision_encoder(CustomEncoder) | ||
| @register_auto_model(arch) | ||
| class CustomQwen3VL: | ||
| pass | ||
|
|
||
| custom_metadata = MultimodalPlaceholderMetadata( | ||
| placeholder_map={"image": "<custom_image>"}) | ||
| MULTIMODAL_PLACEHOLDER_REGISTRY.set_placeholder_metadata( | ||
| "qwen3_vl", custom_metadata, registrant_module=__name__) | ||
|
|
||
| importlib.import_module( | ||
| "tensorrt_llm._torch.models.modeling_qwen3vl") | ||
|
|
||
| assert MODEL_CLASS_VISION_ENCODER_MAPPING[arch][0] is CustomEncoder, ( | ||
| "built-in import clobbered the external vision encoder") | ||
| assert MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder_metadata( | ||
| "qwen3_vl") is custom_metadata, ( | ||
| "built-in import clobbered the external placeholder metadata") | ||
| """) | ||
| ) |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
fd -t f . tests/integration/test_lists | sort
rg -n -C 2 'test_lazy_model_zoo|test_trtllm_serve_endpoints' \
tests/integration/test_lists || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 5373
🏁 Script executed:
git status --short
git diff --stat -- tests/unittest/others/test_lazy_model_zoo.py tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py
git diff --unified=3 -- tests/unittest/others/test_lazy_model_zoo.py tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py
rg -n 'unittest/others|test_lazy_model_zoo|test_trtllm_serve_endpoints' tests/integration/test_listsRepository: NVIDIA/TensorRT-LLM
Length of output: 2915
Add test_lazy_model_zoo.py to a test list.
- Added tests:
test_import_does_not_load_zoo_or_visual_gen,test_lazy_attribute_access_resolves_and_caches,test_placeholder_registry_resolves_in_fresh_process,test_arch_index_matches_decorators,test_models_package_missing_submodule_is_attribute_error,test_resolver_keeps_existing_registration,test_builtin_decorator_does_not_override_external_registration,test_custom_registration_survives_direct_provider_import,test_external_multimodal_override_keeps_provider_importable, andtest_external_sibling_registrations_not_clobbered. test_trtllm_serve_endpoints.pychanges only_create_server; the file is listed intests/integration/test_lists/test-db/l0_b200.yml.test_lazy_model_zoo.pyis absent from alltest-db/andqa/lists. Coverage verdict: insufficient.- Add the file to the appropriate test list and run
pytest tests/unittest/.
📍 Affects 2 files
tests/unittest/others/test_lazy_model_zoo.py#L64-L343(this comment)tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py#L310-L324
🤖 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/others/test_lazy_model_zoo.py` around lines 64 - 343, The new
tests in tests/unittest/others/test_lazy_model_zoo.py are not included in any
test-list configuration, leaving their coverage insufficient. Add
tests/unittest/others/test_lazy_model_zoo.py to the appropriate test-db or qa
test list and run pytest tests/unittest/;
tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py requires no
direct change because its existing l0_b200.yml entry already covers the
_create_server-only modification.
Sources: Coding guidelines, Path instructions
Comment on lines
+183
to
+213
| def test_resolver_keeps_existing_registration(): | ||
| # A registration made by user code (--custom_module_dirs) must win over | ||
| # the built-in module, exactly as it does with the eager import order on | ||
| # main where the zoo loads first and custom code overrides it. | ||
| from tensorrt_llm._torch.models.modeling_utils import ( | ||
| MODEL_CLASS_MAPPING, | ||
| get_registered_model_class, | ||
| ) | ||
|
|
||
| arch = "MistralForCausalLM" | ||
|
|
||
| class _CustomStub: | ||
| pass | ||
|
|
||
| MODEL_CLASS_MAPPING[arch] = _CustomStub | ||
| try: | ||
| assert get_registered_model_class(arch) is _CustomStub, ( | ||
| "resolving an architecture overrode its existing registration" | ||
| ) | ||
| finally: | ||
| # Restore the real class rather than deleting the entry: the | ||
| # built-in decorator skips occupied slots, so if the provider was | ||
| # (or gets) imported while the stub held the slot, a bare delete | ||
| # would leave the architecture unresolvable for the rest of the | ||
| # process (module imports are cached). | ||
| from tensorrt_llm._torch.models.modeling_mistral import MistralForCausalLM | ||
|
|
||
| MODEL_CLASS_MAPPING[arch] = MistralForCausalLM | ||
|
|
||
| # Unknown architectures resolve to None, left to the caller's handling. | ||
| assert get_registered_model_class("DefinitelyNotARegisteredArch") is None |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Keep global registry state isolated.
The cleanup does not restore the entry that existed before Line 197. It always imports modeling_mistral and overwrites the mapping with MistralForCausalLM. This can erase a registration installed by an earlier test and permanently cache the provider module.
Run this test body through _run_fresh so its registry and import state cannot affect later tests.
🤖 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/others/test_lazy_model_zoo.py` around lines 183 - 213, Update
test_resolver_keeps_existing_registration to execute its registry assertions and
cleanup through _run_fresh, isolating MODEL_CLASS_MAPPING and imported provider
state from the surrounding test process. Preserve the existing-registration
assertion and unknown-architecture check within that fresh execution, and remove
the unconditional restoration that overwrites prior global state.
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #64049 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64049 [ run ] completed with state
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
tensorrt_llmimports.get_steady_clock_now_in_seconds()to_utils.Dev Engineer Review
__getattr__and__dir__preserve public API behavior.QA Engineer Review
tests/unittest/others/test_lazy_model_zoo.pycovering lazy imports, caching, missing submodules, architecture indexes, and registration precedence.tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py.tests/integration/test_lists/test-db/ortests/integration/test_lists/qa/.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.