jax/.github/workflows/bazel_cpu.yml at main · mattjj/jax · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
# CI - Bazel CPU tests (RBE)
#
# This workflow runs the Bazel CPU tests with wheel dependencies. It can only be triggered by
# other workflows via `workflow_call`. It is used by the `CI - Wheel Tests (Continuous)` and
# `CI - Wheel Tests (Nightly/Release)` workflows to run the Bazel CPU tests.
#
# It consists of the following job:
# run-tests:
#    - Downloads the jax, jaxlib from a GCS bucket if build_jaxlib is false. Otherwise,
#      the artifacts are built from source.
#    - Executes the `run_bazel_test_cpu_rbe.sh` script, which performs the following actions:
#      - `build_jaxlib=wheel`: Runs the Bazel CPU tests with py_import dependencies.
#      - `build_jaxlib=false`: Runs the Bazel CPU tests with downloaded wheel dependencies.
#      - `build_jaxlib=true`: Runs the Bazel CPU tests with individual Bazel target dependencies.

name: CI - Bazel CPU tests with wheel dependencies (RBE)
permissions: {}
on:
  workflow_call:
    inputs:
      runner:
        description: "Which runner should the workflow run on?"
        type: string
        default: "linux-x86-n4-16"
      python:
        description: "Which python version to test?"
        type: string
        default: "3.12"
      enable-x64:
        description: "Should x64 mode be enabled?"
        type: string
        default: "0"
      halt-for-connection:
        description: 'Should this workflow run wait for a remote connection?'
        type: string
        default: 'no'
      build_jaxlib:
        description: 'Should jaxlib be built from source?'
        required: true
        type: string
      build_jax:
        description: 'Should jax be built from source?'
        required: true
        type: string
      bazel_command:
        description: "Whether to build or run test targets?"
        required: false
        type: string
        default: "test"
      gcs_download_uri:
        description: "GCS location prefix from where the artifacts should be downloaded"
        default: 'gs://general-ml-ci-transient/jax-github-actions/jax/${{ github.workflow }}/${{ github.run_number }}/${{ github.run_attempt }}'
        type: string

jobs:
  run-tests:
    defaults:
      run:
        # Explicitly set the shell to bash
        shell: bash
    runs-on: ${{ inputs.runner }}
    container: ${{ (contains(inputs.runner, 'linux-x86') && 'us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest') ||
                   (contains(inputs.runner, 'linux-arm64') && 'us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-arm64:latest') ||
                   (contains(inputs.runner, 'windows-x86') && null) }}
    env:
      JAXCI_HERMETIC_PYTHON_VERSION: ${{ inputs.python }}
      JAXCI_ENABLE_X64: ${{ inputs.enable-x64 }}
      JAXCI_BUILD_JAXLIB: ${{ inputs.build_jaxlib }}
      JAXCI_BUILD_JAX: ${{ inputs.build_jax }}
      JAXCI_BAZEL_CPU_RBE_MODE: "${{ inputs.bazel_command }}"

# Begin Presubmit Naming Check - name modification requires internal check to be updated
    name: "${{ (contains(inputs.runner, 'linux-x86') && 'linux x86') ||
        (contains(inputs.runner, 'linux-arm64') && 'linux arm64') ||
        (contains(inputs.runner, 'windows-x86') && 'windows x86') }}, Python=${{ inputs.python }}, x64=${{ inputs.enable-x64 }}, build_jax=${{ inputs.build_jax }}, build_jaxlib=${{ inputs.build_jaxlib }}"
# End Presubmit Naming Check github-cpu-presubmits

    steps:
      - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3  # v6.0.0
        with:
          persist-credentials: false
      - name: Download JAX CPU wheels
        if: inputs.build_jaxlib == 'false'
        uses: ./.github/actions/download-jax-cpu-wheels
        with:
          runner: ${{ inputs.runner }}
          python: ${{ inputs.python }}
          gcs_download_uri: ${{ inputs.gcs_download_uri }}
      # Halt for testing
      - name: Wait For Connection
        uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
        with:
          halt-dispatch-input: ${{ inputs.halt-for-connection }}
      - name: "Bazel CPU tests with build_jaxlib=${{ format('{0}', inputs.build_jaxlib) }}"
        timeout-minutes: 60
        run: ./ci/run_bazel_test_cpu_rbe.sh