checks for windows where NCCL backend is not supported

Author: apbose

py/torch_tensorrt/dynamo/utils.py

@@ -937,7 +937,7 @@ def download_plugin_lib_path(platform: str) -> Iterator[str]:
         yield plugin_lib_path
 
 
-def load_and_initialize_trtllm_plugin(plugin_lib_path: str) -> bool:
+def load_and_initialize_trtllm_plugin(plugin_lib_path: str, platform: str) -> bool:
     """
     Loads and initializes the TensorRT-LLM plugin from the given shared library path.
 
@@ -947,6 +947,9 @@ def load_and_initialize_trtllm_plugin(plugin_lib_path: str) -> bool:
     Returns:
         bool: True if successful, False otherwise.
     """
+    if "windows" in platform:
+        logger.info("NCCL backend is not supported on Windows")
+        return False
     try:
         handle = ctypes.CDLL(plugin_lib_path)
         logger.info(f"Successfully loaded plugin library: {plugin_lib_path}")
@@ -1002,8 +1005,10 @@ def load_tensorrt_llm() -> bool:
         bool: True if the plugin was successfully loaded and initialized, False otherwise.
     """
     plugin_lib_path = os.environ.get("TRTLLM_PLUGINS_PATH")
+    platform = Platform.current_platform()
+    platform = str(platform).lower()
     if plugin_lib_path:
-        return load_and_initialize_trtllm_plugin(plugin_lib_path)
+        return load_and_initialize_trtllm_plugin(plugin_lib_path, platform)
     else:
         # this option can be used by user if TRTLLM_PLUGINS_PATH is not set by user
         use_trtllm_plugin = os.environ.get("USE_TRTLLM_PLUGINS", "0").lower() in (
@@ -1017,10 +1022,7 @@ def load_tensorrt_llm() -> bool:
                 "Neither TRTLLM_PLUGIN_PATH is set nor is it directed to download the shared library. Please set either of the two to use TRT-LLM libraries in torchTRT"
             )
             return False
-        else:
-            platform = Platform.current_platform()
-            platform = str(platform).lower()
 
         with download_plugin_lib_path(platform) as plugin_lib_path:
-            return load_and_initialize_trtllm_plugin(plugin_lib_path)
+            return load_and_initialize_trtllm_plugin(plugin_lib_path, platform)
     return False

tests/py/dynamo/distributed/test_nccl_ops.py

@@ -6,6 +6,7 @@
 from distributed_utils import set_environment_variables_pytest
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt._enums import Platform
 
 set_environment_variables_pytest()
 dist.init_process_group(backend="nccl", init_method="env://")
@@ -15,7 +16,12 @@
 
 from conversion.harness import DispatchTestCase
 
+platform_str = str(Platform.current_platform()).lower()
 
+
+@unittest.skipIf(
+    "win" in platform_str, "Skipped on Windows: NCCL backend is not supported."
+)
 class TestGatherNcclOpsConverter(DispatchTestCase):
     @parameterized.expand([8])
     def test_nccl_ops(self, linear_layer_dim):