Refactor nb-tester execution process

scripts/nb-tester/pyproject.toml

@@ -9,7 +9,7 @@ description = "Tool for the Qiskit docs team to test their notebooks"
 requires-python = ">=3.9"
 license = "Apache-2.0"
 dependencies = [
-  "nbconvert~=7.16.0",
+  "nbclient~=0.10.0",
   "nbformat~=5.9.2",
   "ipykernel~=6.29.2",
   "squeaky==0.7.0",

scripts/nb-tester/qiskit_docs_notebook_tester/__init__.py

@@ -25,9 +25,9 @@
 from typing import Iterator
 
 import nbclient
-import nbconvert
 import nbformat
 import tomli
+from jupyter_client.manager import start_new_async_kernel
 from qiskit_ibm_runtime import QiskitRuntimeService
 from squeaky import clean_notebook
 
@@ -140,24 +140,6 @@ def extract_warnings(notebook: nbformat.NotebookNode) -> list[NotebookWarning]:
                 )
     return notebook_warnings
 
-@contextmanager
-def patch_runtime(nb: nbformat.NotebookNode, *, should_patch: bool):
-    if should_patch:
-        nb.cells.insert(0, nbformat.v4.new_code_cell(source=MOCKING_CODE))
-    yield
-    if not should_patch:
-         return
-    nb.cells.pop(0)
-    # Reset execution counts (offset by the MOCKING_CODE cell)
-    for cell in nb.cells:
-        if hasattr(cell, "execution_count"):
-            cell.execution_count -= 1
-        if not hasattr(cell, "outputs"):
-            continue
-        for output in cell.outputs:
-            if hasattr(output, "execution_count"):
-                output.execution_count -= 1
-
 async def execute_notebook(path: Path, args: argparse.Namespace, config: Config) -> bool:
     """
     Wrapper function for `_execute_notebook` to print status and write result
@@ -168,11 +150,11 @@ async def execute_notebook(path: Path, args: argparse.Namespace, config: Config)
     else:
         print(f"▶️ Executing {path}")
     possible_exceptions = (
-        nbconvert.preprocessors.CellExecutionError,
+        nbclient.exceptions.CellExecutionError,
         nbclient.exceptions.CellTimeoutError,
     )
     try:
-        nb = await _execute_notebook(path, args)
+        nb = await _execute_notebook(path, args, is_patched)
     except possible_exceptions as err:
         print(f"❌ Problem in {path}:\n{err}")
         return False
@@ -197,26 +179,31 @@ async def execute_notebook(path: Path, args: argparse.Namespace, config: Config)
     print(f"✅ No problems in {path} (written)")
     return True
 
-async def _execute_notebook(filepath: Path, args: argparse.Namespace) -> nbformat.NotebookNode:
+async def _execute_notebook(filepath: Path, args: argparse.Namespace, patch: bool) -> nbformat.NotebookNode:
     """
-    Use nbconvert to execute notebook.
+    Use nbclient to execute notebook.
     """
     nb = nbformat.read(filepath, as_version=4)
 
-    processor = nbconvert.preprocessors.ExecutePreprocessor(
+    kernel_manager, kernel = await start_new_async_kernel(
+        kernel_name="python3",
+        extra_arguments=["--InlineBackend.figure_format='svg'"],
+    )
+
+    if patch:
+        kernel.execute(MOCKING_CODE, store_history=False)
+
+    notebook_client = nbclient.NotebookClient(
+        nb=nb,
+        km=kernel_manager,
+        kc=kernel,
         # If submitting jobs, we want to wait forever (-1 means no timeout)
         timeout=-1 if args.submit_jobs else 300,
-        kernel_name="python3",
-        extra_arguments=["--InlineBackend.figure_format='svg'"]
     )
 
     # This runs the notebook, including possibly submitting jobs. We run it in a
     # new thread to avoid blocking other notebooks from submitting jobs.
-    with patch_runtime(nb, should_patch=not args.submit_jobs):
-        await asyncio.to_thread(processor.preprocess, nb)
-
-    if not args.write:
-        return nb
+    await notebook_client.async_execute()
 
     for cell in nb.cells:
         # Remove execution metadata to avoid noisy diffs.