Handle submitting jobs to IBM Quantum

.github/workflows/notebook-test.yml

@@ -38,6 +38,7 @@ jobs:
           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
           )

README.md

@@ -121,6 +121,11 @@ 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`. Retrieving information
+> is ok.
+
 ## Check for broken links
 
 CI will check for broken links. You can also check locally:

scripts/nb-tester/test-notebook.py

@@ -14,83 +14,142 @@
 from pathlib import Path
 import nbconvert
 import nbformat
+import nbclient
+from datetime import datetime
+from qiskit_ibm_runtime import QiskitRuntimeService
+import argparse
 
-WRITE_FLAG = "--write"
 NOTEBOOKS_GLOB = "docs/**/*.ipynb"
 NOTEBOOKS_EXCLUDE = [
     "docs/api/**",
     "**/.ipynb_checkpoints/**",
     # Following notebooks have code errors
-    "docs/build/circuit-construction.ipynb",
-    "docs/build/pulse.ipynb",
-    "docs/transpile/custom-transpiler-pass.ipynb",
     "docs/transpile/transpiler-stages.ipynb",
-    # Following notebook has some non-python dependency I can't figure out
-    "docs/build/circuit-visualization.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:
+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,
+        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:
+        # To avoid noisy diffs
         cell.metadata.pop("execution", None)
     nbformat.write(nb, filepath)
 
 
-def find_notebooks() -> list[Path]:
+def find_notebooks(submit_jobs=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)
     ]
 
 
-if __name__ == "__main__":
-    args = sys.argv[1:]
-    write = WRITE_FLAG in args
-    if write:
-        args.remove(WRITE_FLAG)
+def cancel_trailing_jobs(start_time: datetime) -> bool:
+    """
+    Cancel any runtime jobs created after `start_time`.
+    Return True if non exist, False otherwise.
+    """
+    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
 
-    notebook_paths = args or find_notebooks()
+
+parser = argparse.ArgumentParser(
+    prog="Notebook executor",
+    description="For testing notebooks and updating their outputs",
+    epilog="For help, contact Frank Harkins",
+)
+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 execution outputs",
+)
+parser.add_argument(
+    "--submit-jobs",
+    action="store_true",
+    help=(
+        "run notebooks that submit jobs to IBM Quantum and wait indefinitely "
+        "for jobs to complete"
+    ),
+)
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    notebook_paths = args.filenames or find_notebooks(args.submit_jobs)
+    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 notebook_paths]
+    print("Checking for trailing jobs...")
+    results.append(cancel_trailing_jobs(start_time))
     if not all(results):
         sys.exit(1)
     sys.exit(0)