From d0835449ba449c6635b7766f9d0e4a76ce04273b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Simon=20H=C3=B8xbro=20Hansen?= <simon.hansen@me.com>
Date: Wed, 16 Oct 2024 15:13:55 +0200
Subject: [PATCH] Remove subprocess calls and add retry evaluating  a notebook
 (#311)

Co-authored-by: maximlt <mliquet@anaconda.com>
---
 nbsite/cmd.py                 | 35 +++++++++++++++++++++++++++++------
 nbsite/gallery/thumbnailer.py |  2 +-
 nbsite/nbbuild.py             | 23 +++++++++++++++--------
 3 files changed, 45 insertions(+), 15 deletions(-)

diff --git a/nbsite/cmd.py b/nbsite/cmd.py
index 866c3134..6cd2ff91 100644
--- a/nbsite/cmd.py
+++ b/nbsite/cmd.py
@@ -7,6 +7,8 @@
 from collections import ChainMap
 from os.path import dirname
 
+from sphinx.application import Sphinx
+
 from .util import copy_files
 
 DEFAULT_SITE_ORDERING = [
@@ -80,8 +82,7 @@ def build(what='html',
         'EXAMPLES_ASSETS':examples_assets,
         'BINDER':binder
     }
-    merged_env = dict(os.environ, **env)
-    none_vals = {k:v for k,v in merged_env.items() if v is None}
+    none_vals = {k:v for k,v in env.items() if v is None}
     if none_vals:
         raise Exception("Missing value for %s" % list(none_vals.keys()))
 
@@ -90,9 +91,27 @@ def build(what='html',
         for path in glob.glob(os.path.join(paths['doc'], '**', '*.ipynb'), recursive=True):
             print('Removing evaluated notebook from {}'.format(path))
             os.remove(path)
-    extras = [] if disable_parallel else ["-j", "auto"]
-    cmd = ["sphinx-build", "-b", what, paths['doc'], output] + extras
-    subprocess.check_call(cmd, env=merged_env)
+
+    parallel = 0 if disable_parallel else os.cpu_count()
+    # Code smell as there should be a way to configure Sphinx/Nbsite without
+    # env vars, but that's how it was done at the time Sphinx was called
+    # via subprocess.
+    _environ = dict(os.environ)
+    os.environ.update(env)
+    try:
+        app = Sphinx(
+            srcdir=paths["doc"],
+            confdir=paths["doc"],
+            outdir=output,
+            doctreedir=os.path.join(output, ".doctrees"),
+            buildername=what,
+            parallel=parallel,
+        )
+        app.build()
+    finally:
+        os.environ.clear()
+        os.environ.update(_environ)
+
     print('Copying json blobs (used for holomaps) from {} to {}'.format(paths['doc'], output))
     copy_files(paths['doc'], output, '**/*.json')
     copy_files(paths['doc'], output, 'json_*')
@@ -101,10 +120,14 @@ def build(what='html',
         print("Copying examples assets from %s to %s"%(paths['examples_assets'],build_assets))
         copy_files(paths['examples_assets'], build_assets)
     fix_links(output, inspect_links)
+
     # create a .nojekyll file in output for github compatibility
-    subprocess.check_call(["touch", os.path.join(output, '.nojekyll')])
+    with open(os.path.join(output, '.nojekyll'), 'w') as f:
+        f.write('')
+
     if not clean_dry_run:
         print("Call `nbsite build` with `--clean-dry-run` to not actually delete files.")
+
     clean(output, clean_dry_run)
 
 def clean(output, dry_run=False):
diff --git a/nbsite/gallery/thumbnailer.py b/nbsite/gallery/thumbnailer.py
index e8663e50..87c31b9d 100644
--- a/nbsite/gallery/thumbnailer.py
+++ b/nbsite/gallery/thumbnailer.py
@@ -155,7 +155,7 @@ def execute(code, cwd, env):
     with tempfile.NamedTemporaryFile('wb', delete=True) as f:
         f.write(code)
         f.flush()
-        proc = subprocess.Popen(['python', f.name], cwd=cwd, env=env)
+        proc = subprocess.Popen([sys.executable, f.name], cwd=cwd, env=env)
         proc.wait()
         proc.kill()
     return proc.returncode
diff --git a/nbsite/nbbuild.py b/nbsite/nbbuild.py
index aa11533c..68dad9c7 100755
--- a/nbsite/nbbuild.py
+++ b/nbsite/nbbuild.py
@@ -612,14 +612,21 @@ def run(self):
         os.makedirs(dest_dir, exist_ok=True)
 
         # Evaluate Notebook and insert into Sphinx doc
-        evaluate_notebook(
-            nb_abs_path, dest_path,
-            skip_exceptions='skip_exceptions' in self.options,
-            skip_execute=self.options.get('skip_execute'),
-            timeout=setup.config.nbbuild_cell_timeout,
-            ipython_startup=setup.config.nbbuild_ipython_startup,
-            patterns_to_take_with_me=setup.config.nbbuild_patterns_to_take_along
-        )
+        import zmq
+        for n in range(1, 6):
+            try:
+                evaluate_notebook(
+                    nb_abs_path, dest_path,
+                    skip_exceptions='skip_exceptions' in self.options,
+                    skip_execute=self.options.get('skip_execute'),
+                    timeout=setup.config.nbbuild_cell_timeout,
+                    ipython_startup=setup.config.nbbuild_ipython_startup,
+                    patterns_to_take_with_me=setup.config.nbbuild_patterns_to_take_along
+                )
+                break
+            except (zmq.error.ZMQError, RuntimeError) as e:
+                # Sometimes the kernel dies
+                print(f"{nb_abs_path} failed with {e}, retrying ({n}/5)...", flush=True)
 
         preprocessors = self.preprocessors(dest_dir)
         rendered_nodes = render_notebook(