ci: fix nccl builds in CI

Author: narendasan

docker/MODULE.bazel.docker

@@ -8,6 +8,7 @@ bazel_dep(name = "googletest", version = "1.16.0")
 bazel_dep(name = "platforms", version = "0.0.11")
 bazel_dep(name = "rules_cc", version = "0.1.1")
 bazel_dep(name = "rules_python", version = "1.3.0")
+bazel_dep(name = "bazel_skylib", version = "1.7.1")
 
 python = use_extension("@rules_python//python/extensions:python.bzl", "python")
 python.toolchain(
@@ -24,9 +25,12 @@ git_override(
 
 local_repository = use_repo_rule("@bazel_tools//tools/build_defs/repo:local.bzl", "local_repository")
 
-
 new_local_repository = use_repo_rule("@bazel_tools//tools/build_defs/repo:local.bzl", "new_local_repository")
 
+torch_nccl_detect = use_repo_rule("//toolchains/torch_nccl:defs.bzl", "torch_nccl_detect")
+
+torch_nccl_detect(name = "torch_nccl")
+
 # CUDA should be installed on the system locally
 new_local_repository(
     name = "cuda",

docker/MODULE.bazel.ngc

@@ -8,6 +8,7 @@ bazel_dep(name = "googletest", version = "1.14.0")
 bazel_dep(name = "platforms", version = "0.0.10")
 bazel_dep(name = "rules_cc", version = "0.0.9")
 bazel_dep(name = "rules_python", version = "0.34.0")
+bazel_dep(name = "bazel_skylib", version = "1.7.1")
 
 python = use_extension("@rules_python//python/extensions:python.bzl", "python")
 python.toolchain(
@@ -24,9 +25,12 @@ git_override(
 
 local_repository = use_repo_rule("@bazel_tools//tools/build_defs/repo:local.bzl", "local_repository")
 
-
 new_local_repository = use_repo_rule("@bazel_tools//tools/build_defs/repo:local.bzl", "new_local_repository")
 
+torch_nccl_detect = use_repo_rule("//toolchains/torch_nccl:defs.bzl", "torch_nccl_detect")
+
+torch_nccl_detect(name = "torch_nccl")
+
 
 # External dependency for torch_tensorrt if you already have precompiled binaries.
 new_local_repository(

py/torch_tensorrt/__init__.py

@@ -99,12 +99,12 @@ def _register_with_torch() -> None:
     from torch_tensorrt.dynamo import backend  # noqa: F401
     from torch_tensorrt import dynamo  # noqa: F401
 
-from torch_tensorrt import distributed  # noqa: F401
 from torch_tensorrt._compile import *  # noqa: F403
 from torch_tensorrt.distributed._distributed import (  # noqa: F401
     distributed_group,
     set_distributed_group,
 )
+from torch_tensorrt import distributed  # noqa: F401
 from torch_tensorrt.dynamo.runtime._MutableTorchTensorRTModule import (
     MutableTorchTensorRTModule,
 )

setup.py

@@ -472,6 +472,8 @@ def run(self):
 
 dynamo_packages = [
     "torch_tensorrt",
+    "torch_tensorrt.distributed",
+    "torch_tensorrt.distributed.run",
     "torch_tensorrt.dynamo",
     "torch_tensorrt.dynamo.backend",
     "torch_tensorrt.dynamo.conversion",
@@ -506,6 +508,8 @@ def run(self):
 
 dynamo_package_dir = {
     "torch_tensorrt": "py/torch_tensorrt",
+    "torch_tensorrt.distributed": "py/torch_tensorrt/distributed",
+    "torch_tensorrt.distributed.run": "py/torch_tensorrt/distributed/run",
     "torch_tensorrt.dynamo": "py/torch_tensorrt/dynamo",
     "torch_tensorrt.dynamo.backend": "py/torch_tensorrt/dynamo/backend",
     "torch_tensorrt.dynamo.conversion": "py/torch_tensorrt/dynamo/conversion",

tests/py/dynamo/backend/test_backend_compiler.py

@@ -78,8 +78,8 @@ def forward(self, x, y):
         unexpected_ops = {torch.ops.aten.add.Tensor}
 
         inputs = [
-            torch.randint(-40, 40, (16, 7, 5), dtype=torch.int).cuda(),
-            torch.randint(1, 40, (16, 7, 5), dtype=torch.int).cuda(),
+            torch.randn(16, 7, 5, dtype=torch.float).cuda(),
+            torch.randn(16, 7, 5, dtype=torch.float).cuda(),
         ]
 
         (

toolchains/ci_workspaces/MODULE.bazel.tmpl

@@ -8,6 +8,7 @@ bazel_dep(name = "googletest", version = "1.16.0")
 bazel_dep(name = "platforms", version = "0.0.11")
 bazel_dep(name = "rules_cc", version = "0.1.1")
 bazel_dep(name = "rules_python", version = "1.3.0")
+bazel_dep(name = "bazel_skylib", version = "1.7.1")
 
 python = use_extension("@rules_python//python/extensions:python.bzl", "python")
 python.toolchain(
@@ -24,6 +25,10 @@ git_override(
 
 local_repository = use_repo_rule("@bazel_tools//tools/build_defs/repo:local.bzl", "local_repository")
 
+torch_nccl_detect = use_repo_rule("//toolchains/torch_nccl:defs.bzl", "torch_nccl_detect")
+
+torch_nccl_detect(name = "torch_nccl")
+
 # External dependency for torch_tensorrt if you already have precompiled binaries.
 local_repository(
     name = "torch_tensorrt",