diff --git a/dependencies.yaml b/dependencies.yaml index 2976281..bf770a0 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -74,7 +74,9 @@ dependencies: common: - output_types: [conda, requirements, pyproject] packages: + - PyYAML - packaging + - rapids-dependency-file-generator<2.0.dev0 - tomli - tomli-w test: diff --git a/pyproject.toml b/pyproject.toml index 2cab352..11c1717 100644 --- 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<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 index 3a521e9..238626c 100644 --- a/rapids_build_backend/config.py +++ b/rapids_build_backend/config.py @@ -15,6 +15,7 @@ class Config: "allow-nightly-deps": (True, True), "build-backend": (None, False), "commit-file": ("", False), + "dependencies-file": ("dependencies.yaml", True), "disable-cuda-suffix": (False, True), "require-cuda": (True, True), "requires": ([], False), diff --git a/rapids_build_backend/impls.py b/rapids_build_backend/impls.py index 4b190c5..c4fe037 100644 --- a/rapids_build_backend/impls.py +++ b/rapids_build_backend/impls.py @@ -9,11 +9,14 @@ from importlib import import_module import tomli_w -from packaging.requirements import Requirement -from packaging.specifiers import SpecifierSet +import yaml +from rapids_dependency_file_generator.rapids_dependency_file_generator import ( + get_requested_output_types, + make_dependency_files, +) +from . import utils from .config import Config -from .utils import _get_pyproject @lru_cache @@ -31,7 +34,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 +63,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.group(1), match.group(2) except Exception: if not require_cuda: return None @@ -86,126 +89,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 +157,39 @@ 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) + + with open(config.dependencies_file) as f: + parsed_config = yaml.load(f, Loader=yaml.FullLoader) + files = {} + for file_key, file_config in parsed_config["files"].items(): + if "pyproject" not in get_requested_output_types(file_config["output"]): + continue + pyproject_dir = os.path.join( + os.path.dirname(config.dependencies_file), + file_config.get("pyproject_dir", "."), + ) + if not os.path.exists(pyproject_dir): + continue + if not os.path.samefile(pyproject_dir, "."): + continue + file_config["output"] = ["pyproject"] + if cuda_version is not None: + file_config.setdefault("matrix", {})["cuda"] = [ + f"{cuda_version[0]}.{cuda_version[1]}" + ] + files[file_key] = file_config + parsed_config["files"] = files + try: - shutil.move(pyproject_file, bkp_pyproject_file) + shutil.copyfile(pyproject_file, bkp_pyproject_file) + make_dependency_files(parsed_config, config.dependencies_file, 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 +212,7 @@ 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) + requires = [] if hasattr( backend := _get_backend(config.build_backend), @@ -328,7 +226,7 @@ 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) + requires = [] if hasattr( backend := _get_backend(config.build_backend), @@ -342,7 +240,7 @@ 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) + requires = [] if hasattr( backend := _get_backend(config.build_backend), diff --git a/tests/conftest.py b/tests/conftest.py index 0b1e032..035bd05 100644 --- a/tests/conftest.py +++ b/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()[0] != nvcc_version: nvcc = _create_nvcc(nvcc_version) os.environ["PATH"] = os.pathsep.join( [os.path.dirname(nvcc), os.environ["PATH"]] diff --git a/tests/test_impls.py b/tests/test_impls.py index 0f4d223..c6e83b9 100644 --- a/tests/test_impls.py +++ b/tests/test_impls.py @@ -2,15 +2,27 @@ import os.path import tempfile +from contextlib import contextmanager from unittest.mock import Mock, patch import pytest from rapids_build_backend.impls import ( - _VERSIONED_RAPIDS_WHEELS, - _add_cuda_suffix, _edit_git_commit, + _edit_pyproject, + _get_cuda_suffix, ) +from rapids_build_backend.utils import _get_pyproject + + +@contextmanager +def set_cwd(cwd): + old_cwd = os.getcwd() + os.chdir(cwd) + try: + yield + finally: + os.chdir(old_cwd) @pytest.mark.parametrize( @@ -50,18 +62,120 @@ def check_initial_contents(filename): @pytest.mark.parametrize( - ["initial_req", "cuda_suffix", "cuda_major", "expected_req"], + ["cuda_version", "cuda_suffix", "cuda_python_requirement"], [ - ("cupy", "", "11", "cupy-cuda11x"), - ("cupy", "", "12", "cupy-cuda12x"), - ("cupy", "", None, "cupy"), - *((name, "-cu11", "11", f"{name}-cu11") for name in _VERSIONED_RAPIDS_WHEELS), - *((name, "-cu12", "12", f"{name}-cu12") for name in _VERSIONED_RAPIDS_WHEELS), - *((name, "", None, name) for name in _VERSIONED_RAPIDS_WHEELS), - ("poetry", "-cu11", "11", "poetry"), - ("poetry", "-cu12", "12", "poetry"), - ("poetry", "", None, "poetry"), + (("11", "5"), "-cu11", "cuda-python>=11.5,<11.6.dev0"), + (("12", "1"), "-cu12", "cuda-python>=12.1,<12.2.dev0"), ], ) -def tests_add_cuda_suffix(initial_req, cuda_suffix, cuda_major, expected_req): - assert _add_cuda_suffix(initial_req, cuda_suffix, cuda_major) == expected_req +def test_edit_pyproject(cuda_version, cuda_suffix, cuda_python_requirement): + with tempfile.TemporaryDirectory() as d: + original_contents = """[project] +name = "test-project" +dependencies = [] + +[build-system] +requires = [] +""" + + with set_cwd(d): + with open("pyproject.toml", "w") as f: + f.write(original_contents) + + with open("dependencies.yaml", "w") as f: + f.write(""" +files: + project: + output: pyproject + includes: + - project + pyproject_dir: . + matrix: + cuda: ["11.5", "12.1"] + arch: ["x86_64"] + extras: + table: project + build_system: + output: pyproject + includes: + - build_system + pyproject_dir: . + extras: + table: build-system + other_project: + output: pyproject + includes: + - bad + pyproject_dir: python + extras: + table: project + conda: + output: conda + includes: + - bad +dependencies: + project: + common: + - output_types: [pyproject] + packages: + - tomli + specific: + - output_types: [pyproject] + matrices: + - matrix: + arch: "x86_64" + cuda: "11.5" + packages: + - cuda-python>=11.5,<11.6.dev0 + - matrix: + arch: "x86_64" + cuda: "12.1" + packages: + - cuda-python>=12.1,<12.2.dev0 + build_system: + common: + - output_types: [pyproject] + packages: + - scikit-build-core + bad: + common: + - output_types: [pyproject, conda] + packages: + - bad-package +""") + config = Mock( + require_cuda=False, + dependencies_file="dependencies.yaml", + ) + + with patch( + "rapids_build_backend.impls._get_cuda_version", + Mock(return_value=cuda_version), + ), patch( + "rapids_build_backend.impls._get_cuda_suffix", + _get_cuda_suffix.__wrapped__, + ), patch( + "rapids_build_backend.utils._get_pyproject", _get_pyproject.__wrapped__ + ): + with _edit_pyproject(config): + with open("pyproject.toml") as f: + assert ( + f.read() + == f"""[project] +name = "test-project{cuda_suffix}" +dependencies = [ + "{cuda_python_requirement}", + "tomli", +] + +[build-system] +requires = [ + "scikit-build-core", +] +""" + ) + with open(".pyproject.toml.rapids-build-backend.bak") as f: + assert f.read() == original_contents + + with open("pyproject.toml") as f: + assert f.read() == original_contents