lance-cuvs/.github/workflows/python-build.yml at main · lance-format/lance-cuvs · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
name: Python Wheels

on:
  workflow_dispatch:
  pull_request:
  push:
    branches:
      - main

concurrency:
  group: python-wheels-${{ github.event.pull_request.number || github.ref }}
  cancel-in-progress: true

env:
  UV_EXTRA_INDEX_URL: https://pypi.nvidia.com

jobs:
  build-wheels:
    runs-on: ubuntu-24.04
    container:
      image: quay.io/pypa/manylinux_2_28_x86_64

    steps:
      - uses: actions/checkout@v4

      - name: Install system dependencies
        run: ./tools/bootstrap_build_env.sh

      - uses: dtolnay/rust-toolchain@stable

      - uses: astral-sh/setup-uv@v6
        with:
          enable-cache: true

      - name: Use bundled Python 3.12
        run: echo "/opt/python/cp312-cp312/bin" >> "$GITHUB_PATH"

      - uses: Swatinem/rust-cache@v2
        with:
          prefix-key: manylinux-2.28-v1
          workspaces: |
            backends/cuvs_26_02 -> target

      - name: Build release wheels
        run: just build-wheels

      - uses: actions/upload-artifact@v4
        with:
          name: python-wheels-dist
          path: dist/*

  test-loader-wheel:
    needs: build-wheels
    runs-on: ubuntu-24.04

    steps:
      - uses: actions/checkout@v4

      - name: Install system dependencies
        run: |
          sudo apt-get update
          sudo DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
            just

      - uses: astral-sh/setup-uv@v6
        with:
          enable-cache: true

      - name: Install Python 3.12
        run: uv python install 3.12

      - uses: actions/download-artifact@v4
        with:
          name: python-wheels-dist
          path: dist

      - name: Test root wheel in a clean environment
        run: just test-loader-wheel

  test-gpu-wheel:
    needs: build-wheels
    runs-on: ubuntu-24.04-4x-gpu

    steps:
      - uses: actions/checkout@v4

      - name: Verify GPU runner
        run: |
          if ! command -v nvidia-smi >/dev/null 2>&1; then
            echo "This workflow requires a GPU-capable runner with nvidia-smi available." >&2
            exit 1
          fi
          nvidia-smi

      - name: Ensure CUDA toolkit
        run: |
          if ! command -v nvcc >/dev/null 2>&1; then
            curl -fsSL -o /tmp/cuda-keyring.deb \
              https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
            sudo dpkg -i /tmp/cuda-keyring.deb
            sudo apt-get update
            sudo DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
              cuda-toolkit-12-9
          fi
          export CUDA_HOME=/usr/local/cuda-12.9
          export PATH="$CUDA_HOME/bin:$PATH"
          echo "/usr/local/cuda-12.9/bin" >> "$GITHUB_PATH"
          echo "CUDA_HOME=/usr/local/cuda-12.9" >> "$GITHUB_ENV"
          nvcc --version

      - name: Install system dependencies
        run: |
          sudo apt-get update
          sudo DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
            just

      - uses: astral-sh/setup-uv@v6
        with:
          enable-cache: true

      - name: Install Python 3.12
        run: uv python install 3.12

      - uses: actions/download-artifact@v4
        with:
          name: python-wheels-dist
          path: dist

      - name: Test published wheels on a GPU runner
        env:
          LANCE_CUVS_REQUIRE_GPU: "1"
        run: just test-gpu-wheel