120 lines
4.4 KiB
Python
120 lines
4.4 KiB
Python
#!/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)
|
|
env["PYTHONPATH"] = os.pathsep.join(
|
|
filter(None, [str(ROOT), env.get("PYTHONPATH")])
|
|
)
|
|
for name in args.notebooks or NOTEBOOKS:
|
|
path = (ROOT / name).resolve()
|
|
relative = path.relative_to(ROOT)
|
|
# Related notebooks share the library's default local storage.
|
|
work_dir = Path(workspace) / relative.parent
|
|
work_dir.mkdir(parents=True, exist_ok=True)
|
|
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(work_dir), 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=work_dir,
|
|
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()
|