addressing review comments

Author: apbose

examples/distributed_inference/tensor_parallel_initialize_dist.py

@@ -48,7 +48,9 @@ def initialize_logger(
 
     # console handler
     ch = logging.StreamHandler()
-    ch.setLevel(console_level)  # Console handler controls what's printed in console output
+    ch.setLevel(
+        console_level
+    )  # Console handler controls what's printed in console output
     ch.setFormatter(logging.Formatter(f"[Rank {rank}] %(levelname)s: %(message)s"))
     logger.addHandler(ch)
 
@@ -123,21 +125,12 @@ def initialize_distributed_env(
     torch.cuda.set_device(device_id)
 
     # Set C++ TensorRT runtime log level based on most verbose handler
-    # this is similar to set_log_level()
+    # Use the most verbose level to ensure all important logs are captured
     cpp_level = min(file_level_int, console_level_int)
     try:
-        import tensorrt as trt
-        from torch_tensorrt._features import ENABLED_FEATURES
-
-        if ENABLED_FEATURES.torch_tensorrt_runtime:
-            if cpp_level == logging.DEBUG:
-                torch.ops.tensorrt.set_logging_level(int(trt.ILogger.Severity.VERBOSE))
-            elif cpp_level == logging.INFO:
-                torch.ops.tensorrt.set_logging_level(int(trt.ILogger.Severity.INFO))
-            elif cpp_level == logging.WARNING:
-                torch.ops.tensorrt.set_logging_level(int(trt.ILogger.Severity.WARNING))
-            else:
-                torch.ops.tensorrt.set_logging_level(int(trt.ILogger.Severity.ERROR))
+        import torch_tensorrt.logging as torchtrt_logging
+
+        torchtrt_logging.set_level(cpp_level)
     except Exception as e:
         logger.warning(f"Could not set C++ TensorRT log level: {e}")
 

py/torch_tensorrt/_features.py

@@ -78,6 +78,25 @@ def _enabled_features_str() -> str:
     return out_str
 
 
+# Inline helper functions for checking feature availability
+def has_torch_tensorrt_runtime() -> bool:
+    """Check if Torch-TensorRT C++ runtime is available.
+
+    Returns:
+        bool: True if libtorchtrt_runtime.so or libtorchtrt.so is available
+    """
+    return bool(ENABLED_FEATURES.torch_tensorrt_runtime)
+
+
+def has_torchscript_frontend() -> bool:
+    """Check if TorchScript frontend is available.
+
+    Returns:
+        bool: True if libtorchtrt.so is available
+    """
+    return bool(ENABLED_FEATURES.torchscript_frontend)
+
+
 def needs_tensorrt_rtx(f: Callable[..., Any]) -> Callable[..., Any]:
     def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
         if ENABLED_FEATURES.tensorrt_rtx:
@@ -180,6 +199,7 @@ def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
         if ENABLED_FEATURES.trtllm_for_nccl:
             return f(*args, **kwargs)
         else:
+
             def not_implemented(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
                 raise NotImplementedError(
                     "TensorRT-LLM plugin for NCCL is not available"

py/torch_tensorrt/dynamo/backend/backends.py

@@ -7,6 +7,7 @@
 
 import torch
 import torch._dynamo as td
+import torch_tensorrt.logging as torchtrt_logging
 from torch._dynamo.backends.common import aot_autograd
 from torch._dynamo.utils import detect_fake_mode
 from torch._functorch.aot_autograd import aot_export_joint_simple
@@ -23,7 +24,6 @@
 from torch_tensorrt.dynamo.utils import (
     parse_dynamo_kwargs,
     prepare_inputs,
-    set_log_level,
 )
 
 logger = logging.getLogger(__name__)
@@ -40,7 +40,7 @@ def torch_tensorrt_backend(
         and "debug" in kwargs["options"]
         and kwargs["options"]["debug"]
     ) or ("debug" in kwargs and kwargs["debug"]):
-        set_log_level(logger.parent, logging.DEBUG)
+        torchtrt_logging.set_level(logging.DEBUG)
 
     DEFAULT_BACKEND = aot_torch_tensorrt_aten_backend
 

py/torch_tensorrt/dynamo/utils.py

@@ -28,7 +28,6 @@
 from torch.fx.experimental.proxy_tensor import unset_fake_temporarily
 from torch_tensorrt._Device import Device
 from torch_tensorrt._enums import dtype
-from torch_tensorrt._features import ENABLED_FEATURES
 from torch_tensorrt._Input import Input
 from torch_tensorrt._utils import is_tensorrt_version_supported
 from torch_tensorrt.dynamo import _defaults
@@ -270,33 +269,6 @@ def get_model_device(module: torch.fx.GraphModule) -> torch.device:
     return device
 
 
-def set_log_level(parent_logger: Any, level: Any) -> None:
-    """
-    Sets the log level to the user provided level.
-    This is used to set debug logging at a global level
-    at entry points of tracing, dynamo and torch_compile compilation.
-    And set log level for c++ torch trt logger if runtime is available.
-    """
-    if parent_logger:
-        parent_logger.setLevel(level)
-
-    if ENABLED_FEATURES.torch_tensorrt_runtime:
-        if level == logging.DEBUG:
-            log_level = trt.ILogger.Severity.VERBOSE
-        elif level == logging.INFO:
-            log_level = trt.ILogger.Severity.INFO
-        elif level == logging.WARNING:
-            log_level = trt.ILogger.Severity.WARNING
-        elif level == logging.ERROR:
-            log_level = trt.ILogger.Severity.ERROR
-        elif level == logging.CRITICAL:
-            log_level = trt.ILogger.Severity.INTERNAL_ERROR
-        else:
-            raise AssertionError(f"{level} is not valid log level")
-
-        torch.ops.tensorrt.set_logging_level(int(log_level))
-
-
 def prepare_inputs(
     inputs: Input | torch.Tensor | Sequence[Any] | Dict[Any, Any],
     disable_memory_format_check: bool = False,