tokamax/.github/workflows/ci-build.yml at 795635dedf29119c28e9030ddb0f725ee667c00b · openxla/tokamax · 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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
name: CI

on:
  # Trigger the workflow on push or pull request,
  # but only for the main branch
  push:
    branches:
      - main
  pull_request:
    branches:
      - main
  # Event manually triggered (in the Github UI)
  workflow_dispatch:

# Constants for the workflow that are shared across jobs.
env:
  JAX_ENABLE_X64: 0
  JAX_TRACEBACK_FILTERING: off

permissions:
  contents: read
# We should cancel the previous run if it's still running and there is a new commit.
concurrency:
  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
  cancel-in-progress: true

# These jobs should build Tokamax and run the tests on Google Cloud resources.
jobs:
  attention-tests:
    strategy:
      matrix:
        runner: ['linux-x86-a3-8g-h100-1gpu', 'linux-x86-ct6e-44-1tpu']
        pytest_command: [
                         'base_test.py',
                         'api_test.py',
                         'pallas_mosaic_gpu_test.py',
                         'pallas_triton_test.py',
                         'jax_nn_test.py',
                         'xla_chunked_test.py',
                        ]
        include:
          - runner: 'linux-x86-a3-8g-h100-1gpu'
            device: 'cuda'
          - runner: 'linux-x86-ct6e-44-1tpu'
            device: 'tpu'
    runs-on: ${{ matrix.runner }}
    defaults:
      run:
        shell: bash
    container: 'us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cuda13.2-cudnn9.15'
    timeout-minutes: 300
    steps:
      - name: Checkout Tokamax
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
        with:
          persist-credentials: false

      - name: Install dependencies
        run: |
          python3.12 -m uv pip install --upgrade pip
          python3.12 -m uv pip --version
          python3.12 -m uv pip install -e .[${{ matrix.device }},test]
          python3.12 -m uv pip freeze
      - name: Test Tokamax ops with unittest
        run: |
          pytest -s -m "not long" --ignore-glob="*/test_base.py" tokamax/_src/ops/attention/${{ matrix.pytest_command }}

  tokamax-presubmit-ops-tests:
    strategy:
      matrix:
        runner: ['linux-x86-a3-8g-h100-1gpu', 'linux-x86-ct6e-44-1tpu']
        include:
          - runner: 'linux-x86-a3-8g-h100-1gpu'
            device: 'cuda'
          - runner: 'linux-x86-ct6e-44-1tpu'
            device: 'tpu'
    runs-on: ${{ matrix.runner }}
    defaults:
      run:
        shell: bash
    container:
    # Docker image will include CUDA and cuDNN for Tokamax.
      image: 'us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cuda13.2-cudnn9.15'
    timeout-minutes: 300
    steps:
      - name: Checkout Tokamax
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
        with:
          persist-credentials: false

      - name: Install dependencies
        run: |
          python3.12 -m uv pip install --upgrade pip
          python3.12 -m uv pip --version
          python3.12 -m uv pip install -e .[${{ matrix.device }},test]
          python3.12 -m uv pip freeze
      - name: Test Tokamax ops with unittest
        run: |
          pytest -s -m "not long" --ignore-glob="*/test_base.py" --ignore=tokamax/_src/ops/attention/ --ignore-glob="*/splash_attention" tokamax/_src/ops

  tokamax-presubmit:
    strategy:
      matrix:
        runner: ['linux-x86-a3-8g-h100-1gpu', 'linux-x86-ct6e-44-1tpu']
        include:
          - runner: 'linux-x86-a3-8g-h100-1gpu'
            device: 'cuda'
          - runner: 'linux-x86-ct6e-44-1tpu'
            device: 'tpu'
    runs-on: ${{ matrix.runner }}
    defaults:
      run:
        shell: bash
    container:
    # Docker image will include CUDA and cuDNN for Tokamax.
      image: 'us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cuda13.2-cudnn9.15'
    timeout-minutes: 300
    steps:
      - name: Checkout Tokamax
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
        with:
          persist-credentials: false

      - name: Install dependencies
        run: |
          echo "Installing dependencies on device ${{ matrix.device }}"
          python3.12 -m uv pip install --upgrade pip
          python3.12 -m uv pip --version
          python3.12 -m uv pip install -e .[${{ matrix.device }},test]
          python3.12 -m uv pip freeze
      - name: Test Tokamax ops with unittest
        # TODO: Disable benchmarking tests for presubmit until container is fixed
        run: |
          pytest -s -vv -m "not long" --ignore-glob="tokamax/_src/ops/*" --ignore-glob="tokamax/_src/benchmarking_test*" tokamax