Added profile_format and added docstring to the debugger

Author: cehongwang

py/torch_tensorrt/dynamo/_compiler.py

@@ -924,7 +924,7 @@ def contains_metadata(gm: torch.fx.GraphModule) -> bool:
                     os.makedirs(path, exist_ok=True)
                     trt_module.enable_profiling(
                         profiling_results_dir=path,
-                        profile_format="trex",
+                        profile_format=_debugger_settings.profile_format,
                     )
 
     # Parse the graph I/O and store it in dryrun tracker

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py

@@ -719,7 +719,7 @@ def run(
                 if self.compilation_settings.reuse_cached_engines:
                     interpreter_result = self._pull_cached_engine(hash_val)
                     if interpreter_result is not None:  # hit the cache
-                        return interpreter_result
+                        return interpreter_result  # type: ignore[no-any-return]
 
         self._construct_trt_network_def()
 

py/torch_tensorrt/dynamo/debug/_Debugger.py

@@ -30,16 +30,38 @@ def __init__(
         capture_fx_graph_before: Optional[List[str]] = None,
         capture_fx_graph_after: Optional[List[str]] = None,
         save_engine_profile: bool = False,
+        profile_format: str = "perfetto",
         engine_builder_monitor: bool = True,
         logging_dir: str = tempfile.gettempdir(),
     ):
+        """Initialize a debugger for TensorRT conversion.
+
+        Args:
+            log_level (str): Logging level to use. Valid options are:
+                'debug', 'info', 'warning', 'error', 'internal_errors', 'graphs'.
+                Defaults to 'debug'.
+            capture_fx_graph_before (List[str], optional): List of pass names to visualize FX graph
+                before execution of a lowering pass. Defaults to None.
+            capture_fx_graph_after (List[str], optional): List of pass names to visualize FX graph
+                after execution of a lowering pass. Defaults to None.
+            save_engine_profile (bool): Whether to save TensorRT engine profiling information.
+                Defaults to False.
+            profile_format (str): Format for profiling data. Can be either 'perfetto' or 'trex'.
+                If you need to generate engine graph using the profiling files, set it to 'trex' .
+                Defaults to 'perfetto'.
+            engine_builder_monitor (bool): Whether to monitor TensorRT engine building process.
+                Defaults to True.
+            logging_dir (str): Directory to save debug logs and profiles.
+                Defaults to system temp directory.
+        """
 
         os.makedirs(logging_dir, exist_ok=True)
         self.cfg = DebuggerConfig(
             log_level=log_level,
             save_engine_profile=save_engine_profile,
             engine_builder_monitor=engine_builder_monitor,
             logging_dir=logging_dir,
+            profile_format=profile_format,
         )
 
         if log_level == "debug":

py/torch_tensorrt/dynamo/debug/_DebuggerConfig.py

@@ -8,3 +8,4 @@ class DebuggerConfig:
     save_engine_profile: bool = False
     engine_builder_monitor: bool = True
     logging_dir: str = tempfile.gettempdir()
+    profile_format: str = "perfetto"

py/torch_tensorrt/dynamo/debug/_supports_debugger.py

@@ -1,17 +1,15 @@
-from typing import Any, Callable, Type, TypeVar
+from typing import Any, Callable, Type, TypeVar, cast
 
 T = TypeVar("T")
-
+F = TypeVar("F", bound=Callable[..., Any])
 
 _DEBUG_ENABLED_FUNCS = []
 _DEBUG_ENABLED_CLS = []
 
-F = TypeVar("F", bound=Callable[..., Any])
-
 
 def fn_supports_debugger(func: F) -> F:
     _DEBUG_ENABLED_FUNCS.append(func)
-    return func
+    return cast(F, func)
 
 
 def cls_supports_debugger(cls: Type[T]) -> Type[T]: