wehub/nvidia--tensorrt
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/Polygraphy/polygraphy/mod/__init__.py
+3
View File
@@ -0,0 +1,3 @@
from polygraphy.mod.importer import *
from polygraphy.mod.exporter import *
from polygraphy.mod.util import version
tools/Polygraphy/polygraphy/mod/exporter.py
+362
View File
@@ -0,0 +1,362 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import inspect
import sys
import warnings
from textwrap import dedent
import polygraphy
from polygraphy import config
from polygraphy.logger import G_LOGGER
from polygraphy.mod.util import version
def _add_to_all(symbol, module):
if hasattr(module, "__all__"):
module.__all__.append(symbol)
else:
module.__all__ = [symbol]
def _define_in_module(name, symbol, module):
assert name not in vars(module), "This symbol is already defined!"
vars(module)[name] = symbol
_add_to_all(name, module)
def export(funcify=False, func_name=None):
"""
Decorator that exports a symbol into the ``__all__`` attribute of
the caller's module. This makes the symbol visible in a ``*`` import
(e.g. ``from module import *``) and hides other symbols unless they are
also present in ``__all__``.
Args:
funcify (bool):
Whether to create and export a function that will call a decorated Polygraphy loader.
The decorated type *must* be a subclass of ``BaseLoader`` if ``funcify=True``.
This is useful to provide convenient short-hands to immediately evaluate loaders.
For example:
::
@mod.export(funcify=True)
class SuperCoolModelFromPath(BaseLoader):
def __init__(self, init_params):
...
def call_impl(self, call_params):
...
# We can now magically access an immediately evaluated functional
# variant of the loader:
model = super_cool_model_from_path(init_params, call_params)
# Which is equivalent to:
load_model = SuperCoolModelFromPath(init_params)
model = load_model(call_params)
The signature of the generated function is a combination of the signatures
of ``__init__`` and ``call_impl``. Specifically, parameters without defaults will
precede those with defaults, and ``__init__`` parameters will precede ``call_impl``
parameters. Special parameters like ``*args`` and ``**kwargs`` will always be the last
parameters in the generated signature if they are present in the loader method signatures.
The return value(s) will always come from ``call_impl``.
For example:
::
# With __init__ signature:
def __init__(a, b=0) -> None:
# And call_impl signature:
def call_impl(c, d=0) -> z:
# The generated function will have a signature:
def generated(a, c, b=0, d=0) -> z:
func_name (str):
If funcify is True, this controls the name of the generated function.
By default, the exported function will use the same name as the loader, but
``snake_case`` instead of ``PascalCase``.
"""
module = inspect.getmodule(sys._getframe(1))
# Find a method by wallking the inheritance hierarchy of a type:
def find_method(symbol, method):
hierarchy = inspect.getmro(symbol)
for ancestor in hierarchy:
if method in vars(ancestor):
return vars(ancestor)[method]
assert (
False
), f"Could not find method: {method} in the inheritance hierarcy of: {symbol}"
def export_impl(func_or_cls):
_add_to_all(func_or_cls.__name__, module)
if funcify:
# We only support funcify-ing BaseLoaders, and only if __init__ and call_impl
# have no overlapping parameters.
from polygraphy.backend.base import BaseLoader
assert inspect.isclass(
func_or_cls
), "Decorated type must be a loader to use funcify=True"
assert BaseLoader in inspect.getmro(
func_or_cls
), "Decorated type must derive from BaseLoader to use funcify=True"
def get_params(method):
return list(
inspect.signature(
find_method(func_or_cls, method)
).parameters.values()
)[1:]
def is_variadic(param):
return param.kind in [param.VAR_POSITIONAL, param.VAR_KEYWORD]
def has_default(param):
return param.default != param.empty
def get_param_name(p):
# For variadic arguments, p.name will drop the *, **
return str(p) if is_variadic(p) else p.name
def param_names(params):
return [get_param_name(p) for p in params]
loader = func_or_cls
init_params = get_params("__init__")
call_impl_params = get_params("call_impl")
assert (
set(param_names(call_impl_params)) - set(param_names(init_params))
) == set(
param_names(call_impl_params)
), "Cannot funcify a type where call_impl and __init__ have the same argument names!"
# Dynamically generate a function with the right signature.
# To generate the signature, we use the init and call_impl arguments,
# but move required arguments (i.e. without default values) to the front.
def build_arg_list(should_include):
def str_from_param(p):
return get_param_name(p) + (
f"={p.default}" if has_default(p) else ""
)
arg_list = [str_from_param(p) for p in init_params if should_include(p)]
arg_list += [
str_from_param(p) for p in call_impl_params if should_include(p)
]
return arg_list
non_default_args = build_arg_list(
should_include=lambda p: not is_variadic(p) and not has_default(p)
)
default_args = build_arg_list(
should_include=lambda p: not is_variadic(p) and has_default(p)
)
special_args = build_arg_list(should_include=is_variadic)
signature = ", ".join(non_default_args + default_args + special_args)
init_args = ", ".join(param_names(init_params))
call_impl_args = ", ".join(param_names(call_impl_params))
def pascal_to_snake(name):
return "".join(
f"_{c.lower()}" if c.isupper() else c for c in name
).lstrip("_")
nonlocal func_name
func_name = func_name or pascal_to_snake(loader.__name__)
func_code = dedent(
f"""
def {func_name}({signature}):
return loader_binding({init_args})({call_impl_args})
func_var = {func_name}
"""
)
new_locals = {}
exec(
func_code,
# Need to bind the loader this way, or it won't be accesible from func_code.
{"loader_binding": loader},
new_locals,
)
func = new_locals["func_var"]
# Next we setup the docstring so that it is a combination of the __init__
# and call_impl docstrings.
func.__doc__ = f"Immediately evaluated functional variant of :class:`{loader.__name__}` .\n"
def try_add_method_doc(method):
call_impl = find_method(loader, method)
if call_impl.__doc__:
func.__doc__ += dedent(call_impl.__doc__)
try_add_method_doc("__init__")
try_add_method_doc("call_impl")
# Now that the function has been defined, we just need to add it into the module's
# __dict__ so it is accessible like a normal symbol.
_define_in_module(func_name, func, module)
# We don't actually want to modify the decorated object.
return func_or_cls
return export_impl
def warn_deprecated(
name, use_instead, remove_in, module_name=None, always_show_warning=False
):
if version(polygraphy.__version__) >= version(remove_in):
G_LOGGER.internal_error(
f"{name} should have been removed in version: {remove_in}"
)
full_obj_name = f"{module_name}.{name}" if module_name else name
msg = (
f"{full_obj_name} is deprecated and will be removed in Polygraphy {remove_in}."
)
if use_instead is not None:
msg += f" Use {use_instead} instead."
warnings.warn(msg, DeprecationWarning, stacklevel=3)
if always_show_warning:
G_LOGGER.warning(msg)
def deprecate(remove_in, use_instead, module_name=None, name=None):
"""
Decorator that marks a function or class as deprecated.
When the function or class is used, a warning will be issued.
Args:
remove_in (str):
The version in which the decorated type will be removed.
use_instead (str):
The function or class to use instead.
module_name (str):
The name of the containing module. This will be used to
generate more informative warnings.
Defaults to None.
name (str):
The name of the object being deprecated.
If not provided, this is automatically determined based on the decorated type.
Defaults to None.
"""
def deprecate_impl(obj):
if config.INTERNAL_CORRECTNESS_CHECKS and version(
polygraphy.__version__
) >= version(remove_in):
G_LOGGER.internal_error(
f"{obj} should have been removed in version: {remove_in}"
)
nonlocal name
name = name or obj.__name__
if inspect.ismodule(obj):
class DeprecatedModule:
def __getattr__(self, attr_name):
warn_deprecated(name, use_instead, remove_in, module_name)
self = obj
return getattr(self, attr_name)
def __setattr__(self, attr_name, value):
warn_deprecated(name, use_instead, remove_in, module_name)
self = obj
return setattr(self, attr_name, value)
DeprecatedModule.__doc__ = f"Deprecated: Use {use_instead} instead"
return DeprecatedModule()
elif inspect.isclass(obj):
class Deprecated(obj):
def __init__(self, *args, **kwargs):
warn_deprecated(name, use_instead, remove_in, module_name)
super().__init__(*args, **kwargs)
Deprecated.__doc__ = f"Deprecated: Use {use_instead} instead"
return Deprecated
elif inspect.isfunction(obj):
def wrapped(*args, **kwargs):
warn_deprecated(name, use_instead, remove_in, module_name)
return obj(*args, **kwargs)
wrapped.__doc__ = f"Deprecated: Use {use_instead} instead"
return wrapped
else:
G_LOGGER.internal_error(f"deprecate is not implemented for: {obj}")
return deprecate_impl
def export_deprecated_alias(name, remove_in, use_instead=None):
"""
Decorator that creates and exports a deprecated alias for
the decorated class or function.
The alias will behave like the decorated type, except it will
issue a deprecation warning when used.
To create a deprecated alias for an entire module, invoke the
function manually within the module like so:
::
mod.export_deprecated_alias("old_mod_name", remove_in="0.0.0")(sys.modules[__name__])
Args:
name (str):
The name of the deprecated alias.
remove_in (str):
The version, as a string, in which the deprecated alias will be removed.
use_instead (str):
The name of the function, class, or module to use instead.
If this is ``None``, the new name will be automatically determined.
Defaults to None.
"""
module = inspect.getmodule(sys._getframe(1))
def export_deprecated_alias_impl(obj):
new_obj = deprecate(
remove_in,
use_instead=use_instead or obj.__name__,
module_name=module.__name__,
name=name,
)(obj)
_define_in_module(name, new_obj, module)
return obj
return export_deprecated_alias_impl
tools/Polygraphy/polygraphy/mod/importer.py
+377
View File
@@ -0,0 +1,377 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import contextlib
import importlib
import importlib.util
import os
import subprocess as sp
import sys
from typing import List
try:
# Available in Python 3.8+
import importlib.metadata
except ModuleNotFoundError:
pass
from polygraphy import constants
from polygraphy.mod import util as mod_util
# Tracks all of Polygraphy's lazy imports, excluding internal ones.
_all_external_lazy_imports = set()
# Sometimes the Python package name differs from the module name.
_PKG_NAME_FROM_MODULE = {}
# Some packages need additional flags to install correctly.
_EXTRA_FLAGS_FOR_MODULE = {
"onnx_graphsurgeon": ["--extra-index-url=https://pypi.ngc.nvidia.com"],
}
LATEST_VERSION = "==latest"
"""Indicates that the latest version of the package is preferred in lazy_import"""
def _version_ok(ver, preferred):
if preferred == LATEST_VERSION:
return False
pref_ver = preferred.lstrip("<=>").strip()
cond = preferred.rstrip(pref_ver).strip()
check = {
"==": lambda x, y: x == y,
">=": lambda x, y: x >= y,
">": lambda x, y: x > y,
"<=": lambda x, y: x <= y,
"<": lambda x, y: x < y,
}[cond]
return check(mod_util.version(ver), mod_util.version(pref_ver))
def lazy_import(
name: str,
log: bool = None,
pkg_name: str = None,
install_flags: List[str] = None,
requires: List[str] = None,
):
"""
Lazily import a module.
If config.AUTOINSTALL_DEPS is set to 1,
missing modules are automatically installed, and existing modules may be
upgraded if newer versions are required.
Args:
name (str):
The name of the module and optionally the preferred version of the package,
formatted as a version string. For example, ``'example_module>=0.5.0'`` or ``'example_module==1.8.0'``.
log (bool):
Whether to log information about the module.
Defaults to True.
pkg_name (str):
The name of the package that provides this module, if it is different from the module name.
Used only if automatic installation of dependencies is enabled.
install_flags (List[str]):
Additional flags to provide to the installation command.
Used only if automatic installation of dependencies is enabled.
requires (List[str]):
Additional dependencies required by the module which are *not* specified as dependencies.
This parameter should only be required when a module does not correctly specify dependencies.
Defaults to [].
Returns:
LazyModule:
A lazily loaded module. When an attribute is first accessed,
the module will be imported.
"""
def issue_wrong_version_error(installed_version, version):
from polygraphy.logger import G_LOGGER, LogMode
G_LOGGER.error(
f"Module: '{name}' version '{installed_version}' is installed, but version '{version}' is required.\n"
f"Please install the required version or set POLYGRAPHY_AUTOINSTALL_DEPS=1 in your environment variables "
f"to allow Polygraphy to do so automatically.\n"
f"Attempting to continue with the currently installed version of this module, but note that this may cause errors!",
mode=LogMode.ONCE,
)
VERSION_CHARS = ["=", ">", "<"]
log = True if log is None else log
requires = [] if requires is None else requires
def split_name_version(inp):
version_char_indices = [
inp.index(char) for char in VERSION_CHARS if char in inp
]
if not version_char_indices:
return inp, None
min_index = min(version_char_indices)
return inp[:min_index], inp[min_index:]
name, version = split_name_version(name)
all_required_mods = list(map(split_name_version, requires)) + [(name, version)]
if "polygraphy" not in name:
_all_external_lazy_imports.add(name)
def import_mod():
from polygraphy import config
from polygraphy.logger import G_LOGGER, LogMode
def install_mod(install_name, install_version, raise_error=True):
modname = install_name.split(".")[0]
pkg = (
pkg_name
if pkg_name is not None
else _PKG_NAME_FROM_MODULE.get(modname, modname)
)
extra_flags = (
install_flags
if install_flags is not None
else _EXTRA_FLAGS_FOR_MODULE.get(modname, [])
)
def fail():
log_func = G_LOGGER.critical if raise_error else G_LOGGER.warning
log_func(
f"Could not automatically install required module: {pkg}. Please install it manually."
)
if config.ASK_BEFORE_INSTALL:
res = None
while res not in ["y", "n"]:
res = input(
f"Automatically install '{pkg}' (version: {install_version or 'any'}) ([Y]/n)? "
)
res = res.strip()[:1].lower() or "y"
if res == "n":
fail()
if install_version == LATEST_VERSION:
extra_flags.append("--upgrade")
elif install_version is not None:
pkg += install_version
cmd = config.INSTALL_CMD + [pkg] + extra_flags
G_LOGGER.info(f"Running installation command: {' '.join(cmd)}")
status = sp.run(cmd, stdout=sp.PIPE, stderr=sp.PIPE)
if status.returncode != 0:
G_LOGGER.error(
f"Error during installation:\n{constants.TAB}{status.stderr.decode()}"
)
fail()
mod = importlib.import_module(install_name)
return mod
mod = None
try:
mod = importlib.import_module(name)
except ImportError as err:
if config.AUTOINSTALL_DEPS:
for install_name, install_version in all_required_mods:
G_LOGGER.info(
f"Module: '{install_name}' is required, but not installed. Attempting to install now."
)
mod = install_mod(install_name, install_version)
else:
G_LOGGER.critical(
f"Module: '{name}' is required but could not be imported.\nNote: Error was: {err}\n"
f"You can set POLYGRAPHY_AUTOINSTALL_DEPS=1 in your environment variables to allow "
f"Polygraphy to automatically install missing modules.\n"
)
# Auto-upgrade if necessary
for install_name, install_version in all_required_mods:
installed_mod = importlib.import_module(install_name)
if (
install_version is not None
and hasattr(installed_mod, "__version__")
and not _version_ok(installed_mod.__version__, install_version)
):
if config.AUTOINSTALL_DEPS:
G_LOGGER.info(
f"Note: Module: '{install_name}' version '{installed_mod.__version__}' is installed, but version '{install_version}' is required.\n"
f"Attempting to upgrade now."
)
# We can try to use the other version if install fails, so this is non-fatal.
installed_mod = install_mod(
install_name, install_version, raise_error=False
)
if install_name == name:
mod = installed_mod
elif install_version != LATEST_VERSION:
issue_wrong_version_error(
installed_mod.__version__, install_version
)
if log:
G_LOGGER.module_info(mod)
return mod
MODULE_VAR_NAME = "module"
class LazyModule:
def __init__(self):
super().__setattr__(MODULE_VAR_NAME, None)
def __polygraphy_import_mod(self):
if self.module is None:
super().__setattr__(MODULE_VAR_NAME, import_mod())
return self.module
def __getattr__(self, name):
module = self.__polygraphy_import_mod()
return getattr(module, name)
def __setattr__(self, name, value):
module = self.__polygraphy_import_mod()
return setattr(module, name, value)
def is_installed(self):
"""
Checks whether any version of this module is installed.
The module will not be imported by this method.
Returns:
bool: Whether the module is installed.
"""
global importlib
try:
return name in sys.modules or (
importlib.util.find_spec(name) is not None
)
except:
return False
def is_importable(self):
"""
Checks whether this module is importable. Note that a module may be installed but not importable.
Returns:
bool: Whether the module is importable.
"""
try:
importlib.import_module(name)
return True
except:
return False
return LazyModule()
def has_mod(modname):
"""
Checks whether a module is installed without importing the module.
Args:
modname (str): The name of the module to check.
Returns:
bool: Whether the module is installed.
"""
import warnings
import polygraphy
from polygraphy.logger import G_LOGGER
remove_in = "0.50.0"
if mod_util.version(polygraphy.__version__) >= mod_util.version(remove_in):
G_LOGGER.internal_error(
f"has_mod should have been removed in version: {remove_in}"
)
warnings.warn(
f"has_mod is deprecated and will be removed in Polygraphy {remove_in}",
DeprecationWarning,
stacklevel=3,
)
try:
return modname in sys.modules or (importlib.util.find_spec(modname) is not None)
except ValueError:
return False
def autoinstall(lazy_mod):
"""
If the config.AUTOINSTALL_DEPS is set to 1, automatically install or upgrade a module.
Does nothing if autoinstallation is disabled.
Args:
lazy_mod (LazyModule):
A lazy module, like that returned by ``lazy_import``.
"""
from polygraphy import config
if not config.AUTOINSTALL_DEPS:
return
try:
# It doesn't matter which attribute we try to get as any call to `__getattr__` will
# trigger the automatic installation.
getattr(lazy_mod, "__fake_polygraphy_autoinstall_attr")
except:
pass
def import_from_script(path, name):
"""
Imports a specified symbol from a Python script.
Args:
path (str): A path to the Python script. The path must include a '.py' extension.
name (str): The name of the symbol to import from the script.
Returns:
object: The loaded symbol.
"""
from polygraphy.logger import G_LOGGER
dir = os.path.dirname(path)
modname = os.path.splitext(os.path.basename(path))[0]
sys.path.insert(0, dir)
with contextlib.ExitStack() as stack:
def reset_sys_path():
del sys.path[0]
del sys.modules[modname]
stack.callback(reset_sys_path)
try:
importlib.invalidate_caches()
mod = importlib.import_module(modname)
return getattr(mod, name)
except Exception as err:
ext = os.path.splitext(path)[1]
err_msg = f"Could not import symbol: {name} from script: {path}"
if ext != ".py":
err_msg += f"\nThis could be because the extension of the file is not '.py'. Note: The extension is: {ext}"
err_msg += f"\nNote: Error was: {err}"
G_LOGGER.critical(err_msg)
tools/Polygraphy/polygraphy/mod/trt_importer.py
+39
View File
@@ -0,0 +1,39 @@
#
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from polygraphy import config, mod
def tensorrt_module_and_version_string():
"""
Returns the name of the TensorRT module to import. This selects between
TensorRT and TensorRT-RTX based on the value of config.USE_TENSORRT_RTX,
and ensures that a consistent version of the module is imported.
"""
if config.USE_TENSORRT_RTX:
return "tensorrt_rtx>=1.0"
else:
return "tensorrt>=8.5"
def lazy_import_trt():
"""
Returns either tensorrt or tensorrt_rtx based on config.USE_TENSORRT_RTX.
Prefer to use this function instead of mod.lazy_import("tensorrt>=8.5") to
import TensorRT, so that your code can use TensorRT-RTX.
"""
return mod.lazy_import(tensorrt_module_and_version_string())
tools/Polygraphy/polygraphy/mod/util.py
+41
View File
@@ -0,0 +1,41 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
def version(version_str):
def process_version_part(num):
suffix = None
if "+" in num:
num, suffix = num.split("+")
try:
num = int(num)
except ValueError:
VERSION_SUFFIXES = ["a", "b", "rc", "post", "dev"]
# One version part can only contain one of the above suffixes
for version_suffix in VERSION_SUFFIXES:
if version_suffix in num:
num = num.partition(version_suffix)
break
return [num, suffix] if suffix is not None else [num]
ver_list = []
for num in version_str.split("."):
ver_list.extend(process_version_part(num))
return tuple(ver_list)