Migrate kd-tree operations to numpy arrays with RowMajor storage fix #2215
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| # (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. | |
| # | |
| name: PyPI Wheels | |
| on: | |
| # Build wheels on every push to test the build process | |
| push: | |
| branches: | |
| - "**" | |
| paths-ignore: | |
| - "**/website/**" | |
| tags: | |
| - 'v*' | |
| # Build on pull requests | |
| pull_request: | |
| branches: | |
| - "**" | |
| paths-ignore: | |
| - "**/website/**" | |
| # Allow manual trigger for testing builds or publishing | |
| workflow_dispatch: | |
| inputs: | |
| publish: | |
| description: 'Publish to PyPI (otherwise just build and test)' | |
| required: false | |
| type: boolean | |
| default: false | |
| test_pypi: | |
| description: 'Publish to TestPyPI instead of PyPI (for testing)' | |
| required: false | |
| type: boolean | |
| default: false | |
| concurrency: | |
| group: ${{ github.workflow }}-${{ github.ref }} | |
| cancel-in-progress: true | |
| permissions: | |
| contents: read | |
| jobs: | |
| # Linux CPU wheels are built with cibuildwheel (for manylinux_2_28 compatibility) | |
| build_cpu_wheels_linux: | |
| name: pypi-cpu-linux (cibuildwheel) | |
| runs-on: ubuntu-latest | |
| steps: | |
| - uses: actions/checkout@v6 | |
| with: | |
| submodules: recursive | |
| fetch-depth: 0 # Fetch full history for setuptools_scm version detection | |
| fetch-tags: true # Ensure all tags are fetched | |
| - name: Set up Python | |
| uses: actions/setup-python@v6 | |
| with: | |
| python-version: "3.12" | |
| - name: Install Jinja2 | |
| run: pip install jinja2 | |
| - name: Generate pyproject.toml | |
| run: | | |
| python scripts/generate_pyproject.py | |
| # Set a default pyproject.toml for the initial setup | |
| cp pyproject-pypi-cpu-py312.toml pyproject.toml | |
| - name: Build wheels | |
| uses: pypa/cibuildwheel@v3.3.1 | |
| - uses: actions/upload-artifact@v6 | |
| with: | |
| name: wheels-cpu-linux | |
| path: ./wheelhouse/*.whl | |
| # CPU wheels are built with pixi (non-Linux platforms) | |
| # Linux CPU wheels are handled by build_cpu_wheels_linux above using cibuildwheel | |
| build_cpu_wheels: | |
| name: pypi-cpu-py${{ matrix.python-version }}-${{ matrix.os-short }} | |
| runs-on: ${{ matrix.os }} | |
| strategy: | |
| fail-fast: true | |
| matrix: | |
| os: [macos-14] | |
| python-version: ['3.12', '3.13'] | |
| include: | |
| - python-version: '3.12' | |
| pixi-environment: py312 | |
| - python-version: '3.13' | |
| pixi-environment: py313 | |
| - os: macos-14 | |
| os-short: mac-arm64 | |
| env: | |
| # Enable compiler caching (ccache on Linux/macOS) | |
| CMAKE_C_COMPILER_LAUNCHER: ccache | |
| CMAKE_CXX_COMPILER_LAUNCHER: ccache | |
| steps: | |
| # Skip Python 3.13 builds on regular commits to reduce CI cost | |
| # Python 3.13 builds will still run when publishing (tags or manual workflow_dispatch with publish=true) | |
| - name: Check if should run | |
| id: should_run | |
| run: | | |
| if [[ "${{ matrix.python-version }}" == "3.13" ]]; then | |
| if [[ "${{ startsWith(github.ref, 'refs/tags/v') }}" == "true" ]] || \ | |
| [[ "${{ github.event_name == 'workflow_dispatch' && github.event.inputs.publish == 'true' }}" == "true" ]]; then | |
| echo "run=true" >> $GITHUB_OUTPUT | |
| else | |
| echo "run=false" >> $GITHUB_OUTPUT | |
| fi | |
| else | |
| echo "run=true" >> $GITHUB_OUTPUT | |
| fi | |
| shell: bash | |
| - uses: actions/checkout@v6 | |
| if: steps.should_run.outputs.run == 'true' | |
| with: | |
| submodules: recursive | |
| fetch-depth: 0 # Fetch full history for setuptools_scm version detection | |
| fetch-tags: true # Ensure all tags are fetched | |
| - name: Set up ccache | |
| if: steps.should_run.outputs.run == 'true' | |
| uses: hendrikmuhs/ccache-action@v1.2 | |
| with: | |
| key: ${{ github.workflow }}-${{ matrix.os }}-py${{ matrix.python-version }} | |
| - name: Set up pixi | |
| if: steps.should_run.outputs.run == 'true' | |
| uses: prefix-dev/setup-pixi@v0.9.3 | |
| with: | |
| pixi-version: latest | |
| cache: true | |
| cache-write: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }} | |
| environments: ${{ matrix.pixi-environment }} | |
| - name: Generate pyproject configs | |
| if: steps.should_run.outputs.run == 'true' | |
| run: pixi run -e ${{ matrix.pixi-environment }} generate_pyproject | |
| - name: Clean distribution artifacts | |
| if: steps.should_run.outputs.run == 'true' | |
| run: pixi run -e ${{ matrix.pixi-environment }} wheel_clean | |
| - name: Build CPU wheel | |
| if: steps.should_run.outputs.run == 'true' | |
| run: pixi run -e ${{ matrix.pixi-environment }} wheel_build | |
| - name: Print ccache stats | |
| if: steps.should_run.outputs.run == 'true' | |
| run: ccache -s | |
| - name: Upload wheel artifacts | |
| if: steps.should_run.outputs.run == 'true' | |
| uses: actions/upload-artifact@v6 | |
| with: | |
| name: wheels-cpu-${{ matrix.os }}-py${{ matrix.python-version }} | |
| path: dist/*.whl | |
| retention-days: 7 | |
| build_gpu_wheels: | |
| name: pypi-gpu-py${{ matrix.python-version }}-${{ matrix.os-short }} | |
| runs-on: ${{ matrix.os }} | |
| strategy: | |
| fail-fast: true | |
| # Run GPU builds sequentially to avoid disk space exhaustion | |
| # Each PyTorch+CUDA install needs ~3GB and Docker overlay shares host disk | |
| max-parallel: 1 | |
| matrix: | |
| os: [ubuntu-latest] | |
| python-version: ['3.12', '3.13'] | |
| include: | |
| - python-version: '3.12' | |
| pixi-environment: gpu-wheel-build-py312 | |
| cuda-version: "12.9.0" | |
| - python-version: '3.13' | |
| pixi-environment: gpu-wheel-build-py313 | |
| cuda-version: "12.9.0" | |
| - os: ubuntu-latest | |
| os-short: ubuntu | |
| env: | |
| FULL_CUDA_VERSION: ${{ matrix.cuda-version }} | |
| # Use docker for CI builds (default for cibuildwheel on GitHub Actions) | |
| CIBW_CONTAINER_ENGINE: docker | |
| # Enable compiler caching (ccache on Linux) | |
| CMAKE_C_COMPILER_LAUNCHER: ccache | |
| CMAKE_CXX_COMPILER_LAUNCHER: ccache | |
| steps: | |
| # Skip Python 3.13 builds on regular commits to reduce CI cost | |
| # Python 3.13 builds will still run when publishing (tags or manual workflow_dispatch with publish=true) | |
| - name: Check if should run | |
| id: should_run | |
| run: | | |
| if [[ "${{ matrix.python-version }}" == "3.13" ]]; then | |
| if [[ "${{ startsWith(github.ref, 'refs/tags/v') }}" == "true" ]] || \ | |
| [[ "${{ github.event_name == 'workflow_dispatch' && github.event.inputs.publish == 'true' }}" == "true" ]]; then | |
| echo "run=true" >> $GITHUB_OUTPUT | |
| else | |
| echo "run=false" >> $GITHUB_OUTPUT | |
| fi | |
| else | |
| echo "run=true" >> $GITHUB_OUTPUT | |
| fi | |
| shell: bash | |
| - name: Free disk space | |
| if: steps.should_run.outputs.run == 'true' && matrix.os == 'ubuntu-latest' | |
| uses: jlumbroso/free-disk-space@main | |
| with: | |
| # Remove all default tools and applications | |
| tool-cache: false # Keep tool cache for Python | |
| android: true | |
| dotnet: true | |
| haskell: true | |
| large-packages: true | |
| docker-images: false # Keep Docker for cibuildwheel | |
| swap-storage: true | |
| - name: Maximize build space (Ubuntu only) | |
| # DISABLED: This action restructures the filesystem and breaks Docker functionality | |
| # The "Free disk space" step above provides enough space for the build | |
| if: false && steps.should_run.outputs.run == 'true' && matrix.os == 'ubuntu-latest' | |
| uses: easimon/maximize-build-space@master | |
| with: | |
| root-reserve-mb: 30720 | |
| swap-size-mb: 1024 | |
| remove-dotnet: true | |
| remove-android: true | |
| remove-haskell: true | |
| remove-codeql: true | |
| remove-docker-images: false # Keep Docker for cibuildwheel | |
| - uses: actions/checkout@v6 | |
| if: steps.should_run.outputs.run == 'true' | |
| with: | |
| submodules: recursive | |
| fetch-depth: 0 # Fetch full history for setuptools_scm version detection | |
| fetch-tags: true # Ensure all tags are fetched | |
| - name: Set up ccache | |
| if: steps.should_run.outputs.run == 'true' | |
| uses: hendrikmuhs/ccache-action@v1.2 | |
| with: | |
| key: ${{ github.workflow }}-${{ matrix.os }}-py${{ matrix.python-version }}-gpu | |
| - name: Set up pixi | |
| if: steps.should_run.outputs.run == 'true' | |
| uses: prefix-dev/setup-pixi@v0.9.3 | |
| with: | |
| pixi-version: latest | |
| cache: true | |
| cache-write: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }} | |
| environments: ${{ matrix.pixi-environment }} | |
| - name: Generate pyproject configs | |
| if: steps.should_run.outputs.run == 'true' | |
| run: pixi run -e ${{ matrix.pixi-environment }} generate_pyproject | |
| - name: Clean distribution artifacts | |
| if: steps.should_run.outputs.run == 'true' | |
| run: pixi run -e ${{ matrix.pixi-environment }} wheel_clean | |
| - name: Determine version for Docker builds | |
| if: steps.should_run.outputs.run == 'true' | |
| id: get_version | |
| run: | | |
| # Install setuptools_scm to determine version from git tags | |
| pip install setuptools_scm | |
| # Get the version that setuptools_scm would generate | |
| FULL_VERSION=$(python -c "from setuptools_scm import get_version; print(get_version())") | |
| echo "Detected version: $FULL_VERSION" | |
| # Strip local version part (e.g., +gad8997da8) - PyPI doesn't allow local versions | |
| # This converts "0.1.101.dev26+gad8997da8" to "0.1.101.dev26" | |
| VERSION=$(echo "$FULL_VERSION" | sed 's/+.*//') | |
| echo "Version for PyPI: $VERSION" | |
| echo "version=$VERSION" >> $GITHUB_OUTPUT | |
| - name: Build GPU wheel | |
| if: steps.should_run.outputs.run == 'true' | |
| timeout-minutes: 60 | |
| env: | |
| # Pass version to cibuildwheel so Docker containers can use it | |
| # Without this, setuptools_scm inside Docker would generate 0.0.post1 | |
| CIBW_ENVIRONMENT_PASS_LINUX: SETUPTOOLS_SCM_PRETEND_VERSION | |
| SETUPTOOLS_SCM_PRETEND_VERSION: ${{ steps.get_version.outputs.version }} | |
| run: | | |
| echo "=== Starting GPU wheel build ===" | |
| echo "Environment: ${{ matrix.pixi-environment }}" | |
| echo "Version: $SETUPTOOLS_SCM_PRETEND_VERSION" | |
| echo "Docker version:" | |
| docker --version || echo "Docker not available" | |
| echo "=== Checking Docker status ===" | |
| docker info || echo "Docker info failed" | |
| echo "=== Running wheel_build ===" | |
| pixi run -e ${{ matrix.pixi-environment }} wheel_build | |
| echo "=== Build completed ===" | |
| - name: Repair GPU wheel (Linux only) | |
| if: steps.should_run.outputs.run == 'true' && matrix.os == 'ubuntu-latest' | |
| run: pixi run -e ${{ matrix.pixi-environment }} wheel_repair | |
| - name: Print ccache stats | |
| if: steps.should_run.outputs.run == 'true' | |
| run: ccache -s | |
| - name: Upload wheel artifacts | |
| if: steps.should_run.outputs.run == 'true' | |
| uses: actions/upload-artifact@v6 | |
| with: | |
| name: wheels-gpu-${{ matrix.os }}-py${{ matrix.python-version }} | |
| path: dist/*.whl | |
| retention-days: 7 | |
| publish_cpu: | |
| name: Publish pymomentum-cpu to PyPI | |
| needs: [build_cpu_wheels_linux, build_cpu_wheels] | |
| runs-on: ubuntu-latest | |
| environment: | |
| name: pypi-cpu | |
| url: https://pypi.org/p/pymomentum-cpu | |
| permissions: | |
| id-token: write # IMPORTANT: mandatory for trusted publishing | |
| # Only publish on tag push or manual workflow dispatch with publish=true | |
| if: | | |
| (github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')) || | |
| (github.event_name == 'workflow_dispatch' && github.event.inputs.publish == 'true') | |
| steps: | |
| - name: Download CPU wheel artifacts | |
| uses: actions/download-artifact@v7 | |
| with: | |
| path: dist | |
| pattern: wheels-cpu-* | |
| merge-multiple: true | |
| - name: List CPU distributions | |
| run: ls -lh dist/ | |
| - name: Publish CPU to TestPyPI | |
| if: | | |
| github.event_name == 'workflow_dispatch' && github.event.inputs.publish == 'true' && github.event.inputs.test_pypi == 'true' | |
| uses: pypa/gh-action-pypi-publish@release/v1 | |
| with: | |
| repository-url: https://test.pypi.org/legacy/ | |
| - name: Publish CPU to PyPI | |
| if: | | |
| (github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')) || | |
| (github.event_name == 'workflow_dispatch' && github.event.inputs.publish == 'true' && github.event.inputs.test_pypi != 'true') | |
| uses: pypa/gh-action-pypi-publish@release/v1 | |
| publish_gpu: | |
| name: Publish pymomentum-gpu to PyPI | |
| needs: [build_gpu_wheels] | |
| runs-on: ubuntu-latest | |
| environment: | |
| name: pypi-gpu | |
| url: https://pypi.org/p/pymomentum-gpu | |
| permissions: | |
| id-token: write # IMPORTANT: mandatory for trusted publishing | |
| # Only publish on tag push or manual workflow dispatch with publish=true | |
| if: | | |
| (github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')) || | |
| (github.event_name == 'workflow_dispatch' && github.event.inputs.publish == 'true') | |
| steps: | |
| - name: Download GPU wheel artifacts | |
| uses: actions/download-artifact@v7 | |
| with: | |
| path: dist | |
| pattern: wheels-gpu-* | |
| merge-multiple: true | |
| - name: List GPU distributions | |
| run: ls -lh dist/ | |
| - name: Publish GPU to TestPyPI | |
| if: | | |
| github.event_name == 'workflow_dispatch' && github.event.inputs.publish == 'true' && github.event.inputs.test_pypi == 'true' | |
| uses: pypa/gh-action-pypi-publish@release/v1 | |
| with: | |
| repository-url: https://test.pypi.org/legacy/ | |
| - name: Publish GPU to PyPI | |
| if: | | |
| (github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')) || | |
| (github.event_name == 'workflow_dispatch' && github.event.inputs.publish == 'true' && github.event.inputs.test_pypi != 'true') | |
| uses: pypa/gh-action-pypi-publish@release/v1 |