Add jupyter notebook tests to Github CI tests

Author: SurbhiJainUSC

.github/workflows/build_and_test_maxtext.yml

@@ -41,7 +41,8 @@ jobs:
     name: Check for Documentation-Only Changes
     runs-on: ubuntu-latest
     outputs:
-      skip_tests: ${{ steps.check.outputs.skip_tests }}
+      run_tests: ${{ steps.check.outputs.run_tests }}
+      run_notebooks: ${{ steps.check.outputs.run_notebooks }}
     steps:
       - uses: actions/checkout@v4
         with:
@@ -51,7 +52,8 @@ jobs:
         run: |
           if [ "${{ github.event_name }}" != "pull_request" ]; then
             echo "Not a pull request, running all tests"
-            echo "skip_tests=false" >> $GITHUB_OUTPUT
+            echo "run_tests=true" >> $GITHUB_OUTPUT
+            echo "run_notebooks=true" >> $GITHUB_OUTPUT
             exit 0
           fi
 
@@ -63,26 +65,59 @@ jobs:
 
           if [ -z "$CHANGED_FILES" ]; then
             echo "No files changed"
-            echo "skip_tests=false" >> $GITHUB_OUTPUT
-          elif echo "$CHANGED_FILES" | grep -v -E '\.(md|ipynb)$' > /dev/null; then
-            echo "Non-documentation files changed, running tests"
-            echo "skip_tests=false" >> $GITHUB_OUTPUT
+            echo "run_tests=true" >> $GITHUB_OUTPUT
+            echo "run_notebooks=true" >> $GITHUB_OUTPUT
+          fi
+
+          # Check for source code changes (anything not .md and not .ipynb)
+          if echo "$CHANGED_FILES" | grep -v -E '\.(md|ipynb)$' > /dev/null; then
+            echo "Source code files changed, enabling unit tests."
+            echo "run_tests=true" >> $GITHUB_OUTPUT
+          else
+            echo "No source code changes, skipping unit tests."
+            echo "run_tests=false" >> $GITHUB_OUTPUT
+          fi
+
+          # Check for notebook (.ipynb) changes specifically
+          if echo "$CHANGED_FILES" | grep '\.ipynb$' > /dev/null; then
+            echo "Notebook files changed, enabling notebook run."
+            echo "run_notebooks=true" >> $GITHUB_OUTPUT
           else
-            echo "Only documentation (.md|ipynb) files changed, skipping tests"
-            echo "skip_tests=true" >> $GITHUB_OUTPUT
+            echo "No notebook changes, skipping notebook run."
+            echo "run_notebooks=false" >> $GITHUB_OUTPUT
           fi
 
   build_and_upload_maxtext_package:
     needs: doc_only_check
-    if: needs.doc_only_check.outputs.skip_tests != 'true'
+    # Run if either tests or notebooks need to run
+    if: |
+      needs.doc_only_check.outputs.run_tests == 'true' || 
+      needs.doc_only_check.outputs.run_notebooks == 'true'
     uses: ./.github/workflows/build_package.yml
     with:
       device_type: tpu
       device_name: v4-8
       cloud_runner: linux-x86-n2-16-buildkit
 
+  maxtext_jupyter_notebooks:
+    needs: build_and_upload_maxtext_package
+    if: needs.doc_only_check.outputs.run_notebooks == 'true'
+    uses: ./.github/workflows/run_jupyter_notebooks.yml
+    strategy:
+        fail-fast: false
+        matrix:
+          image_type: ["py312"]
+    with:
+      device_type: tpu
+      device_name: v6e-4
+      image_type: ${{ matrix.image_type }}
+      cloud_runner: linux-x86-ct6e-180-4tpu
+    secrets:
+      HF_TOKEN: ${{ secrets.HF_TOKEN }}
+
   maxtext_cpu_unit_tests:
     needs: build_and_upload_maxtext_package
+    if: needs.doc_only_check.outputs.run_tests == 'true'
     uses: ./.github/workflows/run_tests_against_package.yml
     strategy:
         fail-fast: false # don't cancel all jobs on failure
@@ -104,6 +139,7 @@ jobs:
 
   maxtext_tpu_unit_tests:
     needs: build_and_upload_maxtext_package
+    if: needs.doc_only_check.outputs.run_tests == 'true'
     uses: ./.github/workflows/run_tests_against_package.yml
     strategy:
         fail-fast: false
@@ -122,6 +158,7 @@ jobs:
 
   maxtext_tpu_integration_tests:
     needs: build_and_upload_maxtext_package
+    if: needs.doc_only_check.outputs.run_tests == 'true'
     uses: ./.github/workflows/run_tests_against_package.yml
     strategy:
         fail-fast: false
@@ -140,6 +177,7 @@ jobs:
 
   maxtext_tpu_pathways_unit_tests:
     needs: build_and_upload_maxtext_package
+    if: needs.doc_only_check.outputs.run_tests == 'true'
     uses: ./.github/workflows/run_pathways_tests.yml
     strategy:
         fail-fast: false
@@ -158,6 +196,7 @@ jobs:
 
   maxtext_tpu_pathways_integration_tests:
     needs: build_and_upload_maxtext_package
+    if: needs.doc_only_check.outputs.run_tests == 'true'
     uses: ./.github/workflows/run_pathways_tests.yml
     strategy:
         fail-fast: false
@@ -176,6 +215,7 @@ jobs:
 
   maxtext_gpu_unit_tests:
     needs: build_and_upload_maxtext_package
+    if: needs.doc_only_check.outputs.run_tests == 'true'
     uses: ./.github/workflows/run_tests_against_package.yml
     strategy:
         fail-fast: false
@@ -196,6 +236,7 @@ jobs:
 
   maxtext_gpu_integration_tests:
     needs: build_and_upload_maxtext_package
+    if: needs.doc_only_check.outputs.run_tests == 'true'
     uses: ./.github/workflows/run_tests_against_package.yml
     strategy:
         fail-fast: false
@@ -223,7 +264,7 @@ jobs:
       - name: Check test results
         run: |
           # If doc-only, all tests should be skipped
-          if [ "${{ needs.doc_only_check.outputs.skip_tests }}" == "true" ]; then
+          if [ "${{ needs.doc_only_check.outputs.run_tests }}" == "false" ]; then
             echo "Documentation-only changes detected, tests were skipped"
             exit 0
           fi
@@ -246,9 +287,34 @@ jobs:
 
           echo "All required tests passed successfully"
 
+  all_notebooks_passed:
+    name: All Notebooks Passed
+    needs: [doc_only_check, build_and_upload_maxtext_package, maxtext_jupyter_notebooks]
+    if: always()
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check notebooks results
+        run: |
+          if [ "${{ needs.doc_only_check.outputs.run_notebooks }}" == "false" ]; then
+            echo "Non-notebook changes detected, runs were skipped"
+            exit 0
+          fi
+
+          # Otherwise, check that build and notebooks run passed or were skipped
+          echo "Build result: ${{ needs.build_and_upload_maxtext_package.result }}"
+          echo "Jupyter Notebooks result: ${{ needs.maxtext_jupyter_notebooks.result }}"
+
+          # Fail only if any job failed or was cancelled (skipped is OK)
+          if [ "${{ contains(needs.*.result, 'failure') }}" == "true" ] || [ "${{ contains(needs.*.result, 'cancelled') }}" == "true" ]; then
+            echo "One or more jobs failed or were cancelled"
+            exit 1
+          fi
+
+          echo "All required notebooks passed successfully"
+
   notify_failure:
     name: Notify failed build # creates an issue or modifies last open existing issue for failed build
-    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]
+    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]
     if: ${{ always() }}
     runs-on: ubuntu-latest
     permissions:

.github/workflows/run_jupyter_notebooks.yml

@@ -0,0 +1,103 @@
+# Copyright 2026 Google LLC
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This file defines a module for running jupyter notebooks against the built maxtext package.
+
+name: Run Jupyter Notebooks
+
+on:
+  workflow_call:
+    inputs:
+      device_type:
+        required: true
+        type: string
+      device_name:
+        required: true
+        type: string
+      image_type:
+        required: false
+        type: string
+      cloud_runner:
+        required: false
+        type: string
+    secrets:
+      HF_TOKEN:
+        required: true
+
+permissions:
+  contents: read
+jobs:
+  run:
+    runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }}
+    container:
+      image: gcr.io/tpu-prod-env-multipod/maxtext-unit-test-${{ inputs.device_type == 'cpu' && 'tpu' || inputs.device_type }}:${{ inputs.image_type != '' && inputs.image_type }}
+    steps:
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
+      - name: Download the MaxText wheel
+        uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0
+        with:
+          name: maxtext-wheel
+      - name: Install MaxText and Dependencies
+        shell: bash
+        run: |
+          python3 -m uv venv --seed
+          source .venv/bin/activate
+
+          # Install MaxText package
+          maxtext_wheel=$(ls maxtext-*-py3-none-any.whl 2>/dev/null)
+          uv pip install ${maxtext_wheel}[${MAXTEXT_PACKAGE_EXTRA}] --resolution=lowest
+          uv pip install -r src/install_maxtext_extra_deps/extra_deps_from_github.txt
+
+          # Install dependencies for running notebooks
+          uv pip install papermill ipykernel ipywidgets
+          .venv/bin/python3 -m ipykernel install --user --name maxtext_venv
+
+          # Install Tunix for post-training notebooks
+          uv pip install git+https://github.com/google/tunix
+          
+          # Install vllm for post-training notebooks
+          git clone https://github.com/vllm-project/vllm.git
+          VLLM_TARGET_DEVICE="tpu" uv pip install ./vllm
+
+          # Install tpu-inference for post-training notebooks
+          git clone https://github.com/vllm-project/tpu-inference.git
+          uv pip install ./tpu-inference
+
+          uv pip install --no-deps qwix==0.1.4
+          uv pip install --no-deps protobuf==5.29.5
+          python3 -m pip freeze
+      - name: Run Post-Training Notebooks
+        shell: bash
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: |
+          MAXTEXT_REPO_ROOT=$(pwd)
+          MAXTEXT_NOTEBOOKS_ROOT="$MAXTEXT_REPO_ROOT/src/MaxText/examples"
+
+          for notebook in "$MAXTEXT_NOTEBOOKS_ROOT"/{sft,rl}*.ipynb; do
+            filename=$(basename "$notebook")
+            output_name="${filename%.ipynb}_output.ipynb"
+            
+            echo "------------------------------------------------------"
+            echo "Running $filename ..."
+            echo "------------------------------------------------------"
+
+            .venv/bin/papermill "$notebook" "$output_name" -k maxtext_venv
+          done
+      - name: Upload Outputs
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: notebook-outputs-${{ inputs.device_name }}
+          path: ./*_output.ipynb

docs/guides/run_python_notebook.md

@@ -89,7 +89,7 @@ gcloud compute tpus tpu-vm ssh maxtext-tpu-node --zone=YOUR_ZONE -- -L 8888:loca
 
 > **Note**: If you get a "bind: Address already in use" error, it means port 8888 is busy on your local computer. Change the first number to a different port, e.g., -L 9999:localhost:8888. You will then access Jupyter at localhost:9999.
 
-### Step 3: Install Dependencies
+### Step 3: Install JupyterLab
 
 Run the following commands on your TPU-VM:
 
@@ -99,21 +99,40 @@ sudo apt install python3-pip python3-dev git -y
 pip3 install jupyterlab
 ```
 
-### Step 4: Start Jupyter Lab
+### Step 4: Install MaxText and Dependencies
+
+To execute post-training notebooks on your TPU-VM, follow the official [MaxText installation guides](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/rl.html#create-virtual-environment-and-install-maxtext-dependencies) to install MaxText and its dependencies inside a dedicated virtual environment.
+
+### Step 5: Register virtual environment as a Jupyter Kernel
+
+Once the environment is set up, you need to register it so JupyterLab can see it as an available kernel. Run the following command on your TPU-VM, replacing <virtual env name> with the actual name to your virtual environment:
+
+```bash
+python3 -m ipykernel install --user --name <virtual env name> --display-name "Python3 (MaxText Venv)"
+```
+
+### Step 6: Start JupyterLab
+
+Start the Jupyter server on your TPU-VM by running:
 
 ```bash
 jupyter lab --ip=0.0.0.0 --port=8888 --no-browser --allow-root
 ```
 
-### Step 5: Access the Notebook
-5.a. Look at the terminal output for a URL that looks like: `http://127.0.0.1:8888/lab?token=...`
+### Step 7: Access the Notebook
+7.a. Look at the terminal output for a URL that looks like: `http://127.0.0.1:8888/lab?token=...`.
 
-5.b. Copy that URL.
+7.b. Copy that URL.
 
-5.c. Paste it into your **local computer's browser**.
+7.c. Paste it into your **local computer's browser**.
    * **Important:** If you changed the port in Step 2 (e.g., to `9999`), you must manually replace `8888` in the URL with `9999`.
    * *Example:* `http://127.0.0.1:9999/lab?token=...`
 
+7.d. Once the interface opens in your browser, Click on the current kernel name (e.g., `Python 3 (ipykernel)`).
+
+7.e. In the dropdown menu, select the new kernel you just created: `Python3 (MaxText Venv)`.
+
+7.f. Run the jupyter notebook cells.
 
 ## Available Examples