Run notebooks in multiple processes

devtools/conda-envs/examples_env.yml

@@ -13,8 +13,8 @@ dependencies:
     - requests
     - packaging
     # Examples
-    - openff-toolkit-examples>=0.15.2
-    - openff-interchange>=0.3.18
+    - openff-toolkit-examples>=0.16.0
+    - openff-interchange>=0.3.26
     - openff-nagl
     # - openff-fragmenter
     # - openff-qcsubmit
@@ -23,4 +23,4 @@ dependencies:
     - rich
     - jax
     - parmed<4
-    - nglview>=3.0.6
+    - nglview>=3.0.6,<3.1.0

source/_ext/cookbook/utils.py

@@ -2,6 +2,7 @@
 from types import FunctionType
 from typing import (
     Callable,
+    Generic,
     Iterable,
     Iterator,
     Optional,
@@ -10,6 +11,7 @@
     ParamSpec,
     Union,
 )
+import abc
 import contextlib
 import os
 
@@ -32,7 +34,9 @@ def next_or_none(iterator: Iterator[T]) -> Optional[T]:
     return ret
 
 
-def _result_inner(fn, exception, *args, **kwargs):
+def result(
+    fn: Callable[P, T], exception: type[E], *args: P.args, **kwargs: P.kwargs
+) -> Union[T, E]:
     try:
         return fn(*args, **kwargs)
     except exception as e:
@@ -42,7 +46,8 @@ def _result_inner(fn, exception, *args, **kwargs):
 def to_result(
     fn: Callable[P, T], exception: type[E] = Exception
 ) -> Callable[P, Union[T, E]]:
-    return partial(_result_inner, fn, exception)
+    # partial's type stub is not precise enough to work here
+    return partial(result, fn, exception)  # type: ignore [reportReturnType]
 
 
 @contextlib.contextmanager

source/_ext/proc_examples.py

@@ -31,7 +31,7 @@
 )
 from cookbook.github import download_dir, get_tag_matching_installed_version
 from cookbook.globals_ import *
-from cookbook.utils import set_env, to_result
+from cookbook.utils import set_env, result, to_result
 
 
 class NotebookExceptionError(ValueError):
@@ -310,26 +310,24 @@ def main(
     # Execute notebooks in parallel for rendering as HTML
     if do_exec:
         shutil.rmtree(EXEC_IPYNB_ROOT, ignore_errors=True)
-        # # Context manager ensures the pool is correctly terminated if there's
-        # # an exception
-        # with Pool(processes=processes) as pool:
-        #     # Wait a second between launching subprocesses
-        #     # Workaround https://github.com/jupyter/nbconvert/issues/1066
-        #     exec_results = [
-        #         *pool.imap(
-        #             to_result(partial(execute_notebook, cache_branch=cache_branch)),
-        #             delay_iterator(notebooks),
-        #         )
-        #     ]
-        # for result in exec_results:
-        #     if isinstance(result, Exception):
-        #         traceback.print_exception(result)
-        exceptions: list[NotebookExceptionError] = []
-        for notebook in notebooks:
-            try:
-                execute_notebook(notebook, cache_branch=cache_branch)
-            except NotebookExceptionError as e:
-                exceptions.append(e)
+        # Context manager ensures the pool is correctly terminated if there's
+        # an exception
+        with Pool(processes=processes) as pool:
+            # Wait a second between launching subprocesses
+            # Workaround https://github.com/jupyter/nbconvert/issues/1066
+            exec_results = [
+                *pool.imap(
+                    to_result(
+                        partial(execute_notebook, cache_branch=cache_branch),
+                        NotebookExceptionError,
+                    ),
+                    delay_iterator(notebooks),
+                )
+            ]
+
+        exceptions: list[NotebookExceptionError] = [
+            result for result in exec_results if isinstance(result, Exception)
+        ]
 
         for exception in exceptions:
             print(