Skip to content

Commit

Permalink
[nnx] add flaxlib
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 680942447
  • Loading branch information
Cristian Garcia authored and Flax Authors committed Oct 1, 2024
1 parent 62a6a25 commit dfc86d5
Show file tree
Hide file tree
Showing 14 changed files with 875 additions and 19 deletions.
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
# This workflows will upload a Python Package using Twine when a release is created
# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries

name: Upload Python Package
name: Flax - Build and upload to PyPI

on:
release:
types: [created]
types: [published]

jobs:
deploy:
Expand Down
28 changes: 11 additions & 17 deletions .github/workflows/build.yml → .github/workflows/flax_test.yml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
# This workflow will install Python dependencies, run tests and lint with a variety of Python versions
# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions

name: Build
name: Flax - Test

on:
push:
Expand Down Expand Up @@ -70,7 +70,7 @@ jobs:
uses: actions/setup-python@v4
with:
python-version: ${{ matrix.python-version }}
- uses: yezz123/setup-uv@v4
- uses: astral-sh/setup-uv@v2
with:
uv-version: "0.3.0"
- name: Install standalone dependencies only
Expand Down Expand Up @@ -104,23 +104,17 @@ jobs:
uses: actions/setup-python@v4
with:
python-version: ${{ matrix.python-version }}
- uses: yezz123/setup-uv@v4
- name: Setup uv
uses: astral-sh/setup-uv@v2
with:
uv-version: "0.3.0"
- name: Cached virtual environment
id: venv_cache
uses: actions/cache@v3
with:
path: .venv
key: pip-${{ steps.setup_python.outputs.python-version }}-${{ hashFiles('uv.lock') }}
- name: Install Dependencies for cache
if: steps.venv_cache.outputs.cache-hit != 'true'
run: |
if [ -d ".venv" ]; then rm -rf .venv; fi
uv sync --locked --all-extras
- name: Check lockfile
version: "0.3.0"
- name: Setup Rust (flaxlib)
uses: actions-rust-lang/setup-rust-toolchain@v1

- name: Install dependencies
run: |
uv sync --locked --all-extras
uv sync --locked --extra all --extra testing --extra docs
uv pip install ./flaxlib
- name: Install JAX
run: |
if [[ "${{ matrix.jax-version }}" == "newest" ]]; then
Expand Down
94 changes: 94 additions & 0 deletions .github/workflows/flaxlib_publish.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
name: Flaxlib - Build and upload to PyPI

# for testing only:
on:
push:
branches: [main]
paths: ['flaxlib/**']
release:
types: [published]

jobs:
build_wheels:
if: github.event_name == 'push' && contains(github.event.head_commit.modified, 'flaxlib/')
name: Build wheels on ${{ matrix.os }}
runs-on: ${{ matrix.os }}
strategy:
matrix:
# macos-13 is an intel runner, macos-14 is apple silicon
os: [ubuntu-latest, windows-latest, macos-13, macos-14]

steps:
- uses: actions/checkout@v4

- uses: actions/setup-python@v5

- name: Setup Rust
uses: actions-rust-lang/setup-rust-toolchain@v1

- name: Install cibuildwheel
run: python -m pip install cibuildwheel==2.21.0

- name: Build wheels
run: python -m cibuildwheel --output-dir ./flaxlib/wheelhouse ./flaxlib
env:
# rust doesn't seem to be available for musl linux on i686
CIBW_SKIP: "*-musllinux_i686"
CIBW_ENVIRONMENT: 'PATH="$HOME/.cargo/bin:$PATH" CARGO_TERM_COLOR="always"'
CIBW_ENVIRONMENT_WINDOWS: 'PATH="$UserProfile\.cargo\bin;$PATH"'
CIBW_BEFORE_BUILD: rustup show
CIBW_BEFORE_BUILD_LINUX: |
curl https://sh.rustup.rs -sSf | sh -s -- --default-toolchain=stable --profile=minimal -y &&
rustup show
- uses: actions/upload-artifact@v4
with:
name: cibw-wheels-${{ matrix.os }}-${{ strategy.job-index }}
path: ./flaxlib/wheelhouse/*.whl

build_sdist:
if: github.event_name == 'push' && contains(github.event.head_commit.modified, 'flaxlib/')
name: Build source distribution
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4

- name: Setup Rust
uses: actions-rust-lang/setup-rust-toolchain@v1

- name: Build sdist
run: pipx run build --sdist flaxlib

- uses: actions/upload-artifact@v4
with:
name: cibw-sdist
path: ./flaxlib/dist/*.tar.gz

upload_pypi:
if: github.event_name == 'push' && contains(github.event.head_commit.modified, 'flaxlib/')
name: Upload to PyPI
needs: [build_wheels, build_sdist]
runs-on: ubuntu-latest
permissions:
id-token: write
steps:
- uses: actions/setup-python@v1
with:
python-version: '3.x'
- name: Install dependencies
run: |
python -m pip install --upgrade pip
pip install setuptools build wheel twine
- uses: actions/download-artifact@v4
with:
# unpacks all CIBW artifacts into dist/
pattern: cibw-*
path: ./flaxlib/dist
merge-multiple: true

- name: Build and publish
env:
TWINE_USERNAME: __token__
TWINE_PASSWORD: ${{ secrets.FLAXLIB_PYPI_TOKEN }}
run: |
twine upload flaxlib/dist/*
75 changes: 75 additions & 0 deletions flaxlib/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
/target

# Byte-compiled / optimized / DLL files
__pycache__/
.pytest_cache/
*.py[cod]

# C extensions
*.so

# Distribution / packaging
.Python
.venv/
env/
bin/
build/
develop-eggs/
dist/
eggs/
lib/
lib64/
parts/
sdist/
var/
include/
man/
venv/
*.egg-info/
.installed.cfg
*.egg

# Installer logs
pip-log.txt
pip-delete-this-directory.txt
pip-selfcheck.json

# Unit test / coverage reports
htmlcov/
.tox/
.coverage
.cache
nosetests.xml
coverage.xml

# Translations
*.mo

# Mr Developer
.mr.developer.cfg
.project
.pydevproject

# Rope
.ropeproject

# Django stuff:
*.log
*.pot

.DS_Store

# Sphinx documentation
docs/_build/

# PyCharm
.idea/

# VSCode
.vscode/

# Pyenv
.python-version

# cibuildwheel
/wheelhouse
Loading

0 comments on commit dfc86d5

Please sign in to comment.