Skip to content

chore: Fix Logging in torch_compile path #2238

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Aug 17, 2023
11 changes: 11 additions & 0 deletions py/torch_tensorrt/dynamo/backend/backends.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
3 changes: 3 additions & 0 deletions py/torch_tensorrt/dynamo/compile.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 1 addition & 5 deletions py/torch_tensorrt/dynamo/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand All @@ -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

Expand Down