Skip to content

Commit 6cdb451

Browse files
SurbhiJainUSCSurbhiJainUSC
committed
Add jupyter notebook tests to Github CI tests
1 parent 6801c95 commit 6cdb451

7 files changed

Lines changed: 731 additions & 278 deletions

File tree

‎.github/workflows/build_and_test_maxtext.yml‎

Lines changed: 77 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,8 @@ jobs:
4141
name: Check for Documentation-Only Changes
4242
runs-on: ubuntu-latest
4343
outputs:
44-
skip_tests: ${{ steps.check.outputs.skip_tests }}
44+
run_tests: ${{ steps.check.outputs.run_tests }}
45+
run_notebooks: ${{ steps.check.outputs.run_notebooks }}
4546
steps:
4647
- uses: actions/checkout@v4
4748
with:
@@ -51,7 +52,8 @@ jobs:
5152
run: |
5253
if [ "${{ github.event_name }}" != "pull_request" ]; then
5354
echo "Not a pull request, running all tests"
54-
echo "skip_tests=false" >> $GITHUB_OUTPUT
55+
echo "run_tests=true" >> $GITHUB_OUTPUT
56+
echo "run_notebooks=true" >> $GITHUB_OUTPUT
5557
exit 0
5658
fi
5759
@@ -63,26 +65,59 @@ jobs:
6365
6466
if [ -z "$CHANGED_FILES" ]; then
6567
echo "No files changed"
66-
echo "skip_tests=false" >> $GITHUB_OUTPUT
67-
elif echo "$CHANGED_FILES" | grep -v -E '\.(md|ipynb)$' > /dev/null; then
68-
echo "Non-documentation files changed, running tests"
69-
echo "skip_tests=false" >> $GITHUB_OUTPUT
68+
echo "run_tests=true" >> $GITHUB_OUTPUT
69+
echo "run_notebooks=true" >> $GITHUB_OUTPUT
70+
fi
71+
72+
# Check for source code changes (anything not .md and not .ipynb)
73+
if echo "$CHANGED_FILES" | grep -v -E '\.(md|ipynb)$' > /dev/null; then
74+
echo "Source code files changed, enabling unit tests."
75+
echo "run_tests=true" >> $GITHUB_OUTPUT
76+
else
77+
echo "No source code changes, skipping unit tests."
78+
echo "run_tests=false" >> $GITHUB_OUTPUT
79+
fi
80+
81+
# Check for notebook (.ipynb) changes specifically
82+
if echo "$CHANGED_FILES" | grep '\.ipynb$' > /dev/null; then
83+
echo "Notebook files changed, enabling notebook run."
84+
echo "run_notebooks=true" >> $GITHUB_OUTPUT
7085
else
71-
echo "Only documentation (.md|ipynb) files changed, skipping tests"
72-
echo "skip_tests=true" >> $GITHUB_OUTPUT
86+
echo "No notebook changes, skipping notebook run."
87+
echo "run_notebooks=false" >> $GITHUB_OUTPUT
7388
fi
7489
7590
build_and_upload_maxtext_package:
7691
needs: doc_only_check
77-
if: needs.doc_only_check.outputs.skip_tests != 'true'
92+
# Run if either tests or notebooks need to run
93+
if: |
94+
needs.doc_only_check.outputs.run_tests == 'true' ||
95+
needs.doc_only_check.outputs.run_notebooks == 'true'
7896
uses: ./.github/workflows/build_package.yml
7997
with:
8098
device_type: tpu
8199
device_name: v4-8
82100
cloud_runner: linux-x86-n2-16-buildkit
83101

102+
maxtext_jupyter_notebooks:
103+
needs: build_and_upload_maxtext_package
104+
if: needs.doc_only_check.outputs.run_notebooks == 'true'
105+
uses: ./.github/workflows/run_jupyter_notebooks.yml
106+
strategy:
107+
fail-fast: false
108+
matrix:
109+
image_type: ["py312"]
110+
with:
111+
device_type: tpu
112+
device_name: v6e-4
113+
image_type: ${{ matrix.image_type }}
114+
cloud_runner: linux-x86-ct6e-180-4tpu
115+
secrets:
116+
HF_TOKEN: ${{ secrets.HF_TOKEN }}
117+
84118
maxtext_cpu_unit_tests:
85119
needs: build_and_upload_maxtext_package
120+
if: needs.doc_only_check.outputs.run_tests == 'true'
86121
uses: ./.github/workflows/run_tests_against_package.yml
87122
strategy:
88123
fail-fast: false # don't cancel all jobs on failure
@@ -104,6 +139,7 @@ jobs:
104139

105140
maxtext_tpu_unit_tests:
106141
needs: build_and_upload_maxtext_package
142+
if: needs.doc_only_check.outputs.run_tests == 'true'
107143
uses: ./.github/workflows/run_tests_against_package.yml
108144
strategy:
109145
fail-fast: false
@@ -122,6 +158,7 @@ jobs:
122158

123159
maxtext_tpu_integration_tests:
124160
needs: build_and_upload_maxtext_package
161+
if: needs.doc_only_check.outputs.run_tests == 'true'
125162
uses: ./.github/workflows/run_tests_against_package.yml
126163
strategy:
127164
fail-fast: false
@@ -140,6 +177,7 @@ jobs:
140177

141178
maxtext_tpu_pathways_unit_tests:
142179
needs: build_and_upload_maxtext_package
180+
if: needs.doc_only_check.outputs.run_tests == 'true'
143181
uses: ./.github/workflows/run_pathways_tests.yml
144182
strategy:
145183
fail-fast: false
@@ -158,6 +196,7 @@ jobs:
158196

159197
maxtext_tpu_pathways_integration_tests:
160198
needs: build_and_upload_maxtext_package
199+
if: needs.doc_only_check.outputs.run_tests == 'true'
161200
uses: ./.github/workflows/run_pathways_tests.yml
162201
strategy:
163202
fail-fast: false
@@ -176,6 +215,7 @@ jobs:
176215

177216
maxtext_gpu_unit_tests:
178217
needs: build_and_upload_maxtext_package
218+
if: needs.doc_only_check.outputs.run_tests == 'true'
179219
uses: ./.github/workflows/run_tests_against_package.yml
180220
strategy:
181221
fail-fast: false
@@ -196,6 +236,7 @@ jobs:
196236

197237
maxtext_gpu_integration_tests:
198238
needs: build_and_upload_maxtext_package
239+
if: needs.doc_only_check.outputs.run_tests == 'true'
199240
uses: ./.github/workflows/run_tests_against_package.yml
200241
strategy:
201242
fail-fast: false
@@ -223,7 +264,7 @@ jobs:
223264
- name: Check test results
224265
run: |
225266
# If doc-only, all tests should be skipped
226-
if [ "${{ needs.doc_only_check.outputs.skip_tests }}" == "true" ]; then
267+
if [ "${{ needs.doc_only_check.outputs.run_tests }}" == "false" ]; then
227268
echo "Documentation-only changes detected, tests were skipped"
228269
exit 0
229270
fi
@@ -246,9 +287,34 @@ jobs:
246287
247288
echo "All required tests passed successfully"
248289
290+
all_notebooks_passed:
291+
name: All Notebooks Passed
292+
needs: [doc_only_check, build_and_upload_maxtext_package, maxtext_jupyter_notebooks]
293+
if: always()
294+
runs-on: ubuntu-latest
295+
steps:
296+
- name: Check notebooks results
297+
run: |
298+
if [ "${{ needs.doc_only_check.outputs.run_notebooks }}" == "false" ]; then
299+
echo "Non-notebook changes detected, runs were skipped"
300+
exit 0
301+
fi
302+
303+
# Otherwise, check that build and notebooks run passed or were skipped
304+
echo "Build result: ${{ needs.build_and_upload_maxtext_package.result }}"
305+
echo "Jupyter Notebooks result: ${{ needs.maxtext_jupyter_notebooks.result }}"
306+
307+
# Fail only if any job failed or was cancelled (skipped is OK)
308+
if [ "${{ contains(needs.*.result, 'failure') }}" == "true" ] || [ "${{ contains(needs.*.result, 'cancelled') }}" == "true" ]; then
309+
echo "One or more jobs failed or were cancelled"
310+
exit 1
311+
fi
312+
313+
echo "All required notebooks passed successfully"
314+
249315
notify_failure:
250316
name: Notify failed build # creates an issue or modifies last open existing issue for failed build
251-
needs: [maxtext_cpu_unit_tests, maxtext_tpu_unit_tests, maxtext_tpu_integration_tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, maxtext_gpu_unit_tests, maxtext_gpu_integration_tests]
317+
needs: [maxtext_jupyter_notebooks, maxtext_cpu_unit_tests, maxtext_tpu_unit_tests, maxtext_tpu_integration_tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, maxtext_gpu_unit_tests, maxtext_gpu_integration_tests]
252318
if: ${{ always() }}
253319
runs-on: ubuntu-latest
254320
permissions:

‎.github/workflows/run_jupyter_notebooks.yml‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
# Copyright 2026 Google LLC
2+
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# This file defines a module for running jupyter notebooks against the built maxtext package.
16+
17+
name: Run Jupyter Notebooks
18+
19+
on:
20+
workflow_call:
21+
inputs:
22+
device_type:
23+
required: true
24+
type: string
25+
device_name:
26+
required: true
27+
type: string
28+
image_type:
29+
required: false
30+
type: string
31+
cloud_runner:
32+
required: false
33+
type: string
34+
secrets:
35+
HF_TOKEN:
36+
required: true
37+
38+
permissions:
39+
contents: read
40+
jobs:
41+
run:
42+
runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }}
43+
container:
44+
image: gcr.io/tpu-prod-env-multipod/maxtext-unit-test-${{ inputs.device_type == 'cpu' && 'tpu' || inputs.device_type }}:${{ inputs.image_type != '' && inputs.image_type }}
45+
steps:
46+
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
47+
- name: Download the MaxText wheel
48+
uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0
49+
with:
50+
name: maxtext-wheel
51+
- name: Install MaxText and Dependencies
52+
shell: bash
53+
run: |
54+
python3 -m uv venv --seed
55+
source .venv/bin/activate
56+
57+
# Install MaxText package
58+
maxtext_wheel=$(ls maxtext-*-py3-none-any.whl 2>/dev/null)
59+
uv pip install ${maxtext_wheel}[${MAXTEXT_PACKAGE_EXTRA}] --resolution=lowest
60+
uv pip install -r src/install_maxtext_extra_deps/extra_deps_from_github.txt
61+
62+
# Install dependencies for running notebooks
63+
uv pip install papermill ipykernel ipywidgets
64+
.venv/bin/python3 -m ipykernel install --user --name maxtext_venv
65+
66+
# Install Tunix for post-training notebooks
67+
uv pip install git+https://github.com/google/tunix
68+
69+
# Install vllm for post-training notebooks
70+
git clone https://github.com/vllm-project/vllm.git
71+
VLLM_TARGET_DEVICE="tpu" uv pip install ./vllm
72+
73+
# Install tpu-inference for post-training notebooks
74+
git clone https://github.com/vllm-project/tpu-inference.git
75+
uv pip install ./tpu-inference
76+
77+
uv pip install --no-deps qwix==0.1.4
78+
uv pip install --no-deps protobuf==5.29.5
79+
python3 -m pip freeze
80+
- name: Run Post-Training Notebooks
81+
shell: bash
82+
env:
83+
HF_TOKEN: ${{ secrets.HF_TOKEN }}
84+
run: |
85+
MAXTEXT_REPO_ROOT=$(pwd)
86+
MAXTEXT_NOTEBOOKS_ROOT="$MAXTEXT_REPO_ROOT/src/MaxText/examples"
87+
88+
for notebook in "$MAXTEXT_NOTEBOOKS_ROOT"/{sft,rl}*.ipynb; do
89+
filename=$(basename "$notebook")
90+
output_name="${filename%.ipynb}_output.ipynb"
91+
92+
echo "------------------------------------------------------"
93+
echo "Running $filename ..."
94+
echo "------------------------------------------------------"
95+
96+
.venv/bin/papermill "$notebook" "$output_name" -k maxtext_venv
97+
done
98+
- name: Upload Outputs
99+
if: always()
100+
uses: actions/upload-artifact@v4
101+
with:
102+
name: notebook-outputs-${{ inputs.device_name }}
103+
path: ./*_output.ipynb

0 commit comments

Comments
 (0)