diff --git a/.github/workflows/build-wheel.yml b/.github/workflows/build-wheel.yml index 68f2801fb60..d93dd834539 100644 --- a/.github/workflows/build-wheel.yml +++ b/.github/workflows/build-wheel.yml @@ -150,6 +150,18 @@ jobs: run: | twine check --strict cuda_pathfinder/*.whl + - name: Resolve cuda.bindings build dependencies + run: | + container_args=() + if [[ "${{ inputs.host-platform }}" == linux* ]]; then + container_args+=(--container-mount /host) + fi + mkdir -p wheel-constraints + python ci/tools/write_wheel_constraints.py \ + --output wheel-constraints/cuda-bindings.txt \ + "${container_args[@]}" \ + --wheel cuda-pathfinder cuda_pathfinder + - name: Upload cuda.pathfinder build artifacts if: ${{ strategy.job-index == 0 && inputs.host-platform == 'linux-64' }} uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 @@ -172,6 +184,8 @@ jobs: output-dir: ${{ env.CUDA_BINDINGS_ARTIFACTS_DIR }} env: CIBW_BUILD: ${{ env.CIBW_BUILD }} + CIBW_BEFORE_BUILD_LINUX: 'python -m pip install --upgrade "pip>=25.3"' + CIBW_BEFORE_BUILD_WINDOWS: 'python -m pip install --upgrade "pip>=25.3" delvewheel' # TODO: remove cpython-prerelease once 3.15 is officially supported # Allow CPython pre-release builds (currently 3.15 / 3.15t). This is a # no-op for stable Python versions because CIBW_BUILD still filters @@ -181,6 +195,8 @@ jobs: CIBW_ENVIRONMENT_LINUX: > CUDA_PATH=/host/${{ env.CUDA_PATH }} CUDA_PYTHON_PARALLEL_LEVEL=${{ env.CUDA_PYTHON_PARALLEL_LEVEL }} + PIP_BUILD_CONSTRAINT=/host/${{ github.workspace }}/wheel-constraints/cuda-bindings.txt + PIP_CONSTRAINT=/host/${{ github.workspace }}/wheel-constraints/cuda-bindings.txt CC="/host/${{ env.SCCACHE_PATH }} cc" CXX="/host/${{ env.SCCACHE_PATH }} c++" SCCACHE_GHA_ENABLED=true @@ -194,6 +210,8 @@ jobs: CIBW_ENVIRONMENT_WINDOWS: > CUDA_PATH="$(cygpath -w ${{ env.CUDA_PATH }})" CUDA_PYTHON_PARALLEL_LEVEL=${{ env.CUDA_PYTHON_PARALLEL_LEVEL }} + PIP_BUILD_CONSTRAINT="$(cygpath -w ./wheel-constraints/cuda-bindings.txt)" + PIP_CONSTRAINT="$(cygpath -w ./wheel-constraints/cuda-bindings.txt)" # check cache stats before leaving cibuildwheel CIBW_BEFORE_TEST_LINUX: > "/host/${{ env.SCCACHE_PATH }}" --show-adv-stats && @@ -226,6 +244,20 @@ jobs: run: | twine check --strict ${{ env.CUDA_BINDINGS_ARTIFACTS_DIR }}/*.whl + - name: Resolve cuda.core build dependencies + run: | + container_args=() + if [[ "${{ inputs.host-platform }}" == linux* ]]; then + container_args+=(--container-mount /host) + fi + mkdir -p wheel-constraints + python ci/tools/write_wheel_constraints.py \ + --output wheel-constraints/cuda-core.txt \ + "${container_args[@]}" \ + --wheel cuda-pathfinder cuda_pathfinder \ + --wheel cuda-bindings "${CUDA_BINDINGS_ARTIFACTS_DIR}" \ + --expected-major cuda-bindings "${BUILD_CUDA_MAJOR}" + - name: Upload cuda.bindings build artifacts uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: @@ -240,6 +272,8 @@ jobs: output-dir: ${{ env.CUDA_CORE_ARTIFACTS_DIR }} env: CIBW_BUILD: ${{ env.CIBW_BUILD }} + CIBW_BEFORE_BUILD_LINUX: 'python -m pip install --upgrade "pip>=25.3"' + CIBW_BEFORE_BUILD_WINDOWS: 'python -m pip install --upgrade "pip>=25.3" delvewheel' # TODO: remove cpython-prerelease once 3.15 is officially supported # Allow CPython pre-release builds (currently 3.15 / 3.15t). This is a # no-op for stable Python versions because CIBW_BUILD still filters @@ -250,7 +284,8 @@ jobs: CUDA_PATH=/host/${{ env.CUDA_PATH }} CUDA_PYTHON_PARALLEL_LEVEL=${{ env.CUDA_PYTHON_PARALLEL_LEVEL }} CUDA_CORE_BUILD_MAJOR=${{ env.BUILD_CUDA_MAJOR }} - PIP_FIND_LINKS=/host/${{ env.CUDA_BINDINGS_ARTIFACTS_DIR }} + PIP_BUILD_CONSTRAINT=/host/${{ github.workspace }}/wheel-constraints/cuda-core.txt + PIP_CONSTRAINT=/host/${{ github.workspace }}/wheel-constraints/cuda-core.txt CC="/host/${{ env.SCCACHE_PATH }} cc" CXX="/host/${{ env.SCCACHE_PATH }} c++" SCCACHE_GHA_ENABLED=true @@ -265,7 +300,8 @@ jobs: CUDA_PATH="$(cygpath -w ${{ env.CUDA_PATH }})" CUDA_PYTHON_PARALLEL_LEVEL=${{ env.CUDA_PYTHON_PARALLEL_LEVEL }} CUDA_CORE_BUILD_MAJOR=${{ env.BUILD_CUDA_MAJOR }} - PIP_FIND_LINKS="$(cygpath -w ${{ env.CUDA_BINDINGS_ARTIFACTS_DIR }})" + PIP_BUILD_CONSTRAINT="$(cygpath -w ./wheel-constraints/cuda-core.txt)" + PIP_CONSTRAINT="$(cygpath -w ./wheel-constraints/cuda-core.txt)" # check cache stats before leaving cibuildwheel CIBW_BEFORE_TEST_LINUX: > "/host${{ env.SCCACHE_PATH }}" --show-adv-stats && @@ -444,14 +480,29 @@ jobs: OLD_BRANCH=$(yq '.backport_branch' ci/versions.yml) OLD_BASENAME="cuda-bindings-python${PYTHON_VERSION_FORMATTED}-cuda*-${{ inputs.host-platform }}*" LATEST_PRIOR_RUN_ID=$(./ci/tools/lookup-run-id --branch "${OLD_BRANCH}" NVIDIA/cuda-python "CI") + PREV_BINDINGS_DIR="cuda_bindings/dist-prev" gh run download $LATEST_PRIOR_RUN_ID -p ${OLD_BASENAME} -R NVIDIA/cuda-python rm -rf ${OLD_BASENAME}-tests # exclude cython test artifacts ls -al $OLD_BASENAME - mkdir -p "${{ env.CUDA_BINDINGS_ARTIFACTS_DIR }}" - mv $OLD_BASENAME/*.whl "${{ env.CUDA_BINDINGS_ARTIFACTS_DIR }}" + mkdir -p "${PREV_BINDINGS_DIR}" + mv $OLD_BASENAME/*.whl "${PREV_BINDINGS_DIR}" rmdir $OLD_BASENAME + - name: Resolve previous cuda.core build dependencies + run: | + container_args=() + if [[ "${{ inputs.host-platform }}" == linux* ]]; then + container_args+=(--container-mount /host) + fi + mkdir -p wheel-constraints + python ci/tools/write_wheel_constraints.py \ + --output wheel-constraints/cuda-core-prev.txt \ + "${container_args[@]}" \ + --wheel cuda-pathfinder cuda_pathfinder \ + --wheel cuda-bindings cuda_bindings/dist-prev \ + --expected-major cuda-bindings "${BUILD_PREV_CUDA_MAJOR}" + - name: Build cuda.core wheel uses: pypa/cibuildwheel@4726cd35bb13f7bde50cf2761f2499ac7b3aa32c # v4.1.1 with: @@ -459,6 +510,8 @@ jobs: output-dir: ${{ env.CUDA_CORE_ARTIFACTS_DIR }} env: CIBW_BUILD: ${{ env.CIBW_BUILD }} + CIBW_BEFORE_BUILD_LINUX: 'python -m pip install --upgrade "pip>=25.3"' + CIBW_BEFORE_BUILD_WINDOWS: 'python -m pip install --upgrade "pip>=25.3" delvewheel' # TODO: remove cpython-prerelease once 3.15 is officially supported # Allow CPython pre-release builds (currently 3.15 / 3.15t). This is a # no-op for stable Python versions because CIBW_BUILD still filters @@ -469,7 +522,8 @@ jobs: CUDA_PATH=/host/${{ env.CUDA_PATH }} CUDA_PYTHON_PARALLEL_LEVEL=${{ env.CUDA_PYTHON_PARALLEL_LEVEL }} CUDA_CORE_BUILD_MAJOR=${{ env.BUILD_PREV_CUDA_MAJOR }} - PIP_FIND_LINKS=/host/${{ env.CUDA_BINDINGS_ARTIFACTS_DIR }} + PIP_BUILD_CONSTRAINT=/host/${{ github.workspace }}/wheel-constraints/cuda-core-prev.txt + PIP_CONSTRAINT=/host/${{ github.workspace }}/wheel-constraints/cuda-core-prev.txt CC="/host/${{ env.SCCACHE_PATH }} cc" CXX="/host/${{ env.SCCACHE_PATH }} c++" SCCACHE_GHA_ENABLED=true @@ -484,7 +538,8 @@ jobs: CUDA_PATH="$(cygpath -w ${{ env.CUDA_PATH }})" CUDA_PYTHON_PARALLEL_LEVEL=${{ env.CUDA_PYTHON_PARALLEL_LEVEL }} CUDA_CORE_BUILD_MAJOR=${{ env.BUILD_PREV_CUDA_MAJOR }} - PIP_FIND_LINKS="$(cygpath -w ${{ env.CUDA_BINDINGS_ARTIFACTS_DIR }})" + PIP_BUILD_CONSTRAINT="$(cygpath -w ./wheel-constraints/cuda-core-prev.txt)" + PIP_CONSTRAINT="$(cygpath -w ./wheel-constraints/cuda-core-prev.txt)" # check cache stats before leaving cibuildwheel CIBW_BEFORE_TEST_LINUX: > "/host${{ env.SCCACHE_PATH }}" --show-adv-stats && diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index e354a204d02..a7e1d7b90d8 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -118,18 +118,41 @@ jobs: run: | python -m venv .venv - - name: Build cuda-pathfinder + - name: Install pip with build-constraint support + run: .venv/bin/python -m pip install "pip>=25.3" + + - name: Build and install cuda-pathfinder wheel run: | - cd cuda_pathfinder - ../.venv/bin/pip install -v . --group test + .venv/bin/pip wheel -v --no-deps ./cuda_pathfinder -w ./wheels/ + .venv/bin/pip install -v ./wheels/cuda_pathfinder*.whl --group ./cuda_pathfinder/pyproject.toml:test - - name: Build cuda-bindings + - name: Resolve cuda-bindings build dependencies run: | - cd cuda_bindings - ../.venv/bin/pip install -v . --group test + mkdir -p wheel-constraints + .venv/bin/python ci/tools/write_wheel_constraints.py \ + --output wheel-constraints/cuda-bindings.txt \ + --wheel cuda-pathfinder wheels + + - name: Build and install cuda-bindings wheel + run: | + export PIP_BUILD_CONSTRAINT="$(pwd)/wheel-constraints/cuda-bindings.txt" + export PIP_CONSTRAINT="${PIP_BUILD_CONSTRAINT}" + .venv/bin/pip wheel -v --no-deps ./cuda_bindings -w ./wheels/ + .venv/bin/pip install -v ./wheels/cuda_bindings*.whl --group ./cuda_bindings/pyproject.toml:test - - name: Build cuda-core + - name: Resolve cuda-core build dependencies run: | + CUDA_MAJOR="${CUDA_VER%%.*}" + .venv/bin/python ci/tools/write_wheel_constraints.py \ + --output wheel-constraints/cuda-core.txt \ + --wheel cuda-pathfinder wheels \ + --wheel cuda-bindings wheels \ + --expected-major cuda-bindings "${CUDA_MAJOR}" + + - name: Build and install cuda-core + run: | + export PIP_BUILD_CONSTRAINT="$(pwd)/wheel-constraints/cuda-core.txt" + export PIP_CONSTRAINT="${PIP_BUILD_CONSTRAINT}" cd cuda_core ../.venv/bin/pip install -v . --group test @@ -225,20 +248,38 @@ jobs: run: | python -m venv .venv - - name: Build and install cuda.pathfinder + - name: Build cuda.pathfinder wheel run: | - .venv/Scripts/pip install wheel setuptools Cython + .venv/Scripts/python -m pip install "pip>=25.3" wheel setuptools Cython .venv/Scripts/pip wheel -v --no-deps ./cuda_pathfinder -w ./wheels/ + - name: Resolve cuda.bindings build dependencies + run: | + mkdir -p wheel-constraints + .venv/Scripts/python ci/tools/write_wheel_constraints.py \ + --output wheel-constraints/cuda-bindings.txt \ + --wheel cuda-pathfinder wheels + - name: Build cuda.bindings wheel run: | + export PIP_BUILD_CONSTRAINT="$(cygpath -w "$(pwd)/wheel-constraints/cuda-bindings.txt")" + export PIP_CONSTRAINT="${PIP_BUILD_CONSTRAINT}" cd cuda_bindings ../.venv/Scripts/pip wheel -v --no-deps . -w ../wheels/ + - name: Resolve cuda.core build dependencies + run: | + CUDA_MAJOR="${CUDA_VER%%.*}" + .venv/Scripts/python ci/tools/write_wheel_constraints.py \ + --output wheel-constraints/cuda-core.txt \ + --wheel cuda-pathfinder wheels \ + --wheel cuda-bindings wheels \ + --expected-major cuda-bindings "${CUDA_MAJOR}" + - name: Build cuda.core wheel run: | - export PIP_FIND_LINKS="$(pwd)/wheels" - export PIP_PRE=1 + export PIP_BUILD_CONSTRAINT="$(cygpath -w "$(pwd)/wheel-constraints/cuda-core.txt")" + export PIP_CONSTRAINT="${PIP_BUILD_CONSTRAINT}" cd cuda_core ../.venv/Scripts/pip wheel -v --no-deps . -w ../wheels/ diff --git a/.github/workflows/test-sdist-linux.yml b/.github/workflows/test-sdist-linux.yml index 9d077912f3c..52b89961939 100644 --- a/.github/workflows/test-sdist-linux.yml +++ b/.github/workflows/test-sdist-linux.yml @@ -39,7 +39,7 @@ jobs: python-version: "3.12" - name: Install build tools - run: pip install build + run: python -m pip install "pip>=25.3" build # Pure Python packages -- no CTK needed. - name: Build cuda.pathfinder sdist and wheel-from-sdist @@ -52,6 +52,13 @@ jobs: python -m build --sdist cuda_python/ pip wheel --no-deps --wheel-dir cuda_python/dist cuda_python/dist/*.tar.gz + - name: Resolve cuda.bindings build dependencies + run: | + mkdir -p wheel-constraints + python ci/tools/write_wheel_constraints.py \ + --output wheel-constraints/cuda-bindings.txt \ + --wheel cuda-pathfinder cuda_pathfinder/dist + # Cython packages need CTK + sccache. # The env vars ACTIONS_CACHE_SERVICE_V2, ACTIONS_RESULTS_URL, and ACTIONS_RUNTIME_TOKEN # are exposed by this action. @@ -88,10 +95,20 @@ jobs: export CUDA_PYTHON_PARALLEL_LEVEL=$(nproc) export CC="sccache cc" export CXX="sccache c++" - export PIP_FIND_LINKS="$(pwd)/cuda_pathfinder/dist" + export PIP_BUILD_CONSTRAINT="$(pwd)/wheel-constraints/cuda-bindings.txt" + export PIP_CONSTRAINT="${PIP_BUILD_CONSTRAINT}" python -m build --sdist cuda_bindings/ pip wheel --no-deps --wheel-dir cuda_bindings/dist cuda_bindings/dist/*.tar.gz + - name: Resolve cuda.core build dependencies + run: | + CUDA_MAJOR="$(echo '${{ inputs.cuda-version }}' | cut -d. -f1)" + python ci/tools/write_wheel_constraints.py \ + --output wheel-constraints/cuda-core.txt \ + --wheel cuda-pathfinder cuda_pathfinder/dist \ + --wheel cuda-bindings cuda_bindings/dist \ + --expected-major cuda-bindings "${CUDA_MAJOR}" + # cuda_core sdist delegates to setuptools (no CTK needed), but # wheel-from-sdist needs CTK and cuda-bindings (dynamic build dep via # get_requires_for_build_wheel in build_hooks.py). @@ -101,7 +118,8 @@ jobs: export CUDA_CORE_BUILD_MAJOR="$(echo '${{ inputs.cuda-version }}' | cut -d. -f1)" export CC="sccache cc" export CXX="sccache c++" - export PIP_FIND_LINKS="$(pwd)/cuda_bindings/dist $(pwd)/cuda_pathfinder/dist" + export PIP_BUILD_CONSTRAINT="$(pwd)/wheel-constraints/cuda-core.txt" + export PIP_CONSTRAINT="${PIP_BUILD_CONSTRAINT}" python -m build --sdist cuda_core/ pip wheel --no-deps --wheel-dir cuda_core/dist cuda_core/dist/*.tar.gz diff --git a/.github/workflows/test-sdist-windows.yml b/.github/workflows/test-sdist-windows.yml index 043bacc1cad..f9673deec58 100644 --- a/.github/workflows/test-sdist-windows.yml +++ b/.github/workflows/test-sdist-windows.yml @@ -48,7 +48,7 @@ jobs: uses: step-security/msvc-dev-cmd@22c98154b708dbd743e6f27a933cf6ceba3305c4 # v1.13.1 - name: Install build tools - run: pip install build + run: python -m pip install "pip>=25.3" build # Pure Python packages -- no CTK needed. - name: Build cuda.pathfinder sdist and wheel-from-sdist @@ -61,6 +61,13 @@ jobs: python -m build --sdist cuda_python/ pip wheel --no-deps --wheel-dir cuda_python/dist cuda_python/dist/*.tar.gz + - name: Resolve cuda.bindings build dependencies + run: | + mkdir -p wheel-constraints + python ci/tools/write_wheel_constraints.py \ + --output wheel-constraints/cuda-bindings.txt \ + --wheel cuda-pathfinder cuda_pathfinder/dist + # Cython packages need CTK. No sccache on Windows (this is a correctness # smoke test, not a production build; see build-wheel.yml which also # limits sccache to Linux). @@ -73,17 +80,25 @@ jobs: # cuda_bindings/setup.py parses CUDA headers at import time, so CUDA_PATH # (set by fetch_ctk) must be available for both sdist and wheel builds. - # PIP_FIND_LINKS is passed as a native Windows path via cygpath because - # pip on Windows treats space-separated entries as separators and is - # picky about mixed path styles (see build-wheel.yml for the same - # convention). + # Constraint paths are passed as native Windows paths because the pip + # subprocesses run outside Git Bash. - name: Build cuda.bindings sdist and wheel-from-sdist run: | export CUDA_PYTHON_PARALLEL_LEVEL=$(nproc) - export PIP_FIND_LINKS="$(cygpath -w "$(pwd)/cuda_pathfinder/dist")" + export PIP_BUILD_CONSTRAINT="$(cygpath -w "$(pwd)/wheel-constraints/cuda-bindings.txt")" + export PIP_CONSTRAINT="${PIP_BUILD_CONSTRAINT}" python -m build --sdist cuda_bindings/ pip wheel --no-deps --wheel-dir cuda_bindings/dist cuda_bindings/dist/*.tar.gz + - name: Resolve cuda.core build dependencies + run: | + CUDA_MAJOR="$(echo '${{ inputs.cuda-version }}' | cut -d. -f1)" + python ci/tools/write_wheel_constraints.py \ + --output wheel-constraints/cuda-core.txt \ + --wheel cuda-pathfinder cuda_pathfinder/dist \ + --wheel cuda-bindings cuda_bindings/dist \ + --expected-major cuda-bindings "${CUDA_MAJOR}" + # cuda_core sdist delegates to setuptools (no CTK needed), but # wheel-from-sdist needs CTK and cuda-bindings (dynamic build dep via # get_requires_for_build_wheel in build_hooks.py). @@ -91,6 +106,7 @@ jobs: run: | export CUDA_PYTHON_PARALLEL_LEVEL=$(nproc) export CUDA_CORE_BUILD_MAJOR="$(echo '${{ inputs.cuda-version }}' | cut -d. -f1)" - export PIP_FIND_LINKS="$(cygpath -w "$(pwd)/cuda_bindings/dist") $(cygpath -w "$(pwd)/cuda_pathfinder/dist")" + export PIP_BUILD_CONSTRAINT="$(cygpath -w "$(pwd)/wheel-constraints/cuda-core.txt")" + export PIP_CONSTRAINT="${PIP_BUILD_CONSTRAINT}" python -m build --sdist cuda_core/ pip wheel --no-deps --wheel-dir cuda_core/dist cuda_core/dist/*.tar.gz diff --git a/ci/tools/tests/test_write_wheel_constraints.py b/ci/tools/tests/test_write_wheel_constraints.py new file mode 100644 index 00000000000..8af0e38becb --- /dev/null +++ b/ci/tools/tests/test_write_wheel_constraints.py @@ -0,0 +1,297 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +import stat +import sys +import zipfile +from pathlib import Path + +import pytest + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) +from write_wheel_constraints import WheelConstraintError, WheelRequirement, main, write_constraints + + +def _write_wheel( + directory: Path, + filename: str, + *, + name: str, + version: str, + metadata_entries: int = 1, +) -> Path: + directory.mkdir(parents=True, exist_ok=True) + wheel_path = directory / filename + with zipfile.ZipFile(wheel_path, "w") as wheel: + for index in range(metadata_entries): + suffix = f"-{index}" if metadata_entries > 1 else "" + wheel.writestr( + f"{name.replace('-', '_')}-{version}{suffix}.dist-info/METADATA", + f"Metadata-Version: 2.1\nName: {name}\nVersion: {version}\n", + ) + return wheel_path + + +@pytest.mark.agent_authored(model="gpt-5.6-sol") +class TestWriteConstraints: + def test_writes_canonical_direct_references_and_escapes_spaces(self, tmp_path): + wheel_dir = tmp_path / "wheel house" + pathfinder = _write_wheel( + wheel_dir, + "cuda_pathfinder-1.5.2-py3-none-any.whl", + name="cuda.pathfinder", + version="1.5.2", + ) + bindings = _write_wheel( + wheel_dir, + "cuda_bindings-13.3.0.dev1-py3-none-any.whl", + name="CUDA_Bindings", + version="13.3.0.dev1", + ) + _write_wheel( + wheel_dir, + "unrelated-1.0-py3-none-any.whl", + name="unrelated", + version="1.0", + ) + output = tmp_path / "constraints.txt" + + write_constraints( + output, + [ + WheelRequirement("cuda-pathfinder", wheel_dir), + WheelRequirement("cuda.bindings", wheel_dir, expected_major="13"), + ], + ) + + assert output.read_text(encoding="utf-8").splitlines() == [ + f"cuda-pathfinder @ {pathfinder.resolve().as_uri()}", + f"cuda-bindings @ {bindings.resolve().as_uri()}", + ] + assert "%20" in output.read_text(encoding="utf-8") + + def test_expected_major_selects_from_coexisting_cuda_majors(self, tmp_path): + wheel_dir = tmp_path / "wheels" + _write_wheel( + wheel_dir, + "cuda_bindings-12.9.1-py3-none-any.whl", + name="cuda-bindings", + version="12.9.1", + ) + selected = _write_wheel( + wheel_dir, + "cuda_bindings-13.3.0.dev2-py3-none-any.whl", + name="cuda-bindings", + version="13.3.0.dev2", + ) + output = tmp_path / "constraints.txt" + + write_constraints(output, [WheelRequirement("cuda-bindings", wheel_dir, expected_major="13")]) + + assert output.read_text(encoding="utf-8") == f"cuda-bindings @ {selected.resolve().as_uri()}\n" + + @pytest.mark.skipif(os.name != "posix", reason="cibuildwheel /host mapping is POSIX-only") + def test_maps_wheel_uri_to_container_mount(self, tmp_path): + wheel = _write_wheel( + tmp_path / "wheels", + "cuda_pathfinder-1.5.2-py3-none-any.whl", + name="cuda-pathfinder", + version="1.5.2", + ).resolve() + output = tmp_path / "constraints.txt" + + write_constraints( + output, + [WheelRequirement("cuda-pathfinder", wheel.parent)], + container_mount=Path("/host"), + ) + + container_wheel = Path("/host") / wheel.relative_to(wheel.anchor) + assert output.read_text(encoding="utf-8") == f"cuda-pathfinder @ {container_wheel.as_uri()}\n" + assert stat.S_IMODE(output.stat().st_mode) == 0o644 + + def test_rejects_relative_container_mount(self, tmp_path): + wheel_dir = tmp_path / "wheels" + _write_wheel( + wheel_dir, + "cuda_pathfinder-1.5.2-py3-none-any.whl", + name="cuda-pathfinder", + version="1.5.2", + ) + + with pytest.raises(WheelConstraintError, match="Container mount must be an absolute path"): + write_constraints( + tmp_path / "constraints.txt", + [WheelRequirement("cuda-pathfinder", wheel_dir)], + container_mount=Path("host"), + ) + + @pytest.mark.parametrize("directory_kind", ["missing", "empty"]) + def test_rejects_missing_or_empty_wheel_directory(self, tmp_path, directory_kind): + wheel_dir = tmp_path / "wheels" + if directory_kind == "empty": + wheel_dir.mkdir() + + with pytest.raises(WheelConstraintError, match="does not exist|contains no .whl files"): + write_constraints(tmp_path / "constraints.txt", [WheelRequirement("cuda-pathfinder", wheel_dir)]) + + def test_uses_metadata_name_instead_of_spoofed_filename(self, tmp_path): + wheel_dir = tmp_path / "wheels" + _write_wheel( + wheel_dir, + "cuda_bindings-13.3.0-py3-none-any.whl", + name="not-cuda-bindings", + version="13.3.0", + ) + + with pytest.raises(WheelConstraintError, match="Found no wheel for cuda-bindings"): + write_constraints(tmp_path / "constraints.txt", [WheelRequirement("cuda-bindings", wheel_dir)]) + + def test_rejects_multiple_wheels_for_expected_major(self, tmp_path): + wheel_dir = tmp_path / "wheels" + for patch in ("0", "1"): + _write_wheel( + wheel_dir, + f"cuda_bindings-13.3.{patch}-py3-none-any.whl", + name="cuda-bindings", + version=f"13.3.{patch}", + ) + + with pytest.raises(WheelConstraintError, match="Found multiple wheels.*release major 13"): + write_constraints( + tmp_path / "constraints.txt", + [WheelRequirement("cuda-bindings", wheel_dir, expected_major="13")], + ) + + def test_wrong_major_reports_available_version(self, tmp_path): + wheel_dir = tmp_path / "wheels" + _write_wheel( + wheel_dir, + "cuda_bindings-12.9.1-py3-none-any.whl", + name="cuda-bindings", + version="12.9.1", + ) + + with pytest.raises(WheelConstraintError, match=r"release major 13.*cuda-bindings==12\.9\.1"): + write_constraints( + tmp_path / "constraints.txt", + [WheelRequirement("cuda-bindings", wheel_dir, expected_major="13")], + ) + + def test_rejects_duplicate_canonical_project_requests(self, tmp_path): + wheel_dir = tmp_path / "wheels" + _write_wheel( + wheel_dir, + "cuda_bindings-13.3.0-py3-none-any.whl", + name="cuda-bindings", + version="13.3.0", + ) + + with pytest.raises(WheelConstraintError, match="requested more than once"): + write_constraints( + tmp_path / "constraints.txt", + [WheelRequirement("cuda-bindings", wheel_dir), WheelRequirement("CUDA_Bindings", wheel_dir)], + ) + + @pytest.mark.parametrize("project", ["", "cuda bindings", "---", "cuda-bindings-"]) + def test_rejects_invalid_project_name(self, tmp_path, project): + with pytest.raises(WheelConstraintError, match="Invalid project name"): + write_constraints(tmp_path / "constraints.txt", [WheelRequirement(project, tmp_path)]) + + def test_rejects_invalid_wheel_archive(self, tmp_path): + wheel_dir = tmp_path / "wheels" + wheel_dir.mkdir() + (wheel_dir / "cuda_bindings-13.3.0-py3-none-any.whl").write_text("not a zip", encoding="utf-8") + + with pytest.raises(WheelConstraintError, match="Cannot read wheel"): + write_constraints(tmp_path / "constraints.txt", [WheelRequirement("cuda-bindings", wheel_dir)]) + + @pytest.mark.parametrize("metadata_entries", [0, 2]) + def test_rejects_missing_or_multiple_metadata_entries(self, tmp_path, metadata_entries): + wheel_dir = tmp_path / "wheels" + _write_wheel( + wheel_dir, + "cuda_bindings-13.3.0-py3-none-any.whl", + name="cuda-bindings", + version="13.3.0", + metadata_entries=metadata_entries, + ) + + with pytest.raises(WheelConstraintError, match="METADATA entries; expected exactly one"): + write_constraints(tmp_path / "constraints.txt", [WheelRequirement("cuda-bindings", wheel_dir)]) + + def test_failed_resolution_removes_stale_output(self, tmp_path): + output = tmp_path / "constraints.txt" + output.write_text("cuda-bindings==0\n", encoding="utf-8") + + with pytest.raises(WheelConstraintError): + write_constraints(output, [WheelRequirement("cuda-bindings", tmp_path / "missing")]) + + assert not output.exists() + + def test_cli_reports_unknown_expected_major_project(self, tmp_path, capsys): + output = tmp_path / "constraints.txt" + output.write_text("stale\n", encoding="utf-8") + rc = main( + [ + "--output", + str(output), + "--wheel", + "cuda-pathfinder", + str(tmp_path), + "--expected-major", + "cuda-bindings", + "13", + ] + ) + + assert rc == 1 + assert "error: Expected major was specified for an unrequested project" in capsys.readouterr().err + assert not output.exists() + + def test_cli_success_reports_selected_wheel(self, tmp_path, capsys): + wheel_dir = tmp_path / "wheels" + wheel = _write_wheel( + wheel_dir, + "cuda_pathfinder-1.5.2-py3-none-any.whl", + name="cuda-pathfinder", + version="1.5.2", + ) + output = tmp_path / "constraints.txt" + + rc = main( + [ + "--output", + str(output), + "--wheel", + "cuda-pathfinder", + str(wheel_dir), + ] + ) + + assert rc == 0 + assert output.read_text(encoding="utf-8") == f"cuda-pathfinder @ {wheel.resolve().as_uri()}\n" + assert "Selected cuda-pathfinder==1.5.2" in capsys.readouterr().out + + def test_cli_reports_corrupt_wheel_without_traceback(self, tmp_path, capsys): + wheel_dir = tmp_path / "wheels" + wheel_dir.mkdir() + (wheel_dir / "cuda_pathfinder-1.5.2-py3-none-any.whl").write_text("not a zip", encoding="utf-8") + + rc = main( + [ + "--output", + str(tmp_path / "constraints.txt"), + "--wheel", + "cuda-pathfinder", + str(wheel_dir), + ] + ) + + assert rc == 1 + captured = capsys.readouterr() + assert "error: Cannot read wheel" in captured.err + assert "Traceback" not in captured.err diff --git a/ci/tools/write_wheel_constraints.py b/ci/tools/write_wheel_constraints.py new file mode 100644 index 00000000000..5579dea820e --- /dev/null +++ b/ci/tools/write_wheel_constraints.py @@ -0,0 +1,279 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Write pip constraints that select exact, locally built wheel artifacts.""" + +from __future__ import annotations + +import argparse +import os +import re +import sys +import tempfile +import zipfile +from collections.abc import Sequence +from dataclasses import dataclass +from email import policy +from email.errors import MessageError +from email.parser import BytesParser +from pathlib import Path + +_CANONICALIZE_PROJECT_RE = re.compile(r"[-_.]+") +_PROJECT_NAME_RE = re.compile(r"^[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9])?$") +_RELEASE_MAJOR_RE = re.compile(r"^([0-9]+)(?:\.|$)") + + +class WheelConstraintError(RuntimeError): + """Raised when an exact local wheel constraint cannot be produced.""" + + +@dataclass(frozen=True) +class WheelRequirement: + """A distribution that must resolve to one wheel in a local directory.""" + + project: str + directory: Path + expected_major: str | None = None + + +@dataclass(frozen=True) +class _WheelMetadata: + path: Path + name: str + version: str + + +def _canonicalize_project(name: str) -> str: + if _PROJECT_NAME_RE.fullmatch(name) is None: + raise WheelConstraintError(f"Invalid project name: {name!r}") + return _CANONICALIZE_PROJECT_RE.sub("-", name).lower() + + +def _read_wheel_metadata(wheel_path: Path) -> _WheelMetadata: + try: + with zipfile.ZipFile(wheel_path) as wheel: + metadata_entries = [name for name in wheel.namelist() if name.endswith(".dist-info/METADATA")] + if len(metadata_entries) != 1: + raise WheelConstraintError( + f"Wheel {wheel_path} contains {len(metadata_entries)} .dist-info/METADATA entries; expected exactly one" + ) + metadata_bytes = wheel.read(metadata_entries[0]) + metadata = BytesParser(policy=policy.compat32).parsebytes(metadata_bytes) + except WheelConstraintError: + raise + except (OSError, zipfile.BadZipFile, RuntimeError, MessageError) as exc: + raise WheelConstraintError(f"Cannot read wheel {wheel_path}: {exc}") from exc + + name = metadata.get("Name", "").strip() + version = metadata.get("Version", "").strip() + if not name or not version: + raise WheelConstraintError(f"Wheel {wheel_path} METADATA must contain non-empty Name and Version fields") + _canonicalize_project(name) + try: + resolved_path = wheel_path.resolve() + except OSError as exc: + raise WheelConstraintError(f"Cannot resolve wheel path {wheel_path}: {exc}") from exc + return _WheelMetadata(path=resolved_path, name=name, version=version) + + +def _release_major(version: str, wheel_path: Path) -> str: + match = _RELEASE_MAJOR_RE.match(version) + if match is None: + raise WheelConstraintError( + f"Wheel {wheel_path} has version {version!r}, which does not start with a numeric release major" + ) + return match.group(1) + + +def _describe_wheels(wheels: Sequence[_WheelMetadata]) -> str: + if not wheels: + return "none" + return ", ".join(f"{wheel.name}=={wheel.version} ({wheel.path.name})" for wheel in wheels) + + +def _select_wheel(requirement: WheelRequirement) -> _WheelMetadata: + project = _canonicalize_project(requirement.project) + try: + directory = requirement.directory.resolve() + except OSError as exc: + raise WheelConstraintError( + f"Cannot resolve wheel directory for {project}: {requirement.directory}: {exc}" + ) from exc + if not directory.is_dir(): + raise WheelConstraintError(f"Wheel directory for {project} does not exist or is not a directory: {directory}") + + try: + wheel_paths = sorted(directory.glob("*.whl")) + except OSError as exc: + raise WheelConstraintError(f"Cannot inspect wheel directory {directory}: {exc}") from exc + if not wheel_paths: + raise WheelConstraintError(f"Wheel directory for {project} contains no .whl files: {directory}") + + inspected = [_read_wheel_metadata(path) for path in wheel_paths] + matching_project = [wheel for wheel in inspected if _canonicalize_project(wheel.name) == project] + + expected_major = requirement.expected_major + if expected_major is not None: + if re.fullmatch(r"[0-9]+", expected_major) is None: + raise WheelConstraintError(f"Expected major for {project} must contain only digits, got {expected_major!r}") + matching = [wheel for wheel in matching_project if _release_major(wheel.version, wheel.path) == expected_major] + else: + matching = matching_project + + if not matching: + qualifier = f" with release major {expected_major}" if expected_major is not None else "" + raise WheelConstraintError( + f"Found no wheel for {project}{qualifier} in {directory}; inspected: {_describe_wheels(inspected)}" + ) + if len(matching) > 1: + qualifier = f" with release major {expected_major}" if expected_major is not None else "" + raise WheelConstraintError( + f"Found multiple wheels for {project}{qualifier} in {directory}: {_describe_wheels(matching)}" + ) + return matching[0] + + +def _consumer_path(wheel_path: Path, container_mount: Path | None) -> Path: + if container_mount is None: + return wheel_path + if not container_mount.is_absolute(): + raise WheelConstraintError(f"Container mount must be an absolute path: {container_mount}") + if os.name != "posix": + raise WheelConstraintError("--container-mount is only supported on POSIX hosts") + return container_mount / wheel_path.relative_to(wheel_path.anchor) + + +def _remove_stale_output(output_path: Path) -> Path: + try: + resolved_output = output_path.resolve() + resolved_output.unlink(missing_ok=True) + except OSError as exc: + raise WheelConstraintError(f"Cannot remove stale constraints file {output_path}: {exc}") from exc + return resolved_output + + +def write_constraints( + output_path: Path, + requirements: Sequence[WheelRequirement], + *, + container_mount: Path | None = None, +) -> None: + """Write exact direct-reference constraints for the requested wheels.""" + output_path = _remove_stale_output(output_path) + + if not requirements: + raise WheelConstraintError("At least one wheel requirement is required") + + seen_projects: set[str] = set() + selected: list[tuple[str, _WheelMetadata, str]] = [] + for requirement in requirements: + project = _canonicalize_project(requirement.project) + if not project: + raise WheelConstraintError(f"Project name must not be empty: {requirement.project!r}") + if project in seen_projects: + raise WheelConstraintError(f"Project {project} was requested more than once") + seen_projects.add(project) + + wheel = _select_wheel(requirement) + consumer_uri = _consumer_path(wheel.path, container_mount).as_uri() + selected.append((project, wheel, consumer_uri)) + + temporary_path: Path | None = None + try: + output_path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + newline="\n", + dir=output_path.parent, + prefix=f".{output_path.name}.", + delete=False, + ) as output: + temporary_path = Path(output.name) + for project, _, consumer_uri in selected: + output.write(f"{project} @ {consumer_uri}\n") + temporary_path.chmod(0o644) + os.replace(temporary_path, output_path) + except OSError as exc: + cleanup_details = "" + if temporary_path is not None: + try: + temporary_path.unlink(missing_ok=True) + except OSError as cleanup_exc: + cleanup_details = f"; temporary-file cleanup also failed: {cleanup_exc}" + raise WheelConstraintError(f"Cannot write constraints file {output_path}: {exc}{cleanup_details}") from exc + + for project, wheel, consumer_uri in selected: + print(f"Selected {project}=={wheel.version} from {wheel.path} as {consumer_uri}") + print(f"Wrote exact wheel constraints to {output_path}") + + +def _parse_requirements( + wheels: Sequence[Sequence[str]], expected_majors: Sequence[Sequence[str]] +) -> list[WheelRequirement]: + majors_by_project: dict[str, str] = {} + for project, major in expected_majors: + canonical_project = _canonicalize_project(project) + if canonical_project in majors_by_project: + raise WheelConstraintError(f"Expected major for {canonical_project} was specified more than once") + majors_by_project[canonical_project] = major + + requested_projects = {_canonicalize_project(project) for project, _ in wheels} + unused_majors = sorted(majors_by_project.keys() - requested_projects) + if unused_majors: + raise WheelConstraintError( + f"Expected major was specified for an unrequested project: {', '.join(unused_majors)}" + ) + + return [ + WheelRequirement( + project=project, + directory=Path(directory), + expected_major=majors_by_project.get(_canonicalize_project(project)), + ) + for project, directory in wheels + ] + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", required=True, type=Path, help="Constraints file to write") + parser.add_argument( + "--wheel", + action="append", + nargs=2, + required=True, + metavar=("PROJECT", "DIRECTORY"), + help="Project and directory containing its locally built wheel", + ) + parser.add_argument( + "--expected-major", + action="append", + nargs=2, + default=[], + metavar=("PROJECT", "MAJOR"), + help="Require the selected project's wheel version to have this release major", + ) + parser.add_argument( + "--container-mount", + type=Path, + help="Map absolute host wheel paths below this container-visible mount", + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + """Write a constraints file from command-line arguments.""" + args = _parser().parse_args(argv) + try: + _remove_stale_output(args.output) + requirements = _parse_requirements(args.wheel, args.expected_major) + write_constraints(args.output, requirements, container_mount=args.container_mount) + except WheelConstraintError as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main())