chore: Fix Logging in torch_compile path (#2238)

Signed-off-by: Dheeraj Peri <[email protected]>
pytorch · Aug 17, 2023 · f70574e · f70574e
1 parent e77e445
commit f70574e
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 5 deletions.
diff --git a/py/torch_tensorrt/dynamo/backend/backends.py b/py/torch_tensorrt/dynamo/backend/backends.py
@@ -20,6 +20,17 @@
 def torch_tensorrt_backend(
     gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor], **kwargs: Any
 ) -> torch.nn.Module:
+    # Set log level at the top of compilation (torch_tensorrt.dynamo)
+    if (
+        (
+            "options" in kwargs
+            and "debug" in kwargs["options"]
+            and kwargs["options"]["debug"]
+        )
+        or ("debug" in kwargs and kwargs["debug"])
+    ) and logger.parent:
+        logger.parent.setLevel(logging.DEBUG)
+
     DEFAULT_BACKEND = aot_torch_tensorrt_aten_backend
 
     compiled_mod: torch.nn.Module = DEFAULT_BACKEND(gm, sample_inputs, **kwargs)

diff --git a/py/torch_tensorrt/dynamo/compile.py b/py/torch_tensorrt/dynamo/compile.py
@@ -197,6 +197,9 @@ def compile_module(
 
         submodule = getattr(partitioned_module, name)
 
+        logger.debug(
+            "Submodule name: " + str(name) + " Graph: \n" + str(submodule.graph)
+        )
         # Get submodule inputs
         submodule_inputs = partitioning.get_submod_inputs(
             partitioned_module, submodule, sample_inputs

diff --git a/py/torch_tensorrt/dynamo/utils.py b/py/torch_tensorrt/dynamo/utils.py
@@ -159,10 +159,6 @@ def parse_dynamo_kwargs(kwargs: Any) -> CompilationSettings:
         valid_kwargs = {k: v for k, v in kwargs.items() if k in valid_attrs}
         settings = replace(settings, **valid_kwargs)
 
-    # Enable debug/verbose mode if requested
-    if settings.debug:
-        logger.setLevel(logging.DEBUG)
-
     # TODO: Remove once Dynamo precisions refactoring is complete
     if "enabled_precisions" in kwargs:
         enabled_precisions = kwargs["enabled_precisions"]
@@ -188,7 +184,7 @@ def parse_dynamo_kwargs(kwargs: Any) -> CompilationSettings:
     # Parse input runtime specification
     settings.use_python_runtime = use_python_runtime_parser(settings.use_python_runtime)
 
-    logger.debug(f"Compiling with Settings:\n{settings}")
+    logger.info("Compilation Settings: %s\n", settings)
 
     return settings