Check generated JSON into source control (#13)

This will allow projects to download the generated JSON directly from
GitHub.

.pre-commit-config.yaml

@@ -5,7 +5,13 @@ repos:
     rev: v4.6.0
     hooks:
       - id: trailing-whitespace
+        exclude: |
+          (?x)
+              rapids-metadata[.]json$
       - id: end-of-file-fixer
+        exclude: |
+          (?x)
+              rapids-metadata[.]json$
   - repo: https://github.com/rapidsai/dependency-file-generator
     rev: v1.13.11
     hooks:
@@ -37,6 +43,13 @@ repos:
         pass_filenames: false
         additional_dependencies:
           - packaging
+  - repo: local
+    hooks:
+      - id: generate-json
+        name: generate-json
+        entry: ./ci/generate_json.py
+        language: python
+        pass_filenames: false
 
 default_language_version:
       python: python3

ci/generate_json.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+import os.path
+import sys
+
+repo_root = os.path.join(os.path.dirname(__file__), "..")
+sys.path.append(os.path.join(repo_root, "src"))
+
+from rapids_metadata import json as rapids_json  # noqa: E402
+
+if __name__ == "__main__":
+    rapids_json.main(
+        [
+            "--output",
+            os.path.join(repo_root, "rapids-metadata.json"),
+            "--pretty",
+            "--all-versions",
+        ]
+    )

pyproject.toml

@@ -26,6 +26,9 @@ classifiers = [
     "Programming Language :: Python :: 3",
 ]
 requires-python = ">=3.9"
+dependencies = [
+    "packaging",
+]
 
 [project.scripts]
 rapids-metadata-json = "rapids_metadata.json:main"

rapids-metadata.json

@@ -0,0 +1,216 @@
+{
+  "versions": {
+    "24.08": {
+      "repositories": {
+        "_nvidia": {
+          "packages": {
+            "cubinlinker": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "cucim": {
+          "packages": {
+            "cucim": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "cudf": {
+          "packages": {
+            "cudf": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            },
+            "dask-cudf": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "cugraph": {
+          "packages": {
+            "cugraph": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            },
+            "cugraph-dgl": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            },
+            "cugraph-equivariant": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            },
+            "cugraph-pyg": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            },
+            "nx-cugraph": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            },
+            "pylibcugraph": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "cugraph-ops": {
+          "packages": {
+            "pylibcugraphops": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "cuml": {
+          "packages": {
+            "cuml": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            },
+            "libcuml": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            },
+            "libcuml-tests": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "cumlprims_mg": {
+          "packages": {
+            "libcumlprims": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "cuproj": {
+          "packages": {
+            "cuproj": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "cuspatial": {
+          "packages": {
+            "cuspatial": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "cuxfilter": {
+          "packages": {
+            "cuxfilter": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "dask-cuda": {
+          "packages": {
+            "dask-cuda": {
+              "has_cuda_suffix": false,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "ptxcompiler": {
+          "packages": {
+            "ptxcompiler": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "pynvjitlink": {
+          "packages": {
+            "pynvjitlink": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "raft": {
+          "packages": {
+            "libraft": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            },
+            "libraft-headers": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            },
+            "pylibraft": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            },
+            "raft-dask": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "rapids-dask-dependency": {
+          "packages": {
+            "rapids-dask-dependency": {
+              "has_cuda_suffix": false,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "rmm": {
+          "packages": {
+            "librmm": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            },
+            "rmm": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "ucx-py": {
+          "packages": {
+            "ucx-py": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "ucxx": {
+          "packages": {
+            "distributed-ucxx": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            },
+            "libucxx": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            },
+            "ucxx": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            }
+          }
+        },
+        "wholegraph": {
+          "packages": {
+            "pylibwholegraph": {
+              "has_cuda_suffix": true,
+              "publishes_prereleases": true
+            }
+          }
+        }
+      }
+    }
+  }
+}

src/rapids_metadata/json.py

@@ -41,11 +41,28 @@ def default(
         return dataclasses.asdict(o)
 
 
-def main():
+def main(argv: Union[list[str], None] = None):
+    if argv is None:
+        argv = sys.argv[1:]
+
     parser = argparse.ArgumentParser()
-    parser.add_argument("--all-versions", action="store_true")
+    parser.description = "Output RAPIDS metadata as a JSON document."
+    parser.add_argument(
+        "--all-versions",
+        action="store_true",
+        help="Output all versions, ignoring local VERSION file",
+    )
+    parser.add_argument(
+        "--pretty", action="store_true", help="Pretty-print JSON output"
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        metavar="<output file>",
+        help="Write to a file instead of stdout",
+    )
 
-    parsed = parser.parse_args()
+    parsed = parser.parse_args(argv)
     metadata = (
         all_metadata
         if parsed.all_versions
@@ -57,7 +74,24 @@ def main():
             }
         )
     )
-    json.dump(metadata, sys.stdout, cls=_RAPIDSMetadataEncoder)
+
+    def write_file(f):
+        json.dump(
+            metadata,
+            f,
+            cls=_RAPIDSMetadataEncoder,
+            sort_keys=True,
+            separators=(",", ": ") if parsed.pretty else (",", ":"),
+            indent="  " if parsed.pretty else None,
+        )
+        if parsed.pretty:
+            f.write("\n")
+
+    if parsed.output:
+        with open(parsed.output, "w") as f:
+            write_file(f)
+    else:
+        write_file(sys.stdout)
 
 
 if __name__ == "__main__":

src/rapids_metadata/metadata.py

@@ -14,8 +14,8 @@
 
 from dataclasses import dataclass, field
 from os import PathLike
+from typing import Union
 
-from packaging.version import Version
 from .rapids_version import get_rapids_version
 
 __all__ = [
@@ -72,7 +72,11 @@ def cuda_suffixed_packages(self) -> set[str]:
 class RAPIDSMetadata:
     versions: dict[str, RAPIDSVersion] = field(default_factory=dict)
 
-    def get_current_version(self, directory: PathLike) -> RAPIDSVersion:
+    def get_current_version(
+        self, directory: Union[str, PathLike[str]]
+    ) -> RAPIDSVersion:
+        from packaging.version import Version
+
         current_version = get_rapids_version(directory)
         try:
             return self.versions[current_version]

src/rapids_metadata/rapids_version.py

@@ -14,14 +14,15 @@
 
 import os.path
 from os import PathLike
-
-from packaging.version import InvalidVersion, Version
+from typing import Union
 
 
 __all__ = ["get_rapids_version"]
 
 
-def get_rapids_version(directory: PathLike) -> str:
+def get_rapids_version(directory: Union[str, PathLike[str]]) -> str:
+    from packaging.version import InvalidVersion, Version
+
     while not os.path.samefile(directory, os.path.dirname(directory)):
         try:
             with open(os.path.join(directory, "VERSION")) as f: