Feat 865

py/torch_tensorrt/logging.py

@@ -96,3 +96,113 @@ def log(level: Level, msg: str):
         msg (str): Actual message text
     """
     _log(Level._to_internal_level(level), msg)
+
+    InternalError = LogLevel.INTERNAL_ERROR
+    Error = LogLevel.ERROR
+    Warning = LogLevel.WARNING
+    Info = LogLevel.INFO
+    Debug = LogLevel.DEBUG
+    Graph = LogLevel.GRAPH
+
+
+class internal_errors:
+    """Context-manager to limit displayed log messages to just internal errors
+
+    Example::
+
+    with torch_tensorrt.logging.internal_errors():
+        outputs = model_torchtrt(inputs)
+    """
+
+    def __enter__(self):
+        self.external_lvl = get_reportable_log_level()
+        set_reportable_log_level(Level.InternalError)
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        set_reportable_log_level(self.external_lvl)
+
+
+class errors:
+    """Context-manager to limit displayed log messages to just errors and above
+
+    Example::
+
+    with torch_tensorrt.logging.errors():
+        outputs = model_torchtrt(inputs)
+    """
+
+    def __enter__(self):
+        self.external_lvl = get_reportable_log_level()
+        set_reportable_log_level(Level.Error)
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        set_reportable_log_level(self.external_lvl)
+
+
+class warnings:
+    """Context-manager to limit displayed log messages to just warnings and above
+
+    Example::
+
+    with torch_tensorrt.logging.warnings():
+        model_trt = torch_tensorrt.compile(model, **spec)
+    """
+
+    def __enter__(self):
+        self.external_lvl = get_reportable_log_level()
+        set_reportable_log_level(Level.Warning)
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        set_reportable_log_level(self.external_lvl)
+
+
+class info:
+    """Context-manager to display all info and greater severity messages
+
+    Example::
+
+    with torch_tensorrt.logging.info():
+        model_trt = torch_tensorrt.compile(model, **spec)
+    """
+
+    def __enter__(self):
+        self.external_lvl = get_reportable_log_level()
+        set_reportable_log_level(Level.Info)
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        set_reportable_log_level(self.external_lvl)
+
+
+class debug:
+    """Context-manager to display full debug information through the logger
+
+    Example::
+
+    with torch_tensorrt.logging.debug():
+        model_trt = torch_tensorrt.compile(model, **spec)
+    """
+
+    def __enter__(self):
+        self.external_lvl = get_reportable_log_level()
+        set_reportable_log_level(Level.Debug)
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        set_reportable_log_level(self.external_lvl)
+
+
+class graphs:
+    """Context-manager to display the results of intermediate lowering passes
+    as well as full debug information through the logger
+
+    Example::
+
+    with torch_tensorrt.logging.graphs():
+        model_trt = torch_tensorrt.compile(model, **spec)
+    """
+
+    def __enter__(self):
+        self.external_lvl = get_reportable_log_level()
+        set_reportable_log_level(Level.Graph)
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        set_reportable_log_level(self.external_lvl)

tests/py/test_api.py

@@ -323,6 +323,50 @@ def test_is_colored_output_on(self):
         color = torchtrt.logging.get_is_colored_output_on()
         self.assertTrue(color)
 
+    def test_context_managers(self):
+        base_lvl = torchtrt.logging.get_reportable_log_level()
+        with torchtrt.logging.internal_errors():
+            lvl = torchtrt.logging.get_reportable_log_level()
+            self.assertEqual(torchtrt.logging.Level.InternalError, lvl)
+
+        lvl = torchtrt.logging.get_reportable_log_level()
+        self.assertEqual(base_lvl, lvl)
+
+        with torchtrt.logging.errors():
+            lvl = torchtrt.logging.get_reportable_log_level()
+            self.assertEqual(torchtrt.logging.Level.Error, lvl)
+
+        lvl = torchtrt.logging.get_reportable_log_level()
+        self.assertEqual(base_lvl, lvl)
+
+        with torchtrt.logging.warnings():
+            lvl = torchtrt.logging.get_reportable_log_level()
+            self.assertEqual(torchtrt.logging.Level.Warning, lvl)
+
+        lvl = torchtrt.logging.get_reportable_log_level()
+        self.assertEqual(base_lvl, lvl)
+
+        with torchtrt.logging.info():
+            lvl = torchtrt.logging.get_reportable_log_level()
+            self.assertEqual(torchtrt.logging.Level.Info, lvl)
+
+        lvl = torchtrt.logging.get_reportable_log_level()
+        self.assertEqual(base_lvl, lvl)
+
+        with torchtrt.logging.debug():
+            lvl = torchtrt.logging.get_reportable_log_level()
+            self.assertEqual(torchtrt.logging.Level.Debug, lvl)
+
+        lvl = torchtrt.logging.get_reportable_log_level()
+        self.assertEqual(base_lvl, lvl)
+
+        with torchtrt.logging.graphs():
+            lvl = torchtrt.logging.get_reportable_log_level()
+            self.assertEqual(torchtrt.logging.Level.Graph, lvl)
+
+        lvl = torchtrt.logging.get_reportable_log_level()
+        self.assertEqual(base_lvl, lvl)
+
 
 class TestDevice(unittest.TestCase):