Skip to content

Commit

Permalink
Add presubmit script to jax-tpu-embedding/sparsecore
Browse files Browse the repository at this point in the history
I added a presubmit job that builds a wheel from head and runs pytests with it. It's non-blocking right now, but we can change that if it's stable.

PiperOrigin-RevId: 715021685
  • Loading branch information
chandrasekhard2 authored and Google-ML-Automation committed Jan 21, 2025
1 parent b91e774 commit 23d54c8
Showing 1 changed file with 69 additions and 0 deletions.
69 changes: 69 additions & 0 deletions .github/workflows/cloud-tpu-ci-presubmit.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
# Cloud TPU CI (presubmit)
#
# This job currently runs as a non-blocking presubmit. It is experimental and is currently being
# tested to get to a stable state before we enable it as a blocking presubmit.
name: CI - Cloud TPU (presubmit)
on:
workflow_dispatch:
inputs:
halt-for-connection:
description: 'Should this workflow run wait for a remote connection?'
type: choice
required: true
default: 'no'
options:
- 'yes'
- 'no'
pull_request:
branches:
- main

# This should also be set to read-only in the project settings, but it's nice to
# document and enforce the permissions here.
permissions:
contents: read

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

jobs:
cloud-tpu-test:
if: github.event.repository.fork == false
strategy:
fail-fast: false # don't cancel all jobs on failure
matrix:
tpu: [
{type: "v6e-8", cores: "16", runner: "linux-x86-ct6e-180-8tpu"}
]
python-version: ["3.10"]

name: "TPU test (jax-tpu-embedding/sparsecore=head, ${{ matrix.tpu.type }})"

env:
JAXCI_PYTHON: python${{ matrix.python-version }}
JAXCI_TPU_CORES: ${{ matrix.tpu.cores }}
HERMETIC_PYTHON_VERSION: ${{ matrix.python-version }}

runs-on: ${{ matrix.tpu.runner }}
container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest"

timeout-minutes: 60

defaults:
run:
shell: bash -ex {0}
steps:
- name: Checkout jax-tpu-embedding at head
uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
# We need to mark the GitHub workspace as safe as otherwise git commands will fail.
- name: Mark GitHub workspace as safe
run: |
git config --global --add safe.directory "$GITHUB_WORKSPACE"
# Halt for testing
- name: Wait For Connection
uses: google-ml-infra/actions/ci_connection@main
with:
halt-dispatch-input: ${{ inputs.halt-for-connection }}
- name: Build and Run tests
run: bazel test //...

0 comments on commit 23d54c8

Please sign in to comment.