Repository: google/jax

Issue:
RTD build failure executing thinking_in_jax.ipynb
#5426 added a new notebook to our docs, and ever since there have been unpredictable failures in the readthedocs build. An example is here: https://readthedocs.org/projects/jax/builds/12905004/

This is the error:
```pytb
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-1-1cefe89b730a> in <module>
      1 x = np.random.randn(3, 4)
      2 y = np.random.randn(4)
----> 3 f(x, y)

NameError: name 'f' is not defined
```

<details><summary>(Show full traceback)</summary>
<p>

```pytb
Traceback (most recent call last):
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/nbsphinx.py", line 1021, in parse
    rststring, resources = exporter.from_notebook_node(nb, resources)
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/nbsphinx.py", line 832, in from_notebook_node
    nb, resources = pp.preprocess(nb, resources)
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/nbconvert/preprocessors/execute.py", line 79, in preprocess
    self.execute()
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/nbclient/util.py", line 74, in wrapped
    return just_run(coro(*args, **kwargs))
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/nbclient/util.py", line 53, in just_run
    return loop.run_until_complete(coro)
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/nest_asyncio.py", line 70, in run_until_complete
    return f.result()
  File "/home/docs/.pyenv/versions/3.7.9/lib/python3.7/asyncio/futures.py", line 181, in result
    raise self._exception
  File "/home/docs/.pyenv/versions/3.7.9/lib/python3.7/asyncio/tasks.py", line 249, in __step
    result = coro.send(None)
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/nbclient/client.py", line 541, in async_execute
    cell, index, execution_count=self.code_cells_executed + 1
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/nbconvert/preprocessors/execute.py", line 123, in async_execute_cell
    cell, resources = self.preprocess_cell(cell, self.resources, cell_index)
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/nbconvert/preprocessors/execute.py", line 146, in preprocess_cell
    cell = run_sync(NotebookClient.async_execute_cell)(self, cell, index, store_history=self.store_history)
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/nbclient/util.py", line 74, in wrapped
    return just_run(coro(*args, **kwargs))
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/nbclient/util.py", line 53, in just_run
    return loop.run_until_complete(coro)
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/nest_asyncio.py", line 70, in run_until_complete
    return f.result()
  File "/home/docs/.pyenv/versions/3.7.9/lib/python3.7/asyncio/futures.py", line 181, in result
    raise self._exception
  File "/home/docs/.pyenv/versions/3.7.9/lib/python3.7/asyncio/tasks.py", line 249, in __step
    result = coro.send(None)
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/nbclient/client.py", line 832, in async_execute_cell
    self._check_raise_for_error(cell, exec_reply)
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/nbclient/client.py", line 740, in _check_raise_for_error
    raise CellExecutionError.from_cell_and_msg(cell, exec_reply['content'])
nbclient.exceptions.CellExecutionError: An error occurred while executing the following cell:
------------------
x = np.random.randn(3, 4)
y = np.random.randn(4)
f(x, y)
------------------

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-1-1cefe89b730a> in <module>
      1 x = np.random.randn(3, 4)
      2 y = np.random.randn(4)
----> 3 f(x, y)

NameError: name 'f' is not defined


During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/sphinx/cmd/build.py", line 276, in build_main
    app.build(args.force_all, filenames)
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/sphinx/application.py", line 349, in build
    self.builder.build_update()
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/sphinx/builders/__init__.py", line 299, in build_update
    len(to_build))
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/sphinx/builders/__init__.py", line 311, in build
    updated_docnames = set(self.read())
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/sphinx/builders/__init__.py", line 418, in read
    self._read_serial(docnames)
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/sphinx/builders/__init__.py", line 439, in _read_serial
    self.read_doc(docname)
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/sphinx/builders/__init__.py", line 479, in read_doc
    doctree = read_doc(self.app, self.env, self.env.doc2path(docname))
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/sphinx/io.py", line 316, in read_doc
    pub.publish()
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/docutils/core.py", line 218, in publish
    self.settings)
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/sphinx/io.py", line 130, in read
    self.parse()
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/docutils/readers/__init__.py", line 77, in parse
    self.parser.parse(self.input, document)
  File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/5573/lib/python3.7/site-packages/nbsphinx.py", line 1028, in parse
    raise NotebookError('\n'.join(lines))
nbsphinx.NotebookError: CellExecutionError in notebooks/thinking_in_jax.ipynb:
------------------
x = np.random.randn(3, 4)
y = np.random.randn(4)
f(x, y)
------------------

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-1-1cefe89b730a> in <module>
      1 x = np.random.randn(3, 4)
      2 y = np.random.randn(4)
----> 3 f(x, y)

NameError: name 'f' is not defined

You can ignore this error by setting the following in conf.py:

    nbsphinx_allow_errors = True


Notebook error:
CellExecutionError in notebooks/thinking_in_jax.ipynb:
------------------
x = np.random.randn(3, 4)
y = np.random.randn(4)
f(x, y)
------------------

[0;31m---------------------------------------------------------------------------[0m
[0;31mNameError[0m                                 Traceback (most recent call last)
[0;32m<ipython-input-1-1cefe89b730a>[0m in [0;36m<module>[0;34m[0m
[1;32m      1[0m [0mx[0m [0;34m=[0m [0mnp[0m[0;34m.[0m[0mrandom[0m[0;34m.[0m[0mrandn[0m[0;34m([0m[0;36m3[0m[0;34m,[0m [0;36m4[0m[0;34m)[0m[0;34m[0m[0;34m[0m[0m
[1;32m      2[0m [0my[0m [0;34m=[0m [0mnp[0m[0;34m.[0m[0mrandom[0m[0;34m.[0m[0mrandn[0m[0;34m([0m[0;36m4[0m[0;34m)[0m[0;34m[0m[0;34m[0m[0m
[0;32m----> 3[0;31m [0mf[0m[0;34m([0m[0mx[0m[0;34m,[0m [0my[0m[0;34m)[0m[0;34m[0m[0;34m[0m[0m
[0m
[0;31mNameError[0m: name 'f' is not defined
NameError: name 'f' is not defined

You can ignore this error by setting the following in conf.py:

    nbsphinx_allow_errors = True
```

</p>
</details>

The notebook in question is here: https://github.com/google/jax/blob/967f3ac7a335c09739469a18e7f2ad45bbc193df/docs/notebooks/thinking_in_jax.ipynb

Looking at the notebook, it's clear that the name `f` is defined in the previous cell. This appears to be some sort of race condition in cell execution: the cells are nominally executed synchronously, but under the hood nbclient is running cells using [asynchronous calls](https://github.com/jupyter/nbclient/blob/c3898b6df19b1f2fbbf88a6dcf7262c898d71df5/nbclient/client.py#L497-L547) that are synchronized via this utility: https://github.com/jupyter/nbclient/blob/c3898b6df19b1f2fbbf88a6dcf7262c898d71df5/nbclient/util.py#L56-L76

Something about the RTD environment is making this particular cell start running before the previous cell has populated the namespace.


Discussion:
#5635 side-steps the issue by not executing the problematic notebook. I'm going to leave this issue open, though, to track re-enabling the execution if possible.

Resolving patch:
diff --git a/docs/conf.py b/docs/conf.py
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -85,6 +85,8 @@
     'ref.citation',  # Many duplicated citations in numpy/scipy docstrings.
     'ref.footnote',  # Many unreferenced footnotes in numpy/scipy docstrings
     'myst.header',
+    # TODO(jakevdp): remove this suppression once issue is fixed.
+    'misc.highlighting_failure', # https://github.com/ipython/ipython/issues/14142
 ]
 
 # Add any paths that contain templates here, relative to this directory.
@@ -199,8 +201,6 @@
     'notebooks/neural_network_with_tfds_data.*',
     # Slow notebook
     'notebooks/Neural_Network_and_Data_Loading.*',
-    # Strange error apparently due to asynchronous cell execution
-    'notebooks/thinking_in_jax.*',
     # Has extra requirements: networkx, pandas, pytorch, tensorflow, etc.
     'jep/9407-type-promotion.*',
     # TODO(jakevdp): enable execution on the following if possible:


Test patch: