From b52307f86bc8609d082f1591df8c4baebd4057a6 Mon Sep 17 00:00:00 2001 From: Michael Droettboom Date: Tue, 14 Oct 2025 10:12:00 -0400 Subject: [PATCH 1/3] Add a script to help update to a new CTK version --- toolshed/update_ctk.py | 209 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 209 insertions(+) create mode 100644 toolshed/update_ctk.py diff --git a/toolshed/update_ctk.py b/toolshed/update_ctk.py new file mode 100644 index 0000000000..f6c1737bcf --- /dev/null +++ b/toolshed/update_ctk.py @@ -0,0 +1,209 @@ +#!/usr/bin/env python + +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE + +import argparse +import json +import re +import subprocess +import sys +import tarfile +import tempfile +import venv +from pathlib import Path +from urllib.request import urlopen + +# Example URL of an HTML directory listing +CONTENT_URL = "https://developer.download.nvidia.com/compute/cuda/redist" + + +CYBIND_GENERATED_LIBRARIES = [ + ("cufile", "libcufile", "cufile"), + ("nvvm", "libnvvm", "nvvm"), + ("nvjitlink", "libnvjitlink", "nvJitLink"), +] + + +def fetch_headers(version: str, library_name: str, dest_dir: Path): + def tar_filter(members): + for tarinfo in members: + name = Path(tarinfo.name) + parts = name.parts + try: + idx = parts.index("include") + except ValueError: + continue + tarinfo.name = str(Path(*parts[idx + 1 :])) + yield tarinfo + + output_dir = dest_dir / Path(version) + if output_dir.exists(): + print(f"Skipping header download for {library_name} {version}, already exists") + return + + output_dir.mkdir() + + json_url = f"{CONTENT_URL}/redistrib_{version}.json" + with urlopen(json_url) as resp: # noqa: S310 + content = json.loads(resp.read().decode("utf-8")) + if library := content.get(library_name): + archive_url = f"{CONTENT_URL}/{library['linux-x86_64']['relative_path']}" + print(f"Fetching package {archive_url}") + with tempfile.NamedTemporaryFile() as tmp: + tmppath = Path(tmp.name) + + with tmppath.open("wb") as f, urlopen(archive_url) as resp: # noqa: S310 + f.write(resp.read()) + + with tarfile.open(tmppath, "r:xz") as tar: + tar.extractall( # noqa: S202 + members=tar_filter(tar.getmembers()), + path=output_dir, + filter="fully_trusted", + ) + else: + print(f"No {library_name} in version {version}") + + +def update_config(version: str, config_path: Path) -> None: + # This is pretty brittle, but will be better when/if we move all the config to YAML + + out = [] + in_version_section = False + with config_path.open() as f: + for line in f: + if line.strip() == "'versions': [": + in_version_section = True + if in_version_section and line.strip() == "],": + out.append(f" ('{version}', ),\n") + in_version_section = False + out.append(line) + + with config_path.open("w") as f: + f.write("".join(out)) + + +def run_cybind(cybind_repo: Path, cuda_python_repo: Path, libraries: list[str]) -> None: + with tempfile.TemporaryDirectory() as tempdir: + tempdir_path = Path(tempdir) + + venv.create(tempdir_path, with_pip=True) + subprocess.check_call( # noqa: S603 + [ + str(tempdir_path / "bin" / "python"), + "-m", + "pip", + "install", + str(cybind_repo), + ] + ) + try: + subprocess.check_call( # noqa: S603 + [ + str(tempdir_path / "bin" / "python"), + "-m", + "cybind", + "--generate", + *libraries, + "--output-dir", + str(cuda_python_repo / "cuda_bindings"), + ] + ) + except subprocess.CalledProcessError: + print("Error running cybind.") + print("This probably indicates an issue introduced with the new headers.") + print("If necessary, you can edit the headers and re-run this script.") + return 1 + + +def update_version_file(version: str, version_path: Path, is_prev: bool) -> str: + if is_prev: + key = "prev_build" + else: + key = "build" + + with version_path.open() as f: + content = json.load(f) + existing_version = content["cuda"][key]["version"] + content["cuda"][key]["version"] = version + + with version_path.open("w") as f: + content = json.dump(content, f, indent=2) + # json.dump doesn't add a trailing newline + f.write("\n") + + return existing_version + + +def update_matrix(existing_version: str, new_version: str, matrix_path: Path) -> None: + # It would be less brittle to update using JSON here, but that messes up the formatting + + with matrix_path.open() as f: + content = f.read() + + content = re.sub(rf'"CUDA_VER": "{existing_version}"', f'"CUDA_VER": "{new_version}"', content) + + with matrix_path.open("w") as f: + f.write(content) + + +def main(version: str, cuda_python_repo: Path, cybind_repo: Path, is_prev: bool): + cybind_headers_path = cybind_repo / "assets" / "headers" + cybind_config_path = cybind_repo / "assets" / "configs" + + for libname, distname, subdir in CYBIND_GENERATED_LIBRARIES: + fetch_headers(version, distname, cybind_headers_path / subdir) + update_config(version, cybind_config_path / f"config_{libname}.py") + + existing_version = update_version_file(version, cuda_python_repo / "ci" / "versions.json", is_prev) + update_matrix(existing_version, version, cuda_python_repo / "ci" / "test-matrix.json") + + # Do this last, because, if anything, it's the thing that's likely to fail + if run_cybind(cybind_repo, cuda_python_repo, [x[0] for x in CYBIND_GENERATED_LIBRARIES]): + sys.exit(1) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Update cuda-python for a new version of the CTK") + parser.add_argument( + "--cybind-repo", + type=Path, + help="Path to a checkout of cybind (default: ../cybind relative to cuda-python)", + ) + parser.add_argument( + "--is-prev", + action="store_true", + help="When given, update the previous, not latest version", + ) + parser.add_argument( + "version", + type=str, + help="Version to move to", + ) + args = parser.parse_args() + + cuda_python_repo = Path(__file__).parents[1] + + if args.cybind_repo is None: + args.cybind_repo = cuda_python_repo.parent / "cybind" + + print("Before running this script, you need to:") + print(" - Create a new branch in this repo based on upstream/main") + print(f" - Create a new branch in a cybind checkout at {args.cybind_repo} based on upstream/main") + print() + print(f"This will add CTK {args.version} as the {'previous' if args.is_prev else 'latest'} version.") + print("Proceed? [y/N]") + resp = input().strip().lower() + if resp != "y": + print("Aborting") + + main(args.version, cuda_python_repo, args.cybind_repo, args.is_prev) + + print("Remaining manual steps:") + print("- Add a changelog entry:") + print( + f"* Updated the ``cuda.bindings.runtime`` module to statically link " + f"against the CUDA Runtime library from CUDA Toolkit {args.version}." + ) + print("- Inspect the changes to this repo and cybind, commit and submit PRs.") From 7c8cd522bb526ed219f229a14b018fe049ad295e Mon Sep 17 00:00:00 2001 From: Michael Droettboom Date: Tue, 14 Oct 2025 10:16:26 -0400 Subject: [PATCH 2/3] Address copilot's suggestions --- toolshed/update_ctk.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/toolshed/update_ctk.py b/toolshed/update_ctk.py index f6c1737bcf..382fb09cf2 100644 --- a/toolshed/update_ctk.py +++ b/toolshed/update_ctk.py @@ -84,7 +84,7 @@ def update_config(version: str, config_path: Path) -> None: f.write("".join(out)) -def run_cybind(cybind_repo: Path, cuda_python_repo: Path, libraries: list[str]) -> None: +def run_cybind(cybind_repo: Path, cuda_python_repo: Path, libraries: list[str]) -> int: with tempfile.TemporaryDirectory() as tempdir: tempdir_path = Path(tempdir) @@ -116,6 +116,8 @@ def run_cybind(cybind_repo: Path, cuda_python_repo: Path, libraries: list[str]) print("If necessary, you can edit the headers and re-run this script.") return 1 + return 0 + def update_version_file(version: str, version_path: Path, is_prev: bool) -> str: if is_prev: @@ -129,7 +131,7 @@ def update_version_file(version: str, version_path: Path, is_prev: bool) -> str: content["cuda"][key]["version"] = version with version_path.open("w") as f: - content = json.dump(content, f, indent=2) + json.dump(content, f, indent=2) # json.dump doesn't add a trailing newline f.write("\n") @@ -197,6 +199,7 @@ def main(version: str, cuda_python_repo: Path, cybind_repo: Path, is_prev: bool) resp = input().strip().lower() if resp != "y": print("Aborting") + sys.exit(0) main(args.version, cuda_python_repo, args.cybind_repo, args.is_prev) From ab05ab4953c9b3dc484afff7263f7b9c1e604ed8 Mon Sep 17 00:00:00 2001 From: Michael Droettboom Date: Tue, 14 Oct 2025 12:05:08 -0400 Subject: [PATCH 3/3] Include cython-gen in this workflow --- toolshed/update_ctk.py | 71 +++++++++++++++++++++++++++++++++++++----- 1 file changed, 64 insertions(+), 7 deletions(-) diff --git a/toolshed/update_ctk.py b/toolshed/update_ctk.py index 382fb09cf2..27f47b77e4 100644 --- a/toolshed/update_ctk.py +++ b/toolshed/update_ctk.py @@ -150,17 +150,59 @@ def update_matrix(existing_version: str, new_version: str, matrix_path: Path) -> f.write(content) -def main(version: str, cuda_python_repo: Path, cybind_repo: Path, is_prev: bool): +def regenerate_cython_gen(cuda_python_private_repo: Path, cuda_python_repo: Path) -> int: + with tempfile.TemporaryDirectory() as tempdir: + tempdir_path = Path(tempdir) + + venv.create(tempdir_path, with_pip=True) + subprocess.check_call( # noqa: S603 + [ + str(tempdir_path / "bin" / "python"), + "-m", + "pip", + "install", + "-r", + str(cuda_python_private_repo / "requirements.txt"), + ] + ) + try: + subprocess.check_call( # noqa: S603 + [ + str(tempdir_path / "bin" / "python"), + str(cuda_python_private_repo / "regenerate.py"), + "-t", + "driver", + "-t", + "runtime", + "-t", + "nvrtc", + "-o", + str(cuda_python_repo), + ], + cwd=cuda_python_private_repo, + ) + except subprocess.CalledProcessError: + print("Error running cython-gen.") + print("This probably indicates an issue introduced with the new headers.") + return 1 + + return 0 + + +def main(version: str, cuda_python_repo: Path, cybind_repo: Path, cuda_python_private_repo: Path, is_prev: bool): cybind_headers_path = cybind_repo / "assets" / "headers" cybind_config_path = cybind_repo / "assets" / "configs" + existing_version = update_version_file(version, cuda_python_repo / "ci" / "versions.json", is_prev) + update_matrix(existing_version, version, cuda_python_repo / "ci" / "test-matrix.json") + + if regenerate_cython_gen(cuda_python_private_repo, cuda_python_repo): + sys.exit(1) + for libname, distname, subdir in CYBIND_GENERATED_LIBRARIES: fetch_headers(version, distname, cybind_headers_path / subdir) update_config(version, cybind_config_path / f"config_{libname}.py") - existing_version = update_version_file(version, cuda_python_repo / "ci" / "versions.json", is_prev) - update_matrix(existing_version, version, cuda_python_repo / "ci" / "test-matrix.json") - # Do this last, because, if anything, it's the thing that's likely to fail if run_cybind(cybind_repo, cuda_python_repo, [x[0] for x in CYBIND_GENERATED_LIBRARIES]): sys.exit(1) @@ -173,6 +215,11 @@ def main(version: str, cuda_python_repo: Path, cybind_repo: Path, is_prev: bool) type=Path, help="Path to a checkout of cybind (default: ../cybind relative to cuda-python)", ) + parser.add_argument( + "--cuda-python-private-repo", + type=Path, + help="Path to a checkout of cuda-python-private (default: ../cuda-python-private relative to cuda-python)", + ) parser.add_argument( "--is-prev", action="store_true", @@ -190,9 +237,16 @@ def main(version: str, cuda_python_repo: Path, cybind_repo: Path, is_prev: bool) if args.cybind_repo is None: args.cybind_repo = cuda_python_repo.parent / "cybind" + if args.cuda_python_private_repo is None: + args.cuda_python_private_repo = cuda_python_repo.parent / "cuda-python-private" + print("Before running this script, you need to:") print(" - Create a new branch in this repo based on upstream/main") print(f" - Create a new branch in a cybind checkout at {args.cybind_repo} based on upstream/main") + print( + f" - Create a new branch in a cuda-python-private checkout at {cuda_python_repo} based on upstream/cython-gen" + ) + print(" - Install the version of CTK you are adding and make sure CUDA_HOME is pointing to it.") print() print(f"This will add CTK {args.version} as the {'previous' if args.is_prev else 'latest'} version.") print("Proceed? [y/N]") @@ -201,12 +255,15 @@ def main(version: str, cuda_python_repo: Path, cybind_repo: Path, is_prev: bool) print("Aborting") sys.exit(0) - main(args.version, cuda_python_repo, args.cybind_repo, args.is_prev) + main(args.version, cuda_python_repo, args.cybind_repo, args.cuda_python_private_repo, args.is_prev) print("Remaining manual steps:") - print("- Add a changelog entry:") + print("- Add a changelog entry, for example:") + print() print( f"* Updated the ``cuda.bindings.runtime`` module to statically link " f"against the CUDA Runtime library from CUDA Toolkit {args.version}." ) - print("- Inspect the changes to this repo and cybind, commit and submit PRs.") + print() + print("- Inspect the changes to this repo, cuda-python-private and cybind, ") + print(" commit and submit PRs.")