Build wheels and deploy #5

Workflow file for this run

.github/workflows/publish.yml at d12ba9c

	# This workflow will:
	# - Create a new Github release
	# - Build wheels for supported architectures
	# - Deploy the wheels to the Github release
	# - Release the static code to PyPi
	# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries

	name: Build wheels and deploy

	on:
	create:
	tags:
	- v*

	jobs:

	setup_release:
	name: Create Release
	runs-on: ubuntu-latest
	steps:
	- name: Get the tag version
	id: extract_branch
	run: echo ::set-output name=branch::${GITHUB_REF#refs/tags/}
	shell: bash

	- name: Create Release
	id: create_release
	uses: actions/create-release@v1
	env:
	GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
	with:
	tag_name: ${{ steps.extract_branch.outputs.branch }}
	release_name: ${{ steps.extract_branch.outputs.branch }}

	build_wheels:
	name: Build Wheel
	needs: setup_release
	runs-on: ${{ matrix.os }}

	strategy:
	fail-fast: false
	matrix:
	# Using ubuntu-20.04 instead of 22.04 for more compatibility (glibc). Ideally we'd use the
	# manylinux docker image, but I haven't figured out how to install CUDA on manylinux.
	os: [ubuntu-20.04]
	python-version: ['3.9', '3.10', '3.11', '3.12']
	include:
	# 0.4.24 support 11.8 and 12.3
	- jax-version: '0.4.24'
	cuda-version: '11.8.0'
	- jax-version: '0.4.24'
	cuda-version: '12.3.1'
	# 0.4.20 support 11.8, 12.2, 12.3
	- jax-version: '0.4.20'
	cuda-version: '11.8.0'
	- jax-version: '0.4.20'
	cuda-version: '12.2.2'
	- jax-version: '0.4.20'
	cuda-version: '12.3.1'


	steps:
	- name: Checkout
	uses: actions/checkout@v3

	- name: Set up Python
	uses: actions/setup-python@v4
	with:
	python-version: ${{ matrix.python-version }}

	- name: Set CUDA and PyTorch versions
	run: \|
	echo "MATRIX_CUDA_VERSION=$(echo ${{ matrix.cuda-version }} \| awk -F \. {'print $1 $2'})" >> $GITHUB_ENV
	echo "MATRIX_JAX_VERSION=$(echo ${{ matrix.jax-version }} \| awk -F \. {'print $1 "." $2'})" >> $GITHUB_ENV

	- name: Free up disk space
	if: ${{ runner.os == 'Linux' }}
	# https://github.com/easimon/maximize-build-space/blob/master/action.yml
	# https://github.com/easimon/maximize-build-space/tree/test-report
	run: \|
	sudo rm -rf /usr/share/dotnet
	sudo rm -rf /opt/ghc
	sudo rm -rf /opt/hostedtoolcache/CodeQL

	- name: Set up swap space
	if: runner.os == 'Linux'
	uses: pierotofy/set-swap-space@v1.0
	with:
	swap-size-gb: 10

	- name: Install CUDA ${{ matrix.cuda-version }}
	if: ${{ matrix.cuda-version != 'cpu' }}
	uses: Jimver/cuda-toolkit@v0.2.14
	id: cuda-toolkit
	with:
	cuda: ${{ matrix.cuda-version }}
	linux-local-args: '["--toolkit"]'
	# default method is "local", and we're hitting some error with caching for CUDA 11.8 and 12.1
	# method: ${{ (matrix.cuda-version == '11.8.0' \|\| matrix.cuda-version == '12.1.0') && 'network' \|\| 'local' }}
	method: 'network'
	# We need the cuda libraries (e.g. cuSparse, cuSolver) for compiling PyTorch extensions,
	# not just nvcc
	# sub-packages: '["nvcc"]'

	- name: Install Jax ${{ matrix.torch-version }}+cu${{ matrix.cuda-version }}
	run: \|
	pip install --upgrade pip
	pip install --upgrade 'jax[cuda12_local] == ${{ matrix.jax-version }}' -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
	shell:
	bash

	- name: Build wheel
	run: \|
	# We want setuptools >= 49.6.0 otherwise we can't compile the extension if system CUDA version is 11.7 and pytorch cuda version is 11.6
	# https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/torch/utils/cpp_extension.py#L810
	# However this still fails so I'm using a newer version of setuptools
	pip install setuptools==68.0.0
	# setuptools-cuda-cpp on pypi has a bug that breaks ninja
	pip install git+https://github.com/nshepperd/setuptools-cuda-cpp
	pip install ninja packaging wheel pybind11
	export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
	export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/cuda/lib64:$LD_LIBRARY_PATH
	# Set MAX_JOBS to allocate 8GB per job, which should be enough to build comfortably
	RAM=$(free -m \| grep Mem \| awk '{print $4}')
	export MAX_JOBS="$(( RAM/8000 < 1 ? 1 : RAM/8000 ))"
	echo "Building with ${MAX_JOBS} jobs to fit within ${RAM}MB"
	python setup.py bdist_wheel --dist-dir=dist
	tmpname=cu${MATRIX_CUDA_VERSION}jax${{ matrix.jax-version }}
	wheel_name=$(ls dist/*whl \| xargs -n 1 basename \| sed "s/-/+$tmpname-/2")
	ls dist/*whl \|xargs -I {} mv {} dist/${wheel_name}
	echo "wheel_name=${wheel_name}" >> $GITHUB_ENV
	shell:
	bash

	- name: Log Built Wheels
	run: \|
	ls dist

	- name: Get the tag version
	id: extract_branch
	run: echo ::set-output name=branch::${GITHUB_REF#refs/tags/}

	- name: Get Release with tag
	id: get_current_release
	uses: joutvhu/get-release@v1
	with:
	tag_name: ${{ steps.extract_branch.outputs.branch }}
	env:
	GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

	- name: Upload Release Asset
	id: upload_release_asset
	uses: actions/upload-release-asset@v1
	env:
	GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
	with:
	upload_url: ${{ steps.get_current_release.outputs.upload_url }}
	asset_path: ./dist/${{env.wheel_name}}
	asset_name: ${{env.wheel_name}}
	asset_content_type: application/*

	# publish_package:
	# name: Publish package
	# needs: [build_wheels]

	# runs-on: ubuntu-latest

	# steps:
	# - uses: actions/checkout@v3

	# - uses: actions/setup-python@v4
	# with:
	# python-version: '3.10'

	# - name: Install dependencies
	# run: \|
	# pip install ninja packaging setuptools wheel twine
	# # We don't want to download anything CUDA-related here
	# pip install torch --index-url https://download.pytorch.org/whl/cpu

	# - name: Build core package
	# env:
	# FLASH_ATTENTION_SKIP_CUDA_BUILD: "TRUE"
	# run: \|
	# python setup.py sdist --dist-dir=dist

	# - name: Deploy
	# env:
	# TWINE_USERNAME: "__token__"
	# TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
	# run: \|
	# python -m twine upload dist/*

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Build wheels and deploy #5

Workflow file

Build wheels and deploy #5

Jobs

Run details

Workflow file for this run