chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user