Add notebook CI testing and a test mode flag.

.github/workflows/notebooks.yaml

@@ -0,0 +1,47 @@
+name: Notebooks
+
+on:
+  # Trigger the workflow on push or pull request, but only on dev branch
+  push:
+    branches:
+      - revamp-2025
+  pull_request:
+    branches:
+      - revamp-2025
+
+permissions:
+  contents: read  # to fetch code
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  documentation:
+    name: Set up Python ${{ matrix.python-version }}
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    strategy:
+      matrix:
+        python-version: ["3.12"]
+      fail-fast: false
+
+    steps:
+    - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8  # v5.0.0
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c  # v6.0.0
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install .[dev]
+        pip install ipykernel pytest-notebook
+        pip freeze
+    - name: Register Python3 kernel
+      run: |
+        python -m ipykernel install --user --name "python3" --display-name "Python 3"
+    - name: Turn on test mode and test notebooks
+      run: |
+        export AI_STACK_TEST_MODE="True"
+        pytest -n auto --tb=short --nb-force-regen --nb-test-files docs/source/ --ignore=docs/source/JAX_for_LLM_pretraining.ipynb --ignore=docs/source/data_loaders_on_cpu_with_jax.ipynb --ignore=docs/source/data_loaders_on_gpu_with_jax.ipynb --ignore=docs/source/JAX_for_PyTorch_users.ipynb --ignore=docs/source/JAX_porting_PyTorch_model.ipynb

docs/source/conf.py

@@ -67,6 +67,9 @@
     'JAX_transformer_text_classification.md',
     'data_loaders_on_cpu_with_jax.md',
     'data_loaders_on_gpu_with_jax.md',
+    # Ignored test files.
+    'test_mode_test.md',
+    'test_mode_test.ipynb',
 ]
 
 suppress_warnings = [
@@ -104,4 +107,5 @@
     'JAX_transformer_text_classification.ipynb',
     'data_loaders_on_cpu_with_jax.ipynb',
     'data_loaders_on_gpu_with_jax.ipynb',
+    'test_mode_test.ipynb',
 ]

docs/source/test_mode_test.ipynb

@@ -0,0 +1,38 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4765c6f3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Tests that CI sets the test mode environment variable\n",
+    "import os\n",
+    "test_mode = os.environ.get('AI_STACK_TEST_MODE', 'False')\n",
+    "assert test_mode == 'True', test_mode"
+   ]
+  }
+ ],
+ "metadata": {
+  "jupytext": {
+   "formats": "ipynb,md",
+   "main_language": "python"
+  },
+  "language_info": {
+   "name": "python"
+  },
+  "nbreg": {
+   "diff_ignore": [
+    "/cells/0/execution_count",
+    "/cells/0/outputs",
+    "/metadata"
+   ]
+  },
+  "nbsphinx": {
+   "orphan": true
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/source/test_mode_test.md

@@ -0,0 +1,18 @@
+---
+jupyter:
+  jupytext:
+    formats: ipynb,md
+    main_language: python
+    text_representation:
+      extension: .md
+      format_name: markdown
+      format_version: '1.3'
+      jupytext_version: 1.15.2
+---
+
+```python
+# Tests that CI sets the test mode environment variable
+import os
+test_mode = os.environ.get('AI_STACK_TEST_MODE', 'False')
+assert test_mode == 'True', test_mode
+```