diff --git a/.github/workflows/build_flashmm.yaml b/.github/workflows/build_flashmm.yaml new file mode 100644 index 0000000..7c9567f --- /dev/null +++ b/.github/workflows/build_flashmm.yaml @@ -0,0 +1,15 @@ +name: Build Wheels + +on: + push: + # push on github tags that start with "wheel" + tags: + - "wheel*" + +jobs: + flashmm_wheels: + uses: + ./.github/workflows/workflow_build_cuda_wheels.yaml + with: # ${{ github.workspace }} + setuppypath: "csrc/flashmm" + secrets: inherit \ No newline at end of file diff --git a/.github/workflows/workflow_build_cuda_wheels.yaml b/.github/workflows/workflow_build_cuda_wheels.yaml new file mode 100644 index 0000000..f73d04d --- /dev/null +++ b/.github/workflows/workflow_build_cuda_wheels.yaml @@ -0,0 +1,120 @@ +name: install_wheels_cuda + +on: + workflow_call: + inputs: + setuppypath: + required: true + type: string + description: "path to the dir of the setup.py that should be built." + +env: + SETUPPYPATH: ${{ inputs.setuppypath }} + MAX_JOBS: "2" + +jobs: + build_wheels: + name: Build Cuda Wheels for torch + runs-on: ${{ matrix.os }} + + strategy: + matrix: + os: [ubuntu-20.04, windows-latest] + pyver: ["3.8", "3.9", "3.10", "3.11"] + cuda: ["11.8.0", "12.1.1"] + defaults: + run: + shell: pwsh + env: + CUDA_VERSION: ${{ matrix.cuda }} + # define where the setup.py is located + + steps: + - name: Free Disk Space + uses: jlumbroso/free-disk-space@v1.3.0 + if: runner.os == 'Linux' + with: + tool-cache: false + android: true + dotnet: true + haskell: true + large-packages: false + docker-images: true + swap-storage: false + + - uses: actions/checkout@v3 + + - uses: actions/setup-python@v3 + with: + python-version: ${{ matrix.pyver }} + + - name: Setup Mamba + uses: conda-incubator/setup-miniconda@v2.2.0 + with: + activate-environment: "build" + python-version: ${{ matrix.pyver }} + miniforge-variant: Mambaforge + miniforge-version: latest + use-mamba: true + add-pip-as-python-dependency: true + auto-activate-base: false + + - name: Install Dependencies + run: | + # Install CUDA toolkit + mamba install -y 'cuda' -c "nvidia/label/cuda-${env:CUDA_VERSION}" + + # Env variables + $env:CUDA_PATH = $env:CONDA_PREFIX + $env:CUDA_HOME = $env:CONDA_PREFIX + + # Install torch + $cudaVersion = $env:CUDA_VERSION.Replace('.', '') + $cudaVersionPytorch = $cudaVersion.Substring(0, $cudaVersion.Length - 1) + if ([int]$cudaVersionPytorch -gt 118) { $pytorchVersion = "torch==2.1.0" } else {$pytorchVersion = "torch==2.0.1"} + python3 -m pip install --upgrade --no-cache-dir $pytorchVersion+cu$cudaVersionPytorch --index-url https://download.pytorch.org/whl/cu$cudaVersionPytorch + python3 -m pip install build setuptools wheel ninja numpy + + # Print version information + python3 --version + python3 -c "import torch; print('PyTorch:', torch.__version__)" + python3 -c "import torch; print('CUDA:', torch.version.cuda)" + python3 -c "import os; print('CUDA_HOME:', os.getenv('CUDA_HOME', None))" + python3 -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)" + + - name: Build Wheel + run: | + $env:CUDA_PATH = $env:CONDA_PREFIX + $env:CUDA_HOME = $env:CONDA_PREFIX + + # build the wheel + cd ${{ github.workspace }}/$env:SETUPPYPATH + python3 setup.py sdist bdist_wheel + + # rename wheel to reflect cuda version and torch version bash: dist/*.whl + $WHEEL_PATH = python3 -c "if True: \ + import os, shutil, glob, re, torch;\ + torch_name = '+torch' + (torch.__version__.split('+')[0] + 'cu' + torch.version.cuda).replace('.', '');\ + file = glob.glob(os.path.join('dist', '*.whl'))[0];\ + current_dir = os.getcwd();\ + insert_pos = re.search('-cp[0-9]+-cp[0-9]+-', file);\ + torch_name = '' if insert_pos is None else torch_name;\ + insert_pos = re.search('-py3', file).start() if insert_pos is None else insert_pos.start();\ + new_name = f'{file[:insert_pos]}{torch_name}{file[insert_pos:]}';\ + print(os.path.join(os.getcwd(), new_name));\ + shutil.move(file, new_name);" + + # append to powershell env + echo "WHEEL_PATH=$WHEEL_PATH" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append + + - name: Release with Notes + uses: softprops/action-gh-release@v1 + if: startsWith(github.ref, 'refs/tags/') + with: + files: | + ${{ env.WHEEL_PATH }} + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # from previous step + WHEEL_PATH: ${{ env.WHEEL_PATH }} +