Run notebooks in CI

scripts/run-notebooks.py

@@ -0,0 +1,114 @@
+#!/usr/bin/env python3
+"""Execute documentation notebooks in order, with the current Python interpreter."""
+
+import argparse
+import os
+import subprocess
+import sys
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import nbformat
+from jupyter_client import KernelManager
+from nbclient import NotebookClient
+
+ROOT = Path(__file__).resolve().parents[1]
+NOTEBOOKS = [
+    "docs/tutorial/train.ipynb",
+    "docs/tutorial/deploy.ipynb",
+    "docs/examples/simple/data.ipynb",
+    "docs/examples/simple/train.ipynb",
+    "docs/examples/simple/deploy.ipynb",
+    "docs/examples/simple/stacked-bars.ipynb",
+    "docs/examples/simple-mag/train.ipynb",
+    "docs/examples/scibert/data.ipynb",
+    "docs/examples/scibert/train.ipynb",
+    "docs/examples/scibert/deploy.ipynb",
+    "docs/surveys/surveys.ipynb",
+]
+
+CHECK_DEPLOYMENT = """
+import sys
+from importlib import import_module
+from pathlib import Path
+
+from fastapi.testclient import TestClient
+from great_ai.__main__ import find_app, get_script_name, get_script_name_of_notebook
+
+notebook = sys.argv[1]
+try:
+    module_name = get_script_name(notebook)
+    app_name = find_app(module_name)
+    assert app_name is not None, "The CLI could not import the deployment notebook"
+    module = import_module(module_name)
+    function_name = app_name.split(":")[1].split(".")[0]
+    service = getattr(module, function_name)
+    argument = "text" if sys.argv[2] == "simple" else "sentence"
+    with TestClient(service.app) as client:
+        response = client.post("/predict", json={argument: "Neural networks learn useful features."})
+        response.raise_for_status()
+        assert response.json()["output"] is not None
+finally:
+    Path(get_script_name_of_notebook(notebook)).unlink(missing_ok=True)
+"""
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("notebooks", nargs="*", help="Paths relative to the repository")
+    parser.add_argument("--timeout", type=int, default=600, help="Seconds per cell")
+    args = parser.parse_args()
+    output_dir = (
+        ROOT
+        / ".cache/executed-notebooks"
+        / f"python{sys.version_info.major}.{sys.version_info.minor}"
+    )
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # A fresh workspace detects missing dependencies between notebook stages and
+    # avoids silently passing because an old model or database happens to exist.
+    with TemporaryDirectory(prefix="great-ai-notebooks-") as workspace:
+        env = dict(os.environ, GREAT_AI_NOTEBOOK_WORKDIR=workspace)
+        for name in args.notebooks or NOTEBOOKS:
+            path = (ROOT / name).resolve()
+            relative = path.relative_to(ROOT)
+            print(f"Running {relative}", flush=True)
+            notebook = nbformat.read(path, as_version=4)
+            nbformat.validate(notebook)
+            manager = KernelManager(kernel_name="python3")
+            # Do not depend on a user-installed kernelspec pointing at this venv.
+            manager.kernel_spec.argv = [
+                sys.executable,
+                "-m",
+                "ipykernel_launcher",
+                "-f",
+                "{connection_file}",
+            ]
+            client = NotebookClient(notebook, km=manager, timeout=args.timeout)
+            try:
+                client.execute(cwd=str(path.parent), env=env, cleanup_kc=True)
+            finally:
+                destination = output_dir / relative
+                destination.parent.mkdir(parents=True, exist_ok=True)
+                nbformat.write(notebook, destination)
+            if path.stem == "deploy":
+                subprocess.run(
+                    [
+                        sys.executable,
+                        "-c",
+                        CHECK_DEPLOYMENT,
+                        str(destination),
+                        path.parent.name,
+                    ],
+                    cwd=ROOT,
+                    env=env,
+                    check=True,
+                    timeout=args.timeout,
+                    stdout=subprocess.DEVNULL,
+                )
+            print(f"Passed {relative}", flush=True)
+    print(f"Executed notebooks saved to {output_dir}", flush=True)
+
+
+if __name__ == "__main__":
+    main()