jax/.github/workflows/bazel_cuda_presubmit.yml at main · jax-ml/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
name: CI - Bazel CUDA tests

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
  push:
    branches:
      - main
      - 'release/**'

concurrency:
  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
  # Don't cancel in-progress jobs for main/release branches.
  cancel-in-progress: ${{ !contains(github.ref, 'release/') && github.ref != 'main' }}
permissions: {}
jobs:
  build-cuda-artifacts:
    uses: ./.github/workflows/build_artifacts.yml
    strategy:
        fail-fast: false # don't cancel all jobs on failure
        matrix:
          # Python values need to match the matrix stategy in the CUDA tests job below
          runner: ["linux-x86-n4-16"]
          artifact: ["jax-cuda-plugin", "jax-cuda-pjrt"]
          python: ["3.14",]
          cuda-version: ["13"]
    name: "Build ${{ format('{0}', 'CUDA') }} artifacts"
    with:
      runner: ${{ matrix.runner }}
      artifact: ${{ matrix.artifact }}
      python: ${{ matrix.python }}
      cuda-version: ${{ matrix.cuda-version }}
      clone_main_xla: 1
      upload_artifacts_to_gcs: false

  run-tests:
    if: github.event.repository.fork == false
    uses: ./.github/workflows/bazel_cuda.yml
# Begin Presubmit Naming Check - name modification requires internal check to be updated
    strategy:
        fail-fast: false # don't cancel all jobs on failure
        matrix:
          # Note: For testing external dependencies updates, change the python values
          # add 3.13 and 3.13-nogil here when debugging.
          python: ["3.11", "3.14"]
          runner: ["linux-x86-n4-16"]
          jaxlib-version: ["head", "pypi_latest"]
          enable-x64: [1, 0]
          cuda-version: ["12", "13"]
          exclude:
            # Exclude x64=1 on the oldest Python and x64=0 on the newest Python. As long as we have
            # coverage for one of each, we don't need to run both.
            - python: "3.11"
              enable-x64: 1
            - python: "3.11"
              enable-x64: 0
              jaxlib-version: "head"
            - python: "3.11"
              enable-x64: 0
              jaxlib-version: "pypi_latest"
              cuda-version: "13"
            - python: "3.14"
              enable-x64: 0
            - python: "3.14"
              enable-x64: 1
              jaxlib-version: "pypi_latest"
            # Exclude CUDA 12 on jaxlib head because it's too slow.
            - cuda-version: "12"
              jaxlib-version: "head"
    name: "Bazel single accelerator ${{ format('{0}', 'CUDA tests') }}"
# End Presubmit Naming Check github-cuda-presubmits
    with:
      runner: ${{ matrix.runner }}
      python: ${{ matrix.python }}
      download-jax-from-gcs: 0
      skip-download-jaxlib-and-plugins-from-gcs: 0
      jaxlib-version: ${{ matrix.jaxlib-version }}
      halt-for-connection: ${{ inputs.halt-for-connection }}
      cuda-version: ${{ matrix.cuda-version }}
      enable-x64:  ${{ matrix.enable-x64 }}
      # NOTE: For debugging wheel test change build_jaxlib and build_jax to 'wheel'.
      build_jaxlib: ${{ (matrix.jaxlib-version == 'head' && true) || false }}
      build_jax: 'true'
      run_multiaccelerator_tests: 'false'
      clone_main_xla: 1