build

Test multi-host runner #3

	name: build

	on:
	pull_request:
	branches:
	- main

	permissions:
	contents: read
	actions: write # to cancel previous workflows

	concurrency:
	group: ${{ github.workflow }}-${{ github.head_ref \|\| github.ref }}
	cancel-in-progress: true

	jobs:
	build-checkpoint:
	name: "build-checkpoint multi-host"
	runs-on: linux-g2-16-l4-1gpu-x4
	container: python:3.11
	defaults:
	run:
	working-directory: checkpoint
	steps:
	- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
	- name: Extract branch name
	shell: bash
	run: echo "branch=${GITHUB_HEAD_REF:-${GITHUB_REF#refs/heads/}}" >> $GITHUB_OUTPUT
	id: extract_branch
	- name: Install dependencies
	run: \|
	sudo apt-get update
	sudo apt-get install -y protobuf-compiler

	pip install tensorflow

	protoc -I=. --python_out=. $(find orbax/experimental/model/ -name "*.proto")

	pip install -e .

	pip install .[testing] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

	pip install -U jax jaxlib
	- name: Test with pytest
	run: \|
	pytest orbax/experimental/model/core/python/*_test.py

	pytest orbax/experimental/model/tf2obm/*_test.py

	pytest orbax/experimental/model/jax2obm/ \
	--ignore=orbax/experimental/model/jax2obm/main_lib_test.py \
	--ignore=orbax/experimental/model/jax2obm/sharding_test.py \
	--ignore=orbax/experimental/model/jax2obm/jax_to_polymorphic_function_test.py

Provide feedback