refactor!: change for hermetic with minor user interface breaking (#300)

* refactor!: rules_cuda_dependencies only pull in bazel deps for rules_cuda and split the local_cuda part into newly added rules_cuda_toolchains * refactor!: rename local_toolchain as toolkit for bzlmod
bazel-contrib · Feb 22, 2025 · fe377e2 · fe377e2
1 parent 458eccf
commit fe377e2
Show file tree

Hide file tree

Showing 13 changed files with 65 additions and 38 deletions.
diff --git a/.github/release_notes.template b/.github/release_notes.template
@@ -8,7 +8,7 @@ http_archive(
     urls = ["https://github.com/bazel-contrib/rules_cuda/releases/download/{version}/rules_cuda-{version}.tar.gz"],
 )
 
-load("@rules_cuda//cuda:repositories.bzl", "register_detected_cuda_toolchains", "rules_cuda_dependencies")
+load("@rules_cuda//cuda:repositories.bzl", "rules_cuda_dependencies", "rules_cuda_toolchains")
 rules_cuda_dependencies()
-register_detected_cuda_toolchains()
-```
+rules_cuda_toolchains(register_toolchains = True)
+```
diff --git a/MODULE.bazel b/MODULE.bazel
@@ -8,7 +8,7 @@ bazel_dep(name = "bazel_skylib", version = "1.4.2")
 bazel_dep(name = "platforms", version = "0.0.6")
 
 cuda = use_extension("@rules_cuda//cuda:extensions.bzl", "toolchain")
-cuda.local_toolchain(
+cuda.toolkit(
     name = "local_cuda",
     toolkit_path = "",
 )

diff --git a/README.md b/README.md
@@ -18,13 +18,14 @@ http_archive(
     strip_prefix = "rules_cuda-{git_commit_hash}",
     urls = ["https://github.com/bazel-contrib/rules_cuda/archive/{git_commit_hash}.tar.gz"],
 )
-load("@rules_cuda//cuda:repositories.bzl", "register_detected_cuda_toolchains", "rules_cuda_dependencies")
+load("@rules_cuda//cuda:repositories.bzl", "rules_cuda_dependencies", "rules_cuda_toolchains")
 rules_cuda_dependencies()
-register_detected_cuda_toolchains()
+rules_cuda_toolchains(register_toolchains = True)
 ```
 
-**NOTE**: the use of `register_detected_cuda_toolchains` depends on the environment variable `CUDA_PATH`. You must also
-ensure the host compiler is available. On Windows, this means that you will also need to set the environment variable
+**NOTE**: `rules_cuda_toolchains` implicitly calls to `register_detected_cuda_toolchains`, and the use of
+`register_detected_cuda_toolchains` depends on the environment variable `CUDA_PATH`. You must also ensure the
+host compiler is available. On Windows, this means that you will also need to set the environment variable
 `BAZEL_VC` properly.
 
 [`detect_cuda_toolkit`](https://github.com/bazel-contrib/rules_cuda/blob/5633f0c0f7/cuda/private/repositories.bzl#L28-L58)
@@ -47,7 +48,7 @@ archive_override(
 )
 
 cuda = use_extension("@rules_cuda//cuda:extensions.bzl", "toolchain")
-cuda.local_toolchain(
+cuda.toolkit(
     name = "local_cuda",
     toolkit_path = "",
 )

diff --git a/WORKSPACE.bazel b/WORKSPACE.bazel
@@ -5,11 +5,11 @@ local_repository(
     path = "examples",
 )
 
-load("//cuda:repositories.bzl", "register_detected_cuda_toolchains", "rules_cuda_dependencies")
+load("@rules_cuda//cuda:repositories.bzl", "rules_cuda_dependencies", "rules_cuda_toolchains")
 
 rules_cuda_dependencies()
 
-register_detected_cuda_toolchains()
+rules_cuda_toolchains(register_toolchains = True)
 
 load("@bazel_skylib//:workspace.bzl", "bazel_skylib_workspace")
 

diff --git a/cuda/extensions.bzl b/cuda/extensions.bzl
@@ -2,7 +2,7 @@
 
 load("//cuda/private:repositories.bzl", "local_cuda")
 
-cuda_toolkit = tag_class(attrs = {
+cuda_toolkit_tag = tag_class(attrs = {
     "name": attr.string(doc = "Name for the toolchain repository", default = "local_cuda"),
     "toolkit_path": attr.string(doc = "Path to the CUDA SDK, if empty the environment variable CUDA_PATH will be used to deduce this path."),
 })
@@ -22,24 +22,27 @@ def _find_modules(module_ctx):
 
     return root, our_module
 
+def _module_tag_to_dict(t):
+    return {attr: getattr(t, attr) for attr in dir(t)}
+
 def _init(module_ctx):
     # Toolchain configuration is only allowed in the root module, or in rules_cuda.
     root, rules_cuda = _find_modules(module_ctx)
-    toolchains = root.tags.local_toolchain or rules_cuda.tags.local_toolchain
+    toolkits = root.tags.toolkit or rules_cuda.tags.toolkit
 
     registrations = {}
-    for toolchain in toolchains:
-        if toolchain.name in registrations.keys():
-            if toolchain.toolkit_path == registrations[toolchain.name]:
-                # No problem to register a matching toolchain twice
+    for toolkit in toolkits:
+        if toolkit.name in registrations.keys():
+            if toolkit.toolkit_path == registrations[toolkit.name].toolkit_path:
+                # No problem to register a matching toolkit twice
                 continue
-            fail("Multiple conflicting toolchains declared for name {} ({} and {}".format(toolchain.name, toolchain.toolkit_path, registrations[toolchain.name]))
+            fail("Multiple conflicting toolkits declared for name {} ({} and {}".format(toolkit.name, toolkit.toolkit_path, registrations[toolkit.name].toolkit_path))
         else:
-            registrations[toolchain.name] = toolchain.toolkit_path
-    for name, toolkit_path in registrations.items():
-        local_cuda(name = name, toolkit_path = toolkit_path)
+            registrations[toolkit.name] = toolkit
+    for _, toolkit in registrations.items():
+        local_cuda(**_module_tag_to_dict(toolkit))
 
 toolchain = module_extension(
     implementation = _init,
-    tag_classes = {"local_toolchain": cuda_toolkit},
+    tag_classes = {"toolkit": cuda_toolkit_tag},
 )
diff --git a/cuda/private/repositories.bzl b/cuda/private/repositories.bzl
@@ -3,6 +3,7 @@
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 load("@bazel_tools//tools/build_defs/repo:utils.bzl", "maybe")
 load("//cuda/private:template_helper.bzl", "template_helper")
+load("//cuda/private:toolchain.bzl", "register_detected_cuda_toolchains")
 
 def _to_forward_slash(s):
     return s.replace("\\", "/")
@@ -197,12 +198,8 @@ local_cuda = repository_rule(
     # remotable = True,
 )
 
-def rules_cuda_dependencies(toolkit_path = None):
-    """Populate the dependencies for rules_cuda. This will setup workspace dependencies (other bazel rules) and local toolchains.
-
-    Args:
-        toolkit_path: Optionally specify the path to CUDA toolkit. If not specified, it will be detected automatically.
-    """
+def rules_cuda_dependencies():
+    """Populate the dependencies for rules_cuda. This will setup other bazel rules as workspace dependencies"""
     maybe(
         name = "bazel_skylib",
         repo_rule = http_archive,
@@ -223,4 +220,18 @@ def rules_cuda_dependencies(toolkit_path = None):
         ],
     )
 
-    local_cuda(name = "local_cuda", toolkit_path = toolkit_path)
+def rules_cuda_toolchains(toolkit_path = None, register_toolchains = False):
+    """Populate the local_cuda repo.
+
+    Args:
+        toolkit_path: Optionally specify the path to CUDA toolkit. If not specified, it will be detected automatically.
+        register_toolchains: Register the toolchains if enabled.
+    """
+
+    local_cuda(
+        name = "local_cuda",
+        toolkit_path = toolkit_path,
+    )
+
+    if register_toolchains:
+        register_detected_cuda_toolchains()
diff --git a/cuda/repositories.bzl b/cuda/repositories.bzl
@@ -1,6 +1,12 @@
-load("//cuda/private:repositories.bzl", _local_cuda = "local_cuda", _rules_cuda_dependencies = "rules_cuda_dependencies")
+load(
+    "//cuda/private:repositories.bzl",
+    _local_cuda = "local_cuda",
+    _rules_cuda_dependencies = "rules_cuda_dependencies",
+    _rules_cuda_toolchains = "rules_cuda_toolchains",
+)
 load("//cuda/private:toolchain.bzl", _register_detected_cuda_toolchains = "register_detected_cuda_toolchains")
 
-rules_cuda_dependencies = _rules_cuda_dependencies
 local_cuda = _local_cuda
+rules_cuda_dependencies = _rules_cuda_dependencies
+rules_cuda_toolchains = _rules_cuda_toolchains
 register_detected_cuda_toolchains = _register_detected_cuda_toolchains
diff --git a/docs/MODULE.bazel b/docs/MODULE.bazel
@@ -13,7 +13,7 @@ local_path_override(
 )
 
 cuda = use_extension("@rules_cuda//cuda:extensions.bzl", "toolchain")
-cuda.local_toolchain(
+cuda.toolkit(
     name = "local_cuda",
     toolkit_path = "",
 )

diff --git a/docs/user_docs.bzl b/docs/user_docs.bzl
@@ -1,5 +1,10 @@
 load("@rules_cuda//cuda:defs.bzl", _cuda_binary = "cuda_binary", _cuda_library = "cuda_library", _cuda_objects = "cuda_objects", _cuda_test = "cuda_test")
-load("@rules_cuda//cuda:repositories.bzl", _register_detected_cuda_toolchains = "register_detected_cuda_toolchains", _rules_cuda_dependencies = "rules_cuda_dependencies")
+load(
+    "@rules_cuda//cuda:repositories.bzl",
+    _register_detected_cuda_toolchains = "register_detected_cuda_toolchains",
+    _rules_cuda_dependencies = "rules_cuda_dependencies",
+    _rules_cuda_toolchains = "rules_cuda_toolchains",
+)
 load("@rules_cuda//cuda/private:rules/flags.bzl", _cuda_archs_flag = "cuda_archs_flag")
 
 cuda_library = _cuda_library
@@ -12,3 +17,4 @@ cuda_archs = _cuda_archs_flag
 
 register_detected_cuda_toolchains = _register_detected_cuda_toolchains
 rules_cuda_dependencies = _rules_cuda_dependencies
+rules_cuda_toolchains = _rules_cuda_toolchains
diff --git a/examples/MODULE.bazel b/examples/MODULE.bazel
@@ -11,7 +11,7 @@ local_path_override(
 )
 
 cuda = use_extension("@rules_cuda//cuda:extensions.bzl", "toolchain")
-cuda.local_toolchain(
+cuda.toolkit(
     name = "local_cuda",
     toolkit_path = "",
 )

diff --git a/examples/WORKSPACE.bazel b/examples/WORKSPACE.bazel
@@ -12,11 +12,11 @@ local_repository(
 # If you want to have a different version of some dependency,
 # you should fetch it *before* calling this.
 
-load("@rules_cuda//cuda:repositories.bzl", "register_detected_cuda_toolchains", "rules_cuda_dependencies")
+load("@rules_cuda//cuda:repositories.bzl", "rules_cuda_dependencies", "rules_cuda_toolchains")
 
 rules_cuda_dependencies()
 
-register_detected_cuda_toolchains()
+rules_cuda_toolchains(register_toolchains = True)
 
 #################################
 # Dependencies for nccl example #

diff --git a/tests/integration/toolchain_none/MODULE.bazel b/tests/integration/toolchain_none/MODULE.bazel
@@ -7,7 +7,7 @@ local_path_override(
 )
 
 cuda = use_extension("@rules_cuda//cuda:extensions.bzl", "toolchain")
-cuda.local_toolchain(
+cuda.toolkit(
     name = "local_cuda",
     toolkit_path = "/nonexistent/cuda/toolkit/path",
 )

diff --git a/tests/integration/toolchain_root/MODULE.bazel b/tests/integration/toolchain_root/MODULE.bazel
@@ -7,7 +7,7 @@ local_path_override(
 )
 
 cuda = use_extension("@rules_cuda//cuda:extensions.bzl", "toolchain")
-cuda.local_toolchain(
+cuda.toolkit(
     name = "local_cuda",
     toolkit_path = "",
 )