1+ # Cloud TPU CI (presubmit)
2+ #
3+ # This job currently runs as a non-blocking presubmit. It is experimental and is currently being
4+ # tested to get to a stable state before we enable it as a blocking presubmit.
5+ name : CI - Cloud TPU (presubmit)
6+ on :
7+ workflow_dispatch :
8+ inputs :
9+ halt-for-connection :
10+ description : ' Should this workflow run wait for a remote connection?'
11+ type : choice
12+ required : true
13+ default : ' no'
14+ options :
15+ - ' yes'
16+ - ' no'
17+ pull_request :
18+ branches :
19+ - main
20+
21+ # This should also be set to read-only in the project settings, but it's nice to
22+ # document and enforce the permissions here.
23+ permissions :
24+ contents : read
25+
26+ concurrency :
27+ group : ${{ github.workflow }}-${{ github.head_ref || github.ref }}
28+ cancel-in-progress : true
29+
30+ jobs :
31+ cloud-tpu-test :
32+ if : github.event.repository.fork == false
33+ strategy :
34+ fail-fast : false # don't cancel all jobs on failure
35+ matrix :
36+ tpu : [
37+ {type: "v6e-4", cores: "8", runner: "arc-linux-x86-ct6e-180-4tpu"}
38+ ]
39+ python-version : ["3.10"]
40+
41+ name : " TPU test (jax-tpu-embedding/sparsecore=head, ${{ matrix.tpu.type }})"
42+
43+ env :
44+ JAXCI_PYTHON : python${{ matrix.python-version }}
45+ JAXCI_TPU_CORES : ${{ matrix.tpu.cores }}
46+ HERMETIC_PYTHON_VERSION : ${{ matrix.python-version }}
47+
48+ runs-on : ${{ matrix.tpu.runner }}
49+ container : " us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest"
50+
51+ timeout-minutes : 60
52+
53+ defaults :
54+ run :
55+ shell : bash -ex {0}
56+ steps :
57+ - name : Checkout jax-tpu-embedding at head
58+ uses : actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
59+ # We need to mark the GitHub workspace as safe as otherwise git commands will fail.
60+ - name : Mark GitHub workspace as safe
61+ run : |
62+ git config --global --add safe.directory "$GITHUB_WORKSPACE"
63+ # Halt for testing
64+ - name : Wait For Connection
65+ uses : google-ml-infra/actions/ci_connection@main
66+ with :
67+ halt-dispatch-input : ${{ inputs.halt-for-connection }}
68+ - name : Build and Run tests
69+ run : bazel test //...
0 commit comments