Add a script to help update to a new CTK version

toolshed/update_ctk.py

@@ -0,0 +1,269 @@
+#!/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]) -> 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",
+                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
+
+    return 0
+
+
+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:
+        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 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")
+
+    # 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(
+        "--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",
+        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"
+
+    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]")
+    resp = input().strip().lower()
+    if resp != "y":
+        print("Aborting")
+        sys.exit(0)
+
+    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, 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()
+    print("- Inspect the changes to this repo, cuda-python-private and cybind, ")
+    print("  commit and submit PRs.")