Use rapids-dependency-file-generator (#17)

Rather than modifying the dependencies, we should delegate the handling of dependencies to DFG. Call out to DFG's newly-stable Python API to modify `pyproject.toml`.
rapidsai · May 2, 2024 · cdebf61 · cdebf61
1 parent 8cfa39b
commit cdebf61
Show file tree

Hide file tree

Showing 12 changed files with 446 additions and 212 deletions.
diff --git a/README.md b/README.md
@@ -4,11 +4,9 @@
 It currently support `scikit-build-core` and `setuptools` as the wrapped builder.
 The package's primary purpose is to automate the various bits of preprocessing that are typically done to RAPIDS package metadata prior to publishing packages.
 This includes the following notable changes:
+- Running [`rapids-dependency-file-generator`](https://github.com/rapidsai/dependency-file-generator) to get the dependencies for the CUDA version and architecture.
 - Modifying the package name to include CUDA suffixes.
 - Updating the git commit embedded in the importable package.
-- Modifying the package's (build, runtime, etc) dependencies to include CUDA suffixes.
-- Filtering out dependencies based on the CUDA version at build time.
-- Updating dependency specifiers to include an alpha specifier to allow pulling nightly dependencies in nightly builds.
 
 Since some of these modifications are only desirable in certain scenarios (wheel vs conda builds vs editable installs), all of these functions are customizable via the project's configuration in pyproject.toml.
 In cases where more dynamic customization is sensible, suitable environment variables and `config_settings` are supported during builds of distributions.
@@ -17,14 +15,15 @@ In cases where more dynamic customization is sensible, suitable environment vari
 
 Any option without a default is required.
 
-| Option                | Definition                                                          | Type      | Default      | Supports dynamic modification |
-|-----------------------|---------------------------------------------------------------------|-----------|--------------|-------------------------------|
-| `allow-nightly-deps`  | If true, append alpha specifiers to dependencies                    | bool      | true         | Y                             |
-| `build-backend`       | The wrapped build backend (e.g. `setuptools.build_meta`)            | string    |              | N                             |
-| `commit-file`         | The file in which to write the git commit hash                      | string    | "" (No file) | N                             |
-| `disable-cuda-suffix` | If true, don't try to write CUDA suffixes                           | bool      | false        | Y                             |
-| `require-cuda`        | If false, builds will succeed even if nvcc is not available         | bool      | true         | Y                             |
-| `requires`            | List of build requirements (in addition to `build-system.requires`) | list[str] | []           | N                             |
+| Option                | Definition                                                          | Type           | Default             | Supports dynamic modification |
+|-----------------------|---------------------------------------------------------------------|----------------|---------------------|-------------------------------|
+| `build-backend`       | The wrapped build backend (e.g. `setuptools.build_meta`)            | string         |                     | N                             |
+| `commit-file`         | The file in which to write the git commit hash                      | string         | "" (No file)        | N                             |
+| `dependencies-file`   | The path to the `dependencies.yaml` file to use                     | string         | "dependencies.yaml" | Y                             |
+| `disable-cuda-suffix` | If true, don't try to write CUDA suffixes                           | bool           | false               | Y                             |
+| `matrix-entry`        | A `;`-separated list of `=`-delimited key/value pairs               | string         | ""                  | Y                             |
+| `require-cuda`        | If false, builds will succeed even if nvcc is not available         | bool           | true                | Y                             |
+| `requires`            | List of build requirements (in addition to `build-system.requires`) | list[str]      | []                  | N                             |
 
 
 ## Outstanding questions

diff --git a/dependencies.yaml b/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:

diff --git a/pyproject.toml b/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`.

diff --git a/rapids_build_backend/config.py b/rapids_build_backend/config.py
@@ -12,12 +12,13 @@ class Config:
     # required) and whether it may be overridden by an environment variable or a config
     # setting.
     config_options = {
-        "allow-nightly-deps": (True, True),
         "build-backend": (None, False),
         "commit-file": ("", False),
+        "dependencies-file": ("dependencies.yaml", True),
         "disable-cuda-suffix": (False, True),
+        "matrix-entry": ("", True),
         "require-cuda": (True, True),
-        "requires": ([], False),
+        "requires": (lambda: [], False),
     }
 
     def __init__(self, dirname=".", config_settings=None):
@@ -32,6 +33,8 @@ def __getattr__(self, name):
         config_name = name.replace("_", "-")
         if config_name in Config.config_options:
             default_value, allows_override = Config.config_options[config_name]
+            if callable(default_value):
+                default_value = default_value()
 
             # If overrides are allowed environment variables take precedence over the
             # config_settings dict.

diff --git a/rapids_build_backend/impls.py b/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,50 @@ 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)
+
+    try:
+        parsed_config = rapids_dependency_file_generator.load_config_from_file(
+            config.dependencies_file
+        )
+    except FileNotFoundError:
+        parsed_config = None
+
     try:
-        shutil.move(pyproject_file, bkp_pyproject_file)
+        shutil.copyfile(pyproject_file, bkp_pyproject_file)
+        if parsed_config:
+            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)
+                    and os.path.samefile(pyproject_dir, ".")
+                ):
+                    continue
+                matrix = _parse_matrix(config.matrix_entry) 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 +227,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 +243,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 +259,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),

diff --git a/rapids_build_backend/utils.py b/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: