Skip to content

Commit 6bcd411

Browse files
Google-ML-AutomationGoogle-ML-Automation
committed
Merge pull request #2921 from AI-Hypercomputer:jupyter_notebook_tests
PiperOrigin-RevId: 856917509
2 parents 4d99240 + cdd3dd8 commit 6bcd411

7 files changed

Lines changed: 749 additions & 279 deletions

File tree

‎.github/workflows/build_and_test_maxtext.yml‎

Lines changed: 80 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,62 @@ 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+
exit 0
71+
fi
72+
73+
# Check for source code changes (anything not .md and not .ipynb)
74+
if echo "$CHANGED_FILES" | grep -v -E '\.(md|ipynb)$' > /dev/null; then
75+
echo "Source code files changed, enabling unit tests."
76+
echo "run_tests=true" >> $GITHUB_OUTPUT
7077
else
71-
echo "Only documentation (.md|ipynb) files changed, skipping tests"
72-
echo "skip_tests=true" >> $GITHUB_OUTPUT
78+
echo "No source code changes, skipping unit tests."
79+
echo "run_tests=false" >> $GITHUB_OUTPUT
7380
fi
7481
82+
# Check for notebook (.ipynb) changes specifically
83+
if echo "$CHANGED_FILES" | grep '\.ipynb$' > /dev/null; then
84+
echo "Notebook files changed, enabling notebook run."
85+
echo "run_notebooks=true" >> $GITHUB_OUTPUT
86+
else
87+
echo "No notebook changes, skipping notebook run."
88+
echo "run_notebooks=false" >> $GITHUB_OUTPUT
89+
fi
90+
91+
exit 0
92+
7593
build_and_upload_maxtext_package:
7694
needs: doc_only_check
77-
if: needs.doc_only_check.outputs.skip_tests != 'true'
95+
# Run if either tests or notebooks need to run
96+
if: |
97+
needs.doc_only_check.outputs.run_tests == 'true' ||
98+
needs.doc_only_check.outputs.run_notebooks == 'true'
7899
uses: ./.github/workflows/build_package.yml
79100
with:
80101
device_type: tpu
81102
device_name: v4-8
82103
cloud_runner: linux-x86-n2-16-buildkit
83104

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

105143
maxtext_tpu_unit_tests:
106144
needs: build_and_upload_maxtext_package
145+
if: needs.doc_only_check.outputs.run_tests == 'true'
107146
uses: ./.github/workflows/run_tests_against_package.yml
108147
strategy:
109148
fail-fast: false
@@ -122,6 +161,7 @@ jobs:
122161

123162
maxtext_tpu_integration_tests:
124163
needs: build_and_upload_maxtext_package
164+
if: needs.doc_only_check.outputs.run_tests == 'true'
125165
uses: ./.github/workflows/run_tests_against_package.yml
126166
strategy:
127167
fail-fast: false
@@ -140,6 +180,7 @@ jobs:
140180

141181
maxtext_tpu_pathways_unit_tests:
142182
needs: build_and_upload_maxtext_package
183+
if: needs.doc_only_check.outputs.run_tests == 'true'
143184
uses: ./.github/workflows/run_pathways_tests.yml
144185
strategy:
145186
fail-fast: false
@@ -158,6 +199,7 @@ jobs:
158199

159200
maxtext_tpu_pathways_integration_tests:
160201
needs: build_and_upload_maxtext_package
202+
if: needs.doc_only_check.outputs.run_tests == 'true'
161203
uses: ./.github/workflows/run_pathways_tests.yml
162204
strategy:
163205
fail-fast: false
@@ -176,6 +218,7 @@ jobs:
176218

177219
maxtext_gpu_unit_tests:
178220
needs: build_and_upload_maxtext_package
221+
if: needs.doc_only_check.outputs.run_tests == 'true'
179222
uses: ./.github/workflows/run_tests_against_package.yml
180223
strategy:
181224
fail-fast: false
@@ -196,6 +239,7 @@ jobs:
196239

197240
maxtext_gpu_integration_tests:
198241
needs: build_and_upload_maxtext_package
242+
if: needs.doc_only_check.outputs.run_tests == 'true'
199243
uses: ./.github/workflows/run_tests_against_package.yml
200244
strategy:
201245
fail-fast: false
@@ -223,7 +267,7 @@ jobs:
223267
- name: Check test results
224268
run: |
225269
# If doc-only, all tests should be skipped
226-
if [ "${{ needs.doc_only_check.outputs.skip_tests }}" == "true" ]; then
270+
if [ "${{ needs.doc_only_check.outputs.run_tests }}" == "false" ]; then
227271
echo "Documentation-only changes detected, tests were skipped"
228272
exit 0
229273
fi
@@ -246,9 +290,34 @@ jobs:
246290
247291
echo "All required tests passed successfully"
248292
293+
all_notebooks_passed:
294+
name: All Notebooks Passed
295+
needs: [doc_only_check, build_and_upload_maxtext_package, maxtext_jupyter_notebooks]
296+
if: always()
297+
runs-on: ubuntu-latest
298+
steps:
299+
- name: Check notebooks results
300+
run: |
301+
if [ "${{ needs.doc_only_check.outputs.run_notebooks }}" == "false" ]; then
302+
echo "Non-notebook changes detected, runs were skipped"
303+
exit 0
304+
fi
305+
306+
# Otherwise, check that build and notebooks run passed or were skipped
307+
echo "Build result: ${{ needs.build_and_upload_maxtext_package.result }}"
308+
echo "Jupyter Notebooks result: ${{ needs.maxtext_jupyter_notebooks.result }}"
309+
310+
# Fail only if any job failed or was cancelled (skipped is OK)
311+
if [ "${{ contains(needs.*.result, 'failure') }}" == "true" ] || [ "${{ contains(needs.*.result, 'cancelled') }}" == "true" ]; then
312+
echo "One or more jobs failed or were cancelled"
313+
exit 1
314+
fi
315+
316+
echo "All required notebooks passed successfully"
317+
249318
notify_failure:
250319
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]
320+
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]
252321
if: ${{ always() }}
253322
runs-on: ubuntu-latest
254323
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)