Test notebooks on pull request

.github/workflows/notebook-test.yml

@@ -0,0 +1,73 @@
+# This code is a Qiskit project.
+#
+# (C) Copyright IBM 2023.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+name: Test notebooks
+on:
+  pull_request:
+    paths:
+      - "docs/**/*.ipynb"
+      - "!docs/api/**/*"
+  workflow_dispatch:
+jobs:
+  execute:
+    name: Execute notebooks
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v4
+        with:
+          python-version: "3.11"
+          cache: "pip"
+
+      - name: Get all changed files
+        id: all-changed-files
+        uses: tj-actions/changed-files@af2816c65436325c50621100d67f6e853cd1b0f1
+
+      - name: Check for notebooks that require LaTeX
+        id: latex-changed-files
+        uses: tj-actions/changed-files@af2816c65436325c50621100d67f6e853cd1b0f1
+        with:
+          # Add your notebook to this list if it needs latex to run
+          files: |
+            docs/build/circuit-visualization.ipynb
+
+      - name: Install LaTeX dependencies
+        if: steps.latex-changed-files.outputs.any_changed == 'true'
+        run: |
+          sudo apt-get update
+          sudo apt-get install texlive-pictures texlive-latex-extra poppler-utils
+
+      - name: Install Python packages
+        # This is to save our account in the next step. Note that the
+        # package will be re-installed during the "Run tox" step.
+        run: pip install qiskit-ibm-runtime tox
+
+      - name: Save IBM Quantum account
+        if: ${{ github.event.pull_request.head.repo.full_name == github.repository }}
+        shell: python
+        run: |
+          from qiskit_ibm_runtime import QiskitRuntimeService
+          QiskitRuntimeService.save_account(
+              channel="ibm_quantum",
+              instance="ibm-q/open/main",
+              token="${{ secrets.IBM_QUANTUM_TEST_TOKEN }}",
+              set_as_default=True
+          )
+
+      - name: Cache tox environment
+        uses: actions/cache@v3
+        with:
+          path: ".tox"
+          key: ${{ hashFiles('scripts/nb-tester/requirements.txt') }}
+
+      - name: Run tox
+        run: tox -- ${{ steps.all-changed-files.outputs.all_changed_files }}

README.md

@@ -121,6 +121,15 @@ pipx install tox
   tox -- optional/paths/to/notebooks.ipynb --write
   ```
 
+> [!NOTE]
+> If your notebook submits hardware jobs to IBM Quantum, you must add it to the
+> ignore list in `scripts/nb-tester/test-notebooks.py`. This is not needed if
+> you only retrieve information.
+>
+> If your notebook uses the latex circuit drawer (`qc.draw("latex")`), you must
+> add it to the "Check for notebooks that require LaTeX" step in
+> `.github/workflows/notebook-test.yml`.
+
 ## Check for broken links
 
 CI will check for broken links. You can also check locally:

scripts/nb-tester/test-notebook.py

@@ -10,82 +10,166 @@
 # notice, and modified files need to carry a notice indicating that they have
 # been altered from the originals.
 
+import argparse
 import sys
+from dataclasses import dataclass
+from datetime import datetime
 from pathlib import Path
+
+import nbclient
 import nbconvert
 import nbformat
+from qiskit_ibm_runtime import QiskitRuntimeService
 
-WRITE_FLAG = "--write"
 NOTEBOOKS_GLOB = "docs/**/*.ipynb"
 NOTEBOOKS_EXCLUDE = [
     "docs/api/**",
     "**/.ipynb_checkpoints/**",
-    # Following notebooks have code errors
+    # Following notebooks are broken
     "docs/transpile/transpiler-stages.ipynb",
-    # Following notebooks make requests so can't be tested yet
     "docs/run/get-backend-information.ipynb",
+]
+NOTEBOOKS_THAT_SUBMIT_JOBS = [
     "docs/start/hello-world.ipynb",
 ]
 
 
-def execute_notebook(path: Path, write=False) -> bool:
+@dataclass(frozen=True)
+class ExecuteOptions:
+    write: bool
+    submit_jobs: bool
+
+
+def execute_notebook(path: Path, options: ExecuteOptions) -> bool:
     """
     Wrapper function for `_execute_notebook` to print status
     """
     print(f"▶️  {path}", end="", flush=True)
+    possible_exceptions = (
+        nbconvert.preprocessors.CellExecutionError,
+        nbclient.exceptions.CellTimeoutError,
+    )
     try:
-        _execute_notebook(path, write)
-    except nbconvert.preprocessors.CellExecutionError as err:
+        _execute_notebook(path, options)
+    except possible_exceptions as err:
         print("\r❌\n")
         print(err)
         return False
     print("\r✅")
     return True
 
 
-def _execute_notebook(filepath: Path, write=False) -> None:
+def _execute_notebook(filepath: Path, options: ExecuteOptions) -> None:
     """
     Use nbconvert to execute notebook
     """
     nb = nbformat.read(filepath, as_version=4)
 
     processor = nbconvert.preprocessors.ExecutePreprocessor(
-        timeout=100,
+        # If submitting jobs, we want to wait forever (-1 means no timeout)
+        timeout=-1 if options.submit_jobs else 100,
         kernel_name="python3",
     )
 
     processor.preprocess(nb)
 
-    if not write:
+    if not options.write:
         return
 
     for cell in nb.cells:
+        # Remove execution metadata to avoid noisy diffs.
         cell.metadata.pop("execution", None)
     nbformat.write(nb, filepath)
 
 
-def find_notebooks() -> list[Path]:
+def find_notebooks(*, submit_jobs: bool = False) -> list[Path]:
     """
     Get paths to all notebooks in NOTEBOOKS_GLOB that are not excluded by
     NOTEBOOKS_EXCLUDE
     """
     all_notebooks = Path(".").rglob(NOTEBOOKS_GLOB)
+    excluded_notebooks = NOTEBOOKS_EXCLUDE
+    if not submit_jobs:
+        excluded_notebooks += NOTEBOOKS_THAT_SUBMIT_JOBS
     return [
         path
         for path in all_notebooks
-        if not any(path.match(glob) for glob in NOTEBOOKS_EXCLUDE)
+        if not any(path.match(glob) for glob in excluded_notebooks)
     ]
 
 
+def cancel_trailing_jobs(start_time: datetime) -> bool:
+    """
+    Cancel any runtime jobs created after `start_time`.
+    Return True if non exist, False otherwise.
+
+    Notebooks should not submit jobs during a normal test run. If they do, the
+    cell will time out and this function will cancel the job to avoid wasting
+    device time.
+
+    If a notebook submits a job but does not wait for the result, this check
+    will also catch it and cancel the job.
+    """
+    service = QiskitRuntimeService()
+    jobs = [j for j in service.jobs(created_after=start_time) if not j.in_final_state()]
+    if not jobs:
+        return True
+
+    print(
+        f"⚠️ Cancelling {len(jobs)} job(s) created after {start_time}.\n"
+        "Add any notebooks that submit jobs to NOTEBOOKS_EXCLUDE in "
+        "`scripts/nb-tester/test-notebook.py`."
+    )
+    for job in jobs:
+        job.cancel()
+    return False
+
+
+def create_argument_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="Notebook executor",
+        description="For testing notebooks and updating their outputs",
+    )
+    parser.add_argument(
+        "filenames",
+        help=(
+            "Paths to notebooks. If not provided, the script will search for "
+            "notebooks in `docs/`. To exclude a notebook from this process, add it "
+            "to `NOTEBOOKS_EXCLUDE` in the script."
+        ),
+        nargs="*",
+    )
+    parser.add_argument(
+        "-w",
+        "--write",
+        action="store_true",
+        help="Overwrite notebooks with the results of this script's execution.",
+    )
+    parser.add_argument(
+        "--submit-jobs",
+        action="store_true",
+        help=(
+            "Run notebooks that submit jobs to IBM Quantum and wait indefinitely "
+            "for jobs to complete. Warning: this has a real cost because it uses "
+            "quantum resources! Only use this argument occasionally and intentionally." 
+        ),
+    )
+    return parser
+
+
 if __name__ == "__main__":
-    args = sys.argv[1:]
-    write = WRITE_FLAG in args
-    if write:
-        args.remove(WRITE_FLAG)
+    args = create_argument_parser().parse_args()
 
-    notebook_paths = args or find_notebooks()
+    paths = map(Path, args.filenames or find_notebooks(submit_jobs=args.submit_jobs))
+
+    # Execute notebooks
+    start_time = datetime.now()
     print("Executing notebooks:")
-    results = [execute_notebook(path, write) for path in notebook_paths]
+    results = [
+        execute_notebook(path, args) for path in paths if path.suffix == ".ipynb"
+    ]
+    print("Checking for trailing jobs...")
+    results.append(cancel_trailing_jobs(start_time))
     if not all(results):
         sys.exit(1)
     sys.exit(0)