chore: import upstream snapshot with attribution

tools/Polygraphy/polygraphy/logger/__init__.py

@@ -0,0 +1 @@
+from polygraphy.logger.logger import G_LOGGER, LogMode

tools/Polygraphy/polygraphy/logger/logger.py

@@ -0,0 +1,700 @@
+#
+# 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 copy
+import enum
+import functools
+import logging
+import os
+import sys
+import time
+import traceback
+
+COLORED_MODULE_PRESENT = None
+
+
+def has_colors():
+    global COLORED_MODULE_PRESENT
+    if COLORED_MODULE_PRESENT is None:
+        try:
+            import colored
+
+            COLORED_MODULE_PRESENT = True
+        except:
+            COLORED_MODULE_PRESENT = False
+            print(
+                "[W] 'colored' module is not installed, will not use colors when logging. "
+                "To enable colors, please install the 'colored' module: python3 -m pip install colored"
+            )
+    return COLORED_MODULE_PRESENT
+
+
+# Context manager to apply indentation to messages
+class LoggerIndent:
+    def __init__(self, logger, indent):
+        self.logger = logger
+        self.old_indent = self.logger.logging_indent
+        self.indent = indent
+
+    def __enter__(self):
+        self.logger.logging_indent = self.indent
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.logger.logging_indent = self.old_indent
+
+
+# Context manager to temporarily set verbosity
+class LoggerVerbosity:
+    def __init__(self, logger, severity):
+        self.logger = logger
+        self.old_severity = copy.copy(self.logger.module_severity)
+        self.module_severity = severity
+
+    def __enter__(self):
+        self.logger.module_severity = self.module_severity
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.logger.module_severity = self.old_severity
+
+
+class LogMode(enum.IntEnum):
+    """
+    Specifies how messages should be logged.
+    """
+
+    EACH = 0
+    """Log the message each time"""
+    ONCE = 1
+    """Log the message only once. The same message will not be logged again."""
+
+
+class SeverityTrie:
+    """
+    A trie that represents per-path logging verbosities.
+    """
+
+    def _split_path(self, path):
+        # Leading or duplicate slashes can create empty elements in the path components. We ignore those.
+        return list(filter(lambda x: x, path.split(os.path.sep)))
+
+    def __init__(self, severity_dict):
+        assert "" in severity_dict, "severity_dict must include default severity!"
+
+        self.trie = {}
+        for path, severity in severity_dict.items():
+            cur_dict = self.trie
+            for path_component in self._split_path(path):
+                if path_component not in cur_dict:
+                    cur_dict[path_component] = {}
+                cur_dict = cur_dict[path_component]
+            cur_dict[""] = severity
+
+        # Skip path checking if we don't have any path entries.
+        self.has_non_default_entries = len(self.trie) > 1
+
+    def get(self, path=None):
+        """
+        Get the logging verbosity for the given path.
+
+        Args:
+            path (str): The path
+
+        Returns:
+            int: The logging verbosity.
+        """
+        default_severity = self.trie[""]
+        if path is None or not self.has_non_default_entries:
+            return default_severity
+
+        cur_dict = self.trie
+
+        def get_value(dct):
+            return dct.get("", default_severity)
+
+        for path_component in self._split_path(path):
+            if path_component not in cur_dict:
+                return get_value(cur_dict)
+            cur_dict = cur_dict[path_component]
+        return get_value(cur_dict)
+
+    def __str__(self):
+        return str(self.trie)
+
+
+class Logger:
+    """
+    Global logging interface. Do **not** construct a logger manually.
+    Instead, use ``G_LOGGER``, the global logger.
+    """
+
+    ULTRA_VERBOSE = -20  # Cast it into the flames!
+    """Enable unreasonably verbose messages and above"""
+    SUPER_VERBOSE = -10
+    """Enable extremely verbose messages and above"""
+    EXTRA_VERBOSE = 0
+    """Enable extra verbose messages and above"""
+    VERBOSE = 10
+    """Enable verbose messages and above"""
+    INFO = 20
+    """Enable informative messages and above"""
+    START = 22
+    """Enable messages indicating when a task is started and above"""
+    FINISH = 28
+    """Enable messages indicating when a task is finished and above"""
+    WARNING = 30
+    """Enable only warning messages and above"""
+    ERROR = 40
+    """Enable only error messages and above"""
+    CRITICAL = 50
+    """Enable only critical/fatal error messages and above"""
+
+    SEVERITY_LETTER_MAPPING = {
+        ULTRA_VERBOSE: "[U]",
+        SUPER_VERBOSE: "[S]",
+        EXTRA_VERBOSE: "[X]",
+        VERBOSE: "[V]",
+        INFO: "[I]",
+        START: "[I]",
+        FINISH: "[I]",
+        WARNING: "[W]",
+        ERROR: "[E]",
+        CRITICAL: "[!]",
+    }
+
+    SEVERITY_COLOR_MAPPING = {
+        ULTRA_VERBOSE: "dark_gray",
+        SUPER_VERBOSE: "medium_violet_red",
+        EXTRA_VERBOSE: "medium_purple",
+        VERBOSE: "light_magenta",
+        INFO: None,
+        START: "light_cyan",
+        FINISH: "light_green",
+        WARNING: "light_yellow",
+        ERROR: "light_red",
+        CRITICAL: "light_red",
+    }
+
+    def __init__(
+        self, severity=INFO, colors=True, letter=True, timestamp=False, line_info=False
+    ):
+        """
+        Args:
+            severity (Union[int, Dict[str, int]]):
+                    Severity below which messages will be ignored.
+                    This can be specified on a per-submodule/file basis by providing a dictionary of paths to
+                    logging severities. In this case, use the ``""`` to indicate the default severity.
+                    Paths should be relative to the `polygraphy/` directory.
+                    For example, `polygraphy/backend` can be specified with just `backend/`.
+                    For example: ``{"": G_LOGGER.INFO, "backend/trt": G_LOGGER.VERBOSE}``
+                    This is converted to a ``SeverityTrie`` on assignment.
+                    Defaults to G_LOGGER.INFO.
+            colors (bool):
+                    Whether to use colored output.
+                    Defaults to True.
+            letter (bool):
+                    Whether to prepend each logging message with a letter indicating it's severity.
+                    Defaults to True.
+            timestamp (bool):
+                    Whether to include a timestamp in the logging output.
+                    Defaults to False.
+            line_info (bool):
+                    Whether to include file and line number information in the logging output.
+                    Defaults to False.
+        """
+        self.logging_indent = 0
+        self.once_logged = set()
+        self.colors = colors
+        self.letter = letter
+        self.timestamp = timestamp
+        self.line_info = line_info
+        self.logger_callbacks = []
+        self.log_file = None
+        """
+        Path to a log file to write logging output from Polygraphy.
+        This will not include logging messages from libraries used by Polygraphy, like
+        TensorRT or ONNX-Runtime.
+        """
+        self.module_severity = severity
+        """
+        Severity below which messages will be ignored.
+        This can be specified on a per-submodule/file basis by providing a dictionary of paths to
+        logging severities. In this case, use the ``""`` to indicate the default severity.
+        Paths should be relative to the `polygraphy/` directory.
+        For example, `polygraphy/backend` can be specified with just `backend/`.
+        For example: ``{"": G_LOGGER.INFO, "backend/trt": G_LOGGER.VERBOSE}``
+        This is converted to a ``SeverityTrie`` on assignment.
+        Defaults to G_LOGGER.INFO.
+        """
+        self._use_python_logging_system = False
+        """
+        A flag indicating whether to use the Python `logging` module for log emission.
+        By default, logs are emitted to `stdout` or `stderr`. When this flag is set to `True`,
+        the logger uses the Python `logging` module instead of `stdout`/`stderr`.
+        This allows logs to be integrated into a unified logging system which can be helpful
+        in advanced cases like using multiprocessing or when the user wants one logging system
+        to manage Polygraphy and other libraries logs.
+        """
+
+    @property
+    def log_file(self):
+        return self._log_path
+
+    @log_file.setter
+    def log_file(self, value):
+        self._log_path = value
+        self._log_file = None
+        if self._log_path:
+            dir_path = os.path.dirname(self._log_path)
+            if dir_path:
+                dir_path = os.path.realpath(dir_path)
+                os.makedirs(dir_path, exist_ok=True)
+            self._log_file = open(self._log_path, "w")
+
+    @property
+    def module_severity(self):
+        return self._module_severity
+
+    @module_severity.setter
+    def module_severity(self, value):
+        if isinstance(value, SeverityTrie):
+            self._module_severity = value
+        else:
+            if not isinstance(value, dict):
+                value = {"": value}
+            if "" not in value:
+                value[""] = Logger.INFO
+            self._module_severity = SeverityTrie(value)
+
+        self._run_callbacks()
+
+    @property
+    def severity(self):
+        print(
+            "Warning: Accessing the `severity` property of G_LOGGER is deprecated and will be removed in v0.50.0. Use `module_severity` instead"
+        )
+        return self._module_severity.get()
+
+    @severity.setter
+    def severity(self, value):
+        print(
+            "Warning: Accessing the `severity` property of G_LOGGER is deprecated and will be removed in v0.50.0. Use `module_severity` instead"
+        )
+        self.module_severity = value
+
+    @property
+    def use_python_logging_system(self):
+        return self._use_python_logging_system
+
+    @use_python_logging_system.setter
+    def use_python_logging_system(self, value):
+        self._use_python_logging_system = value
+        if value:
+            self.python_logger = logging.getLogger("Polygraphy")
+            self.severity_level_mapping = {
+                Logger.ULTRA_VERBOSE: 2,
+                Logger.SUPER_VERBOSE: 4,
+                Logger.EXTRA_VERBOSE: 6,
+            }
+
+    def module_path(self, path):
+        """
+        Converts a given path to a path relative to the Polygraphy root module.
+        If the path is not part of the Polygraphy module, returns a path relative to the common prefix.
+
+        Args:
+            path (str): The path
+
+        Returns:
+            str: The path relative to the Polygraphy root module or common prefix.
+        """
+        import polygraphy
+
+        module_root_dir = polygraphy.__path__[0]
+        file_path = os.path.relpath(path, module_root_dir)
+        if os.pardir in file_path:
+            common_path_len = len(os.path.commonpath([module_root_dir, path]))
+            file_path = path[common_path_len:].lstrip(os.path.sep)
+        return file_path
+
+    def _run_callbacks(self):
+        for callback in self.logger_callbacks:
+            callback(self._module_severity)
+
+    def register_callback(self, callback):
+        """
+        Registers a callback with the logger, which will be invoked when the logging severity is modified.
+        The callback is guaranteed to be called at least once in the register_callback function.
+
+        Args:
+            callback (Callable(SeverityTrie)):
+                    A callback that accepts the current logger severity trie.
+        """
+        callback(self._module_severity)
+        self.logger_callbacks.append(callback)
+
+    def indent(self, level=1):
+        """
+        Returns a context manager that indents all strings logged by the specified amount.
+
+        Args:
+            level (int): The indentation level
+        """
+        return LoggerIndent(self, level + self.logging_indent)
+
+    def verbosity(self, severity=CRITICAL):
+        """
+        Returns a context manager that temporarily changes the severity of the logger for its duration.
+
+        Args:
+            severity (Union[int, Dict[str, int]]):
+                    Severity below which messages will be ignored.
+                    This can be specified on a per-submodule/file basis by providing a dictionary of paths to
+                    logging severities. In this case, use the ``""`` to indicate the default severity.
+                    For example: ``{"": G_LOGGER.INFO, "backend/trt": G_LOGGER.VERBOSE}``
+                    Defaults to Logger.CRITICAL, which will suppress all messages.
+        """
+        return LoggerVerbosity(self, severity)
+
+    def log(self, message, severity, mode=LogMode.EACH, stack_depth=2, error_ok=False):
+        """
+        Logs a message to stdout.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            severity (int):
+                    The severity with which to log this message. If the severity is less than
+                    the logger's current severity, the message is suppressed. Provided callables
+                    will not be called in that case.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+            stack_depth (int):
+                    The stack depth to use to determine file and line information.
+                    Defaults to 2.
+            error_ok (bool):
+                    Whether to suppress errors encountered while logging.
+                    When this is True, in the event of an error, the message will not be
+                    logged, but the logger will recover and resume execution.
+                    When False, the logger will re-raise the exception.
+        """
+        from polygraphy import config, constants
+
+        def get_rel_file_path_and_lineno():
+            file_path = sys._getframe(stack_depth).f_code.co_filename
+            line_no = sys._getframe(stack_depth).f_lineno
+            # If we can't get a valid path, keep walking the stack until we can.
+            new_stack_depth = stack_depth
+            while not os.path.exists(file_path) and new_stack_depth > 0:
+                new_stack_depth -= 1
+                file_path = sys._getframe(new_stack_depth).f_code.co_filename
+                line_no = sys._getframe(new_stack_depth).f_lineno
+
+            return self.module_path(file_path), line_no
+
+        def process_message(message, file_path, line_no):
+            def get_prefix():
+                prefix = ""
+                if self.letter:
+                    prefix += Logger.SEVERITY_LETTER_MAPPING[severity] + " "
+                if self.timestamp:
+                    prefix += f"({time.strftime('%X')}) "
+                if self.line_info:
+                    prefix += f"[{file_path}:{line_no}] "
+                return prefix
+
+            def apply_indentation(prefix, message):
+                message_lines = str(message).splitlines()
+                tab = constants.TAB * self.logging_indent
+                newline_tab = "\n" + tab + " " * len(prefix)
+                return tab + newline_tab.join([line for line in message_lines])
+
+            def apply_color(message):
+                if self.colors and has_colors():
+                    import colored
+
+                    color = Logger.SEVERITY_COLOR_MAPPING[severity]
+
+                    if color:
+                        try:
+                            color = colored.fore(color)
+                        except:
+                            color = [colored.fg(color)]
+
+                        return colored.stylize(message, color)
+                return message
+
+            prefix = get_prefix()
+            message = apply_indentation(prefix, message)
+            return apply_color(f"{prefix}{message}")
+
+        file_path, line_no = None, None
+        if self.line_info or self.module_severity.has_non_default_entries:
+            file_path, line_no = get_rel_file_path_and_lineno()
+
+        def should_log(message):
+            if severity < self.module_severity.get(file_path):
+                return False
+
+            if mode == LogMode.ONCE:
+                message_hash = hash(message)
+                if message_hash in self.once_logged:
+                    return False
+                self.once_logged.add(message_hash)
+            return True
+
+        if not should_log(message):
+            return
+
+        if callable(message):
+            try:
+                message = message()
+            except Exception as err:
+                if not error_ok or config.INTERNAL_CORRECTNESS_CHECKS:
+                    raise
+                message = f"<Error while logging this message: {str(err)}>"
+
+        message = str(message)
+
+        # Use the warnings module in correctness checking mode so all warnings are
+        # visible in the test result summary.
+        if config.INTERNAL_CORRECTNESS_CHECKS and severity == Logger.WARNING:
+            import warnings
+
+            warnings.warn(message)
+
+        message = process_message(message, file_path, line_no)
+
+        if self._log_file is not None:
+            self._log_file.write(message + "\n")
+            self._log_file.flush()
+
+        if self._use_python_logging_system:
+            # python logging system does not handle negative levels, map them to positive values
+            level = severity if severity > 0 else self.severity_level_mapping[severity]
+            self.python_logger.log(level=level, msg=message)
+        else:
+            print(
+                message, file=sys.stdout if severity < Logger.CRITICAL else sys.stderr
+            )
+
+    def backtrace(self, depth=0, limit=None, severity=ERROR):
+        limit = (
+            limit if limit is not None else (3 - self.module_severity.get() // 10) * 2
+        )  # Info provides 1 stack frame
+        limit = max(limit, 0)
+        frame = sys._getframe(depth + 2)
+        self.log(
+            " ".join(traceback.format_stack(f=frame, limit=limit)), severity=severity
+        )
+
+    def ultra_verbose(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with ULTRA_VERBOSE severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.ULTRA_VERBOSE, mode=mode, stack_depth=3, error_ok=True)
+
+    def super_verbose(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with SUPER_VERBOSE severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.SUPER_VERBOSE, mode=mode, stack_depth=3, error_ok=True)
+
+    def extra_verbose(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with EXTRA_VERBOSE severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.EXTRA_VERBOSE, mode=mode, stack_depth=3, error_ok=True)
+
+    def verbose(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with VERBOSE severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.VERBOSE, mode=mode, stack_depth=3, error_ok=True)
+
+    def info(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with INFO severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.INFO, mode=mode, stack_depth=3)
+
+    def start(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with START severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.START, mode=mode, stack_depth=3)
+
+    def finish(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with FINISH severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.FINISH, mode=mode, stack_depth=3)
+
+    def warning(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with WARNING severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.WARNING, mode=mode, stack_depth=3)
+
+    def error(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with ERROR severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.ERROR, mode=mode, stack_depth=3)
+
+    def critical(self, message, ExceptionType=None):
+        """
+        Logs a message to stdout with CRITICAL severity and raises an exception.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            ExceptionType (type):
+                    The type of exception to raise.
+                    Defaults to PolygraphyException.
+
+        Raises:
+            ExceptionType
+        """
+        self.log(message, Logger.CRITICAL, stack_depth=3)
+        from polygraphy.exception import PolygraphyException
+
+        ExceptionType = ExceptionType or PolygraphyException
+        raise ExceptionType(message) from None
+
+    def internal_error(self, message):
+        from polygraphy import config
+
+        if not config.INTERNAL_CORRECTNESS_CHECKS:
+            return
+
+        self.log(message, Logger.CRITICAL, stack_depth=3)
+        from polygraphy.exception import PolygraphyInternalException
+
+        raise PolygraphyInternalException(message)
+
+    def _str_from_module_info(self, module, name=None):
+        ret = ""
+
+        def try_append(func):
+            nonlocal ret
+            try:
+                ret += func()
+            except:
+                pass
+
+        try_append(lambda: name or f"Loaded Module: {module.__name__}")
+        try_append(lambda: f" | Version: {module.__version__}")
+        try_append(lambda: f" | Path: {list(map(os.path.realpath, module.__path__))}")
+        return ret
+
+    def module_info(self, module, name=None, severity=VERBOSE):
+        message = self._str_from_module_info(module, name)
+        self.log(message, severity=severity, stack_depth=3, mode=LogMode.ONCE)
+
+    def log_exception(self, func):
+        """
+        Decorator that causes exceptions in a function to be logged.
+        This is useful in cases where the exception is caught by a caller, but should
+        still be logged.
+        """
+
+        @functools.wraps(func)
+        def wrapped(*args, **kwargs):
+            from polygraphy.exception import PolygraphyException
+
+            try:
+                return func(*args, **kwargs)
+            except PolygraphyException:
+                # `PolygraphyException`s are always logged.
+                raise
+            except Exception as err:
+                G_LOGGER.error(err)
+                raise
+
+        return wrapped
+
+
+G_LOGGER = Logger()
+"""The global logger. Use this instead of constructing a logger"""
+
+# For backwards compatibility
+G_LOGGER.exit = G_LOGGER.critical