Use rapids-dependency-file-generator

dependencies.yaml

@@ -74,7 +74,9 @@ dependencies:
     common:
       - output_types: [conda, requirements, pyproject]
         packages:
+          - PyYAML
           - packaging
+          - rapids-dependency-file-generator>=1.13.3,<2.0.dev0
           - tomli
           - tomli-w
   test:

pyproject.toml

@@ -12,7 +12,9 @@ name = "rapids-build-backend"
 version = "0.0.1"
 description = "Custom PEP517 builder for RAPIDS"
 dependencies = [
+    "PyYAML",
     "packaging",
+    "rapids-dependency-file-generator>=1.13.3,<2.0.dev0",
     "tomli",
     "tomli-w",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit dependencies.yaml and run `rapids-dependency-file-generator`.

rapids_build_backend/config.py

@@ -15,7 +15,9 @@ class Config:
         "allow-nightly-deps": (True, True),
         "build-backend": (None, False),
         "commit-file": ("", False),
+        "dependencies-file": ("dependencies.yaml", True),
         "disable-cuda-suffix": (False, True),
+        "matrix": ("", True),
         "require-cuda": (True, True),
         "requires": ([], False),
     }

rapids_build_backend/impls.py

@@ -8,12 +8,19 @@
 from functools import lru_cache
 from importlib import import_module
 
+import rapids_dependency_file_generator
 import tomli_w
-from packaging.requirements import Requirement
-from packaging.specifiers import SpecifierSet
 
+from . import utils
 from .config import Config
-from .utils import _get_pyproject
+
+
+def _parse_matrix(matrix):
+    if not matrix:
+        return None
+    return {
+        key: [value] for key, value in (item.split("=") for item in matrix.split(";"))
+    }
 
 
 @lru_cache
@@ -31,7 +38,7 @@ def _get_backend(build_backend):
 
 
 @lru_cache
-def _get_cuda_major(require_cuda=False):
+def _get_cuda_version(require_cuda=False):
     """Get the CUDA suffix based on nvcc.
 
     Parameters
@@ -60,10 +67,10 @@ def _get_cuda_major(require_cuda=False):
 
         output_lines = process_output.stdout.decode().splitlines()
 
-        match = re.search(r"release (\d+)", output_lines[3])
+        match = re.search(r"release (\d+)\.(\d+)", output_lines[3])
         if match is None:
             raise ValueError("Failed to parse CUDA version from nvcc output.")
-        return match.group(1)
+        return match.groups()
     except Exception:
         if not require_cuda:
             return None
@@ -86,126 +93,9 @@ def _get_cuda_suffix(require_cuda=False):
         The CUDA suffix (e.g., "-cu11") or an empty string if CUDA could not be
         detected.
     """
-    if (major := _get_cuda_major(require_cuda)) is None:
+    if (version := _get_cuda_version(require_cuda)) is None:
         return ""
-    return f"-cu{major}"
-
-
-# Wheels with a CUDA suffix.
-_VERSIONED_RAPIDS_WHEELS = [
-    "rmm",
-    "pylibcugraphops",
-    "pylibcugraph",
-    "nx-cugraph",
-    "dask-cudf",
-    "cuspatial",
-    "cuproj",
-    "cuml",
-    "cugraph",
-    "cudf",
-    "ptxcompiler",
-    "cubinlinker",
-    "cugraph-dgl",
-    "cugraph-pyg",
-    "cugraph-equivariant",
-    "raft-dask",
-    "pylibwholegraph",
-    "pylibraft",
-    "cuxfilter",
-    "cucim",
-    "ucx-py",
-    "ucxx",
-    "pynvjitlink",
-    "distributed-ucxx",
-]
-
-# Wheels without a CUDA suffix.
-_UNVERSIONED_RAPIDS_WHEELS = [
-    "dask-cuda",
-    "rapids-dask-dependency",
-]
-
-# Wheels that don't release regular alpha versions
-_CUDA_11_ONLY_WHEELS = (
-    "ptxcompiler",
-    "cubinlinker",
-)
-
-
-def _add_cuda_suffix(req, cuda_suffix, cuda_major):
-    req = Requirement(req)
-    if req.name == "cupy" and cuda_major is not None:
-        req.name += f"-cuda{cuda_major}x"
-    elif req.name in _VERSIONED_RAPIDS_WHEELS:
-        req.name += cuda_suffix
-
-    return str(req)
-
-
-def _add_alpha_specifier(req):
-    req = Requirement(req)
-    if (
-        req.name in _VERSIONED_RAPIDS_WHEELS or req.name in _UNVERSIONED_RAPIDS_WHEELS
-    ) and req.name not in _CUDA_11_ONLY_WHEELS:
-        req.specifier &= SpecifierSet(">=0.0.0a0")
-    return str(req)
-
-
-def _process_dependencies(config, dependencies=None):
-    """Add the CUDA suffix to any versioned RAPIDS wheels in dependencies.
-
-    If dependencies is None, then config.requires is used.
-
-    Parameters
-    ----------
-    config : Config
-        The project's configuration.
-    dependencies : list of str, optional
-        The dependencies to suffix. If None, then config.requires is used.
-
-    Returns
-    -------
-    list of str
-        The dependencies with the CUDA suffix added to any versioned RAPIDS wheels.
-    """
-    # Note that this implementation is currently suboptimal, in each step to allow the
-    # steps to be more freely composable based on options. We could optimize by using a
-    # single loop with multiple nested conditionals, but that would make the logic
-    # harder to understand and modify. The performance of this code should be negligible
-    # in the context of a build anyway.
-    dependencies = dependencies or config.requires
-
-    # Step 1: Filter out CUDA 11-only wheels if we're not in a CUDA 11 build. Skip this
-    # step if if we were unable to detect a CUDA version.
-    major = _get_cuda_major(config.require_cuda)
-    if major is not None and major != "11":
-        dependencies = filter(
-            lambda dep: dep not in _CUDA_11_ONLY_WHEELS,
-            dependencies,
-        )
-
-    # Step 2: Allow nightlies of RAPIDS packages except in release builds. Do this
-    # before suffixing the names so that lookups in _add_alpha_specifier are accurate.
-    if config.allow_nightly_deps:
-        dependencies = map(
-            _add_alpha_specifier,
-            dependencies,
-        )
-
-    # Step 3: Add the CUDA suffix to any versioned RAPIDS wheels. This step may be
-    # explicitly skipped by setting the disable_cuda_suffix option to True, or
-    # implicitly skipped if we were unable to detect a CUDA version and require_cuda was
-    # False.
-    if not config.disable_cuda_suffix:
-        suffix = _get_cuda_suffix(config.require_cuda)
-        # If we can't determine the local CUDA version then we can skip this step
-        if suffix:
-            dependencies = map(
-                lambda dep: _add_cuda_suffix(dep, suffix, major),
-                dependencies,
-            )
-
-    return list(dependencies)
+    return f"-cu{version[0]}"
 
 
 @lru_cache
@@ -271,27 +161,45 @@ def _edit_pyproject(config):
     being built. This is useful for projects that want to build wheels
     with a different name than the package name.
     """
-    pyproject = _get_pyproject()
-    project_data = pyproject["project"]
-    project_data["name"] += _get_cuda_suffix(config.require_cuda)
-
-    dependencies = pyproject["project"].get("dependencies")
-    if dependencies is not None:
-        project_data["dependencies"] = _process_dependencies(
-            config, project_data["dependencies"]
-        )
-
-    optional_dependencies = pyproject["project"].get("optional-dependencies")
-    if optional_dependencies is not None:
-        project_data["optional-dependencies"] = {
-            extra: _process_dependencies(config, deps)
-            for extra, deps in optional_dependencies.items()
-        }
-
     pyproject_file = "pyproject.toml"
     bkp_pyproject_file = ".pyproject.toml.rapids-build-backend.bak"
+
+    cuda_version = _get_cuda_version(config.require_cuda)
+
+    parsed_config = rapids_dependency_file_generator.load_config_from_file(
+        config.dependencies_file
+    )
+
     try:
-        shutil.move(pyproject_file, bkp_pyproject_file)
+        shutil.copyfile(pyproject_file, bkp_pyproject_file)
+        for file_key, file_config in parsed_config.files.items():
+            if (
+                rapids_dependency_file_generator.Output.PYPROJECT
+                not in file_config.output
+            ):
+                continue
+            pyproject_dir = os.path.join(
+                os.path.dirname(config.dependencies_file),
+                file_config.pyproject_dir,
+            )
+            if not os.path.exists(pyproject_dir):
+                continue
+            if not os.path.samefile(pyproject_dir, "."):
+                continue
+            matrix = _parse_matrix(config.matrix) or dict(file_config.matrix)
+            if cuda_version is not None:
+                matrix["cuda"] = [f"{cuda_version[0]}.{cuda_version[1]}"]
+            rapids_dependency_file_generator.make_dependency_files(
+                parsed_config=parsed_config,
+                file_keys=[file_key],
+                output={rapids_dependency_file_generator.Output.PYPROJECT},
+                matrix=matrix,
+                prepend_channels=[],
+                to_stdout=False,
+            )
+        pyproject = utils._get_pyproject()
+        project_data = pyproject["project"]
+        project_data["name"] += _get_cuda_suffix(config.require_cuda)
         with open(pyproject_file, "wb") as f:
             tomli_w.dump(pyproject, f)
         yield
@@ -314,7 +222,9 @@ def _edit_pyproject(config):
 def get_requires_for_build_wheel(config_settings):
     config = Config(config_settings=config_settings)
     with _edit_pyproject(config), _edit_git_commit(config):
-        requires = _process_dependencies(config)
+        # Reload the config for a possibly updated tool.rapids-build-backend.requires
+        reloaded_config = Config(config_settings=config_settings)
+        requires = list(reloaded_config.requires)
 
         if hasattr(
             backend := _get_backend(config.build_backend),
@@ -328,7 +238,9 @@ def get_requires_for_build_wheel(config_settings):
 def get_requires_for_build_sdist(config_settings):
     config = Config(config_settings=config_settings)
     with _edit_pyproject(config), _edit_git_commit(config):
-        requires = _process_dependencies(config)
+        # Reload the config for a possibly updated tool.rapids-build-backend.requires
+        reloaded_config = Config(config_settings=config_settings)
+        requires = list(reloaded_config.requires)
 
         if hasattr(
             backend := _get_backend(config.build_backend),
@@ -342,7 +254,9 @@ def get_requires_for_build_sdist(config_settings):
 def get_requires_for_build_editable(config_settings):
     config = Config(config_settings=config_settings)
     with _edit_pyproject(config):
-        requires = _process_dependencies(config)
+        # Reload the config for a possibly updated tool.rapids-build-backend.requires
+        reloaded_config = Config(config_settings=config_settings)
+        requires = list(reloaded_config.requires)
 
         if hasattr(
             backend := _get_backend(config.build_backend),

rapids_build_backend/utils.py

@@ -1,13 +1,10 @@
 # Copyright (c) 2024, NVIDIA CORPORATION.
 
 import os
-from functools import lru_cache
 
 import tomli
 
 
-# Avoid unnecessary I/O by caching.
-@lru_cache
 def _get_pyproject(dirname="."):
     """Parse and return the pyproject.toml file."""
     with open(os.path.join(dirname, "pyproject.toml"), "rb") as f:

tests/conftest.py

@@ -16,7 +16,7 @@
 from jinja2 import Environment, FileSystemLoader
 from packaging.version import parse as parse_version
 
-from rapids_build_backend.impls import _get_cuda_major
+from rapids_build_backend.impls import _get_cuda_version
 
 DIR = Path(__file__).parent.parent.resolve()
 
@@ -69,7 +69,7 @@ def patch_nvcc_if_needed(nvcc_version):
     try:
         # Only create a patch if one is required. In addition to reducing overhead, this
         # also ensures that we test the real nvcc and don't mask any relevant errors.
-        if _get_cuda_major() != nvcc_version:
+        if _get_cuda_version() is None or _get_cuda_version()[0] != nvcc_version:
             nvcc = _create_nvcc(nvcc_version)
             os.environ["PATH"] = os.pathsep.join(
                 [os.path.dirname(nvcc), os.environ["PATH"]]
@@ -235,9 +235,9 @@ def jinja_environment():
 _TEMPLATE_DEFAULTS = {
     "pyproject.toml": {
         "name": "pkg",
-        "dependencies": [],
+        "dependencies": {},
         "extras": {},
-        "build_requires": [],
+        "build_requires": {},
         "flags": {},
         "build_backend": "setuptools.build_meta",
     }