From 935b4fcb35a3974c505fbaecce8e67a021860e48 Mon Sep 17 00:00:00 2001 From: Cloud Han Date: Sun, 29 Dec 2024 23:55:00 +0800 Subject: [PATCH 1/4] feat: add cuda_redist_json repo rule to pull in remote redistrib.json file, to simplify hermetic repo setup --- cuda/extensions.bzl | 33 +++++++++++- cuda/private/repositories.bzl | 68 ++++++++++++++++++++++++ cuda/private/template_helper.bzl | 45 ++++++++++++++++ cuda/private/templates/BUILD.redist_json | 18 +++++++ cuda/private/templates/redist.bzl.tpl | 15 ++++++ cuda/repositories.bzl | 2 + 6 files changed, 180 insertions(+), 1 deletion(-) create mode 100644 cuda/private/templates/BUILD.redist_json create mode 100644 cuda/private/templates/redist.bzl.tpl diff --git a/cuda/extensions.bzl b/cuda/extensions.bzl index 35368f4b..f1d32c86 100644 --- a/cuda/extensions.bzl +++ b/cuda/extensions.bzl @@ -1,7 +1,7 @@ """Entry point for extensions used by bzlmod.""" load("//cuda/private:compat.bzl", "components_mapping_compat") -load("//cuda/private:repositories.bzl", "cuda_component", "local_cuda") +load("//cuda/private:repositories.bzl", "cuda_component", "cuda_redist_json", "local_cuda") cuda_component_tag = tag_class(attrs = { "name": attr.string(mandatory = True, doc = "Repo name for the deliverable cuda_component"), @@ -33,6 +33,30 @@ cuda_component_tag = tag_class(attrs = { ), }) +cuda_redist_json_tag = tag_class(attrs = { + "name": attr.string(mandatory = True, doc = "Repo name for the cuda_redist_json"), + "components": attr.string_list(mandatory = True, doc = "components to be used"), + "integrity": attr.string( + doc = "Expected checksum in Subresource Integrity format of the file downloaded. " + + "This must match the checksum of the file downloaded.", + ), + "sha256": attr.string( + doc = "The expected SHA-256 of the file downloaded. " + + "This must match the SHA-256 of the file downloaded.", + ), + "urls": attr.string_list( + doc = "A list of URLs to a file that will be made available to Bazel. " + + "Each entry must be a file, http or https URL. Redirections are followed. " + + "Authentication is not supported. " + + "URLs are tried in order until one succeeds, so you should list local mirrors first. " + + "If all downloads fail, the rule will fail.", + ), + "version": attr.string( + doc = "Generate a URL by using the specified version." + + "This URL will be tried after all URLs specified in the `urls` attribute.", + ), +}) + cuda_toolkit_tag = tag_class(attrs = { "name": attr.string(mandatory = True, doc = "Name for the toolchain repository", default = "local_cuda"), "toolkit_path": attr.string( @@ -70,17 +94,23 @@ def _impl(module_ctx): # Toolchain configuration is only allowed in the root module, or in rules_cuda. root, rules_cuda = _find_modules(module_ctx) components = None + redist_jsons = None toolkits = None if root.tags.toolkit: components = root.tags.component + redist_jsons = root.tags.redist_json toolkits = root.tags.toolkit else: components = rules_cuda.tags.component + redist_jsons = rules_cuda.tags.redist_json toolkits = rules_cuda.tags.toolkit for component in components: cuda_component(**_module_tag_to_dict(component)) + for redist_json in redist_jsons: + cuda_redist_json(**_module_tag_to_dict(redist_json)) + registrations = {} for toolkit in toolkits: if toolkit.name in registrations.keys(): @@ -97,6 +127,7 @@ toolchain = module_extension( implementation = _impl, tag_classes = { "component": cuda_component_tag, + "redist_json": cuda_redist_json_tag, "toolkit": cuda_toolkit_tag, }, ) diff --git a/cuda/private/repositories.bzl b/cuda/private/repositories.bzl index 68590c3a..117a3030 100644 --- a/cuda/private/repositories.bzl +++ b/cuda/private/repositories.bzl @@ -337,6 +337,74 @@ def default_components_mapping(components): """ return {c: "@local_cuda_" + c for c in components} +def _cuda_redist_json_impl(repository_ctx): + the_url = None # the url that successfully fetch redist json, we then use it to fetch deliverables + urls = [u for u in repository_ctx.attr.urls] + + redist_ver = repository_ctx.attr.version + if redist_ver: + urls.append("https://developer.download.nvidia.com/compute/cuda/redist/redistrib_{}.json".format(redist_ver)) + + if len(urls) == 0: + fail("`urls` or `version` must be specified.") + + for url in urls: + ret = repository_ctx.download( + output = "redist.json", + integrity = repository_ctx.attr.integrity, + sha256 = repository_ctx.attr.sha256, + url = url, + ) + if ret.success: + the_url = url + break + + if the_url == None: + fail("Failed to retrieve the redist json file.") + + # convert redist.json to list of spec (list of dicts with cuda_components attrs) + specs = [] + redist = json.decode(repository_ctx.read("redist.json")) + if not redist_ver: + redist_ver = redist["release_label"] + for c in repository_ctx.attr.components: + c_full = FULL_COMPONENT_NAME[c] + os = None + if _is_linux(repository_ctx): + os = "linux" + elif _is_windows(repository_ctx): + os = "windows" + + arch = "x86_64" # TODO: support cross compiling + platform = "{os}-{arch}".format(os = os, arch = arch) + + payload = redist[c_full][platform] + payload_relative_path = payload["relative_path"] + payload_url = the_url.rsplit("/", 1)[0] + "/" + payload_relative_path + archive_name = payload_relative_path.rsplit("/", 1)[1].split("-archive.")[0] + "-archive" + + specs.append({ + "component_name": c, + "urls": [payload_url], + "sha256": payload["sha256"], + "strip_prefix": archive_name, + "version": redist[c_full]["version"], + }) + + template_helper.generate_redist_bzl(repository_ctx, specs, redist_ver) + repository_ctx.symlink(Label("//cuda/private:templates/BUILD.redist_json"), "BUILD") + +cuda_redist_json = repository_rule( + implementation = _cuda_redist_json_impl, + attrs = { + "components": attr.string_list(mandatory = True), + "integrity": attr.string(mandatory = False), + "sha256": attr.string(mandatory = False), + "urls": attr.string_list(mandatory = False), + "version": attr.string(mandatory = False), + }, +) + def rules_cuda_dependencies(): """Populate the dependencies for rules_cuda. This will setup other bazel rules as workspace dependencies""" maybe( diff --git a/cuda/private/template_helper.bzl b/cuda/private/template_helper.bzl index 9590b7ba..ffa5966c 100644 --- a/cuda/private/template_helper.bzl +++ b/cuda/private/template_helper.bzl @@ -89,6 +89,50 @@ def _generate_defs_bzl(repository_ctx, is_local_ctk): } repository_ctx.template("defs.bzl", tpl_label, substitutions = substitutions, executable = False) +def _generate_redist_bzl(repository_ctx, component_specs, redist_version): + """Generate `@rules_cuda_redist_json//:redist.bzl` + + Args: + repository_ctx: repository_ctx + component_specs: list of dict, dict keys are component_name, urls, sha256, strip_prefix and version + """ + + rules_cuda_components_body = [] + mapping = {} + + component_tpl = """cuda_component( + name = "{repo_name}", + component_name = "{component_name}", + sha256 = {sha256}, + strip_prefix = {strip_prefix}, + urls = {urls}, + )""" + + for spec in component_specs: + repo_name = "local_cuda_" + spec["component_name"] + version = spec.get("version", None) + if version != None: + repo_name = repo_name + "_v" + version + + rules_cuda_components_body.append( + component_tpl.format( + repo_name = repo_name, + component_name = spec["component_name"], + sha256 = repr(spec["sha256"]), + strip_prefix = repr(spec["strip_prefix"]), + urls = repr(spec["urls"]), + ), + ) + mapping[spec["component_name"]] = "@" + repo_name + + tpl_label = Label("//cuda/private:templates/redist.bzl.tpl") + substitutions = { + "%{rules_cuda_components_body}": "\n\n ".join(rules_cuda_components_body), + "%{components_mapping}": repr(mapping), + "%{version}": redist_version, + } + repository_ctx.template("redist.bzl", tpl_label, substitutions = substitutions, executable = False) + def _generate_toolchain_build(repository_ctx, cuda): tpl_label = Label( "//cuda/private:templates/BUILD.local_toolchain_" + @@ -127,6 +171,7 @@ def _generate_toolchain_clang_build(repository_ctx, cuda, clang_path): template_helper = struct( generate_build = _generate_build, generate_defs_bzl = _generate_defs_bzl, + generate_redist_bzl = _generate_redist_bzl, generate_toolchain_build = _generate_toolchain_build, generate_toolchain_clang_build = _generate_toolchain_clang_build, ) diff --git a/cuda/private/templates/BUILD.redist_json b/cuda/private/templates/BUILD.redist_json new file mode 100644 index 00000000..746e8daa --- /dev/null +++ b/cuda/private/templates/BUILD.redist_json @@ -0,0 +1,18 @@ +package( + default_visibility = ["//visibility:public"], +) + +filegroup( + name = "redist_bzl", + srcs = [":redist.bzl"], +) + +filegroup( + name = "redist_json", + srcs = [":redist.json"], +) + +exports_files([ + "redist.bzl", + "redist.json", +]) diff --git a/cuda/private/templates/redist.bzl.tpl b/cuda/private/templates/redist.bzl.tpl new file mode 100644 index 00000000..eeb6df80 --- /dev/null +++ b/cuda/private/templates/redist.bzl.tpl @@ -0,0 +1,15 @@ +load("@rules_cuda//cuda:repositories.bzl", "cuda_component", "rules_cuda_toolchains") + +def rules_cuda_components(): + # See template_helper.generate_redist_bzl(...) for body generation logic + %{rules_cuda_components_body} + + return %{components_mapping} + +def rules_cuda_components_and_toolchains(register_toolchains = False): + components_mapping = rules_cuda_components() + rules_cuda_toolchains( + components_mapping= components_mapping, + register_toolchains = register_toolchains, + version = "%{version}", + ) diff --git a/cuda/repositories.bzl b/cuda/repositories.bzl index 9635d0db..04a8a493 100644 --- a/cuda/repositories.bzl +++ b/cuda/repositories.bzl @@ -1,6 +1,7 @@ load( "//cuda/private:repositories.bzl", _cuda_component = "cuda_component", + _cuda_redist_json = "cuda_redist_json", _default_components_mapping = "default_components_mapping", _local_cuda = "local_cuda", _rules_cuda_dependencies = "rules_cuda_dependencies", @@ -10,6 +11,7 @@ load("//cuda/private:toolchain.bzl", _register_detected_cuda_toolchains = "regis # rules cuda_component = _cuda_component +cuda_redist_json = _cuda_redist_json local_cuda = _local_cuda # macros From 92c346136ceedb107ed7bbfb5f7dc2d113a2295b Mon Sep 17 00:00:00 2001 From: Cloud Han Date: Mon, 24 Feb 2025 23:48:08 +0800 Subject: [PATCH 2/4] test: add cuda_redist_json integration tests for workspace --- tests/integration/test_all.sh | 11 +++++++++ .../toolchain_redist_json/BUILD.bazel | 1 + .../toolchain_redist_json/WORKSPACE.bazel | 23 +++++++++++++++++++ 3 files changed, 35 insertions(+) create mode 120000 tests/integration/toolchain_redist_json/BUILD.bazel create mode 100644 tests/integration/toolchain_redist_json/WORKSPACE.bazel diff --git a/tests/integration/test_all.sh b/tests/integration/test_all.sh index 4d8aec70..62c30e54 100755 --- a/tests/integration/test_all.sh +++ b/tests/integration/test_all.sh @@ -71,3 +71,14 @@ pushd "$this_dir/toolchain_components" bazel build --enable_bzlmod //:use_rule bazel clean && bazel shutdown popd + +# toolchain configured with deliverables (redistrib.json with workspace) +pushd "$this_dir/toolchain_redist_json" + bazel build --enable_workspace //... --@rules_cuda//cuda:enable=False + bazel build --enable_workspace //... --@rules_cuda//cuda:enable=True + bazel build --enable_workspace //:optinally_use_rule --@rules_cuda//cuda:enable=False + bazel build --enable_workspace //:optinally_use_rule --@rules_cuda//cuda:enable=True + bazel build --enable_workspace //:use_library + bazel build --enable_workspace //:use_rule + bazel clean && bazel shutdown +popd diff --git a/tests/integration/toolchain_redist_json/BUILD.bazel b/tests/integration/toolchain_redist_json/BUILD.bazel new file mode 120000 index 00000000..60a18fe7 --- /dev/null +++ b/tests/integration/toolchain_redist_json/BUILD.bazel @@ -0,0 +1 @@ +../BUILD.to_symlink \ No newline at end of file diff --git a/tests/integration/toolchain_redist_json/WORKSPACE.bazel b/tests/integration/toolchain_redist_json/WORKSPACE.bazel new file mode 100644 index 00000000..0cc935f2 --- /dev/null +++ b/tests/integration/toolchain_redist_json/WORKSPACE.bazel @@ -0,0 +1,23 @@ +local_repository( + name = "rules_cuda", + path = "../../..", +) + +# buildifier: disable=load-on-top +load("@rules_cuda//cuda:repositories.bzl", "cuda_redist_json", "rules_cuda_dependencies") + +rules_cuda_dependencies() + +cuda_redist_json( + name = "rules_cuda_redist_json", + components = [ + "cccl", + "cudart", + "nvcc", + ], + version = "12.6.3", +) + +load("@rules_cuda_redist_json//:redist.bzl", "rules_cuda_components_and_toolchains") + +rules_cuda_components_and_toolchains(register_toolchains = True) From 7a5951237194212cee7ef918f1e4692180987c81 Mon Sep 17 00:00:00 2001 From: Cloud Han Date: Tue, 25 Feb 2025 04:32:05 +0000 Subject: [PATCH 3/4] test: add maually specified components as a fallback for tests with bzlmod --- .../toolchain_redist_json/MODULE.bazel | 53 +++++++++++++++++++ 1 file changed, 53 insertions(+) create mode 100644 tests/integration/toolchain_redist_json/MODULE.bazel diff --git a/tests/integration/toolchain_redist_json/MODULE.bazel b/tests/integration/toolchain_redist_json/MODULE.bazel new file mode 100644 index 00000000..ae57f4f4 --- /dev/null +++ b/tests/integration/toolchain_redist_json/MODULE.bazel @@ -0,0 +1,53 @@ +module(name = "bzlmod_components") + +# FIXME: cuda_redist_json is not exposed in bzlmod now. Fallback to manually specified components for tests +bazel_dep(name = "rules_cuda", version = "0.0.0") +local_path_override( + module_name = "rules_cuda", + path = "../../..", +) + +cuda = use_extension("@rules_cuda//cuda:extensions.bzl", "toolchain") +cuda.component( + name = "local_cuda_cccl", + component_name = "cccl", + sha256 = "9c3145ef01f73e50c0f5fcf923f0899c847f487c529817daa8f8b1a3ecf20925", + strip_prefix = "cuda_cccl-linux-x86_64-12.6.77-archive", + urls = ["https://developer.download.nvidia.com/compute/cuda/redist/cuda_cccl/linux-x86_64/cuda_cccl-linux-x86_64-12.6.77-archive.tar.xz"], +) +cuda.component( + name = "local_cuda_cudart", + component_name = "cudart", + sha256 = "f74689258a60fd9c5bdfa7679458527a55e22442691ba678dcfaeffbf4391ef9", + strip_prefix = "cuda_cudart-linux-x86_64-12.6.77-archive", + urls = ["https://developer.download.nvidia.com/compute/cuda/redist/cuda_cudart/linux-x86_64/cuda_cudart-linux-x86_64-12.6.77-archive.tar.xz"], +) +cuda.component( + name = "local_cuda_nvcc", + component_name = "nvcc", + sha256 = "840deff234d9bef20d6856439c49881cb4f29423b214f9ecd2fa59b7ac323817", + strip_prefix = "cuda_nvcc-linux-x86_64-12.6.85-archive", + urls = ["https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/linux-x86_64/cuda_nvcc-linux-x86_64-12.6.85-archive.tar.xz"], +) +cuda.toolkit( + name = "local_cuda", + components_mapping = { + "cccl": "@local_cuda_cccl", + "cudart": "@local_cuda_cudart", + "nvcc": "@local_cuda_nvcc", + }, + version = "12.6", +) +use_repo( + cuda, + "local_cuda", + "local_cuda_cccl", + "local_cuda_cudart", + "local_cuda_nvcc", +) + +bazel_dep(name = "rules_cuda_examples", version = "0.0.0") +local_path_override( + module_name = "rules_cuda_examples", + path = "../../../examples", +) From 998736d96774344535172b9bfd1c51be0212476f Mon Sep 17 00:00:00 2001 From: Cloud Han Date: Tue, 25 Feb 2025 21:04:28 +0800 Subject: [PATCH 4/4] fix: add WORKSPACE.bzlmod --- tests/integration/toolchain_redist_json/WORKSPACE.bzlmod | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/integration/toolchain_redist_json/WORKSPACE.bzlmod diff --git a/tests/integration/toolchain_redist_json/WORKSPACE.bzlmod b/tests/integration/toolchain_redist_json/WORKSPACE.bzlmod new file mode 100644 index 00000000..e69de29b